> This page is for version Nightly (default).
> 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.7.0 · 26.10: https://docs.nvidia.com/nemo/automodel/v0.7/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.

# nemo_automodel.components.models.inkling2.parallelization

Inkling precision policy on Automodel's shared FSDP and expert parallelizers.

## Module Contents

### Classes

| Name                                                                                                      | Description                                                             |
| --------------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------- |
| [`Inkling2Parallelizer`](#nemo_automodel-components-models-inkling2-parallelization-Inkling2Parallelizer) | Keep model-owned casts while reusing shared sharding and communication. |

### Functions

| Name                                                                                                      | Description                                                                 |
| --------------------------------------------------------------------------------------------------------- | --------------------------------------------------------------------------- |
| [`_is_fp32_module_name`](#nemo_automodel-components-models-inkling2-parallelization-_is_fp32_module_name) | Return whether a dotted module name is one of the model-owned FP32 modules. |
| [`_shard`](#nemo_automodel-components-models-inkling2-parallelization-_shard)                             | Make `module` one FSDP unit that leaves already-owned parameters alone.     |
| [`iter_fp32_islands`](#nemo_automodel-components-models-inkling2-parallelization-iter_fp32_islands)       | Yield the FP32 modules below `module` that need their own FSDP unit.        |

### Data

[`FP32_MODULES`](#nemo_automodel-components-models-inkling2-parallelization-FP32_MODULES)

### API

```python
class nemo_automodel.components.models.inkling2.parallelization.Inkling2Parallelizer()
```

**Bases:** [ModelParallelizer](/nemo/automodel/nemo-automodel/nemo_automodel/components/distributed/parallelizer#nemo_automodel-components-distributed-parallelizer-ModelParallelizer)

Keep model-owned casts while reusing shared sharding and communication.

```python
nemo_automodel.components.models.inkling2.parallelization.Inkling2Parallelizer._fully_shard_module(
    module: torch.nn.Module,
    mesh: torch.distributed.device_mesh.DeviceMesh,
    mp_policy: torch.distributed.fsdp.MixedPrecisionPolicy | None = None,
    offload_policy: torch.distributed.fsdp.OffloadPolicy | None = None,
    reshard_after_forward: bool | int | None = None,
    ignored_params: set[torch.nn.Parameter] | None = None
) -> torch.nn.Module
```

Shard the FP32 modules as their own units, then the module around them.

**Parameters:**

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

Owner of matrices \[output, input], expert tensors \[experts,
input, output], and vectors \[features]. Nested FSDP units keep ownership.

---

**`mesh`** `DeviceMesh`

One-dimensional DP shard mesh or two-dimensional HSDP mesh.

---

**`mp_policy`** `MixedPrecisionPolicy | None` — default: None

Requested compute policy; model code owns input/output casts.

---

**`offload_policy`** `OffloadPolicy | None` — default: None

Shared parameter storage/offload policy.

---

**`reshard_after_forward`** `bool | int | None` — default: None

Shared FSDP reshard setting.

---

**`ignored_params`** `set[nn.Parameter] | None` — default: None

Parameters with the module's tensor layouts already
owned by another parallelism unit; their placements are preserved.

---

**Returns:** `nn.Module`

The same module with global-shaped DTensor parameters. Dense parameters

```python
nemo_automodel.components.models.inkling2.parallelization.Inkling2Parallelizer._prepare_model(
    model: torch.nn.Module,
    mesh_context: nemo_automodel.components.distributed.mesh.MeshContext
) -> None
```

Validate qualified text topology before shared parallelization mutates ownership.

**Parameters:**

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

Inkling text model with shared grouped experts and an untied head.

---

**`mesh_context`** `MeshContext`

Shared meshes and strategy/compute policy.

---

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

Residual histories are explicit immutable inputs to each decoder layer.

```python
nemo_automodel.components.models.inkling2.parallelization._is_fp32_module_name(
    name: str
) -> bool
```

Return whether a dotted module name is one of the model-owned FP32 modules.

**Parameters:**

**`name`** `str`

Module name relative to the module being sharded; activation
checkpoint wrappers add a `_checkpoint_wrapped_module` segment.

---

**Returns:** `bool`

True when the name equals or ends with an `FP32_MODULES` entry.

```python
nemo_automodel.components.models.inkling2.parallelization._shard(
    module: torch.nn.Module,
    mesh: torch.distributed.device_mesh.DeviceMesh,
    mp_policy: torch.distributed.fsdp.MixedPrecisionPolicy | None,
    offload_policy: torch.distributed.fsdp.OffloadPolicy | None,
    reshard_after_forward: bool | int | None,
    owned: set[torch.nn.Parameter]
) -> None
```

Make `module` one FSDP unit that leaves already-owned parameters alone.

**Parameters:**

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

Module whose remaining parameters become one FSDP unit.

---

**`mesh`** `DeviceMesh`

One-dimensional DP shard mesh or two-dimensional HSDP mesh.

---

**`mp_policy`** `MixedPrecisionPolicy | None`

Compute policy of this unit.

---

**`offload_policy`** `OffloadPolicy | None`

Shared parameter storage/offload policy.

---

**`reshard_after_forward`** `bool | int | None`

Shared FSDP reshard setting.

---

**`owned`** `set[nn.Parameter]`

Parameters owned elsewhere; only those inside `module` are passed on.

---

```python
nemo_automodel.components.models.inkling2.parallelization.iter_fp32_islands(
    module: torch.nn.Module,
    owned: set[torch.nn.Parameter]
) -> collections.abc.Iterator[torch.nn.Module]
```

Yield the FP32 modules below `module` that need their own FSDP unit.

**Parameters:**

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

Module about to become an FSDP unit; only its descendants are searched.

---

**`owned`** `set[nn.Parameter]`

Parameters already owned by nested FSDP units or by expert parallelism.

---

```python
nemo_automodel.components.models.inkling2.parallelization.FP32_MODULES = ('attn_query', 'mlp_query', 'output_query', 'q_l2_clip', 'rel_l2_clip', 'mlp.gat...
```