> 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.diffusion_gemma.fsdp

FSDP2 sharding for `diffusion_gemma` under pure FSDP (`ep_size=1`).

At `ep_size=1` there is no MoE mesh, so the model is sharded by the generic
`ModelParallelizer`
(via `ModelParallelizer.parallelize`), which applies `fully_shard` per decoder
layer and to the root.  The generic `fully_shard` flattens *all* of a decoder
layer's parameters into one FSDP unit, which folds each layer's grouped-expert
tensors (`moe.experts.{gate_and_up_projs,down_projs}`, the bulk of the 26B
parameters) into the layer's single all-gather.  That gathers every expert of a
layer at once on each forward — a large activation-memory spike for a model that
runs the shared stack twice (causal encode + bidirectional decode, plus an
optional self-conditioning pass).

`fully_shard_diffusion_gemma` mirrors `deepseek_v4`'s
`fully_shard_deepseek_v4`: it makes `moe.experts` its **own** FSDP unit
(sharded dim-0 on the dp mesh) *before* wrapping the rest of the decoder layer.
Consequences:

* The grouped-expert parameters become global-`[n_experts]` `Shard(0)`
  DTensors on the dp mesh, so DCP sees the checkpoint's global expert shape and
  each rank reads only its shard (no `[128] vs [16]` size mismatch).
* During the experts' forward, FSDP all-gathers their parameters back to the
  full `[n_experts, ...]` tensor, so `GroupedExperts` sees a plain
  (non-DTensor) tensor and runs with `ep_size == 1` — all experts local, no
  expert-parallel token shuffle.  This is **pure FSDP, not EP**.
* Experts gather/reshard independently of the rest of the layer, bounding peak
  memory across the double (encode + decode) pass.
* `moe.experts` becomes a distinct `FSDPModule` that
  `MoEFSDPSyncMixin._iter_fsdp_modules` discovers (`block.moe.experts`).

No expert parallelism is introduced; this is the `ep_size=1` path only.

## Module Contents

### Classes

| Name                                                                                                                        | Description                                                         |
| --------------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------- |
| [`DiffusionGemmaModelParallelizer`](#nemo_automodel-components-models-diffusion_gemma-fsdp-DiffusionGemmaModelParallelizer) | Pure-FSDP2 strategy that shards grouped experts as their own units. |

### Functions

| Name                                                                                                                | Description                                                             |
| ------------------------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------- |
| [`_fully_shard_once`](#nemo_automodel-components-models-diffusion_gemma-fsdp-_fully_shard_once)                     | Apply `fully_shard` to `module` unless it is already an FSDP unit.      |
| [`_has_fsdp_state`](#nemo_automodel-components-models-diffusion_gemma-fsdp-_has_fsdp_state)                         | Return True if `module` has already been wrapped by `fully_shard`.      |
| [`fully_shard_diffusion_gemma`](#nemo_automodel-components-models-diffusion_gemma-fsdp-fully_shard_diffusion_gemma) | Apply FSDP2 to a `diffusion_gemma` decoder layer (or any other module). |

### Data

[`PARALLELIZER`](#nemo_automodel-components-models-diffusion_gemma-fsdp-PARALLELIZER)

### API

```python
class nemo_automodel.components.models.diffusion_gemma.fsdp.DiffusionGemmaModelParallelizer()
```

**Bases:** `ModelParallelizer`

Pure-FSDP2 strategy that shards grouped experts as their own units.

```python
nemo_automodel.components.models.diffusion_gemma.fsdp.DiffusionGemmaModelParallelizer._fully_shard_module(
    module,
    kwargs = {}
)
```

```python
nemo_automodel.components.models.diffusion_gemma.fsdp._fully_shard_once(
    module: torch.nn.Module,
    mesh,
    mp_policy,
    offload_policy,
    fsdp_kwargs = {}
) -> torch.nn.Module
```

Apply `fully_shard` to `module` unless it is already an FSDP unit.

```python
nemo_automodel.components.models.diffusion_gemma.fsdp._has_fsdp_state(
    module: torch.nn.Module
) -> bool
```

Return True if `module` has already been wrapped by `fully_shard`.

```python
nemo_automodel.components.models.diffusion_gemma.fsdp.fully_shard_diffusion_gemma(
    module: torch.nn.Module,
    mesh,
    mp_policy,
    offload_policy = None,
    fsdp_kwargs = {}
) -> torch.nn.Module
```

Apply FSDP2 to a `diffusion_gemma` decoder layer (or any other module).

For a `DiffusionGemmaMoEDecoderLayer`, shard its grouped experts
(`moe.experts`) as a separate FSDP unit first, then shard the rest of the
layer.  All other modules (embeddings, final norm, self-conditioning, the
root model) are sharded as a single unit.

**Parameters:**

**`module`** `nn.Module`

The module to shard (a decoder layer or the root model).

---

**`mesh`**

The (1-D) data-parallel device mesh to shard across.

---

**`mp_policy`**

FSDP2 mixed-precision policy.

---

**`offload_policy`** — default: None

Optional FSDP2 CPU-offload policy.

---

**`**fsdp_kwargs`** — default: \{}

Forwarded to `fully_shard` (e.g. `reshard_after_forward`).

---

**Returns:** `nn.Module`

The sharded module.

```python
nemo_automodel.components.models.diffusion_gemma.fsdp.PARALLELIZER = DiffusionGemmaModelParallelizer()
```