> 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.models.gemma4_moe.loss

Loss adapters for Gemma4 tensor parallelism.

## Module Contents

### Classes

| Name                                                                                                                                           | Description                                                    |
| ---------------------------------------------------------------------------------------------------------------------------------------------- | -------------------------------------------------------------- |
| [`Gemma4TensorParallelFusedLinearCrossEntropy`](#nemo_automodel-components-models-gemma4_moe-loss-Gemma4TensorParallelFusedLinearCrossEntropy) | Fused linear CE for Gemma4's tied, vocabulary-sharded LM head. |

### API

```python
class nemo_automodel.components.models.gemma4_moe.loss.Gemma4TensorParallelFusedLinearCrossEntropy()
```

**Bases:** [FusedLinearCrossEntropy](/nemo-automodel/nemo_automodel/components/loss/linear_ce#nemo_automodel-components-loss-linear_ce-FusedLinearCrossEntropy)

Fused linear CE for Gemma4's tied, vocabulary-sharded LM head.

Gemma4 E-series TP keeps the tied embedding/LM-head weight sharded on the
TP vocabulary axis. Cut cross entropy needs a rank-local full vocabulary
weight, but the generic loss intentionally rejects a DTensor mesh that does
not match the DP/CP loss-reduction group. This adapter gathers only the
Gemma4 TP layout and marks its backward gradient replicated: TP peers see
the same tokens and therefore must not sum duplicate weight gradients.
FSDP continues to reduce independent DP/CP contributions normally.

```python
nemo_automodel.components.models.gemma4_moe.loss.Gemma4TensorParallelFusedLinearCrossEntropy.materialize_lm_weight(
    lm_weight: torch.Tensor,
    grad_reduce_group: torch.distributed.ProcessGroup | None = None
) -> torch.Tensor
```

staticmethod

Materialize a Gemma4 TP LM head with correct TP gradient semantics.

**Parameters:**

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

Regular or DTensor LM-head weight shaped
`[vocab, hidden]`.

---

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

DP/CP group whose ranks own independent tokens.
It remains owned by FSDP when the visible DTensor mesh contains
only TP, as it does while an FSDP unit is unsharded.

---

**Returns:** `torch.Tensor`

A rank-local full-vocabulary tensor. Backward slices, rather than