nemo_automodel.components.models.flint.parallelization

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

Module Contents

Classes

Name	Description
`FlintModelParallelizer`	Preserve FP32 compute units within the existing generic FSDP2 flow.

API

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

Bases: ModelParallelizer

Preserve FP32 compute units within the existing generic FSDP2 flow.

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

AnyDefaults to {}

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

Returns: nn.Module

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

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 | NoneDefaults to None

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

offload_policy

OffloadPolicy | NoneDefaults to None

Existing FSDP storage-offload policy.

reshard_after_forward

bool | int | NoneDefaults to None

Existing FSDP reshard policy, or its default.

ignored_params

set[nn.Parameter] | NoneDefaults to 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