> For clean Markdown of any page, append .md to the page URL. > For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/v0.5/nemo-automodel/nemo_automodel/components/loss/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... ``` ## Docs - [nemo_automodel.components.loss.chunked_ce](https://docs.nvidia.com/nemo/automodel/v0.5/nemo-automodel/nemo_automodel/components/loss/chunked_ce.md) - [nemo_automodel.components.loss.dist_utils](https://docs.nvidia.com/nemo/automodel/v0.5/nemo-automodel/nemo_automodel/components/loss/dist_utils.md) - [nemo_automodel.components.loss.dllm_loss](https://docs.nvidia.com/nemo/automodel/v0.5/nemo-automodel/nemo_automodel/components/loss/dllm_loss.md) - [nemo_automodel.components.loss.embedding_distill](https://docs.nvidia.com/nemo/automodel/v0.5/nemo-automodel/nemo_automodel/components/loss/embedding_distill.md) - [nemo_automodel.components.loss.infonce](https://docs.nvidia.com/nemo/automodel/v0.5/nemo-automodel/nemo_automodel/components/loss/infonce.md) - [nemo_automodel.components.loss.intermediate_distill](https://docs.nvidia.com/nemo/automodel/v0.5/nemo-automodel/nemo_automodel/components/loss/intermediate_distill.md) - [nemo_automodel.components.loss.kd_loss](https://docs.nvidia.com/nemo/automodel/v0.5/nemo-automodel/nemo_automodel/components/loss/kd_loss.md) - [nemo_automodel.components.loss.linear_ce](https://docs.nvidia.com/nemo/automodel/v0.5/nemo-automodel/nemo_automodel/components/loss/linear_ce.md) - [nemo_automodel.components.loss.linear_ce_base](https://docs.nvidia.com/nemo/automodel/v0.5/nemo-automodel/nemo_automodel/components/loss/linear_ce_base.md) - [nemo_automodel.components.loss.listmle](https://docs.nvidia.com/nemo/automodel/v0.5/nemo-automodel/nemo_automodel/components/loss/listmle.md) - [nemo_automodel.components.loss.loss](https://docs.nvidia.com/nemo/automodel/v0.5/nemo-automodel/nemo_automodel/components/loss/loss.md) - [nemo_automodel.components.loss.masked_ce](https://docs.nvidia.com/nemo/automodel/v0.5/nemo-automodel/nemo_automodel/components/loss/masked_ce.md) - [nemo_automodel.components.loss.mtp](https://docs.nvidia.com/nemo/automodel/v0.5/nemo-automodel/nemo_automodel/components/loss/mtp.md) - [nemo_automodel.components.loss.soft_ce](https://docs.nvidia.com/nemo/automodel/v0.5/nemo-automodel/nemo_automodel/components/loss/soft_ce.md) - [nemo_automodel.components.loss.te_parallel_ce](https://docs.nvidia.com/nemo/automodel/v0.5/nemo-automodel/nemo_automodel/components/loss/te_parallel_ce.md) - [nemo_automodel.components.loss.utils](https://docs.nvidia.com/nemo/automodel/v0.5/nemo-automodel/nemo_automodel/components/loss/utils.md)