> 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

## Submodules

* **[`nemo_automodel.components.loss.chunked_ce`](/nemo-automodel/nemo_automodel/components/loss/chunked_ce)**
* **[`nemo_automodel.components.loss.dist_utils`](/nemo-automodel/nemo_automodel/components/loss/dist_utils)**
* **[`nemo_automodel.components.loss.dllm_loss`](/nemo-automodel/nemo_automodel/components/loss/dllm_loss)**
* **[`nemo_automodel.components.loss.embedding_distill`](/nemo-automodel/nemo_automodel/components/loss/embedding_distill)**
* **[`nemo_automodel.components.loss.infonce`](/nemo-automodel/nemo_automodel/components/loss/infonce)**
* **[`nemo_automodel.components.loss.intermediate_distill`](/nemo-automodel/nemo_automodel/components/loss/intermediate_distill)**
* **[`nemo_automodel.components.loss.kd_loss`](/nemo-automodel/nemo_automodel/components/loss/kd_loss)**
* **[`nemo_automodel.components.loss.linear_ce`](/nemo-automodel/nemo_automodel/components/loss/linear_ce)**
* **[`nemo_automodel.components.loss.linear_ce_base`](/nemo-automodel/nemo_automodel/components/loss/linear_ce_base)**
* **[`nemo_automodel.components.loss.listmle`](/nemo-automodel/nemo_automodel/components/loss/listmle)**
* **[`nemo_automodel.components.loss.loss`](/nemo-automodel/nemo_automodel/components/loss/loss)**
* **[`nemo_automodel.components.loss.masked_ce`](/nemo-automodel/nemo_automodel/components/loss/masked_ce)**
* **[`nemo_automodel.components.loss.mtp`](/nemo-automodel/nemo_automodel/components/loss/mtp)**
* **[`nemo_automodel.components.loss.soft_ce`](/nemo-automodel/nemo_automodel/components/loss/soft_ce)**
* **[`nemo_automodel.components.loss.te_parallel_ce`](/nemo-automodel/nemo_automodel/components/loss/te_parallel_ce)**
* **[`nemo_automodel.components.loss.utils`](/nemo-automodel/nemo_automodel/components/loss/utils)**

## Package Contents

### Classes

| Name                                                                                        | Description                                                                  |
| ------------------------------------------------------------------------------------------- | ---------------------------------------------------------------------------- |
| [`ChunkedCEConfig`](#nemo_automodel-components-loss-loss-ChunkedCEConfig)                   | Chunked output projection and cross entropy.                                 |
| [`FusedLinearCEConfig`](#nemo_automodel-components-loss-loss-FusedLinearCEConfig)           | `FusedLinearCrossEntropy`.                                                   |
| [`KDLossConfig`](#nemo_automodel-components-loss-loss-KDLossConfig)                         | `KDLoss` (knowledge distillation).                                           |
| [`LossConfig`](#nemo_automodel-components-loss-loss-LossConfig)                             | Base loss config.  Subclasses expose their full field surface and            |
| [`LossFromFactoryConfig`](#nemo_automodel-components-loss-loss-LossFromFactoryConfig)       | Escape hatch for external integrations (e.g. veRL) and the YAML recipe path. |
| [`MaskedCrossEntropyConfig`](#nemo_automodel-components-loss-loss-MaskedCrossEntropyConfig) | `MaskedCrossEntropy`.                                                        |
| [`TEParallelCEConfig`](#nemo_automodel-components-loss-loss-TEParallelCEConfig)             | `TEParallelCrossEntropy`.                                                    |

### Functions

| Name                                                                          | Description                                                       |
| ----------------------------------------------------------------------------- | ----------------------------------------------------------------- |
| [`build_loss_config`](#nemo_automodel-components-loss-loss-build_loss_config) | Normalize a loss `loss` target plus `kwargs` into a `LossConfig`. |
| [`build_loss_module`](#nemo_automodel-components-loss-loss-build_loss_module) | Build a loss function.                                            |

### Data

[`LOSS_CONFIG_REGISTRY`](#nemo_automodel-components-loss-loss-LOSS_CONFIG_REGISTRY)

### API

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

Dataclass

**Bases:** [LossConfig](#nemo_automodel-components-loss-loss-LossConfig)

Chunked output projection and cross entropy.

**`chunk_len`** `int = 512`

Maximum token rows per output projection.

---

**`compile`** `bool = True`

Compile each chunk's projection and CE.

---

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

Ignored target value.

---

**`reduction`** `str = 'sum'`

Reduction across all tokens.

---

```python
nemo_automodel.components.loss.loss.ChunkedCEConfig.build() -> torch.nn.Module
```

```python
class nemo_automodel.components.loss.loss.FusedLinearCEConfig(
    ignore_index: int = -100,
    logit_softcapping: float = 0.0,
    reduction: str = 'sum',
    impl: typing.Literal['cce', 'torch_compile'] = 'cce'
)
```

Dataclass

**Bases:** [LossConfig](#nemo_automodel-components-loss-loss-LossConfig)

`FusedLinearCrossEntropy`.

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

Label value marking padding tokens.

---

**`impl`** `Literal['cce', 'torch_compile'] = 'cce'`

`"cce"` (flat memory, slowest) or `"torch_compile"` (faster,
memory grows with supervised tokens x vocab).

---

**`logit_softcapping`** `float = 0.0`

Softcap logits before CE (0 = disabled).

---

**`reduction`** `str = 'sum'`

Reduction mode.

---

```python
nemo_automodel.components.loss.loss.FusedLinearCEConfig.build() -> torch.nn.Module
```

```python
class nemo_automodel.components.loss.loss.KDLossConfig(
    ignore_index: int = -100,
    temperature: float = 1.0,
    fp32_upcast: bool = True,
    tp_group: typing.Any = None,
    chunk_size: int = 0
)
```

Dataclass

**Bases:** [LossConfig](#nemo_automodel-components-loss-loss-LossConfig)

`KDLoss` (knowledge distillation).

**`chunk_size`** `int = 0`

Number of valid tokens per KD loss chunk (0 = no chunking).

---

**`fp32_upcast`** `bool = True`

Cast logits to float32 for numerical stability.

---

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

Label value marking padding tokens.

---

**`temperature`** `float = 1.0`

Softmax temperature T.  Loss is scaled by T².

---

**`tp_group`** `Any = None`

Tensor-parallel process group the loss reduces over (runtime
arg; usually left as `None` in YAML and supplied programmatically).

---

```python
nemo_automodel.components.loss.loss.KDLossConfig.build() -> torch.nn.Module
```

```python
class nemo_automodel.components.loss.loss.LossConfig()
```

Dataclass

Base loss config.  Subclasses expose their full field surface and
implement `build`.

```python
nemo_automodel.components.loss.loss.LossConfig.build() -> torch.nn.Module
```

Construct the loss module.

```python
class nemo_automodel.components.loss.loss.LossFromFactoryConfig(
    factory: collections.abc.Callable[..., torch.nn.Module] | None = None,
    kwargs: dict[str, typing.Any] = dict()
)
```

Dataclass

**Bases:** [LossConfig](#nemo_automodel-components-loss-loss-LossConfig)

Escape hatch for external integrations (e.g. veRL) and the YAML recipe path.

Rather than exposing typed fields, it wraps a loss class/callable and the
`**kwargs` to construct it, keeping the factory path on the same
`config.build()` contract as the typed configs so callers never have to
special-case it.  The factory is called as `factory(**kwargs)`.

**`factory`** `Callable[..., Module] | None = None`

---

**`kwargs`** `dict[str, Any] = field(default_factory=dict)`

---

```python
nemo_automodel.components.loss.loss.LossFromFactoryConfig.build() -> torch.nn.Module
```

```python
class nemo_automodel.components.loss.loss.MaskedCrossEntropyConfig(
    fp32_upcast: bool = True,
    ignore_index: int = -100,
    reduction: str = 'sum'
)
```

Dataclass

**Bases:** [LossConfig](#nemo_automodel-components-loss-loss-LossConfig)

`MaskedCrossEntropy`.

**`fp32_upcast`** `bool = True`

Cast logits to float32 before computing CE.

---

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

Label value marking padding tokens.

---

**`reduction`** `str = 'sum'`

Reduction mode — `"sum"`, `"mean"`, or `"none"`.

---

```python
nemo_automodel.components.loss.loss.MaskedCrossEntropyConfig.build() -> torch.nn.Module
```

```python
class nemo_automodel.components.loss.loss.TEParallelCEConfig(
    ignore_index: int = -100,
    reduction: str = 'sum',
    tp_group: typing.Any = None
)
```

Dataclass

**Bases:** [LossConfig](#nemo_automodel-components-loss-loss-LossConfig)

`TEParallelCrossEntropy`.

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

Label value marking padding tokens.

---

**`reduction`** `str = 'sum'`

Reduction mode.

---

**`tp_group`** `Any = None`

Tensor-parallel process group the loss reduces over (runtime
arg; usually left as `None` in YAML and supplied programmatically).

---

```python
nemo_automodel.components.loss.loss.TEParallelCEConfig.build() -> torch.nn.Module
```

```python
nemo_automodel.components.loss.loss.build_loss_config(
    loss: nemo_automodel.components.loss.loss.LossConfig | str | collections.abc.Callable[..., torch.nn.Module],
    loss_kwargs: typing.Any = {}
) -> nemo_automodel.components.loss.loss.LossConfig
```

Normalize a loss `loss` target plus `kwargs` into a `LossConfig`.

The single normalization entry point shared by the recipe layer (which
resolves a YAML `_target_` to a Python object) and `build_loss_module`.
It dispatches on `loss`:

* a typed `LossConfig` instance — returned as-is; `**loss_kwargs`
  must be empty (hyperparameters live on the config).
* a `LossConfig` subclass (the class) — rejected; pass an instance.
* a registry name (see `LOSS_CONFIG_REGISTRY`, e.g.
  `"MaskedCrossEntropy"`) — built from the matching typed config.
* a loss class / callable plus `**loss_kwargs` — if it is a registered loss
  class and the kwargs fit the config's fields, it is upgraded to that typed
  config; otherwise it is wrapped in a `LossFromFactoryConfig`.  The
  caller resolves any dotted path to a callable; the component never does
  dotted-path resolution.

**Returns:** `LossConfig`

class:`LossConfig` ready to `build()`.

```python
nemo_automodel.components.loss.loss.build_loss_module(
    loss: nemo_automodel.components.loss.loss.LossConfig | collections.abc.Callable[..., torch.nn.Module],
    loss_kwargs: typing.Any = {}
) -> torch.nn.Module
```

Build a loss function.

Thin dispatcher: it normalizes `loss` to a `LossConfig` via
`build_loss_config` and returns `config.build()`.  Dispatches on
`loss`:

* **Typed config** (`LossConfig` instance) — the Automodel-native
  path.  Hyperparameters come from the config; `**loss_kwargs` must be
  empty.
* **Loss class / callable** (e.g. `MaskedCrossEntropy`) plus
  `**loss_kwargs` — the integration / YAML escape hatch.  The caller
  resolves any dotted path to a callable; the component never does string
  resolution.

**Parameters:**

**`loss`** `LossConfig | Callable[..., nn.Module]`

Typed `LossConfig` instance, or a loss class/callable to
construct with `**loss_kwargs`.

---

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

Constructor kwargs for the class/callable form.  Must be
empty when `loss` is a typed config.

---

**Returns:** `nn.Module`

Instantiated loss function.

```python
nemo_automodel.components.loss.loss.LOSS_CONFIG_REGISTRY: dict[str, type[LossConfig]] = {'MaskedCrossEntropy': MaskedCrossEntropyConfig, 'FusedLinearCrossEntropy': Fuse...
```