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

## Module Contents

### Classes

| Name                                                                                           | Description                                            |
| ---------------------------------------------------------------------------------------------- | ------------------------------------------------------ |
| [`FusedLinearCrossEntropy`](#nemo_automodel-components-loss-linear_ce-FusedLinearCrossEntropy) | Fused linear-projection and cross-entropy loss module. |

### Functions

| Name                                                                                                                     | Description                                                                               |
| ------------------------------------------------------------------------------------------------------------------------ | ----------------------------------------------------------------------------------------- |
| [`_get_triton_version`](#nemo_automodel-components-loss-linear_ce-_get_triton_version)                                   | -                                                                                         |
| [`new_is_triton_greater_or_equal`](#nemo_automodel-components-loss-linear_ce-new_is_triton_greater_or_equal)             | Check if pytorch-triton/triton version is greater than or equal to the specified version. |
| [`new_is_triton_greater_or_equal_3_2_0`](#nemo_automodel-components-loss-linear_ce-new_is_triton_greater_or_equal_3_2_0) | Check if pytorch-triton/triton version is greater than or equal to 3.1.0.                 |

### Data

[`HAVE_CUT_CROSS_ENTROPY`](#nemo_automodel-components-loss-linear_ce-HAVE_CUT_CROSS_ENTROPY)

### API

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

**Bases:** [LinearCrossEntropy](/nemo-automodel/nemo_automodel/components/loss/linear_ce_base#nemo_automodel-components-loss-linear_ce_base-LinearCrossEntropy)

Fused linear-projection and cross-entropy loss module.

```python
nemo_automodel.components.loss.linear_ce.FusedLinearCrossEntropy.forward(
    hidden_states: torch.Tensor,
    labels: torch.Tensor,
    lm_weight: torch.Tensor,
    num_label_tokens: int | None = None,
    grad_reduce_group: torch.distributed.ProcessGroup | None = None,
    loss_weights: torch.Tensor | None = None
) -> torch.Tensor
```

Compute fused linear cross entropy matching PyTorch behavior.

**Parameters:**

**`hidden_states`** `torch.Tensor`

Rank-local hidden states with shape
`[batch, sequence, hidden]`.

---

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

Rank-local target token IDs with shape `[batch, sequence]`.

---

**`lm_weight`** `torch.Tensor`

LM-head weight with global shape `[vocab, hidden]`.
It may be a regular tensor or an FSDP-sharded DTensor.

---

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

Global number of non-padding target tokens used
to normalize a sum-reduced loss.

---

**`grad_reduce_group`** `dist.ProcessGroup | None` — default: None

Group that contributes independent loss shards
when `lm_weight` is a sharded DTensor.

---

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

Optional per-token multipliers matching
`labels.shape`. Only supported with `reduction="sum"`.

---

**Returns:** `torch.Tensor`

Scalar loss tensor on the same device as `hidden_states`. The

```python
nemo_automodel.components.loss.linear_ce._get_triton_version()
```

```python
nemo_automodel.components.loss.linear_ce.new_is_triton_greater_or_equal(
    version_str
)
```

Check if pytorch-triton/triton version is greater than or equal to the specified version.

**Parameters:**

**`version_str`**

Version string to check

---

**Returns:**

True if pytorch-triton/triton version >= specified version

```python
nemo_automodel.components.loss.linear_ce.new_is_triton_greater_or_equal_3_2_0()
```

Check if pytorch-triton/triton version is greater than or equal to 3.1.0.

**Returns:**

True if pytorch-triton/triton version >= 3.1.0

```python
nemo_automodel.components.loss.linear_ce.HAVE_CUT_CROSS_ENTROPY = True
```