> 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                                                                |
| -------------------------------------------------------------------------------------------- | -------------------------------------------------------------------------- |
| [`_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. |
| [`calculate_loss`](#nemo_automodel-components-loss-utils-calculate_loss)                     | Calculate a logit-based or fused linear cross-entropy loss.                |

### API

```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.calculate_loss(
    loss_fn: torch.nn.Module,
    kwargs: typing.Any = {}
) -> torch.Tensor
```

Calculate a logit-based or fused linear cross-entropy loss.

**Parameters:**

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

Loss module. `FusedLinearCrossEntropy` consumes
`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
fused-loss shards. The caller's mapping and tensors are not mutated.

---

**Returns:** `torch.Tensor`

Scalar loss tensor that does not alias an input.