nemo_automodel.components.models.mistral_large4.parallelization

Model-owned FSDP precision and activation-checkpointing policy.

Module Contents

Classes

Name	Description
`MistralLarge4Parallelizer`	Use existing FSDP machinery while preserving Large 4 precision and sharing.

API

class nemo_automodel.components.models.mistral_large4.parallelization.MistralLarge4Parallelizer()

Bases: ModelParallelizer

Use existing FSDP machinery while preserving Large 4 precision and sharing.

Expert ownership and dispatch use the common MoE parallelizer.

nemo_automodel.components.models.mistral_large4.parallelization.MistralLarge4Parallelizer._apply(
    model: torch.nn.Module,
    device_mesh: torch.distributed.device_mesh.DeviceMesh,
    kwargs = {}
) -> torch.nn.Module

Apply the inherited FSDP flow and retain the CP token mesh.

Parameters:

model

nn.Module

Native full model or pruned PP stage.

device_mesh

DeviceMesh

Resolved shared mesh, including its named CP axis.

**kwargs

Defaults to {}

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

Returns: nn.Module

Parallelized model with its optional CP mesh.

nemo_automodel.components.models.mistral_large4.parallelization.MistralLarge4Parallelizer._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 existing dtype-aware FSDP grouping to a model-owned boundary.

Parameters:

module

nn.Module

Parameter-owning module; its matrices have shape [output, input], vectors [features]. Existing nested FSDP units retain ownership.

mesh

DeviceMesh

One-dimensional DP shard mesh or [replicate, shard] HSDP mesh.

mp_policy

MixedPrecisionPolicy | NoneDefaults to None

Requested compute policy; sensitive parameters always use FP32.

offload_policy

OffloadPolicy | NoneDefaults to None

Existing FSDP storage policy.

reshard_after_forward

bool | int | NoneDefaults to None

Existing FSDP reshard setting.

ignored_params

set[nn.Parameter] | NoneDefaults to None

Optional parameter set with the module’s tensor layouts; these parameters retain their existing ownership and placements.

Returns: nn.Module

Same module, with global-shaped DTensor parameters sharded on their

nemo_automodel.components.models.mistral_large4.parallelization.MistralLarge4Parallelizer._use_full_layer_activation_checkpointing(
    model: torch.nn.Module
) -> bool

Shared KV and index selections are explicit immutable block inputs.