> For clean Markdown of any page, append .md to the page URL.
> For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/llms.txt.
> For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/nemo/automodel/_mcp/server.

# 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`](#nemo_automodel-components-speculative-dspark-target-DSparkTargetBatch)     | Target-model features needed by the DSpark trainer.                    |
| [`HFDSparkTargetModel`](#nemo_automodel-components-speculative-dspark-target-HFDSparkTargetModel) | Capture intermediate + final hidden states from a frozen HF causal LM. |

### Data

[`__all__`](#nemo_automodel-components-speculative-dspark-target-__all__)

### API

```python
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
)
```

Dataclass

Target-model features needed by the DSpark trainer.

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

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.

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

Return the final norm module whose output feeds `lm_head`.

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

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

```python
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`).

```python
nemo_automodel.components.speculative.dspark.target.HFDSparkTargetModel.generate_batch(
    input_ids: torch.Tensor,
    attention_mask: torch.Tensor,
    loss_mask: 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.

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

Return the target model input embeddings.

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

Return the target model output embeddings (lm\_head).

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