source/src/speculators/train/cli.py

"""Training entrypoint — core logic moved from scripts/train.py."""

import argparse
import gc
import logging
import random
import warnings
from copy import deepcopy
from pathlib import Path

import numpy as np
import torch
import torch.distributed as dist
import transformers
from packaging import version
from transformers import LlamaConfig, PretrainedConfig
from transformers.models.auto.configuration_auto import AutoConfig
from transformers.models.qwen3.configuration_qwen3 import Qwen3Config

from hs_connectors import HiddenStatesBackend
from speculators.model import SpeculatorModel
from speculators.models.eagle3.data import shift_batch
from speculators.models.eagle3.rotary_partial import install_partial_neox_rotary
from speculators.models.mtp.data import shift_batch_mtp
from speculators.models.utils import (
    get_verifier_config,
    resolve_draft_intermediate_size,
)
from speculators.train.config import TrainConfig
from speculators.train.dataloader import create_train_val_loaders
from speculators.train.distributed import (
    get_rank,
    is_distributed,
    maybe_destroy_distributed,
    maybe_setup_distributed,
)
from speculators.train.logger import (
    log_run_config,
    setup_metric_logger,
    setup_root_logger,
)
from speculators.train.trainer import Trainer, TrainerConfig
from speculators.train.utils import resolve_mask_token_id
from speculators.train.vocab_mapping import (
    build_vocab_mappings_from_distribution,
    get_target_vocab_size,
)
from speculators.utils.loading import is_config_only_dir

logger = logging.getLogger(__name__)

DRAFT_ARCH_CONFIGS: dict[str, type] = {
    "llama": LlamaConfig,
    "qwen3": Qwen3Config,
}
MROPE_INVERSE_TOLERANCE = 1e-6


def set_seed(seed: int, deterministic: bool = False):
    """Set random seeds for reproducibility."""
    random.seed(seed)
    np.random.seed(seed)  # noqa: NPY002
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)

    if deterministic:
        # For deterministic behavior (may impact performance)
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False


def _maybe_apply_mrope_full_head_hack(
    rope_params: dict,
    resolved_head_dim: int,
    enabled: bool,
) -> None:
    """Optionally rescale partial MRoPE settings to full-head semantics."""
    if "mrope_section" not in rope_params:
        return

    inherited_partial = float(rope_params.get("partial_rotary_factor", 1.0))
    if enabled and inherited_partial < 1.0:
        old_section = list(rope_params["mrope_section"])
        inv = 1.0 / inherited_partial
        if abs(inv - round(inv)) > MROPE_INVERSE_TOLERANCE:
            raise ValueError(
                "mrope_full_head_hack cannot rescale mrope_section because "
                f"1/partial_rotary_factor={inv} is not an integer."
            )
        scale = int(round(inv))
        new_section = [int(x) * scale for x in old_section]
        if 2 * sum(new_section) != resolved_head_dim:
            raise ValueError(
                "mrope_full_head_hack rescaling produced inconsistent "
                f"mrope_section {new_section}: 2*sum={2 * sum(new_section)} "
                f"but head_dim={resolved_head_dim}."
            )
        rope_params["mrope_section"] = new_section
        rope_params["partial_rotary_factor"] = 1.0
        logger.warning(
            "MRoPE full-head hack applied: partial_rotary_factor "
            f"{inherited_partial} -> 1.0, mrope_section {old_section} -> "
            f"{new_section}."
        )
    elif not enabled and inherited_partial < 1.0:
        logger.warning(
            "mrope_full_head_hack=False with partial_rotary_factor="
            f"{inherited_partial} < 1.0 can cause HF trainer / vLLM "
            "partial-rotation mismatch."
        )


def create_transformer_layer_config(  # noqa: C901
    verifier_name_or_path: str,
    num_layers: int,
    draft_arch: str,
    hidden_act: str | None,
    sliding_window: int,
    full_attention_indices: list[int],
    mrope_full_head_hack: bool = True,
    trust_remote_code: bool = False,
) -> PretrainedConfig:
    if draft_arch not in DRAFT_ARCH_CONFIGS:
        raise ValueError(
            f"Unknown draft architecture: {draft_arch}. "
            f"Available: {list(DRAFT_ARCH_CONFIGS.keys())}"
        )

    if draft_arch not in ("llama", "qwen3"):
        warnings.warn(
            f"Draft architecture '{draft_arch}' is not yet supported in vLLM. "
            "The trained model may not be usable for inference in vLLM. "
            "Consider using 'llama' or 'qwen3' for full vLLM compatibility.",
            stacklevel=2,
        )

    config_class = DRAFT_ARCH_CONFIGS[draft_arch]
    verifier_config = AutoConfig.from_pretrained(
        verifier_name_or_path,
        trust_remote_code=trust_remote_code,
    )

    # For multimodal models (Qwen3VL, etc.), extract text_config
    if hasattr(verifier_config, "text_config"):
        verifier_config = verifier_config.text_config

    hidden_act = (
        hidden_act
        or getattr(verifier_config, "hidden_act", None)
        or getattr(verifier_config, "hidden_activation", None)
    )
    if hidden_act is None:
        raise AttributeError(
            f"{type(verifier_config).__name__} has neither 'hidden_act' "
            "nor 'hidden_activation'"
        )

    head_dim = getattr(verifier_config, "head_dim", None)
    num_attention_heads = verifier_config.num_attention_heads
    num_key_value_heads = verifier_config.num_key_value_heads

    if (
        head_dim
        and verifier_config.hidden_size % num_attention_heads != 0
        and verifier_config.hidden_size % head_dim == 0
    ):
        num_attention_heads = verifier_config.hidden_size // head_dim
        if num_attention_heads % num_key_value_heads != 0:
            num_key_value_heads = num_attention_heads
    resolved_head_dim = head_dim or verifier_config.hidden_size // num_attention_heads

    if full_attention_indices and (
        min(full_attention_indices) < 0 or max(full_attention_indices) >= num_layers
    ):
        raise ValueError(
            "Full attention indices must be valid draft layer ids "
            "in range [0, num_layers)."
        )
    layer_types = [
        "full_attention" if i in full_attention_indices else "sliding_attention"
        for i in range(num_layers)
    ]

    config = config_class(
        vocab_size=verifier_config.vocab_size,
        hidden_size=verifier_config.hidden_size,
        intermediate_size=resolve_draft_intermediate_size(verifier_config),
        num_hidden_layers=num_layers,
        num_attention_heads=num_attention_heads,
        num_key_value_heads=num_key_value_heads,
        hidden_act=hidden_act,
        max_position_embeddings=verifier_config.max_position_embeddings,
        initializer_range=verifier_config.initializer_range,
        rms_norm_eps=verifier_config.rms_norm_eps,
        head_dim=head_dim,
        tie_word_embeddings=False,
        sliding_window=sliding_window,
        use_sliding_window="sliding_attention" in layer_types,
        layer_types=layer_types,
    )

    # New rope parameters definition introduced in transformers 5.0
    if version.parse(transformers.__version__) >= version.parse("5.0.0"):
        if hasattr(verifier_config, "rope_parameters"):
            rope_params = deepcopy(verifier_config.rope_parameters)
            # Some verifiers (e.g. Laguna) use the nested per-layer-type rope format
            # {"full_attention": {...}, "sliding_attention": {...}} with no top-level
            # "rope_theta". The llama-style draft uses a single rope, so collapse it to
            # a flat default rope (preferring the sliding-attention theta).
            if isinstance(rope_params, dict) and "rope_theta" not in rope_params:
                sub = (
                    rope_params.get("sliding_attention")
                    or rope_params.get("full_attention")
                    or {}
                )
                if isinstance(sub, dict):
                    rope_params = dict(sub)
                    rope_params.setdefault("rope_type", "default")
                    rope_params.setdefault("rope_theta", 10000.0)
                else:
                    rope_params = {"rope_type": "default", "rope_theta": 10000.0}

            if isinstance(rope_params, dict):
                _maybe_apply_mrope_full_head_hack(
                    rope_params, resolved_head_dim, mrope_full_head_hack
                )
                # ``type`` is a legacy alias (only "mrope" on VL models) that
                # transformers strips during validation and that breaks vLLM's
                # config checks; drop it while keeping the real MRoPE fields.
                rope_params.pop("type", None)
                rope_params.pop("mrope_interleaved", None)
                # The verifier (e.g. Mistral) may use partial rotary embeddings,
                # but the draft model doesn't support partial_rotary_factor.
                # Only keep it for MRoPE configs that need it.
                if "mrope_section" not in rope_params:
                    rope_params.pop("partial_rotary_factor", None)
            config.rope_parameters = rope_params
    else:
        if hasattr(verifier_config, "rope_scaling"):
            rope_scaling = deepcopy(verifier_config.rope_scaling)
            if isinstance(rope_scaling, dict):
                _maybe_apply_mrope_full_head_hack(
                    rope_scaling, resolved_head_dim, mrope_full_head_hack
                )
                # Strip legacy fields for consistency with rope_parameters path
                rope_scaling.pop("type", None)
                rope_scaling.pop("mrope_interleaved", None)
                # Same partial_rotary_factor guard as the rope_parameters path.
                if "mrope_section" not in rope_scaling:
                    rope_scaling.pop("partial_rotary_factor", None)
            config.rope_scaling = rope_scaling
        config.rope_theta = getattr(verifier_config, "rope_theta", 10000.0)

    return config


def load_draft_transformer_layer_config(
    draft_config: str,
    verifier_name_or_path: str,
    trust_remote_code: bool = False,
) -> PretrainedConfig:
    """Load the draft decoder ``transformer_layer_config`` from a config source.

    ``draft_config`` may be a HF hub id, a local directory containing a
    ``config.json``, or a path to a config JSON file. It is expected to hold a
    plain decoder config (``LlamaConfig`` for eagle3/peagle, ``Qwen3Config`` for
    dflash). If a full speculator config is given instead, its nested
    ``transformer_layer_config`` is extracted as a convenience.

    The decoder is reconciled against the verifier: ``hidden_size`` must match
    (draft/verifier hidden-size mismatch is not yet supported) and ``vocab_size``
    is aligned to the verifier's target vocabulary. The pruned draft vocabulary
    is controlled separately via ``--draft-vocab-size``.
    """
    config_dict, _ = PretrainedConfig.get_config_dict(draft_config)
    if "transformer_layer_config" in config_dict:
        # A full speculator config was passed; use only the decoder definition.
        config_dict = config_dict["transformer_layer_config"]

    model_type = config_dict.get("model_type")
    if not model_type:
        raise ValueError(
            "--draft-config must define a 'model_type' (e.g. 'llama' for "
            "eagle3/peagle, 'qwen3' for dflash); none was found in the config "
            f"loaded from '{draft_config}'."
        )
    config_class: type[PretrainedConfig] = type(AutoConfig.for_model(model_type))
    draft_config_obj = config_class.from_dict(config_dict)

    verifier_config = get_verifier_config(
        verifier_name_or_path,
        trust_remote_code=trust_remote_code,
    )
    if draft_config_obj.hidden_size != verifier_config.hidden_size:
        raise ValueError(
            f"--draft-config hidden_size ({draft_config_obj.hidden_size}) must match "
            f"the verifier hidden_size ({verifier_config.hidden_size}). Draft/verifier "
            "hidden-size mismatch is not yet supported."
        )
    if draft_config_obj.vocab_size != verifier_config.vocab_size:
        logger.warning(
            "Overriding --draft-config vocab_size (%s) with the verifier vocab_size "
            "(%s). Use --draft-vocab-size to control the pruned draft vocabulary.",
            draft_config_obj.vocab_size,
            verifier_config.vocab_size,
        )
        draft_config_obj.vocab_size = verifier_config.vocab_size
    return draft_config_obj


def _load_mappings(d2t_path, t2d_path, expected_draft_vocab_size: int | None):
    logger.info(f"Loading vocab mappings from '{d2t_path}' and '{t2d_path}'")
    # Load d2t and t2d tensors if provided
    d2t = torch.from_numpy(np.load(d2t_path))
    t2d = torch.from_numpy(np.load(t2d_path))
    draft_vocab_size = d2t.shape[0]
    if (
        expected_draft_vocab_size is not None
        and expected_draft_vocab_size != draft_vocab_size
    ):
        raise ValueError(
            f"Vocab mapping d2t has size {draft_vocab_size}, but "
            f"--draft-vocab-size requires {expected_draft_vocab_size}."
        )
    return d2t, t2d, draft_vocab_size


def _generate_vocab_mappings(args: argparse.Namespace, data_path: Path):
    token_freq_path = args.token_freq_path or data_path / "token_freq.pt"
    token_freq_path = Path(token_freq_path)
    if not token_freq_path.exists():
        return None

    logger.info("No vocab mappings provided. Regenerating from token frequencies")
    token_freq_dict = torch.load(token_freq_path, weights_only=True)

    target_vocab_size = get_target_vocab_size(
        None,
        args.verifier_name_or_path,
        trust_remote_code=args.trust_remote_code,
    )

    d2t, t2d = build_vocab_mappings_from_distribution(
        token_freq_dict=token_freq_dict,
        draft_vocab_size=args.draft_vocab_size,
        target_vocab_size=target_vocab_size,
    )
    draft_vocab_size = d2t.shape[0]
    if args.draft_vocab_size != draft_vocab_size:
        raise ValueError(
            f"Generated vocab mapping d2t has size {draft_vocab_size}, but "
            f"--draft-vocab-size requires {args.draft_vocab_size}."
        )

    d2t_cache_path = data_path / f"d2t-{draft_vocab_size}.npy"
    t2d_cache_path = data_path / f"t2d-{draft_vocab_size}.npy"
    logger.info(
        f"Caching vocab mapping files to '{d2t_cache_path}' and '{t2d_cache_path}'"
    )
    np.save(d2t_cache_path, d2t.cpu().numpy())
    np.save(t2d_cache_path, t2d.cpu().numpy())

    return d2t, t2d, draft_vocab_size


def parse_vocab_mappings(args: argparse.Namespace):
    if (args.d2t_path or args.t2d_path) and not (args.d2t_path and args.t2d_path):
        raise ValueError(
            "Both t2d and d2t must be provided together, or both must be omitted. "
            f"Got t2d={'provided' if args.t2d_path is not None else 'not provided'}"
            f"d2t={'provided' if args.d2t_path is not None else 'not provided'}"
        )

    result = None
    if not is_distributed() or get_rank() == 0:
        if args.d2t_path or args.t2d_path:
            result = _load_mappings(args.d2t_path, args.t2d_path, args.draft_vocab_size)
        else:
            data_path = Path(args.data_path)
            if args.draft_vocab_size is not None:
                default_t2d_path = data_path / f"t2d-{args.draft_vocab_size}.npy"
                default_d2t_path = data_path / f"d2t-{args.draft_vocab_size}.npy"
                if default_t2d_path.exists() and default_d2t_path.exists():
                    result = _load_mappings(
                        default_d2t_path, default_t2d_path, args.draft_vocab_size
                    )
                else:
                    result = _generate_vocab_mappings(args, data_path)

            if result is None:
                token_freq_path = args.token_freq_path or data_path / "token_freq.pt"
                token_freq_path = Path(token_freq_path)

                logger.warning(
                    "No vocab mappings found, and can't generate new ones because "
                    f"either token_freq_path='{token_freq_path}' doesn't exist or "
                    "--draft-vocab-size is None. Using full verifier vocab"
                )
                # When vocab mapping is not provided, use the full verifier vocab
                verifier_config = AutoConfig.from_pretrained(
                    args.verifier_name_or_path,
                    trust_remote_code=args.trust_remote_code,
                )
                if hasattr(verifier_config, "text_config"):
                    verifier_config = verifier_config.text_config
                result = None, None, verifier_config.vocab_size

    if is_distributed():
        payload = [result]
        dist.broadcast_object_list(payload, src=0)
        result = payload[0]
    return result


def _build_from_config_only(
    model_class: type[SpeculatorModel],
    path: str,
    t2d: torch.Tensor | None,
    d2t: torch.Tensor | None,
    verifier_name_or_path: str | None = None,
    draft_attn_impl: str | None = None,
) -> SpeculatorModel:
    """Initialize a fresh draft from a saved speculator *config* (no weights).

    Mirrors the tail of ``from_training_args``: build the model from the full
    speculator config, load vocab mappings, and pull verifier weights -- but with
    no trained draft weights to restore (decoder weights are randomly initialized).
    """
    config = model_class.config_class.from_pretrained(path)
    if draft_attn_impl is not None:
        config.transformer_layer_config._attn_implementation = draft_attn_impl  # noqa: SLF001
    speculators_config = getattr(config, "speculators_config", None)
    # Fall back to the CLI --verifier-name-or-path only when the saved config has
    # no verifier path -- either null or blanked to "". A real path in the config
    # takes precedence and the CLI value is ignored.
    if (
        verifier_name_or_path
        and speculators_config is not None
        and not getattr(speculators_config.verifier, "name_or_path", None)
    ):
        speculators_config.verifier.name_or_path = verifier_name_or_path
    model = model_class(config=config)
    if hasattr(model, "load_vocab_mappings"):
        model.load_vocab_mappings(t2d, d2t)  # type: ignore[attr-defined, operator]
    if hasattr(model, "load_verifier_weights"):
        model.load_verifier_weights()  # type: ignore[attr-defined, operator]
    return model


def build_draft_model(
    args: argparse.Namespace,
    model_class: type[SpeculatorModel],
    t2d: torch.Tensor | None,
    d2t: torch.Tensor | None,
    draft_vocab_size: int | None,
) -> SpeculatorModel:
    """Resolve the draft model from one of these sources:

    * ``--from-pretrained``: finetune existing weights, or -- when the path is a
      config-only directory -- initialize fresh weights from a full saved
      speculator config.
    * ``--draft-config``: take the decoder ``transformer_layer_config`` from a
      config file; build the rest of the speculator from the other CLI args.
    * neither: synthesize the decoder from the verifier config + CLI flags.

    MTP is special-cased: when not loading ``--from-pretrained``, it reuses the
    verifier's own decoder config as the draft ``transformer_layer_config`` and
    extracts the native MTP head weights from the verifier, so the decoder-shaping
    flags and ``--draft-config`` do not apply.
    """
    if args.from_pretrained:
        if is_config_only_dir(args.from_pretrained):
            logger.info(
                "--from-pretrained points to a config-only directory ('%s'); "
                "initializing fresh draft weights from the saved speculator config.",
                args.from_pretrained,
            )
            return _build_from_config_only(
                model_class,
                args.from_pretrained,
                t2d=t2d,
                d2t=d2t,
                verifier_name_or_path=args.verifier_name_or_path,
                draft_attn_impl=(
                    args.draft_attn_impl if args.speculator_type != "mtp" else None
                ),
            )
        if args.speculator_type != "mtp":
            # _attn_implementation is never serialized by HF configs, so re-apply
            # the CLI selection before construction -- mirroring from_training_args.
            # MTP is skipped: its from_training_args never sets the field and its
            # __init__ resolves its own default ("sdpa") when it is absent.
            config = model_class.config_class.from_pretrained(args.from_pretrained)
            config.transformer_layer_config._attn_implementation = args.draft_attn_impl  # noqa: SLF001
            return model_class.from_pretrained(
                args.from_pretrained,
                config=config,
                t2d=t2d,
                d2t=d2t,
                verifier=args.verifier_name_or_path,
            )
        return model_class.from_pretrained(
            args.from_pretrained,
            t2d=t2d,
            d2t=d2t,
            verifier=args.verifier_name_or_path,
        )

    if args.speculator_type == "mtp":
        # MTP uses the verifier's own decoder config as the draft
        # transformer_layer_config and extracts the native MTP head weights from
        # the verifier; the decoder-shaping flags and --draft-config do not apply,
        # and there is no draft mask token to resolve.
        transformer_layer_config = get_verifier_config(
            args.verifier_name_or_path,
            trust_remote_code=getattr(args, "trust_remote_code", False),
        )
    else:
        if args.draft_config:
            transformer_layer_config = load_draft_transformer_layer_config(
                args.draft_config,
                args.verifier_name_or_path,
                trust_remote_code=args.trust_remote_code,
            )
        else:
            full_attention_indices = args.full_attention_indices
            if not full_attention_indices:
                logger.info(
                    "All %d draft layers using sliding window attention "
                    "(window=%d). To use full attention on specific layers, "
                    "pass '--full-attention-indices <layer_ids>'.",
                    args.num_layers,
                    args.sliding_window,
                )

            transformer_layer_config = create_transformer_layer_config(
                verifier_name_or_path=args.verifier_name_or_path,
                num_layers=args.num_layers,
                draft_arch=args.draft_arch,
                hidden_act=args.draft_hidden_act,
                sliding_window=args.sliding_window,
                full_attention_indices=full_attention_indices,
                mrope_full_head_hack=args.draft_mrope_full_head_hack,
                trust_remote_code=args.trust_remote_code,
            )

        args.mask_token_id = resolve_mask_token_id(
            args.verifier_name_or_path,
            transformer_layer_config.vocab_size,
            args.mask_token_id,
            trust_remote_code=args.trust_remote_code,
        )

    args.draft_vocab_size = draft_vocab_size
    return model_class.from_training_args(
        verifier_config=transformer_layer_config,
        t2d=t2d,
        d2t=d2t,
        **vars(args),
    )


def main(cfg: TrainConfig):  # noqa: C901
    # Phase-1 adapter: the model layer still consumes a flat vars(args)-shaped
    # dict via **kwargs, so flatten the typed config back into a namespace here.
    # New code should read cfg.<group>.<field> directly and must NOT add new
    # args.* accesses below this line.
    args = argparse.Namespace(**cfg.flatten())

    # Set random seed for reproducibility
    set_seed(args.seed, args.deterministic_cuda)

    # Setup logging
    setup_root_logger()
    setup_metric_logger(
        loggers=args.logger, run_name=args.run_name, output_dir=args.log_dir
    )

    # Setup distributed training
    maybe_setup_distributed()

    # Publish train config to metric backends that support hyperparameter logging
    log_run_config(cfg)

    if args.fsdp_shard and not is_distributed():
        raise ValueError(
            "--fsdp-shard requires launching with torchrun/distributed training; "
            "otherwise parameters are not sharded."
        )

    # Install partial-neox rotary patch if not using full-head hack
    if not args.draft_mrope_full_head_hack:
        install_partial_neox_rotary()
        logger.info(
            "Installed partial-neox rotary patch for HF/vLLM RoPE alignment "
            "(draft_mrope_full_head_hack=False)"
        )
    # Write the reproducibility artifacts (run.yaml + train_command.txt) next to
    # the checkpoints at rank 0 only, so every checkpoint carries the resolved
    # config that produced it.
    if get_rank() == 0:
        cfg.save(args.save_path)

    hidden_states_dtype = getattr(torch, args.hidden_states_dtype)

    if args.speculator_type == "mtp":
        if args.draft_attn_impl != "simple_flex_attention":
            raise ValueError(
                "--draft-attn-impl is not configurable for MTP. "
                "Must be left with the default value ('simple_flex_attention')."
            )
        # MTP reuses the verifier's own decoder as the draft and extracts the
        # native MTP head weights from the verifier, so there are no vocab
        # mappings or draft mask token to resolve from the CLI. This works both
        # with --from-pretrained (a previously converted checkpoint) and without
        # it (weights are extracted from the verifier on the fly). The decoder
        # transformer_layer_config is resolved later in build_draft_model.
        d2t, t2d, draft_vocab_size = None, None, None
        args.mask_token_id = None
    else:
        d2t, t2d, draft_vocab_size = parse_vocab_mappings(args)

        if args.full_attention_indices and args.speculator_type == "mtp":
            raise ValueError(
                "--full-attention-indices is not supported for mtp draft models."
            )

    registry = SpeculatorModel.registry
    if registry is None or args.speculator_type not in registry:
        available = list(registry.keys()) if registry else []
        raise ValueError(
            f"Unknown speculator type: {args.speculator_type}. Available: {available}"
        )

    model_class = registry[args.speculator_type]

    draft_model = build_draft_model(args, model_class, t2d, d2t, draft_vocab_size)

    # Get target layer IDs from the model (resolved at model level)
    num_target_layers = len(draft_model.target_layer_ids)  # type: ignore[arg-type]

    if args.speculator_type == "mtp":
        args.num_speculative_steps = draft_model.config.num_speculative_steps

    # Dry-run: persist an initialized checkpoint and exit before training so the
    # config/weights can be validated (e.g. in vLLM). The saved checkpoint can be
    # fed straight back via --from-pretrained to start training.
    if args.dry_run:
        # Save in hidden_states_dtype (bf16) for compact checkpoints.
        draft_model.to(hidden_states_dtype)
        if get_rank() == 0:
            logger.info(
                "[dry-run] Saving initialized checkpoint (%s) to '%s'",
                hidden_states_dtype,
                args.save_path,
            )
            draft_model.save_pretrained(args.save_path)
            logger.info(
                "[dry-run] Done. Validate this checkpoint, then train with "
                "'--from-pretrained %s'.",
                args.save_path,
            )
        maybe_destroy_distributed()
        return

    hidden_size = draft_model.config.transformer_layer_config.hidden_size

    # Setup dataloaders
    preprocess_fns = {
        "eagle3": shift_batch,
        "peagle": shift_batch,
        "mtp": shift_batch_mtp,
    }
    preprocess = preprocess_fns.get(args.speculator_type)

    backend_registry = HiddenStatesBackend.registry
    backend_cls = backend_registry[args.hidden_states_backend]
    # from_train_args is the live runtime consumer of the backend's (mirrored)
    # train-args, read off the flattened namespace. hs_connectors stays
    # argparse-based and standalone so vLLM can use it without speculators; that
    # is why the backend's train-args are mirrored into the pydantic schema rather
    # than the plugin depending on pydantic. test_backend_reconciliation.py keeps
    # the mirror complete so nothing read here was dropped during resolution.
    transfer = backend_cls.from_train_args(args, args.data_path)

    train_loader, val_loader = create_train_val_loaders(
        data_path=args.data_path,
        total_seq_len=args.total_seq_len,
        hidden_states_dtype=hidden_states_dtype,
        noise_std=args.noise_std,
        transfer=transfer,
        vllm_endpoint=args.vllm_endpoint,
        on_missing=args.on_missing,
        on_generate=args.on_generate,
        verifier_name_or_path=args.verifier_name_or_path,
        request_timeout=args.request_timeout,
        max_retries=args.max_retries,
        generation_validation_retries=args.generation_validation_retries,
        max_consecutive_generation_failures=args.max_consecutive_generation_failures,
        fail_on_hidden_state_error=args.fail_on_hidden_state_error,
        fetch_threads=args.fetch_threads,
        dataloader_in_order=args.dataloader_in_order,
        vllm_http_keepalive=args.vllm_http_keepalive,
        hidden_size=hidden_size,
        num_target_layers=num_target_layers,
        num_workers=args.num_workers,
        prefetch_factor=args.prefetch_factor,
        preprocess=preprocess,
        train_data_ratio=args.train_data_ratio,
    )

    # Get trainer kwargs from model class
    train_call_kwargs, val_call_kwargs = model_class.get_trainer_kwargs(**vars(args))

    trainer_config = TrainerConfig(
        num_epochs=args.epochs,
        save_path=args.save_path,
        lr=args.lr,
        resume_from_checkpoint=not args.no_resume_from_checkpoint,
        train_call_kwargs=train_call_kwargs,
        val_call_kwargs=val_call_kwargs,
        optimizer=args.optimizer,
        weight_decay=args.weight_decay,
        weight_decay_exclude_1d=args.weight_decay_exclude_1d,
        muon_lr=args.muon_lr,
        muon_momentum=args.muon_momentum,
        muon_weight_decay=args.muon_weight_decay,
        muon_ns_steps=args.muon_ns_steps,
        muon_adjust_lr_fn=args.muon_adjust_lr_fn,
        scheduler_type=args.scheduler_type,
        scheduler_warmup_steps=args.scheduler_warmup_steps,
        scheduler_warmup_ratio=args.scheduler_warmup_ratio,
        scheduler_total_steps=args.scheduler_total_steps,
        scheduler_num_cosine_cycles=args.scheduler_num_cosine_cycles,
        checkpoint_freq=args.checkpoint_freq,
        save_best=args.save_best,
        hidden_states_dtype=hidden_states_dtype,
        log_freq=args.log_freq,
        fsdp_shard=args.fsdp_shard,
        max_steps=args.max_steps,
    )
    trainer = Trainer(draft_model, trainer_config, train_loader, val_loader)

    # Run training
    trainer.run_training()

    # Cleanup
    del trainer, draft_model
    gc.collect()
    acc = torch.accelerator.current_accelerator(check_available=True)
    if acc is not None:
        torch.get_device_module(acc.type).empty_cache()
    maybe_destroy_distributed()