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

## Module Contents

### Classes

| Name                                                                                    | Description                                                              |
| --------------------------------------------------------------------------------------- | ------------------------------------------------------------------------ |
| [`ChunkedCrossEntropy`](#nemo_automodel-components-loss-chunked_ce-ChunkedCrossEntropy) | Chunk the LM-head projection and CE, recomputing each chunk in backward. |
| [`_ChunkedLinearCE`](#nemo_automodel-components-loss-chunked_ce-_ChunkedLinearCE)       | -                                                                        |

### Functions

| Name                                                                                        | Description                             |
| ------------------------------------------------------------------------------------------- | --------------------------------------- |
| [`_linear_cross_entropy`](#nemo_automodel-components-loss-chunked_ce-_linear_cross_entropy) | Project one token chunk and compute CE. |
| [`_validate_chunk_len`](#nemo_automodel-components-loss-chunked_ce-_validate_chunk_len)     | Validate that `chunk_len` is positive.  |
| [`compute_cross_entropy`](#nemo_automodel-components-loss-chunked_ce-compute_cross_entropy) | Compute fp32 CE.                        |

### API

```python
class nemo_automodel.components.loss.chunked_ce.ChunkedCrossEntropy(
    chunk_len: int = 512,
    compile: bool = True,
    ignore_index: int = -100,
    reduction: str = 'sum'
)
```

**Bases:** [LinearCrossEntropy](/nemo-automodel/nemo_automodel/components/loss/linear_ce_base#nemo_automodel-components-loss-linear_ce_base-LinearCrossEntropy)

Chunk the LM-head projection and CE, recomputing each chunk in backward.

Unlike a logits-based loss, this consumes final hidden states and the
projection weight. Neither full logits nor a full logits gradient is
allocated. Only one chunk's vocabulary activations are live at a time.
Higher-order gradients are not supported.

**`_compute_loss`** `Callable[[torch.Tensor, torch.Tensor, torch.Tensor, int, torch.dtype | None], Tensor]`

---

```python
nemo_automodel.components.loss.chunked_ce.ChunkedCrossEntropy.forward(
    hidden_states: torch.Tensor,
    labels: torch.Tensor,
    lm_weight: torch.Tensor,
    num_label_tokens: int | None = None,
    grad_reduce_group: torch.distributed.ProcessGroup | None = None,
    loss_weights: torch.Tensor | None = None,
    mask: torch.Tensor | None = None,
    logits_dtype: torch.dtype | None = None
) -> torch.Tensor
```

Compute CE directly from hidden states, with no full-logit allocation.

**Parameters:**

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

Rank-local states shaped \[..., hidden], with arbitrary
leading token dimensions, including packed \[tokens, hidden].

---

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

Target indices shaped \[...] matching the token dimensions,
or \[tokens] for packed hidden states shaped \[1, tokens, hidden].

---

**`lm_weight`** `torch.Tensor`

Weight shaped \[vocab, hidden]; FSDP DTensors use the
LinearCrossEntropy.materialize\_lm\_weight contract. TP-sharded
vocabulary/hidden states are not supported by this loss.

---

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

Global count normalizing a sum-reduced loss.

---

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

DP/CP group contributing independent token losses.

---

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

Optional constant multipliers shaped \[...], sum only.

---

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

Optional mask shaped \[...]; zero positions are ignored.

---

**`logits_dtype`** `torch.dtype | None` — default: None

Optional output cast after projection, matching the
owning model. Projection uses lm\_weight.dtype; mixed-dtype
hidden states are converted per chunk, never the full weight.

---

**Returns:** `torch.Tensor`

FP32 scalar, or losses shaped \[...] for reduction="none". Inputs

```python
nemo_automodel.components.loss.chunked_ce.ChunkedCrossEntropy.prepare_lm_weight(
    lm_head: torch.nn.Module | None,
    model_config: transformers.PretrainedConfig | None = None,
    grad_reduce_group: torch.distributed.ProcessGroup | None = None
) -> torch.Tensor
```

Materialize one head in its effective projection dtype for all losses.

**Parameters:**

**`lm_head`** `nn.Module | None`

Plain linear head with weight of global shape \[vocab, hidden].
Its FSDP policy, when present, owns the compute dtype.

---

**`model_config`** `PretrainedConfig | None` — default: None

Owning model config, used to reject post-projection transforms.

---

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

Group contributing independent token losses.

---

**Returns:** `torch.Tensor`

Dense \[vocab, hidden] weight under `materialize_lm_weight`'s

```python
nemo_automodel.components.loss.chunked_ce.ChunkedCrossEntropy.validate_lm_head(
    lm_head: torch.nn.Module | None,
    model_config: transformers.PretrainedConfig | None = None
) -> torch.nn.Linear
```

staticmethod

Validate the head and the owning model's output-logit configuration.

A linear module alone does not rule out post-projection transformations.
Read the text configuration through the HF config contract, including
multimodal models whose output settings live in a nested text config.

**Parameters:**

**`lm_head`** `nn.Module | None`

Output head with weight of global shape \[vocab, hidden].

---

**`model_config`** `PretrainedConfig | None` — default: None

Owning model's configuration, when available.

---

**Returns:** `nn.Linear`

The validated linear head, unchanged.

```python
class nemo_automodel.components.loss.chunked_ce._ChunkedLinearCE()
```

**Bases:** `Function`

```python
nemo_automodel.components.loss.chunked_ce._ChunkedLinearCE.backward(
    ctx: typing.Any,
    grad_out: torch.Tensor
) -> tuple[torch.Tensor | None, torch.Tensor | None, None, None, None, None, None]
```

staticmethod

Recompute one projection/CE graph at a time.

**Parameters:**

**`ctx`** `Any`

Context holding \[tokens, hidden] states, \[vocab, hidden]
weight and \[tokens] labels.

---

**`grad_out`** `torch.Tensor`

Upstream loss gradients shaped \[tokens].

---

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

Hidden-state and weight gradients matching the original shapes and

```python
nemo_automodel.components.loss.chunked_ce._ChunkedLinearCE.forward(
    ctx: typing.Any,
    hidden: torch.Tensor,
    weight: torch.Tensor,
    labels: torch.Tensor,
    chunk_len: int,
    ignore_index: int,
    logits_dtype: torch.dtype | None,
    compute_loss: collections.abc.Callable[[torch.Tensor, torch.Tensor, torch.Tensor, int, torch.dtype | None], torch.Tensor]
) -> torch.Tensor
```

staticmethod

Compute chunk losses without retaining logits.

**Parameters:**

**`ctx`** `Any`

Autograd context.

---

**`hidden`** `torch.Tensor`

Local hidden states shaped \[tokens, hidden].

---

**`weight`** `torch.Tensor`

Dense projection weight shaped \[vocab, hidden].

---

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

Target indices shaped \[tokens].

---

**`chunk_len`** `int`

Maximum token rows per projection.

---

**`ignore_index`** `int`

Ignored target value.

---

**`logits_dtype`** `torch.dtype | None`

Optional output cast matching the model projection.

---

**`compute_loss`** `Callable[[torch.Tensor, torch.Tensor, torch.Tensor, int, torch.dtype | None], torch.Tensor]`

Eager or compiled per-chunk loss.

---

**Returns:** `torch.Tensor`

FP32 losses shaped \[tokens]; inputs are not mutated.

```python
nemo_automodel.components.loss.chunked_ce._linear_cross_entropy(
    hidden: torch.Tensor,
    weight: torch.Tensor,
    labels: torch.Tensor,
    ignore_index: int,
    logits_dtype: torch.dtype | None
) -> torch.Tensor
```

Project one token chunk and compute CE.

**Parameters:**

**`hidden`** `torch.Tensor`

Local hidden states shaped \[chunk\_tokens, hidden].

---

**`weight`** `torch.Tensor`

Dense LM-head weight shaped \[vocab, hidden].

---

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

Target indices shaped \[chunk\_tokens].

---

**`ignore_index`** `int`

Ignored target value.

---

**`logits_dtype`** `torch.dtype | None`

Optional model-owned output dtype after projection.

---

**Returns:** `torch.Tensor`

FP32 losses shaped \[chunk\_tokens].

```python
nemo_automodel.components.loss.chunked_ce._validate_chunk_len(
    chunk_len: int
) -> int
```

Validate that `chunk_len` is positive.

```python
nemo_automodel.components.loss.chunked_ce.compute_cross_entropy(
    logits: torch.Tensor,
    targets: torch.Tensor,
    ignore_index: int = -100,
    reduction: str = 'sum'
) -> torch.Tensor
```

Compute fp32 CE.

**Parameters:**

**`logits`** `torch.Tensor`

Scores shaped \[tokens, vocab].

---

**`targets`** `torch.Tensor`

Class indices shaped \[tokens].

---

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

Ignored target value.

---

**`reduction`** `str` — default: 'sum'

PyTorch CE reduction.

---

**Returns:** `torch.Tensor`

Scalar loss, or \[tokens] losses for reduction="none".