> 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.                                |
| [`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                                                             |
| ----------------------------------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------- |
| [`_packed_seq_ids_from_batch`](#nemo_automodel-components-models-common-mtp-mtp-_packed_seq_ids_from_batch)                   | Normalize supported packed-boundary metadata to token-aligned IDs.      |
| [`_packed_seq_ids_from_padded_lengths`](#nemo_automodel-components-models-common-mtp-mtp-_packed_seq_ids_from_padded_lengths) | Expand padded packed-sequence lengths into token-aligned sequence IDs.  |
| [`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 = ''`

---

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

---

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

---

**`num_physical_depths`** `int`

---

**`pattern_length`** `int`

---

**`total_sublayers`** `int`

---

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

---

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

---

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

---

**`position_ids_seq_dim`** `int`

---

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

---

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

---

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

**`mtp_config`** `MTPConfig`

:class:`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._packed_seq_ids_from_batch(
    batch: collections.abc.MutableMapping[str, object],
    input_ids: torch.Tensor
) -> torch.LongTensor | None
```

Normalize supported packed-boundary metadata to token-aligned IDs.

**Parameters:**

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

Unsharded batch. Optional seq\_idx or \_packed\_seq\_ids tensors
have shape \[batch, sequence] (or \[sequence] when batch is one).
seq\_lens\_padded has shape \[batch, num\_sequences].
cu\_seqlens\_padded or cu\_seqlens contains flattened cumulative
boundaries of shape \[num\_sequences + 1] with optional negative
sentinels.

---

**`input_ids`** `torch.Tensor`

Global token-ID tensor of shape \[batch, sequence] whose
materialized token layout defines the expected output shape.

---

**Returns:** `torch.LongTensor | None`

Sequence-ID tensor of shape \[batch, sequence] on the input device,

```python
nemo_automodel.components.models.common.mtp.mtp._packed_seq_ids_from_padded_lengths(
    seq_lens_padded: torch.Tensor,
    batch_size: int,
    seq_len: int,
    device: torch.device
) -> torch.LongTensor
```

Expand padded packed-sequence lengths into token-aligned sequence IDs.

**Parameters:**

**`seq_lens_padded`** `torch.Tensor`

Tensor of shape \[batch, num\_sequences] or
\[num\_sequences] for a single batch row. Negative entries are
unused sentinels; nonnegative entries are materialized sequence
lengths including padding.

---

**`batch_size`** `int`

Expected batch dimension of the returned tensor.

---

**`seq_len`** `int`

Expected materialized token count in each batch row.

---

**`device`** `torch.device`

Device for the returned token-aligned IDs.

---

**Returns:** `torch.LongTensor`

Sequence-ID tensor of shape \[batch, sequence] on device.

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