> This page is for version 0.7.0 · 26.10.
> 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.
> 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.flint.parallelization

Flint-owned FSDP precision boundaries for learned sinks and the output head.

## Module Contents

### Classes

| Name                                                                                                       | Description                                                         |
| ---------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------- |
| [`FlintModelParallelizer`](#nemo_automodel-components-models-flint-parallelization-FlintModelParallelizer) | Preserve FP32 compute units within the existing generic FSDP2 flow. |

### API

```python
class nemo_automodel.components.models.flint.parallelization.FlintModelParallelizer()
```

**Bases:** `ModelParallelizer`

Preserve FP32 compute units within the existing generic FSDP2 flow.

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

Apply inherited FSDP policies and retain this stage's CP token mesh.

**Parameters:**

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

Native full model or pruned PP stage; parameter tensor layouts
remain \[out,in] projections and grouped expert matrices.

---

**`device_mesh`** `DeviceMesh`

Resolved shared mesh, including its named CP axis.

---

**`**kwargs`** `Any` — default: \{}

Existing shared FSDP/AC/compile policies, passed unchanged.

---

**Returns:** `nn.Module`

Parallelized native module with its optional CP mesh. This mesh owns

```python
nemo_automodel.components.models.flint.parallelization.FlintModelParallelizer._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
```

Apply the shared primitive with Flint's two protected dtype boundaries.

**Parameters:**

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

Decoder block, optional activation-checkpoint wrapper, input
embedding, output head, or outer model. Block ordinary parameters
use native \[out,in] matrices and fused expert \[experts,hidden,2\*intermediate]
/ \[experts,intermediate,hidden] layouts. Sink holders own \[heads]
parameters; the head owns \[vocab,hidden].

---

**`mesh`** `DeviceMesh`

FSDP data-parallel mesh. Parameters become DTensors sharded by
the shared FSDP primitive; callers must enter on all mesh ranks.

---

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

Bulk compute/reduction policy; FP32 master storage remains
unchanged. Sink holders and the independently wrapped head use FP32.

---

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

Existing FSDP storage-offload policy.

---

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

Existing FSDP reshard policy, or its default.

---

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

Runtime parameters owned by another FSDP unit, if any.
They may have arbitrary global/native shapes or rank-local
shards; their existing DTensor mesh/placements are preserved.

---

**Returns:** `nn.Module`

The same module, mutated in place into FSDP units. Existing shapes