nemo_automodel.components.speculative.dflash.target

Target-model wrapper for DFlash training.

Unlike EAGLE-3 (which captures exactly three aux layers and left-shifts the supervision), DFlash captures an arbitrary set of decoder layers, concatenates them along the feature dim, and feeds the result to the draft as context. No shifting is applied — the DFlash block attention mask handles anchor alignment.

Module Contents

Classes

Name	Description
`DFlashTargetBatch`	Target-model context features needed by the DFlash trainer.
`HFDFlashTargetModel`	Capture a set of decoder-layer hidden states from a frozen HF causal LM.

API

class nemo_automodel.components.speculative.dflash.target.DFlashTargetBatch(
    hidden_states: torch.Tensor,
    input_ids: torch.Tensor,
    attention_mask: torch.Tensor,
    loss_mask: torch.Tensor,
    logits: typing.Optional[torch.Tensor] = None,
    position_ids: typing.Optional[torch.Tensor] = None,
    seq_lens: typing.Optional[torch.Tensor] = None,
    doc_remaining: typing.Optional[torch.Tensor] = None
)

Dataclass

Target-model context features needed by the DFlash trainer.

position_ids / seq_lens / doc_remaining are None off the packing path and carry the (unshifted) packing metadata to the trainer on it.

attention_mask

Tensor

doc_remaining

Optional[Tensor] = None

hidden_states

Tensor

input_ids

Tensor

logits

Optional[Tensor] = None

loss_mask

Tensor

position_ids

Optional[Tensor] = None

seq_lens

Optional[Tensor] = None

class nemo_automodel.components.speculative.dflash.target.HFDFlashTargetModel(
    model: torch.nn.Module,
    target_layer_ids: typing.Sequence[int],
    capture_logits: bool = False,
    cp_mesh = None
)

Capture a set of decoder-layer hidden states from a frozen HF causal LM.

A forward hook on decoder layer i captures that layer’s output, which in HuggingFace’s output_hidden_states convention is hidden_states[i + 1] — matching SpecForge’s extract_context_feature (offset 1).

_cp_size

= cp_mesh.size() if cp_mesh is not None else 1

capture_logits

= bool(capture_logits)

model

= model.eval()

target_layer_ids

= self._validate_layer_ids(target_layer_ids)

nemo_automodel.components.speculative.dflash.target.HFDFlashTargetModel._check_captured(
    captured: dict[int, torch.Tensor]
) -> None

nemo_automodel.components.speculative.dflash.target.HFDFlashTargetModel._get_transformer_layers() -> list[torch.nn.Module]

Return decoder layers as an ordered, integer-indexable list.

nemo_automodel.components.speculative.dflash.target.HFDFlashTargetModel._validate_layer_ids(
    target_layer_ids: typing.Sequence[int]
) -> list[int]

nemo_automodel.components.speculative.dflash.target.HFDFlashTargetModel.generate_batch(
    input_ids: torch.Tensor,
    attention_mask: torch.Tensor,
    loss_mask: torch.Tensor,
    position_ids: typing.Optional[torch.Tensor] = None,
    seq_lens: typing.Optional[torch.Tensor] = None,
    doc_remaining: typing.Optional[torch.Tensor] = None
) -> nemo_automodel.components.speculative.dflash.target.DFlashTargetBatch

Run the target model and capture the selected layers’ hidden states as context.

With seq_lens ([B, max_docs], per-document lengths summing to S) the target runs with a document-level block-causal mask and per-document position_ids so the captured context hidden states do not leak across document boundaries (SDPA/eager consume the [B, 1, S, S] block-causal additive mask; FlashAttention infers boundaries from position_ids at batch size 1). The packing metadata is carried through to the trainer unchanged.

nemo_automodel.components.speculative.dflash.target.HFDFlashTargetModel.get_input_embeddings() -> torch.nn.Embedding

Return the target model input embeddings.