> For clean Markdown of any page, append .md to the page URL. > For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/v0.5/nemo-automodel/nemo_automodel/components/models/common/mtp/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 ## Docs - [nemo_automodel.components.models.common.mtp.mtp](https://docs.nvidia.com/nemo/automodel/v0.5/nemo-automodel/nemo_automodel/components/models/common/mtp/mtp.md)