nemo_automodel.components.distributed.megatron_fsdp

Module Contents

Classes

Name	Description
`MegatronFSDPManager`	Manager for parallelizing models using MegatronFSDP with TP, DP, CP sharding.

Functions

Name	Description
`fully_shard_optimizer`	Register an already-built optimizer with a MegatronFSDP-wrapped model.
`maybe_shard_optimizer`	Shard the optimizer with Megatron-FSDP when the strategy requires it.
`restore_distributed_param_attrs`	Re-apply Megatron-FSDP per-parameter attributes dropped by a post-wrap rebuild.
`snapshot_distributed_param_attrs`	Snapshot the per-parameter attributes Megatron-FSDP stamps on distributed params.

Data

HAS_MEGATRON_FSDP

logger

API

class nemo_automodel.components.distributed.megatron_fsdp.MegatronFSDPManager(
    config: nemo_automodel.components.distributed.config.MegatronFSDPConfig,
    device_mesh: torch.distributed.device_mesh.DeviceMesh
)

Manager for parallelizing models using MegatronFSDP with TP, DP, CP sharding.

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

The device mesh must be created externally and passed in.

Parameters:

config

MegatronFSDPConfig

Configuration for MegatronFSDP distributed training.

device_mesh

DeviceMesh

Device mesh for distributed operations.

activation_checkpointing

= config.activation_checkpointing

average_in_collective

= config.average_in_collective

calculate_per_token_loss

= config.calculate_per_token_loss

check_for_nan_in_grad

= config.check_for_nan_in_grad

disable_bucketing

= config.disable_bucketing

fsdp_double_buffer

= config.fsdp_double_buffer

grad_reduce_in_fp32

= config.grad_reduce_in_fp32

init_fsdp_with_meta_device

= config.init_fsdp_with_meta_device

keep_fp8_transpose_cache

= config.keep_fp8_transpose_cache

megatron_fsdp_unit_modules

= config.megatron_fsdp_unit_modules

nccl_ub

= config.nccl_ub

overlap_grad_reduce

= config.overlap_grad_reduce

overlap_param_gather

= config.overlap_param_gather

preserve_fp32_weights

= config.preserve_fp32_weights

report_nan_in_param_grad

= config.report_nan_in_param_grad

zero_dp_strategy

= config.zero_dp_strategy

nemo_automodel.components.distributed.megatron_fsdp.MegatronFSDPManager.parallelize(
    model,
    optimizer = None
)

Parallelizes the given model using MegatronFSDP and TP sharding strategies.

Parameters:

model

The model to be parallelized.

optimizer

Defaults to None

The optimizer for the model. If None, user needs to call model.finish_grad_sync() before optimizer.step(), model.install_optimized_model_weights() and model.zero_grad_buffer() after optimizer.zero_grad().

Returns:

(parallelized_model, optimizer)

nemo_automodel.components.distributed.megatron_fsdp.fully_shard_optimizer(
    model: torch.nn.Module,
    optimizer: torch.optim.Optimizer,
    preproc_state_dict_for_dcp_ckpt: bool = True
) -> torch.optim.Optimizer

Megatron-FSDP 0.5.0’s fully_shard_optimizer recovers the owning MegatronFSDP from a _megatron_fsdp_model attribute that MegatronFSDP.__init__ stamps onto each distributed Parameter. That attribute is a plain Python attribute and does not survive operations that rebuild Parameter objects (e.g. the dtype/device cast the from_pretrained load path performs after wrapping). The combined fully_shard(model, optimizer) entry point never hits this because it registers the optimizer in the same call, before any such op runs; the recipe’s separate build-model-then-build-optimizer order does, leaving fully_shard_optimizer unable to find the reference and aborting before the first optimizer step. Re-stamp the reference (mirroring the wheel’s own __init__ logic) on the current distributed params right before deferred sharding so the separate sequence matches the combined entry point.

nemo_automodel.components.distributed.megatron_fsdp.maybe_shard_optimizer(
    model_part: torch.nn.Module,
    optimizer: torch.optim.Optimizer,
    distributed_config: nemo_automodel.components.distributed.config.DistributedConfig | None,
    allow: bool = True
) -> torch.optim.Optimizer

Shard the optimizer with Megatron-FSDP when the strategy requires it.

Returns the optimizer unchanged unless distributed_config is a :class:MegatronFSDPConfig running in a distributed (world size > 1) job.

Parameters:

model_part

nn.Module

The (already sharded) model part the optimizer belongs to.

optimizer

torch.optim.Optimizer

The optimizer to (optionally) shard.

distributed_config

DistributedConfig | None

Distributed strategy config; only triggers sharding when it is a :class:MegatronFSDPConfig.

allow

boolDefaults to True

Guard for optimizers incompatible with Megatron-FSDP sharding (e.g. Dion); asserts when sharding would otherwise apply.

nemo_automodel.components.distributed.megatron_fsdp.restore_distributed_param_attrs(
    model: torch.nn.Module,
    snapshot: dict[str, dict] | None
) -> None

Re-apply Megatron-FSDP per-parameter attributes dropped by a post-wrap rebuild.

Companion to :func:snapshot_distributed_param_attrs. For each current parameter (matched by name, since the rebuild replaced the objects) it restores any snapshot attribute the rebuilt parameter is missing, following the fix suggested by the Megatron-FSDP maintainer on NVIDIA/Megatron-LM#5790: only attributes absent on the new parameter are copied, so genuinely re-derived state is never clobbered.

Parameters:

model

nn.Module

The Megatron-FSDP-wrapped model whose parameters were rebuilt.

snapshot

dict[str, dict] | None

The mapping returned by :func:snapshot_distributed_param_attrs, or None (no-op).

nemo_automodel.components.distributed.megatron_fsdp.snapshot_distributed_param_attrs(
    model: torch.nn.Module
) -> dict[str, dict] | None

Snapshot the per-parameter attributes Megatron-FSDP stamps on distributed params.

MegatronFSDP.__init__ decorates each distributed Parameter with plain Python attributes that later training steps depend on: _megatron_fsdp_model (owning-model back-ref used by :func:fully_shard_optimizer), _is_shared (set on tied parameters so _grad_acc routes their gradients through the root hook instead of double-accumulating), orig_param/megatron_fsdp_dist_index/ megatron_fsdp_slice and the reset_attribute closure. These attributes live in the Parameter.__dict__ and are silently dropped by any post-wrap operation that rebuilds Parameter objects — e.g. the from_pretrained checkpoint reload and the lm_head re-tie the recipe performs after wrapping. Capture them here, keyed by parameter name (object identity does not survive the rebuild), so :func:restore_distributed_param_attrs can re-apply them afterwards.

remove_duplicate=False is required so tied parameters (e.g. lm_head.weight aliasing model.embed_tokens.weight) are captured under every name they appear under, including the _is_shared marker Megatron-FSDP places on the tied alias.

Parameters:

model

nn.Module

The (possibly Megatron-FSDP-wrapped) model to snapshot.

Returns: dict[str, dict] | None

A mapping from parameter name to a copy of its __dict__, or None when

nemo_automodel.components.distributed.megatron_fsdp.HAS_MEGATRON_FSDP = True

nemo_automodel.components.distributed.megatron_fsdp.logger = logging.getLogger(__name__)