nemo_automodel.components.speculative.dflash.dflash2_core

DFlash 2 online training wrapper.

DFlash 2 (https://inco.ai/blog/dflash2/) keeps DFlash’s parallel block draft and adds a two-tap in-block convolution to the backbone plus a pairwise path selector over each position’s top-k candidates (see nemo_automodel.components.speculative.dflash.draft_qwen3_dflash2). The convolution needs no new supervision — it is part of the backbone and is trained by the ordinary DFlash objective — so this wrapper differs from DFlashTrainerModule in exactly one place: it adds a second term that trains the selector::

loss = base_loss + selector_loss_weight * selector_loss

base_loss is DFlash’s decay-weighted block CE over the full vocabulary, i.e. what makes each position’s candidate list good. selector_loss is a CE over the selector_top_k candidates of that same position, scored against the ground-truth predecessor (the token the walk would have committed had every earlier position been right) and supervised with the index of the true token inside the candidate list. Positions whose true token missed the candidate list carry no selector signal — there is nothing there to select — and are excluded from the selector term; candidate_recall reports how often that happens.

Both terms use the same block-position decay weights, so a position’s importance is identical in the two objectives.

Module Contents

Classes

API

class nemo_automodel.components.speculative.dflash.dflash2_core.DFlash2StepMetrics(
    loss: torch.Tensor,
    loss_weight: torch.Tensor,
    accuracy: torch.Tensor,
    valid_tokens: torch.Tensor,
    correct_tokens: torch.Tensor,
    accept_len: torch.Tensor,
    accept_len_sum: torch.Tensor,
    valid_blocks: torch.Tensor,
    base_loss: torch.Tensor,
    selector_loss: torch.Tensor,
    base_accuracy: torch.Tensor,
    base_correct_tokens: torch.Tensor,
    base_accept_len: torch.Tensor,
    base_accept_len_sum: torch.Tensor,
    candidate_recall: torch.Tensor
)

class nemo_automodel.components.speculative.dflash.dflash2_core.DFlash2TrainerModule(
    draft_model: nemo_automodel.components.speculative.dflash.draft_qwen3_dflash2.Qwen3DFlash2DraftModel,
    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: float | None = None,
    selector_loss_weight: float = 1.0,
    sliding_window: int | None = None
)

nemo_automodel.components.speculative.dflash.dflash2_core.DFlash2TrainerModule._depth_weights(
    mask: torch.Tensor
) -> torch.Tensor

nemo_automodel.components.speculative.dflash.dflash2_core.DFlash2TrainerModule._selector_scores(
    hidden: torch.Tensor,
    logits: torch.Tensor,
    target_ids: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]

nemo_automodel.components.speculative.dflash.dflash2_core.DFlash2TrainerModule.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.dflash2_core.DFlash2StepMetrics