> 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                                                           |
| --------------------------------------------------------------------------------------------------------------------------------- | --------------------------------------------------------------------- |
| [`_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.              |
| [`register_gemma4_parallel_strategy`](#nemo_automodel-components-models-gemma4_moe-parallelization-register_gemma4_parallel_strategy) | Register Gemma4's model-owned FSDP2 strategy once.                        |

### API

```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.register_gemma4_parallel_strategy() -> None
```

Register Gemma4's model-owned FSDP2 strategy once.