nemo_automodel.recipes.speculative_decoding.train_dspark

DSpark draft-model training recipe (Qwen3, Gemma4, DeepSeek V4, GLM-5.2, and MiniMax M3 VL targets).

DSpark is a semi-autoregressive parallel drafter: a parallel backbone produces a block of tokens per anchor in one pass, a serial Markov head injects intra-block dependency, and a confidence head predicts per-position acceptance. This recipe mirrors the EAGLE / DFlash scaffolding — online target hidden-state capture, gradient accumulation with a trailing-window flush, and the shared checkpointer plumbing — and trains the draft with the three-term DSpark objective.

Module Contents

Classes

Name	Description
`TrainDSparkRecipe`	Recipe for DSpark draft-model training on Qwen3, Gemma4, DeepSeek V4, GLM-5.2, and MiniMax M3 VL targets.
`_DSparkMetricWindow`	Metric sums accumulated between two log points, reduced in one collective.
`_DraftArgs`	Dict with attribute access for the per-architecture draft-config builders.

Functions

Name	Description
`_add_accept_rate_per_position`	Add measured per-position acceptance rates to a metrics dictionary.
`_build_dspark_optimizer`	Build the DSpark trainer’s optimizer from its `optimizer:` config.
`_extract_mm_kwargs`	Return only the multimodal keys present in batch, for `generate_batch(**kwargs)`.
`_resolve_dspark_optimizer_spec`	Normalize the recipe’s `optimizer:` config into a `build_optimizer` spec.
`_validate_cached_dspark_manifest`	Validate that a DSpark offline cache matches the configured target/draft run.
`_validate_draft_packing`	Reject sequence packing for a draft whose forward ignores the packing metadata.
`_validate_remote_target_options`	Reject the recipe options the remote target backend does not support.
`main`	Entrypoint for `TrainDSparkRecipe`.

Data

_DSPARK_MM_KEYS

_DSPARK_WINDOW_SCALARS

logger

API

class nemo_automodel.recipes.speculative_decoding.train_dspark.TrainDSparkRecipe(
    cfg
)

Bases: SpecDecodeRecipeBase

Recipe for DSpark draft-model training on Qwen3, Gemma4, DeepSeek V4, GLM-5.2, and MiniMax M3 VL targets.

min_warmup_steps

= 20

nemo_automodel.recipes.speculative_decoding.train_dspark.TrainDSparkRecipe._build_draft_optimizer(
    opt_cfg
) -> torch.optim.Optimizer

Build the configurable DSpark optimizer (see _build_dspark_optimizer).

nemo_automodel.recipes.speculative_decoding.train_dspark.TrainDSparkRecipe._close_remote_target() -> None

Disconnect from the remote target servers, which then shut down.

Only the main rank disconnects: a server shuts down on the first disconnect it receives, and every rank’s last target request precedes a collective the main rank also joins, so by now no rank still needs the servers.

nemo_automodel.recipes.speculative_decoding.train_dspark.TrainDSparkRecipe._finish_wandb() -> None

nemo_automodel.recipes.speculative_decoding.train_dspark.TrainDSparkRecipe._forward_batch(
    batch
)

Run one batch through live target capture or the offline cache.

nemo_automodel.recipes.speculative_decoding.train_dspark.TrainDSparkRecipe._load_extra_state(
    ckpt_dir: str
) -> None

Restore DSpark meta: global_step and epoch, and validate mask_token_id.

nemo_automodel.recipes.speculative_decoding.train_dspark.TrainDSparkRecipe._maybe_precompute_fp8_scales() -> None

Precompute float8 dynamic scales after an optimizer step (FSDP2 fp8 all-gather only).

nemo_automodel.recipes.speculative_decoding.train_dspark.TrainDSparkRecipe._resolve_mask_token_id(
    recipe_cfg,
    vocab_size: int
) -> int

staticmethod

Resolve and validate the MASK token id filling non-anchor block positions.

The draft’s embed_tokens row at this id is the learned “predict here” signal. It must be a deliberately chosen reserved / unused token id (never a silent fallback to pad, which is commonly aliased to eos), and the inference runtime must fill block slots with the same id.

nemo_automodel.recipes.speculative_decoding.train_dspark.TrainDSparkRecipe._run_eval()

Evaluate the draft on the validation stream.

Reports the loss and the acceptance diagnostics that decide whether the draft is worth serving: the per-position accept_rate@k, its aggregate, the expected accepted block length tau, and the confidence head’s calibration against the measured acceptance. Every batch already computes these (DSparkStepMetrics); training reduces them over a log window and validation over the whole split, both as unreduced numerator/denominator sums so the ratio is formed once, after the data-parallel reduction, rather than averaged over per-rank ratios.

Returns:

The metric dict, or None when no validation dataloader is configured.

nemo_automodel.recipes.speculative_decoding.train_dspark.TrainDSparkRecipe._save_extra_state(
    path: str,
    epoch: int
) -> None

Persist DSpark meta: global_step, epoch, block_size, mask, and target layers.

nemo_automodel.recipes.speculative_decoding.train_dspark.TrainDSparkRecipe._should_shard_dense_target(
    recipe_cfg
) -> bool

Whether to load a frozen dense target FSDP2-sharded via the standard distributed setup.

Opt-in (recipe_args.shard_dense_target, default False). A dense target (Qwen3 / Gemma4) is otherwise loaded whole and replicated on every rank. For a large dense target (e.g. Gemma4-31B) the frozen target is ~62 GiB, leaving no room for the draft’s training activations, so training OOMs at the first backward on 80 GiB GPUs. Loading it through create_distributed_setup_from_config + NeMoAutoModelForCausalLM.from_pretrained(distributed_setup=...) FSDP2-shards it across the mesh, the same path the MoE / VL targets already use.

A small target (e.g. Qwen3-0.6B) stays replicated by default, since sharding a target that already fits is pure all-gather overhead. Requires distributed.strategy='fsdp2' on more than one rank; otherwise the request is ignored with a warning and the target stays replicated.

Raises:

ValueError: if shard_dense_target is requested together with a model-parallel or replication axis (tp_size/pp_size/cp_size/ep_size/ dp_replicate_size > 1). DSpark’s forward-hook hidden-state capture needs one non-pipelined model(...) call per rank (pp_size > 1 builds an AutoPipeline instead of a module), the other model-parallel axes are untested for the frozen dense target here, and HSDP replication re-replicates the target across the replicate dimension, defeating the sharding.

nemo_automodel.recipes.speculative_decoding.train_dspark.TrainDSparkRecipe.run_train_validation_loop()

Run the DSpark training loop.

nemo_automodel.recipes.speculative_decoding.train_dspark.TrainDSparkRecipe.setup()

Build the target model, DSpark draft, data, optimizer, and trainer module.

class nemo_automodel.recipes.speculative_decoding.train_dspark._DSparkMetricWindow(
    block_size: int,
    device: torch.device | None = None,
    loss: float = 0.0,
    ce_loss: float = 0.0,
    l1_loss: float = 0.0,
    confidence_loss: float = 0.0,
    tau_num: float = 0.0,
    tau_den: float = 0.0,
    confidence_abs_error_num: float = 0.0,
    confidence_bias_num: float = 0.0,
    confidence_cumprod_bias_num: float = 0.0,
    confidence_diag_den: float = 0.0,
    num_micro_batches: float = 0.0
)

Dataclass

Metric sums accumulated between two log points, reduced in one collective.

The scalar sums and the two [block_size] per-position accept vectors are concatenated into a single tensor by pack so one all-reduce covers the whole window, and unpack turns the reduced tensor into the metrics to log.

The losses are window means of already normalized per-micro-batch values, so they divide by the micro-batch count. The acceptance diagnostics accumulate as (num, den) sums and divide once after the reduction, which gives the exact global ratio regardless of per-rank token imbalance. A diagnostic whose denominator is zero was not measured this window (e.g. an ablation without the confidence head) and is omitted, so it shows no curve rather than a flat zero that reads like collapsed acceptance.

accept_den

Tensor = field(init=False)

accept_num

Tensor = field(init=False)

block_size

int

ce_loss

float = 0.0

confidence_abs_error_num

float = 0.0

confidence_bias_num

float = 0.0

confidence_cumprod_bias_num

float = 0.0

confidence_diag_den

float = 0.0

confidence_loss

float = 0.0

device

device | None = None

l1_loss

float = 0.0

loss

float = 0.0

num_micro_batches

float = 0.0

tau_den

float = 0.0

tau_num

float = 0.0

nemo_automodel.recipes.speculative_decoding.train_dspark._DSparkMetricWindow.__post_init__() -> None

nemo_automodel.recipes.speculative_decoding.train_dspark._DSparkMetricWindow.add(
    metrics: nemo_automodel.components.speculative.dspark.core.DSparkStepMetrics
) -> None

Accumulate one micro-batch’s outputs.

nemo_automodel.recipes.speculative_decoding.train_dspark._DSparkMetricWindow.pack() -> torch.Tensor

Flatten the window into the 1-D tensor handed to the DP all-reduce.

nemo_automodel.recipes.speculative_decoding.train_dspark._DSparkMetricWindow.reset() -> None

Zero every sum, starting a new window.

nemo_automodel.recipes.speculative_decoding.train_dspark._DSparkMetricWindow.unpack(
    reduced: torch.Tensor
) -> dict[str, float]

Turn the DP-reduced pack tensor into the metrics to log.

class nemo_automodel.recipes.speculative_decoding.train_dspark._DraftArgs()

Bases: dict

Dict with attribute access for the per-architecture draft-config builders.

nemo_automodel.recipes.speculative_decoding.train_dspark._DraftArgs.__getattr__(
    key
)

nemo_automodel.recipes.speculative_decoding.train_dspark._add_accept_rate_per_position(
    metrics: dict[str, float],
    accept_num: torch.Tensor,
    accept_den: torch.Tensor
) -> None

Add measured per-position acceptance rates to a metrics dictionary.

nemo_automodel.recipes.speculative_decoding.train_dspark._build_dspark_optimizer(
    trainer_module,
    opt_cfg,
    device_mesh = None
) -> torch.optim.Optimizer

Build the DSpark trainer’s optimizer from its optimizer: config.

Thin wrapper around build_optimizer so TrainDSparkRecipe._build_draft_optimizer has a single, unit-testable call site (build_optimizer itself needs no distributed environment for a non-pipelined single-part model like the DSpark draft, so this is testable with a plain CPU module).

nemo_automodel.recipes.speculative_decoding.train_dspark._extract_mm_kwargs(
    batch: dict
) -> dict

Return only the multimodal keys present in batch, for generate_batch(**kwargs).

Empty for a text-only batch (Qwen3, Gemma4, or MiniMax M3 without multimodal: true), so the generate_batch call is unchanged in that case.

nemo_automodel.recipes.speculative_decoding.train_dspark._resolve_dspark_optimizer_spec(
    opt_cfg
) -> tuple[str, dict]

Normalize the recipe’s optimizer: config into a build_optimizer spec.

Reads an optional _target_ (a registry short name such as "fused_adam" or a dotted import path, e.g. transformer_engine.pytorch.optimizers.FusedAdam) plus whatever other fields the config carries — lr/betas/weight_decay and any optimizer-specific kwargs (master_weights, master_weight_dtype, exp_avg_dtype, exp_avg_sq_dtype, store_param_remainders, …) — and returns the (target, kwargs) tuple that build_optimizer resolves via its registry / dotted-import-path / OptimizerFromFactoryConfig escape hatch.

Absent an explicit _target_, this defaults to plain torch.optim.AdamW with its prior betas/weight_decay defaults (matching the previous hardcoded behavior, so existing DSpark configs are unaffected). Those two AdamW-shaped defaults are only injected in that no-_target_ case: forcing them onto an arbitrary explicit _target_ would break optimizers that do not accept a betas kwarg (e.g. plain SGD).

nemo_automodel.recipes.speculative_decoding.train_dspark._validate_cached_dspark_manifest(
    cache_dir: str,
    manifest: dict,
    target_config,
    target_layer_ids: list[int],
    target_model: str,
    target_model_type: str,
    seq_length: int,
    compute_dtype: torch.dtype,
    chat_template_sha256: str | None,
    chat_template_kwargs: dict[str, typing.Any] | None
) -> None

Validate that a DSpark offline cache matches the configured target/draft run.

chat_template_sha256 is the hash of the template the run configured, or None when it configured none and so has no framing to compare against the cache’s; a mismatch only warns, since an edit to a template comment changes the hash too. chat_template_kwargs are the template arguments the run configured (None means none, i.e. {}); like the mask options they shape the cached supervision, so a mismatch raises.

nemo_automodel.recipes.speculative_decoding.train_dspark._validate_draft_packing(
    draft_model: torch.nn.Module
) -> None

Reject sequence packing for a draft whose forward ignores the packing metadata.

Such a draft would silently let anchors cross document boundaries.

nemo_automodel.recipes.speculative_decoding.train_dspark._validate_remote_target_options(
    cached_target_path: str | None,
    is_multimodal: bool,
    packed_sequence_size: int,
    cp_size: int
) -> None

Reject the recipe options the remote target backend does not support.

The remote server receives only input_ids / attention_mask / loss_mask and runs one unsharded forward per request, so it cannot serve image inputs, the block-causal packing metadata, or a context-parallel target. A precomputed cache already replaces the target, so it cannot be combined with a remote one.

nemo_automodel.recipes.speculative_decoding.train_dspark.main(
    config_path: str | None = None
)

Entrypoint for TrainDSparkRecipe.

nemo_automodel.recipes.speculative_decoding.train_dspark._DSPARK_MM_KEYS = tuple(k for k in VLM_INPUT_KEYS if k != 'input_ids')

nemo_automodel.recipes.speculative_decoding.train_dspark._DSPARK_WINDOW_SCALARS = ('loss', 'ce_loss', 'l1_loss', 'confidence_loss', 'tau_num', 'tau_den', 'confide...

nemo_automodel.recipes.speculative_decoding.train_dspark.logger = logging.getLogger(__name__)