nemo_automodel.components.speculative.dflash.core

DFlash online training wrapper.

Ported from SpecForge’s specforge/core/dflash.py. DFlashTrainerModule samples a set of anchor positions per sequence, builds one parallel draft block per anchor (the block’s first token is the real anchor token, the rest are MASK), runs the draft model under a bespoke block attention mask, and computes a block-wise cross-entropy loss against the ground-truth continuation of each anchor.

Two training objectives are supported via loss_type:

"dflash" (default): the DFlash paper’s fixed-anchor objective. Only block position 0 is a real token; positions 1..block_size-1 are supervised with the decay-weighted CE of Eq. 4 (w_k = exp(-(k-1)/gamma)).
"variable_prefix": the D2SD VP-Drafter objective (arXiv:2606.04446). Each block draws a visible-prefix length l from a truncated geometric prior (Pr(l) ~ prefix_weight_base ** l), positions < l are filled with the real tokens, and only the masked positions >= l are supervised, with the decay re-anchored at the prefix boundary (w_k = exp(-(k-l)/gamma)). This trains the draft to re-draft block suffixes behind a variable accepted prefix, the regime a variable-prefix drafter sees at inference time.

Module Contents

Classes

Name	Description
`DFlashStepMetrics`	Per-step training outputs for the DFlash draft.
`DFlashTrainerModule`	DFlash online training wrapper with block-wise CE loss.
`NoValidAnchorsError`	Raised when a batch has no sample long enough to form a DFlash block.

Functions

Name	Description
`_context_doc_ids`	Per-context-token document id `[B, S]` from packed `seq_lens` `[B, max_docs]`.
`_to_full_tensor`	Materialise a (possibly tensor-parallel) tensor as a plain local tensor.

Data

_DFLASH_LOSS_TYPES

API

class nemo_automodel.components.speculative.dflash.core.DFlashStepMetrics(
    loss: torch.Tensor,
    accuracy: torch.Tensor,
    valid_tokens: torch.Tensor
)

Dataclass

Per-step training outputs for the DFlash draft.

accuracy

Tensor

loss

Tensor

valid_tokens

Tensor

class nemo_automodel.components.speculative.dflash.core.DFlashTrainerModule(
    draft_model: nemo_automodel.components.speculative.dflash.draft_qwen3.Qwen3DFlashDraftModel,
    target_lm_head: torch.nn.Module,
    target_embed_tokens: torch.nn.Module,
    mask_token_id: int,
    block_size: int = 16,
    attention_backend: str = 'flex_attention',
    num_anchors: int = 512,
    loss_decay_gamma: typing.Optional[float] = None,
    loss_type: str = 'dflash',
    prefix_weight_base: float = 0.9
)

Bases: Module

DFlash online training wrapper with block-wise CE loss.

_min_prefix

= min(2, block_size - 1)

loss_fn

prefix_weight_base

= float(prefix_weight_base)

nemo_automodel.components.speculative.dflash.core.DFlashTrainerModule._build_block_targets(
    input_ids: torch.Tensor,
    loss_mask: torch.Tensor,
    anchor_positions: torch.Tensor,
    block_keep_mask: torch.Tensor,
    seq_len: int,
    label_start: int = 0,
    doc_remaining: torch.Tensor | None = None
) -> typing.Tuple[torch.Tensor, torch.Tensor, torch.Tensor]

Per-block ground-truth tokens and the supervised-position mask.

Returns (label_indices, target_ids, block_mask) each of shape [B, N, block_size]. label_indices[..., k] is the sequence position block position k predicts (anchor + label_start + k); block_mask is the product of block validity, in-bounds, and the gathered loss mask. Shared by the block-wise trainers (DFlash, JetSpec, and Domino, which passes label_start=1 for shift_label) so the label/mask gathering lives in one place. Under packing, doc_remaining [B, S] truncates labels at the anchor’s document boundary: anchor sampling only keeps offsets up to block_size - 1 inside the anchor’s document, and a shifted label window (label_start > 0) reaches one past that guarantee.

nemo_automodel.components.speculative.dflash.core.DFlashTrainerModule._create_noise_embed(
    input_ids,
    anchor_positions,
    block_keep_mask
)

Embed each block as [anchor_token, MASK, MASK, ...] (invalid blocks all MASK).

nemo_automodel.components.speculative.dflash.core.DFlashTrainerModule._create_position_ids(
    anchor_positions: torch.Tensor,
    context_position_ids: typing.Optional[torch.Tensor] = None
) -> torch.Tensor

Position ids for the parallel draft blocks (anchor position + offset).

Without packing the anchor’s position equals its row index, so the block positions are anchor + offset. Under packing context_position_ids [B, S] holds per-document reset positions, so the block’s base position is gathered from it at the anchor (context_position_ids[anchor] + offset) to keep the draft’s RoPE phase document-local.

nemo_automodel.components.speculative.dflash.core.DFlashTrainerModule._create_vp_noise_embed(
    input_ids,
    anchor_positions,
    block_keep_mask,
    prefix_lengths
)

Embed the draft blocks with a visible prefix of real tokens, then MASK.

The variable-prefix analogue of :meth:_create_noise_embed: block positions < prefix_lengths hold the real sequence tokens, the rest (and every position of an invalid block) hold MASK.

Parameters:

input_ids

Long tensor of shape [batch, sequence].

anchor_positions

Long tensor of shape [batch, blocks]; each block’s start position in the sequence.

block_keep_mask

Bool tensor of shape [batch, blocks]; invalid (padding) blocks are embedded as all MASK.

prefix_lengths

Long tensor of shape [batch, blocks]; this block’s visible-prefix length.

Returns:

Tensor of shape [batch, blocks * block_size, hidden].

nemo_automodel.components.speculative.dflash.core.DFlashTrainerModule._prepare_block_inputs(
    input_ids: torch.Tensor,
    loss_mask: torch.Tensor,
    position_ids: torch.Tensor | None = None,
    seq_lens: torch.Tensor | None = None,
    doc_remaining: torch.Tensor | None = None,
    causal: bool = False
) -> typing.Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, 'torch.Tensor | BlockMask', typing.Optional[torch.Tensor]]

Shared block-drafting prologue: anchors, noise embedding, positions, mask.

Centralises the sequence-packing handling for every block-wise trainer (DFlash and its Domino / JetSpec subclasses): anchors are sampled so each block stays inside one document, the block’s context prefix is restricted to its anchor’s document, and the draft RoPE uses per-document positions.

Under loss_type="variable_prefix" the block embedding shows a sampled real-token prefix (_create_vp_noise_embed) instead of the single anchor token, and the sampled prefix_lengths are returned for the loss; every other loss_type returns prefix_lengths=None.

Parameters:

input_ids

torch.Tensor

[B, S] context token ids (long).

loss_mask

torch.Tensor

[B, S] supervised-token mask.

position_ids

torch.Tensor | NoneDefaults to None

[B, S] per-document reset positions (packing), or None.

seq_lens

torch.Tensor | NoneDefaults to None

[B, max_docs] packed document lengths, or None (unpacked).

doc_remaining

torch.Tensor | NoneDefaults to None

[B, S] remaining real tokens of each position’s document.

causal

boolDefaults to False

When True, build the in-block-causal (JetSpec) mask instead of the bidirectional (DFlash / Domino) one.

Returns: torch.Tensor

“(anchor_positions [B, N], block_keep_mask [B, N], noise_embedding

nemo_automodel.components.speculative.dflash.core.DFlashTrainerModule._sample_anchor_positions(
    seq_len: int,
    loss_mask: torch.Tensor,
    device: torch.device,
    doc_remaining: typing.Optional[torch.Tensor] = None
) -> typing.Tuple[torch.Tensor, torch.Tensor]

Randomly sample anchor positions per sample; returns (anchors, keep_mask).

doc_remaining [B, S] (sequence packing) restricts anchors so the whole block stays inside one document (doc_remaining >= block_size - 1), the per-document analogue of the anchor <= seq_len - block_size bound. This is required for correctness — _build_block_targets gathers labels by absolute offset and does not encode document boundaries, so a block that crossed one would be supervised on the next document’s tokens. A side effect is that a packed document shorter than block_size yields no anchors (the unpacked path still supervises such a short sequence’s partial block); pack with documents at least block_size long to avoid dropping their signal.

nemo_automodel.components.speculative.dflash.core.DFlashTrainerModule._sample_prefix_lengths(
    bsz: int,
    n_blocks: int,
    device: torch.device
) -> torch.Tensor

Sample a visible-prefix length per block for variable-prefix training.

A prefix length l means block positions [0, l) are visible real tokens and positions [l, block_size) are masked prediction targets. Lengths follow D2SD’s truncated geometric prior Pr(l) ~ base ** l over [min(2, block_size - 1), block_size - 1]; a base below 1 biases toward short prefixes. The lower bound skips the degenerate fixed-anchor DFlash case, and the upper bound keeps at least one masked target per block.

Returns: torch.Tensor

Long tensor of shape [batch, blocks].

nemo_automodel.components.speculative.dflash.core.DFlashTrainerModule._variable_prefix_loss(
    logits: torch.Tensor,
    target_ids: torch.Tensor,
    block_mask: torch.Tensor,
    prefix_lengths: torch.Tensor
) -> nemo_automodel.components.speculative.dflash.core.DFlashStepMetrics

Decay-weighted CE over each block’s masked suffix (D2SD Eq. for L_VP).

Only positions at or past the sampled visible prefix are supervised, and the exponential decay restarts at the prefix boundary: w_k = exp(-(k - l) / gamma) for block position k >= l (loss_decay_gamma=None disables decay). The loss is the weighted mean sum(nll * w) / sum(w), mirroring the fixed-anchor path’s normalize="mean". Assumes every prefix_lengths entry is at least min(2, block_size - 1) (what :meth:_sample_prefix_lengths produces), so the leading always-visible positions can be sliced off before the CE.

Parameters:

logits

torch.Tensor

Tensor of shape [batch, blocks, block_size, vocab].

target_ids

torch.Tensor

Long tensor of shape [batch, blocks, block_size]; the ground-truth token at anchor + k for block position k.

block_mask

torch.Tensor

Tensor of shape [batch, blocks, block_size]; 0/1 product of block validity, in-bounds, and the loss mask.

prefix_lengths

torch.Tensor

Long tensor of shape [batch, blocks]; this block’s visible-prefix length.

Returns: DFlashStepMetrics

DFlashStepMetrics with the weighted-mean loss, the argmax accuracy

nemo_automodel.components.speculative.dflash.core.DFlashTrainerModule.forward(
    input_ids: torch.Tensor,
    hidden_states: torch.Tensor,
    loss_mask: torch.Tensor,
    position_ids: torch.Tensor | None = None,
    seq_lens: torch.Tensor | None = None,
    doc_remaining: torch.Tensor | None = None
) -> nemo_automodel.components.speculative.dflash.core.DFlashStepMetrics

Parallel block-wise training forward pass.

Sequence packing (position_ids [B, S] per-document reset positions, seq_lens [B, max_docs] document lengths, doc_remaining [B, S]) keeps every block inside one document: anchors are constrained so the block does not cross a boundary, the block’s context prefix attends only within the anchor’s document, and the draft’s RoPE uses the per-document positions.

class nemo_automodel.components.speculative.dflash.core.NoValidAnchorsError()

Bases: ValueError

Raised when a batch has no sample long enough to form a DFlash block.

A DFlash anchor needs at least block_size + 1 supervised tokens (the anchor plus its block). Datasets always contain some short conversations; the training loop catches this and skips the offending micro-batch rather than aborting the run.

nemo_automodel.components.speculative.dflash.core._context_doc_ids(
    seq_lens: torch.Tensor,
    seq_len: int,
    device: torch.device
) -> torch.Tensor

Per-context-token document id [B, S] from packed seq_lens [B, max_docs].

Mirrors the doc_id construction in build_block_causal_additive_mask: a token’s id is the number of document boundaries at or before its position, so 0-length padding entries never split a real document.

nemo_automodel.components.speculative.dflash.core._to_full_tensor(
    tensor: torch.Tensor
) -> torch.Tensor

Materialise a (possibly tensor-parallel) tensor as a plain local tensor.

Under tensor parallelism the target’s column-parallel lm_head and vocab-parallel embed_tokens return DTensor outputs. The draft and the block-wise loss consume plain tensors, so gather the full tensor. A no-op for an already-plain (unsharded / replicated) tensor.

nemo_automodel.components.speculative.dflash.core._DFLASH_LOSS_TYPES = ('dflash', 'variable_prefix')