nemo_automodel.components.models.inkling2.parallelization

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

Module Contents

Classes

Name	Description
`Inkling2Parallelizer`	Keep model-owned casts while reusing shared sharding and communication.

Functions

Name	Description
`_is_fp32_module_name`	Return whether a dotted module name is one of the model-owned FP32 modules.
`_shard`	Make `module` one FSDP unit that leaves already-owned parameters alone.
`iter_fp32_islands`	Yield the FP32 modules below `module` that need their own FSDP unit.

Data

FP32_MODULES

API

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

Bases: ModelParallelizer

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

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

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

offload_policy

OffloadPolicy | NoneDefaults to None

Shared parameter storage/offload policy.

reshard_after_forward

bool | int | NoneDefaults to None

Shared FSDP reshard setting.

ignored_params

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

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.

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.

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.

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.

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.

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