> 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)         | Cross-entropy loss computed over sequence chunks.                 |
| [`_ChunkedCrossEntropySum`](#nemo_automodel-components-loss-chunked_ce-_ChunkedCrossEntropySum) | Sum-reduced cross-entropy with recomputed fp32 chunk activations. |

### Functions

| Name                                                                                        | Description                                                 |
| ------------------------------------------------------------------------------------------- | ----------------------------------------------------------- |
| [`_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) | Computes the cross-entropy loss between logits and targets. |

### Data

[`_compiled_compute_cross_entropy`](#nemo_automodel-components-loss-chunked_ce-_compiled_compute_cross_entropy)

### API

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

**Bases:** `Module`

Cross-entropy loss computed over sequence chunks.

**`chunk_len`** `= _validate_chunk_len(chunk_len)`

---

```python
nemo_automodel.components.loss.chunked_ce.ChunkedCrossEntropy.forward(
    logits: torch.Tensor,
    labels: torch.Tensor,
    mask: torch.Tensor | None = None,
    num_label_tokens: int | None = None
) -> torch.Tensor
```

Computes cross-entropy loss in chunks to handle long sequences more efficiently.

**Parameters:**

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

Model output tensor of shape \[..., vocab], with arbitrary
leading token dimensions.

---

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

Target indices of shape \[...] matching `logits.shape[:-1]`.
When `mask` is provided, ignored positions are replaced with
`ignore_index` in this tensor.

---

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

Optional tensor of shape \[...] matching `labels`. Nonzero
positions contribute to the loss and zero positions are ignored.

---

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

Optional global count used to normalize the
sum-reduced scalar loss.

---

**Returns:** `torch.Tensor`

Scalar tensor containing the reduced cross-entropy loss.

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

**Bases:** `Function`

Sum-reduced cross-entropy with recomputed fp32 chunk activations.

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

staticmethod

Recompute the fp32 softmax chunks and return the logits gradient.

**Parameters:**

**`ctx`**

Autograd context containing the original-dtype tensors saved by
:meth:`forward`.

---

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

Scalar tensor containing the upstream loss gradient.

---

**Returns:** `torch.Tensor`

Tuple whose first element is the logits gradient of shape \[tokens,

```python
nemo_automodel.components.loss.chunked_ce._ChunkedCrossEntropySum.forward(
    ctx,
    logits: torch.Tensor,
    labels: torch.Tensor,
    ignore_index: int,
    chunk_len: int
) -> torch.Tensor
```

staticmethod

Compute sum-reduced cross-entropy one fp32 chunk at a time.

**Parameters:**

**`ctx`**

Autograd context used to save tensors for backward.

---

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

Prediction scores of shape \[tokens, vocab]. The tensor stays
in its original floating-point dtype and is not mutated.

---

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

Target class indices of shape \[tokens]. Positions equal to
`ignore_index` contribute zero loss and gradient.

---

**`ignore_index`** `int`

Target value excluded from the loss.

---

**`chunk_len`** `int`

Maximum number of token rows upcast to fp32 at once.

---

**Returns:** `torch.Tensor`

Scalar fp32 tensor containing the sum-reduced loss.

```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 = -100,
    reduction = 'sum'
)
```

Computes the cross-entropy loss between logits and targets.

**Parameters:**

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

Model predictions of shape (sequence\_length, num\_classes).

---

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

Ground-truth labels of shape (sequence\_length,).

---

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

Target value that is ignored when computing the loss.
Defaults to -100.

---

**Returns:**

torch.Tensor: The sum of cross-entropy losses over the sequence.

```python
nemo_automodel.components.loss.chunked_ce._compiled_compute_cross_entropy = None
```