> 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.utils

## Module Contents

### Functions

| Name                                                                                         | Description                                                                |
| -------------------------------------------------------------------------------------------- | -------------------------------------------------------------------------- |
| [`_count_label_tokens`](#nemo_automodel-components-loss-utils-_count_label_tokens)           | Count supervised entries in an integer target tensor of any shape.         |
| [`_get_final_hidden_states`](#nemo_automodel-components-loss-utils-_get_final_hidden_states) | Return the final hidden-states tensor from an HF-like model output.        |
| [`_get_lm_head_module`](#nemo_automodel-components-loss-utils-_get_lm_head_module)           | Return the model's LM-head module, if one can be found.                    |
| [`_get_lm_head_weight`](#nemo_automodel-components-loss-utils-_get_lm_head_weight)           | Return the model's LM-head weight without changing its distributed layout. |
| [`_get_loss_ignore_index`](#nemo_automodel-components-loss-utils-_get_loss_ignore_index)     | Return the label sentinel consumed by `loss_fn`.                           |
| [`_normalize_kd_labels`](#nemo_automodel-components-loss-utils-_normalize_kd_labels)         | Align KD supervision with the main loss mask.                              |
| [`_normalize_loss_labels`](#nemo_automodel-components-loss-utils-_normalize_loss_labels)     | Map dataset padding to a loss's configured ignore index.                   |
| [`calculate_loss`](#nemo_automodel-components-loss-utils-calculate_loss)                     | Calculate a logits-based or linear-projection cross-entropy loss.          |
| [`prepare_lm_weight`](#nemo_automodel-components-loss-utils-prepare_lm_weight)               | Prepare one shared projection weight for main and auxiliary losses.        |

### Data

[`_DATASET_IGNORE_INDEX`](#nemo_automodel-components-loss-utils-_DATASET_IGNORE_INDEX)

### API

```python
nemo_automodel.components.loss.utils._count_label_tokens(
    labels: torch.Tensor,
    ignore_index: int
) -> int
```

Count supervised entries in an integer target tensor of any shape.

Both the dataset's `-100` padding and the configured `ignore_index`
are excluded. The return value is a Python integer.

```python
nemo_automodel.components.loss.utils._get_final_hidden_states(
    model_output: typing.Any
) -> typing.Any | None
```

Return the final hidden-states tensor from an HF-like model output.

Local copy of `components.training.model_output_utils.get_final_hidden_states`
to keep `components/loss/` import-independent from `components/training/`.

```python
nemo_automodel.components.loss.utils._get_lm_head_module(
    model: torch.nn.Module
) -> torch.nn.Module | None
```

Return the model's LM-head module, if one can be found.

Local copy of `components.utils.model_utils.get_lm_head_module` to keep
`components/loss/` import-independent from `components/utils/` (see the
`Components must not import each other` import-linter contract).

```python
nemo_automodel.components.loss.utils._get_lm_head_weight(
    model: torch.nn.Module
) -> torch.Tensor
```

Return the model's LM-head weight without changing its distributed layout.

```python
nemo_automodel.components.loss.utils._get_loss_ignore_index(
    loss_fn: object
) -> int
```

Return the label sentinel consumed by `loss_fn`.

```python
nemo_automodel.components.loss.utils._normalize_kd_labels(
    labels: torch.Tensor,
    loss_ignore_index: int,
    kd_ignore_index: int
) -> torch.Tensor
```

Align KD supervision with the main loss mask.

**Parameters:**

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

Integer target tensor of any shape.

---

**`loss_ignore_index`** `int`

Label sentinel consumed by the main loss.

---

**`kd_ignore_index`** `int`

Label sentinel consumed by the KD loss.

---

**Returns:** `torch.Tensor`

Target tensor with the same shape, dtype, and device as `labels`.

```python
nemo_automodel.components.loss.utils._normalize_loss_labels(
    labels: torch.Tensor,
    ignore_index: int
) -> torch.Tensor
```

Map dataset padding to a loss's configured ignore index.

**Parameters:**

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

Integer target tensor of any shape.

---

**`ignore_index`** `int`

Label sentinel consumed by the loss.

---

**Returns:** `torch.Tensor`

Target tensor with the same shape, dtype, and device as `labels`.

```python
nemo_automodel.components.loss.utils.calculate_loss(
    loss_fn: torch.nn.Module,
    kwargs: typing.Any = {}
) -> torch.Tensor
```

Calculate a logits-based or linear-projection cross-entropy loss.

**Parameters:**

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

Loss module. `LinearCrossEntropy` implementations consume
`hidden_states` with shape `[batch, sequence, hidden]`, labels
with shape `[batch, sequence]`, and an LM-head weight with global
shape `[vocab, hidden]`. Other loss modules consume logits with
shape `[batch, sequence, vocab]` and labels.

---

**`**kwargs`** `Any` — default: \{}

Loss inputs. Rank-local tensors keep their existing layout;
`grad_reduce_group` describes the ranks contributing independent
linear-loss shards. The caller's mapping and tensors are not mutated.

---

**Returns:** `torch.Tensor`

Scalar loss tensor that does not alias an input.

```python
nemo_automodel.components.loss.utils.prepare_lm_weight(
    loss_fn: nemo_automodel.components.loss.linear_ce_base.LinearCrossEntropy,
    model: torch.nn.Module,
    grad_reduce_group: torch.distributed.ProcessGroup | None = None
) -> torch.Tensor
```

Prepare one shared projection weight for main and auxiliary losses.

**Parameters:**

**`loss_fn`** `LinearCrossEntropy`

Linear-projection loss that owns materialization semantics.

---

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

Model or DDP wrapper whose wrapped module owns the config and
output head of global shape \[vocab, hidden].

---

**`grad_reduce_group`** `torch.distributed.ProcessGroup | None` — default: None

Group contributing independent token losses.

---

**Returns:** `torch.Tensor`

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

```python
nemo_automodel.components.loss.utils._DATASET_IGNORE_INDEX = -100
```