> 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.loss.mtp

## Module Contents

### Classes

| Name                                                                               | Description                                                             |
| ---------------------------------------------------------------------------------- | ----------------------------------------------------------------------- |
| [`MTPLossConfig`](#nemo_automodel-components-loss-mtp-MTPLossConfig)               | Typed config for the Multi-Token-Prediction auxiliary loss.             |
| [`MTPLossOutput`](#nemo_automodel-components-loss-mtp-MTPLossOutput)               | Aggregate MTP loss and its per-depth components.                        |
| [`PipelineCausalLMLoss`](#nemo_automodel-components-loss-mtp-PipelineCausalLMLoss) | Pipeline schedule loss that can add MTP auxiliary CE on the last stage. |

### Functions

| Name                                                                           | Description                                                    |
| ------------------------------------------------------------------------------ | -------------------------------------------------------------- |
| [`calculate_mtp_loss`](#nemo_automodel-components-loss-mtp-calculate_mtp_loss) | Compute the DeepSeek-V3 Multi-Token Prediction auxiliary loss. |

### API

```python
class nemo_automodel.components.loss.mtp.MTPLossConfig(
    scaling_factor: float | None = None,
    ignore_index: int = -100
)
```

Dataclass

Typed config for the Multi-Token-Prediction auxiliary loss.

MTP is gated on the model emitting per-depth outputs; this config only
carries its hyperparameters. `scaling_factor=None` keeps the
model-provided value (`out.mtp_loss_scaling_factor` /
`get_mtp_loss_scaling_factor`); set it to override.

**`ignore_index`** `int = -100`

---

**`scaling_factor`** `float | None = None`

---

```python
nemo_automodel.components.loss.mtp.MTPLossConfig.build(
    loss_fn: torch.nn.Module,
    model: torch.nn.Module,
    grad_reduce_group: torch.distributed.ProcessGroup | None = None
) -> nemo_automodel.components.loss.mtp.PipelineCausalLMLoss
```

Build the pipeline-schedule, MTP-aware loss for `loss_fn`/`model`.

```python
class nemo_automodel.components.loss.mtp.MTPLossOutput(
    loss: torch.Tensor,
    per_depth_losses: list[torch.Tensor]
)
```

Dataclass

Aggregate MTP loss and its per-depth components.

**`loss`** `Tensor`

---

**`per_depth_losses`** `list[Tensor]`

---

```python
class nemo_automodel.components.loss.mtp.PipelineCausalLMLoss(
    loss_fn: torch.nn.Module,
    model: torch.nn.Module,
    scaling_factor: float | None = None,
    ignore_index: int = -100,
    grad_reduce_group: torch.distributed.ProcessGroup | None = None
)
```

**Bases:** `Module`

Pipeline schedule loss that can add MTP auxiliary CE on the last stage.

Per-microbatch `seq_idx` is read from a trailing element of the
last-stage output tuple — the model appends an `[B, S] int32` tail
when MTP is enabled. This binds each microbatch's seq\_idx to its loss
call via the PP runtime's output→loss contract, so the wiring is
schedule-agnostic. Legacy `cu_seqlens` (THD path) is a fallback for
models that don't emit a seq\_idx tail.

**`cu_seqlens`** `Tensor | None = None`

---

```python
nemo_automodel.components.loss.mtp.PipelineCausalLMLoss._extract_seq_idx_tail(
    output
) -> tuple[torch.Tensor | None, object]
```

staticmethod

Detect and strip a trailing per-microbatch seq\_idx from output.

Convention: with MTP enabled the last-stage output is
`(logits, *mtp_per_depth_h, seq_idx)` with an `[B, S] int32`
tail — dtype alone discriminates.

```python
nemo_automodel.components.loss.mtp.PipelineCausalLMLoss.forward(
    output,
    labels: torch.Tensor
) -> torch.Tensor
```

Last-stage pipeline loss (main CE plus optional MTP aux CE).

B=microbatch, S=seq, H=hidden, V=vocab.

**Parameters:**

**`output`**

bare hidden states `[B, S, H]` (FusedLinearCrossEntropy
path), a HF output with logits `[B, S, V]`, or an MTP tuple
`(logits, *mtp_per_depth_h[, seq_idx])` with `seq_idx`
`[B, S]` int32. A tuple with FusedLinearCrossEntropy raises.

---

**`labels`** `torch.Tensor`

target token ids `[B, S]` int64.

---

**Returns:** `torch.Tensor`

Scalar loss tensor.

```python
nemo_automodel.components.loss.mtp.calculate_mtp_loss(
    loss_fn: torch.nn.Module,
    mtp_per_depth_h: list[torch.Tensor] | None = None,
    mtp_per_depth_logits: list[torch.Tensor] | None = None,
    mtp_per_depth_targets: collections.abc.Sequence[torch.Tensor] | None = None,
    labels: torch.Tensor,
    model: torch.nn.Module,
    scaling_factor: float = 0.1,
    num_label_tokens: int | None = None,
    ignore_index: int = -100,
    cu_seqlens: torch.Tensor | None = None,
    seq_idx: torch.Tensor | None = None,
    lm_weight: torch.Tensor | None = None,
    grad_reduce_group: torch.distributed.ProcessGroup | None = None,
    return_per_depth: bool = False
) -> torch.Tensor | nemo_automodel.components.loss.mtp.MTPLossOutput
```

Compute the DeepSeek-V3 Multi-Token Prediction auxiliary loss.

Each depth's CE is dispatched through :func:`calculate_loss` with the
same loss class as the main path, so MTP inherits FusedLinearCrossEntropy
/ MaskedCrossEntropy memory and numerical characteristics.

**Parameters:**

**`loss_fn`** `nn.Module`

Configured per-token loss class (same instance the main
path uses).

---

**`mtp_per_depth_h`** `list[torch.Tensor] | None` — default: None

Per-depth hidden-state tensors of shape
`[batch, sequence, hidden]`, or `[1, tokens, hidden]` for a
flattened THD-packed stream.

---

**`mtp_per_depth_logits`** `list[torch.Tensor] | None` — default: None

Per-depth logit tensors of shape
`[batch, sequence, vocab]`, or `[1, tokens, vocab]` for a
flattened THD-packed stream.

---

**`mtp_per_depth_targets`** `Sequence[torch.Tensor] | None` — default: None

Optional precomputed target tensors, one per
MTP depth, with the same local shape and token layout as `labels`.
Context-parallel callers must shift and boundary-mask them in global
sequence order before applying the model input shard. These targets
are authoritative: this function cannot infer or repair global
packed-sequence boundaries from a rank-local CP shard.

---

**`labels`** `torch.Tensor`

Original unshifted label tensor of shape `[batch, sequence]`
or `[tokens]` for a flattened THD-packed stream.

---

**`model`** `nn.Module`

The wrapped model; used to fetch the shared LM head when the
loss class needs materialized logits (non-FusedLinearCE path).

---

**`scaling_factor`** `float` — default: 0.1

Coefficient applied to the summed per-depth CE.

---

**`num_label_tokens`** `int | None` — default: None

Total non-ignore label-token count used for
sum-reduction normalization.

---

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

Label value masked out of the CE loss for the trailing
`k+1` rolled positions at depth `k`.

---

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

Optional cumulative sequence lengths `[num_seqs+1]`
(THD-pack layout). When supplied and `seq_idx` is not, builds
a per-token sub-sequence index via searchsorted. Without packing
this can be omitted.

---

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

Optional per-token sub-sequence index `[B, S]` (or `[S]`).
Equality classes are what matter; absolute values can be any
ints. Takes precedence over `cu_seqlens`. Used to mask label
rolls whose source position lies in a different sub-sequence.

---

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

Optional LM-head weight tensor of shape `[vocab, hidden]`.
Supplying it lets the main loss and all MTP depths share one DTensor
`full_tensor()` gather on the FusedLinearCrossEntropy path.

---

**`grad_reduce_group`** `dist.ProcessGroup | None` — default: None

Group that contributes independent loss shards when
the shared LM-head weight is a DTensor.

---

**`return_per_depth`** `bool` — default: False

Return the aggregate loss together with the unscaled
loss for each MTP depth. Defaults to `False` to preserve the
scalar return expected by existing callers.

---

**Returns:** `torch.Tensor | MTPLossOutput`

Scalar MTP loss tensor. When `return_per_depth=True`, returns an