nemo_automodel.components.distributed.fsdp2

Module Contents

Classes

Name	Description
`FSDP2Manager`	Deprecated compatibility wrapper for FSDP2 parallelization.

Functions

Name	Description
`_patch_is_packed_sequence_for_training`	Eliminate CPU-GPU sync from flash attention for standard (non-packed) training.
`fsdp2_sharding_enabled`	Report whether `ModelParallelizer` shards the model for this mesh.

API

class nemo_automodel.components.distributed.fsdp2.FSDP2Manager(
    config: nemo_automodel.components.distributed.config.FSDP2Config,
    device_mesh: torch.distributed.device_mesh.DeviceMesh,
    moe_mesh: torch.distributed.device_mesh.DeviceMesh | None = None,
    moe_config: nemo_automodel.components.distributed.config.MoEParallelizerConfig | None = None
)

Deprecated compatibility wrapper for FSDP2 parallelization.

.. deprecated:: 0.7 Pass FSDP2Config through the config-driven infrastructure and provide model-specific behavior with ModelParallelizer. This compatibility class is scheduled for removal in 0.8.

This manager applies parallelization to the model using a prescribed TP sharding plan. It supports mixed precision and CPU offloading options.

The device mesh must be created externally and passed in.

Parameters:

config

FSDP2Config

Configuration for FSDP2 distributed training.

device_mesh

DeviceMesh

Device mesh for distributed operations.

moe_mesh

Optional[DeviceMesh]Defaults to None

Optional device mesh for expert parallelism.

moe_config

MoEParallelizerConfig | NoneDefaults to None

Optional expert-parallel policy included in the model’s MeshContext.

nemo_automodel.components.distributed.fsdp2.FSDP2Manager.__getattr__(
    name
)

nemo_automodel.components.distributed.fsdp2.FSDP2Manager.maybe_compile(
    model
)

Apply per-layer compile after sharding, alongside whole-model compile_model().

nemo_automodel.components.distributed.fsdp2.FSDP2Manager.parallelize(
    model: torch.nn.Module,
    reapply_trainability: collections.abc.Callable[[nn.Module], None] | None = None
) -> torch.nn.Module

Parallelizes the given model using FSDP2 and TP sharding strategies.

Parameters:

model

nn.Module

The model to be parallelized.

reapply_trainability

Callable[[nn.Module], None] | NoneDefaults to None

Optional callback that re-resolves parameter trainability after model surgery and before FSDP construction.

Returns: nn.Module

The parallelized model.

nemo_automodel.components.distributed.fsdp2._patch_is_packed_sequence_for_training() -> None

Eliminate CPU-GPU sync from flash attention for standard (non-packed) training.

transformers._is_packed_sequence() returns a GPU bool scalar when batch_size==1, which causes Python’s if to call aten::is_nonzero — a CPU-GPU sync — once per attention layer per forward pass. With FSDP+TP+gradient-checkpointing this fires hundreds of times per iteration.

For standard (non-packed) training sequences are never packed, so returning the Python False immediately is both correct and avoids the sync. Do NOT apply this patch when using packed-sequence training (multiple sequences concatenated into one tensor with position_ids that reset to 0 mid-sequence).

nemo_automodel.components.distributed.fsdp2.fsdp2_sharding_enabled(
    device_mesh: torch.distributed.device_mesh.DeviceMesh
) -> bool

Report whether ModelParallelizer shards the model for this mesh.

Parallelization is skipped on a single-rank world or a single-element mesh, which also skips every side effect of fully_shard — most importantly the MixedPrecisionPolicy cast of parameters to the compute dtype. Callers that depend on that cast must check this instead of assuming FSDP2 is active.

Parameters:

device_mesh

DeviceMesh

Device mesh from the model’s parallelization context.

Returns: bool

True when fully_shard is applied, False when parallelization is skipped.