> 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.models.common.mtp.mtp

Model-agnostic MTP scaffolding: depth iteration, token rolling, and loss.

## Module Contents

### Classes

| Name                                                                      | Description                              |
| ------------------------------------------------------------------------- | ---------------------------------------- |
| [`MTPConfig`](#nemo_automodel-components-models-common-mtp-mtp-MTPConfig) | Runtime configuration for the MTP block. |
| [`MTPModule`](#nemo_automodel-components-models-common-mtp-mtp-MTPModule) | Multi-Token Prediction block.            |

### Functions

| Name                                                                                                          | Description                                                       |
| ------------------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------- |
| [`get_mtp_loss_scaling_factor`](#nemo_automodel-components-models-common-mtp-mtp-get_mtp_loss_scaling_factor) | Return the model's configured MTP auxiliary-loss scaling factor.  |
| [`roll_tensor`](#nemo_automodel-components-models-common-mtp-mtp-roll_tensor)                                 | Roll a tensor along `dim` by `shifts` and zero the wrapped slice. |

### API

```python
class nemo_automodel.components.models.common.mtp.mtp.MTPConfig(
    num_layers: int = 0,
    layer_pattern: str = '',
    loss_scaling_factor: float = 0.1,
    use_repeated_layer: bool = False
)
```

Dataclass

Runtime configuration for the MTP block.

```python
class nemo_automodel.components.models.common.mtp.mtp.MTPModule(
    mtp_config: nemo_automodel.components.models.common.mtp.mtp.MTPConfig,
    block_types_per_sublayer: list[str],
    sublayer_factory: typing.Callable[..., torch.nn.Module]
)
```

**Bases:** `Module`

Multi-Token Prediction block.

Holds a flat :class:`nn.ModuleList` of sublayers (length
`num_physical_depths * pattern_length`) where the first sublayer of
each physical depth carries the fusion modules (`enorm`, `hnorm`,
`eh_proj`) and the last sublayer of each physical depth carries
`final_layernorm`. This flat layout matches the HuggingFace export
format used by Nemotron-V3 (`mtp.layers.{i}.*`).

The model-specific sublayer construction (which decoder block to use, how
to handle MoE / attention / Mamba) is delegated to the caller via
`sublayer_factory`.

**Parameters:**

:class:`MTPConfig` describing depth and pattern.

List of block-type strings (one per inner
sublayer position), length must equal `mtp_config.pattern_length`.
Caller is responsible for parsing the model-specific symbol
convention; this module does not interpret symbols.

Callable
`factory(global_idx, depth, sublayer_idx, block_type, has_fusion, has_final_norm) -> nn.Module`
constructing one sublayer. The returned module must be callable
as `sublayer(hidden_states, **kwargs) -> Tensor` and, when
`has_fusion=True`, expose attributes `enorm`, `hnorm`,
`eh_proj`. When `has_final_norm=True` it must expose
`final_layernorm`.

```python
nemo_automodel.components.models.common.mtp.mtp.MTPModule.forward(
    hidden_states: torch.Tensor,
    input_ids: torch.LongTensor | None = None,
    embed_fn: typing.Callable[[torch.LongTensor], torch.Tensor] | None = None,
    embed_inputs: tuple[torch.Tensor, ...] | None = None,
    position_ids: torch.LongTensor | None = None,
    block_kwargs = {}
) -> list[torch.Tensor]
```

Iterate over MTP depths and return per-depth hidden states.

Two mutually-exclusive input modes:

* **Single-rank / first-stage PP** (default): pass `input_ids` plus
  `embed_fn`. The module rolls `input_ids` cumulatively left by 1
  per depth and applies `embed_fn` to produce the future-token
  embedding for that depth.
* **Final-stage PP / multimodal**: pass `embed_inputs` (a tuple of
  pre-rolled per-depth embeddings, length `num_depths`). Used when
  the last PP stage no longer owns `embed_tokens`, or for multimodal
  models (e.g. SALM) where some positions carry continuous audio
  embeddings that cannot be recovered by re-embedding an integer token
  id — the caller pre-rolls the fused embedding tensor and passes it
  here.

**Parameters:**

Output of the main model's final norm (`h_0`);
shape matches the model's residual stream.

Token ids `[B, S]` (or `[T]` in THD). Rolled
cumulatively left by 1 per depth. Mutually exclusive with
`embed_inputs`.

Callable applied to rolled `input_ids` to produce the
future-token embedding (typically the model's input embedding
layer). Required when `input_ids` is supplied.

Optional tuple of `num_depths` pre-computed
future-token embeddings, one per depth in MTP order.
Mutually exclusive with `input_ids`/`embed_fn`.

Position ids matching `input_ids`. When supplied,
rolled cumulatively per depth in lockstep with `input_ids`
(so slot `t` carries the original position of the rolled
token) and forwarded to each sublayer via `block_kwargs`.
Required for RoPE-using sublayers; ignored by sublayers that
don't consume it.

Forwarded to each sublayer's `__call__` (e.g.
`attention_mask`).

**Returns:** `list[torch.Tensor]`

List of length `num_depths` containing the hidden state

```python
nemo_automodel.components.models.common.mtp.mtp.get_mtp_loss_scaling_factor(
    model: torch.nn.Module,
    default: float = 0.1
) -> float
```

Return the model's configured MTP auxiliary-loss scaling factor.

```python
nemo_automodel.components.models.common.mtp.mtp.roll_tensor(
    t: torch.Tensor,
    shifts: int = -1,
    dim: int = -1
) -> torch.Tensor
```

Roll a tensor along `dim` by `shifts` and zero the wrapped slice.

Used to shift `input_ids` / `position_ids` / `labels` left by one
position per MTP depth. Single-GPU path only (no CP / packed-sequence
handling).

**Parameters:**

Input tensor.

Number of positions to shift (negative = left shift).

Dimension to roll along.

**Returns:** `torch.Tensor`

New tensor with the trailing `|shifts|` positions along `dim`