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

Gemma4-owned tensor-parallel plan and FSDP2 strategy registration.

## Module Contents

### Classes

| Name                                                                                                                              | Description                                                           |
| --------------------------------------------------------------------------------------------------------------------------------- | --------------------------------------------------------------------- |
| [`Gemma4ModelParallelizer`](#nemo_automodel-components-models-gemma4_moe-parallelization-Gemma4ModelParallelizer)                 | Apply the variant-aware Gemma4 TP plan before standard FSDP2.         |
| [`_Gemma4RowwiseParallel`](#nemo_automodel-components-models-gemma4_moe-parallelization-_Gemma4RowwiseParallel)                   | Shard a Gemma4 linear on input features with a synchronous reduction. |
| [`_Gemma4VocabParallelEmbedding`](#nemo_automodel-components-models-gemma4_moe-parallelization-_Gemma4VocabParallelEmbedding)     | Shard a Gemma4 embedding by vocabulary without DTensor `MaskPartial`. |
| [`_ReduceFromTensorParallelRegion`](#nemo_automodel-components-models-gemma4_moe-parallelization-_ReduceFromTensorParallelRegion) | Sum local TP values in forward and leave replicated gradients local.  |

### Functions

| Name                                                                                                                            | Description                                                               |
| ------------------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------- |
| [`_gemma4_rowwise_linear_forward`](#nemo_automodel-components-models-gemma4_moe-parallelization-_gemma4_rowwise_linear_forward) | Apply one Gemma4 row-parallel linear and synchronously reduce its output. |
| [`_gemma4_tp_plan`](#nemo_automodel-components-models-gemma4_moe-parallelization-_gemma4_tp_plan)                               | Return the TP plan for the concrete Gemma4 checkpoint variant.            |
| [`_gemma4_vocab_parallel_forward`](#nemo_automodel-components-models-gemma4_moe-parallelization-_gemma4_vocab_parallel_forward) | Look up replicated token ids in one Gemma4 vocabulary shard.              |
| [`_get_attention_head_counts`](#nemo_automodel-components-models-gemma4_moe-parallelization-_get_attention_head_counts)         | Return every per-layer attention-head shape used by a Gemma4 config.      |

### Data

[`PARALLELIZER`](#nemo_automodel-components-models-gemma4_moe-parallelization-PARALLELIZER)

### API

```python
class nemo_automodel.components.models.gemma4_moe.parallelization.Gemma4ModelParallelizer()
```

**Bases:** `ModelParallelizer`

Apply the variant-aware Gemma4 TP plan before standard FSDP2.

```python
nemo_automodel.components.models.gemma4_moe.parallelization.Gemma4ModelParallelizer._apply(
    model: torch.nn.Module,
    device_mesh: torch.distributed.device_mesh.DeviceMesh,
    kwargs: typing.Any = {}
) -> torch.nn.Module
```

Apply Gemma4 tensor parallelism.

```python
nemo_automodel.components.models.gemma4_moe.parallelization.Gemma4ModelParallelizer._use_full_layer_activation_checkpointing(
    model: torch.nn.Module
) -> bool
```

Keep Gemma4's replay-safe shared attention inside the checkpoint region.

```python
nemo_automodel.components.models.gemma4_moe.parallelization.Gemma4ModelParallelizer._validate_tp_mesh(
    model: torch.nn.Module,
    tp_mesh: torch.distributed.device_mesh.DeviceMesh
) -> None
```

Validate every heterogeneous Gemma4 attention-head shape.

```python
class nemo_automodel.components.models.gemma4_moe.parallelization._Gemma4RowwiseParallel(
    output_holder: list[torch.Tensor | None]
)
```

**Bases:** `ParallelStyle`

Shard a Gemma4 linear on input features with a synchronous reduction.

```python
nemo_automodel.components.models.gemma4_moe.parallelization._Gemma4RowwiseParallel._apply(
    module: torch.nn.Module,
    device_mesh: torch.distributed.device_mesh.DeviceMesh
) -> torch.nn.Module
```

Partition `[output, input]` weight as `Shard(1)` and install the local linear.

```python
class nemo_automodel.components.models.gemma4_moe.parallelization._Gemma4VocabParallelEmbedding()
```

**Bases:** `ParallelStyle`

Shard a Gemma4 embedding by vocabulary without DTensor `MaskPartial`.

The input ids are replicated `[batch, sequence]`. The weight changes from
`[vocab, embedding]` to a DTensor with `Shard(0)`; the forward returns a
local, replicated `[batch, sequence, embedding]` Tensor.

```python
nemo_automodel.components.models.gemma4_moe.parallelization._Gemma4VocabParallelEmbedding._apply(
    module: torch.nn.Module,
    device_mesh: torch.distributed.device_mesh.DeviceMesh
) -> torch.nn.Module
```

Partition one Gemma4 embedding weight and install its local lookup.

```python
class nemo_automodel.components.models.gemma4_moe.parallelization._ReduceFromTensorParallelRegion()
```

**Bases:** `Function`

Sum local TP values in forward and leave replicated gradients local.

```python
nemo_automodel.components.models.gemma4_moe.parallelization._ReduceFromTensorParallelRegion.backward(
    ctx,
    grad_output: torch.Tensor
)
```

staticmethod

Pass replicated gradients unchanged to each owning TP shard.

```python
nemo_automodel.components.models.gemma4_moe.parallelization._ReduceFromTensorParallelRegion.forward(
    ctx,
    local_output: torch.Tensor,
    process_group
) -> torch.Tensor
```

staticmethod

Reduce a local contribution Tensor into one replicated Tensor.

```python
nemo_automodel.components.models.gemma4_moe.parallelization._gemma4_rowwise_linear_forward(
    module: torch.nn.Linear,
    input_: torch.Tensor
) -> torch.Tensor
```

Apply one Gemma4 row-parallel linear and synchronously reduce its output.

**Parameters:**

**`module`** `nn.Linear`

Linear whose weight is a DTensor sharded on input features.

---

**`input_`** `torch.Tensor`

Local Tensor shaped `[..., input_features / tp_size]`.

---

**Returns:** `torch.Tensor`

Local Tensor shaped `[..., output_features]`, replicated across the

```python
nemo_automodel.components.models.gemma4_moe.parallelization._gemma4_tp_plan(
    model: torch.nn.Module,
    sequence_parallel: bool = False
) -> dict[str, torch.distributed.tensor.parallel.ParallelStyle]
```

Return the TP plan for the concrete Gemma4 checkpoint variant.

E2B/E4B add a packed per-layer embedding table that is absent from 31B.
Row-sharding that table by vocabulary is important: at E4B dimensions it is
roughly 2.8 billion parameters, so leaving it replicated defeats much of
the memory benefit of TP.

**Parameters:**

**`model`** `nn.Module`

Gemma4 conditional-generation model whose text config selects
the dense variant.

---

**`sequence_parallel`** `bool` — default: False

Whether sequence parallelism was requested. Gemma4
does not currently support it, so the request is ignored.

---

**Returns:** `dict[str, ParallelStyle]`

Mapping from Gemma4 module paths to PyTorch parallel styles.

```python
nemo_automodel.components.models.gemma4_moe.parallelization._gemma4_vocab_parallel_forward(
    module: torch.nn.Embedding,
    input_ids: torch.Tensor
) -> torch.Tensor
```

Look up replicated token ids in one Gemma4 vocabulary shard.

**Parameters:**

**`module`** `nn.Embedding`

Gemma4 embedding whose weight is a DTensor sharded on vocabulary.

---

**`input_ids`** `torch.Tensor`

Replicated integer Tensor shaped `[batch, sequence]`.

---

**Returns:** `torch.Tensor`

Local Tensor shaped `[batch, sequence, embedding]`, replicated across

```python
nemo_automodel.components.models.gemma4_moe.parallelization._get_attention_head_counts(
    text_config
) -> set[tuple[int, int]]
```

Return every per-layer attention-head shape used by a Gemma4 config.

```python
nemo_automodel.components.models.gemma4_moe.parallelization.PARALLELIZER = Gemma4ModelParallelizer()
```