> This page is for version 0.5.0 · 26.06.
> For other versions, use one of these documentation indexes:
> - Nightly (default): https://docs.nvidia.com/nemo/automodel/nightly/llms.txt
> - Latest: https://docs.nvidia.com/nemo/automodel/latest/llms.txt
> - 0.5.0 · 26.06: https://docs.nvidia.com/nemo/automodel/v0.5/llms.txt
> - 0.4.0 · 26.04: https://docs.nvidia.com/nemo/automodel/v0.4/llms.txt

> 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

Shared scaffolding for Multi-Token Prediction (MTP) auxiliary heads.

MTP follows the DeepSeek-V3 design (Liu et al., 2024). Each MTP "depth"
predicts one additional future token; per depth the input is rolled left by
one position, fused with the previous-depth hidden state, and passed through
an inner block before producing logits via the shared LM head.

Components in this package are model-agnostic. Model-specific glue (building
the inner block out of the model's own decoder layers, wiring HF state-dict
keys) lives in the model's own package.

## Submodules

* **[`nemo_automodel.components.models.common.mtp.mtp`](/nemo-automodel/nemo_automodel/components/models/common/mtp/mtp)**

## Package Contents

### Classes

| Name                                                                                                    | Description                                                             |
| ------------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------- |
| [`MTPConfig`](#nemo_automodel-components-models-common-mtp-mtp-MTPConfig)                               | Runtime configuration for the MTP block.                                |
| [`MTPContextParallelInputs`](#nemo_automodel-components-models-common-mtp-mtp-MTPContextParallelInputs) | Globally shifted MTP tensors prepared before context-parallel sharding. |
| [`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.        |
| [`prepare_mtp_context_parallel_inputs`](#nemo_automodel-components-models-common-mtp-mtp-prepare_mtp_context_parallel_inputs) | Prepare global future-token tensors before context-parallel sharding.   |
| [`roll_tensor`](#nemo_automodel-components-models-common-mtp-mtp-roll_tensor)                                                 | Roll a tensor along `dim` by `shifts` and zero the wrapped slice.       |
| [`shift_packed_tensor`](#nemo_automodel-components-models-common-mtp-mtp-shift_packed_tensor)                                 | Shift a token-aligned tensor left without crossing sequence boundaries. |

### 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.

**`enabled`** `bool`

---

**`layer_pattern`** `str = ''`

Per-depth inner-block pattern, e.g. `"*E"` for one
attention + one MoE sublayer per depth.

---

**`loss_scaling_factor`** `float = 0.1`

Coefficient applied to the summed per-depth CE
loss (default `0.1`). The effective per-depth weight is
`loss_scaling_factor / num_layers`.

---

**`num_layers`** `int = 0`

Number of MTP forward iterations (D). `0` disables MTP.
Equivalent to Megatron's `--mtp-num-layers`.

---

**`num_physical_depths`** `int`

---

**`pattern_length`** `int`

---

**`total_sublayers`** `int`

---

**`use_repeated_layer`** `bool = False`

When `True`, build a single physical depth's
worth of sublayers and reuse it for all `num_layers` forward
iterations (weight-tied across depths). Equivalent to Megatron's
`--mtp-use-repeated-layer`.

---

```python
class nemo_automodel.components.models.common.mtp.mtp.MTPContextParallelInputs(
    input_ids: tuple[torch.LongTensor, ...],
    position_ids: tuple[torch.LongTensor, ...],
    targets: tuple[torch.LongTensor, ...],
    valid_masks: tuple[torch.BoolTensor, ...],
    position_ids_seq_dim: int
)
```

Dataclass

Globally shifted MTP tensors prepared before context-parallel sharding.

**`input_ids`** `tuple[LongTensor, ...]`

Per-depth token IDs of shape `[batch, sequence]` in global
sequence order.

---

**`position_ids`** `tuple[LongTensor, ...]`

Per-depth position IDs in global sequence order. Standard
positions have shape `[batch, sequence]`; multi-axis RoPE uses
`[axes, batch, sequence]`.

---

**`position_ids_seq_dim`** `int`

Sequence dimension of each position-ID tensor.

---

**`targets`** `tuple[LongTensor, ...]`

Per-depth loss targets of shape `[batch, sequence]` in global
sequence order, with invalid positions set to the loss ignore index.

---

**`valid_masks`** `tuple[BoolTensor, ...]`

Per-depth masks of shape `[batch, sequence]` identifying
positions whose future token remains in the same packed sequence.

---

```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 `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:**

**`mtp_config`** `MTPConfig`

`MTPConfig` describing depth and pattern.

---

**`block_types_per_sublayer`** `list[str]`

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.

---

**`sublayer_factory`** `Callable[..., nn.Module]`

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

---

**`layers`** `= nn.ModuleList(layers)`

---

**`num_depths`** `int`

---

**`pattern_length`** `int`

---

```python
nemo_automodel.components.models.common.mtp.mtp.MTPModule.forward(
    hidden_states: torch.Tensor,
    input_ids: torch.LongTensor | None = None,
    input_ids_per_depth: tuple[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,
    position_ids_per_depth: tuple[torch.LongTensor, ...] | None = None,
    block_kwargs = {}
) -> list[torch.Tensor]
```

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

Three 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.
* **Context parallel**: pass `input_ids_per_depth` plus `embed_fn`.
  Each tensor is globally shifted and then sharded into the local CP
  token layout, so this module embeds it directly without a rank-local
  roll.
* **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:**

**`hidden_states`** `torch.Tensor`

Output of the main model's final norm (`h_0`);
tensor of shape `[batch, sequence, hidden]` or
`[tokens, hidden]` for THD.

---

**`input_ids`** `torch.LongTensor | None` — default: None

Token ids of shape `[batch, sequence]` (or
`[tokens]` in THD). Rolled
cumulatively left by 1 per depth. Mutually exclusive with
`input_ids_per_depth` and `embed_inputs`.

---

**`input_ids_per_depth`** `tuple[torch.LongTensor, ...] | None` — default: None

Optional tuple of `num_depths` pre-shifted
token-ID tensors. Each has local CP shape
`[batch, sequence]` or `[tokens]` and is embedded directly
with `embed_fn`. Requires `position_ids_per_depth`.

---

**`embed_fn`** `Callable[[torch.LongTensor], torch.Tensor] | None` — default: None

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.

---

**`embed_inputs`** `tuple[torch.Tensor, ...] | None` — default: None

Optional tuple of `num_depths` pre-computed
future-token embeddings, each of shape
`[batch, sequence, hidden]` or `[tokens, hidden]`.
Mutually exclusive with `input_ids`/
`input_ids_per_depth`/`embed_fn`.

---

**`position_ids`** `torch.LongTensor | None` — default: None

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.

---

**`position_ids_per_depth`** `tuple[torch.LongTensor, ...] | None` — default: None

Optional tuple of `num_depths`
pre-computed future-token position tensors. Each tensor has
shape `[batch, sequence]`, `[axes, batch, sequence]` for
multi-axis RoPE, or `[tokens]` for THD. When supplied, these
tensors are forwarded directly instead of rolling rank-local
`position_ids`. Use this when context parallelism has already sharded the sequence.
Required with `input_ids_per_depth` and incompatible with
rank-local `input_ids` rolling.

---

**`**block_kwargs`** — default: \{}

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

---

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

List of length `num_depths` containing hidden states of shape

```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.prepare_mtp_context_parallel_inputs(
    batch: collections.abc.MutableMapping[str, object],
    num_depths: int,
    ignore_index: int = -100
) -> nemo_automodel.components.models.common.mtp.mtp.MTPContextParallelInputs
```

Prepare global future-token tensors before context-parallel sharding.

Each MTP depth is shifted in global sequence order before CP partitions the
token axis. Packed boundaries are preserved, so no future token, position,
or target crosses from one document into another. Missing or shared
position IDs are materialized in batch so the main model and MTP heads are
subsequently sharded from the same global source.

**Parameters:**

**`batch`** `MutableMapping[str, object]`

Mutable unsharded batch. input\_ids and labels are tensors of
shape \[batch, sequence]. Optional position\_ids has shape
\[batch, sequence], shared shape \[1, sequence], or multi-axis RoPE
shape \[axes, batch, sequence]. Packed boundaries may use the tensor layouts documented by
\_packed\_seq\_ids\_from\_batch.

---

**`num_depths`** `int`

Number of MTP future-token depths; must be positive.

---

**`ignore_index`** `int` — default: -100

Fill value for invalid targets at trailing and packed
boundary positions.

---

**Returns:** `MTPContextParallelInputs`

Per-depth input IDs, position IDs, targets, and validity masks. Token

```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:**

**`t`** `torch.Tensor`

Input tensor.

---

**`shifts`** `int` — default: -1

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

---

**`dim`** `int` — default: -1

Dimension to roll along.

---

**Returns:** `torch.Tensor`

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

```python
nemo_automodel.components.models.common.mtp.mtp.shift_packed_tensor(
    tensor: torch.Tensor,
    depth: int,
    seq_idx: torch.Tensor | None = None,
    fill_value: float | int = 0,
    batch_dim: int = 0,
    seq_dim: int = 1
) -> torch.Tensor
```

Shift a token-aligned tensor left without crossing sequence boundaries.

**Parameters:**

**`tensor`** `torch.Tensor`

Token-aligned tensor in global sequence order. Its batch and
sequence axes are selected by `batch_dim` and `seq_dim`.

---

**`depth`** `int`

Number of future-token positions to shift; must be positive.

---

**`seq_idx`** `torch.Tensor | None` — default: None

Optional sequence IDs of shape `[batch, sequence]`. Tokens
whose shifted source has a different ID are filled.

---

**`fill_value`** `float | int` — default: 0

Scalar used for trailing and cross-sequence positions.

---

**`batch_dim`** `int` — default: 0

Batch dimension in `tensor`.

---

**`seq_dim`** `int` — default: 1

Sequence dimension in `tensor`.

---

**Returns:** `torch.Tensor`

Tensor with the same shape, dtype, and device as `tensor`. The output