nemo_automodel.components.speculative.dspark.target

Target-model wrapper for DSpark training (online hidden-state capture).

DSpark feeds the draft two things from the frozen target: the concatenation of a configured set of decoder-layer hidden states (the draft fc context), and the final post-norm hidden state (the input the target’s lm_head consumes, used by the TV / confidence losses). Both are captured in a single forward pass via forward hooks, mirroring the DFlash target wrapper.

Module Contents

Classes

Name	Description
`DSparkTargetBatch`	Target-model features needed by the DSpark trainer.
`HFDSparkTargetModel`	Capture intermediate + final hidden states from a frozen HF causal LM.

Data

__all__

API

class nemo_automodel.components.speculative.dspark.target.DSparkTargetBatch(
    target_hidden_states: torch.Tensor,
    target_last_hidden_states: torch.Tensor,
    input_ids: 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
)

Dataclass

Target-model features needed by the DSpark trainer.

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

doc_remaining

Optional[Tensor] = None

input_ids

Tensor

loss_mask

Tensor

position_ids

Optional[Tensor] = None

seq_lens

Optional[Tensor] = None

target_hidden_states

Tensor

target_last_hidden_states

Tensor

class nemo_automodel.components.speculative.dspark.target.HFDSparkTargetModel(
    model: torch.nn.Module,
    target_layer_ids: typing.Sequence[int],
    cp_mesh = None
)

Capture intermediate + final hidden states from a frozen HF causal LM.

A forward hook on decoder layer i captures hidden_states[i + 1] (the HuggingFace output_hidden_states offset-1 convention); a hook on the final norm captures the post-norm last hidden state.

_cp_size

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

_num_layers

= len(self._get_transformer_layers())

model

= model.eval()

target_layer_ids

nemo_automodel.components.speculative.dspark.target.HFDSparkTargetModel._collapse_hc_streams(
    tensor: torch.Tensor
) -> torch.Tensor

Collapse a 4D Hyper-Connection stream [B, S, hc_mult, H] to [B, S, H].

DeepSeek V4 decoder layers emit hc_mult parallel residual copies; only the final-norm output is already collapsed. For an intermediate target-feature layer we reduce the streams with their mean: a simple, in-distribution reduction that the draft’s learnable fc then reprojects. We deliberately avoid the model’s final hc_head here, since it is trained for the last-layer stream distribution, not the intermediate ones. Non-HC targets emit 3D states and pass through unchanged.

nemo_automodel.components.speculative.dspark.target.HFDSparkTargetModel._get_final_norm() -> torch.nn.Module

Return the final norm module whose output feeds lm_head.

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

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

nemo_automodel.components.speculative.dspark.target.HFDSparkTargetModel._inner_model() -> torch.nn.Module

Return the base transformer module that owns layers and norm.

Handles the common nestings: a plain causal LM (model.model), a decoder-only base (model), and multimodal targets whose text stack is under language_model (e.g. Gemma4: model.model.language_model).

nemo_automodel.components.speculative.dspark.target.HFDSparkTargetModel.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,
    mm_kwargs: torch.Tensor = {}
) -> nemo_automodel.components.speculative.dspark.target.DSparkTargetBatch

Run the target model once and capture the DSpark context + last hidden state.

Features follow common.extract_context_feature exactly: -1 is the embedding output, the final layer is the post-norm hidden state (HF output_hidden_states[num_layers]), and any other id is that decoder layer’s output. The final-norm output is also returned as the last hidden state for the TV / confidence losses.

mm_kwargs carries any multimodal inputs present in the training batch (pixel_values, image_grid_thw, …); :func:filter_forward_kwargs drops whatever the target’s own forward signature does not accept, so a text-only target (Qwen3, Gemma4, or MiniMax M3 without multimodal training data) is never passed inputs it cannot handle — this is a no-behavior-change extension for every existing (mm_kwargs-empty) caller. A VLM target (e.g. MiniMax M3) splices the vision features into its embedding sequence internally when they are provided.

nemo_automodel.components.speculative.dspark.target.HFDSparkTargetModel.get_input_embeddings() -> torch.nn.Embedding

Return the target model input embeddings.

nemo_automodel.components.speculative.dspark.target.HFDSparkTargetModel.get_output_embeddings() -> torch.nn.Module

Return the target model output embeddings (lm_head).

nemo_automodel.components.speculative.dspark.target.__all__ = ['HFDSparkTargetModel', 'DSparkTargetBatch']