Is this page helpful?

core.optimizer.fully_sharded_optimizer#

MCore optimizer wrapper for experimental Megatron-FSDP v2.

Module Contents#

Classes#

FullyShardedOptimizer

MCore optimizer wrapper for MFSDP-owned sharded parameters and gradients.

Functions#

count_replication

Return how many ranks hold an identical copy of tensor’s local shard.

API#

core.optimizer.fully_sharded_optimizer.count_replication(tensor: torch.distributed.tensor.DTensor) → int#

Return how many ranks hold an identical copy of tensor’s local shard.

A sharded mesh axis holds disjoint pieces that must all be counted; a replicated axis holds identical copies that must be counted once, so a gradient statistic summed over the grad-stats group has to divide by this.

MFSDP v2 gradients are always DTensors, so this takes one rather than accepting a plain tensor and guessing a layout for it.

class core.optimizer.fully_sharded_optimizer.FullyShardedOptimizer(
optimizer: torch.optim.Optimizer,
config: core.optimizer.optimizer_config.OptimizerConfig,
grad_scaler: Optional[core.optimizer.grad_scaler.MegatronGradScaler],
init_state_fn: Callable,
model_chunks: List[core.transformer.module.MegatronModule],
)#

Bases: core.optimizer.optimizer.MixedPrecisionOptimizer

MCore optimizer wrapper for MFSDP-owned sharded parameters and gradients.

MFSDP v2 owns the optimizer-facing parameter and gradient shards directly. Unlike :class:DistributedOptimizer, this wrapper does not build DDP param-and-grad-buffer range maps or allocate separate main-parameter shards. It preserves MCore’s mixed-precision optimizer step contract while making MFSDP-specific storage operations explicit.

Initialization

Initialize the MFSDP optimizer wrapper.

Parameters:

  • optimizer – Base optimizer such as Adam or SGD.

  • config – Optimizer configuration.

  • grad_scaler – Optional loss scaler. Currently unsupported for MFSDP v2, but accepted to match the MCore optimizer construction contract.

  • init_state_fn – Function used to initialize optimizer state.

  • model_chunks – MFSDP v2 model chunks optimized by this wrapper.

static _validate_config(
config: core.optimizer.optimizer_config.OptimizerConfig,
model_chunks: List[core.transformer.module.MegatronModule],
) → None#

Validate the MFSDP v2 optimizer support contract.

abstractmethod state_dict()#

Return optimizer state.

Deliberately unsupported: MFSDP v2 checkpoints through :meth:sharded_state_dict, and Megatron-FSDP always runs with --ckpt-format fsdp_dtensor (see validate_args), so nothing reaches this accessor. Forwarding the wrapped optimizer’s state here would only produce a checkpoint that silently drops the cross-rank resharding metadata.

_trainable_parameters() → collections.abc.Iterator[torch.nn.Parameter]#

Yield the parameters the base optimizer owns before the empty-shard filter.

_get_param_groups keeps every requires_grad parameter, and under TE < 2.18 MFSDP v2 then drops the ones whose local shard is empty, working around a TE FusedAdam bug (see :func:get_megatron_optimizer). That filter is applied per rank, so on an older TE the surviving set differs across ranks; this iterator reconstructs the rank-invariant superset, which is the same thing on either TE.

_shard_process_groups() → list[torch.distributed.ProcessGroup]#

Return the process groups across which the parameters are sharded.

Every sharded dimension of the parameter device mesh contributes one group. Replicated dimensions are skipped: replicas hold identical shards, hence an identical keyspace. These are the groups preprocess_state_dict_for_uneven_dtensor gathers over, so a collective run here reaches the same ranks the save and load collectives do.

Raises:

NotImplementedError – If the parameters do not all share one device mesh, which is what expert parallelism produces. Not yet supported rather than refused on principle; training under expert parallelism is unaffected.

_gather_state_keys_by_fqn() → dict[str, list[str]]#

Map every parameter’s FQN to the keys of its DTensor optimizer state entries.

The keys are read off live state (exp_avg/exp_avg_sq for Adam) rather than hard-coded, so any base optimizer works. They are gathered because a parameter that the TE < 2.18 empty-shard filter dropped from this rank’s optimizer is only described by the rank that owns a non-empty shard of it, which is what lets every rank synthesize a placeholder with exactly the owning rank’s keys. A parameter that has no state anywhere (nothing has stepped yet) is simply absent.

Which key set a parameter has is a property of the base optimizer, not of the shard, so ranks that both describe an FQN describe it identically and the merge below can let any of them win. Sorting is what makes that hold as stated rather than by luck: it turns each entry into a canonical list, so two ranks whose state dicts happened to insert the keys in different orders still produce equal values.

_param_to_group_meta() → dict[str, Any]#

Map each locally-owned parameter’s FQN to its param-group hyperparameters.

Only the parameters this rank’s optimizer owns appear, which is exactly what both ends need: on save DCP unions the (non-tensor) entries written by all ranks, and on load each rank reads back the entries of the parameters it owns. The base optimizer (TE FusedAdam) tracks step per group rather than per parameter, so step round-trips here rather than in the per-parameter state.

_param_groups_from_group_meta(
param_to_group_meta: dict[str, Any],
) → list[dict[str, Any]]#

Rebuild FQN-keyed torch param groups matching this rank’s current optimizer.

Iterating self.optimizer.param_groups preserves the current parameter order within each group, which is what :meth:torch.optim.Optimizer.load_state_dict uses to map the checkpoint’s FQN-keyed state onto parameter tensors.

load_state_dict(
state_dict: core.dist_checkpointing.mapping.ShardedStateDict,
) → None#

Load optimizer state produced by :meth:sharded_state_dict.

By the time this runs DCP has already written the checkpoint’s tensors into the resting optimizer-state DTensors in place, since sharded_state_dict(is_loading=True) exposed them as the load destinations. What is left is to restore the param-group hyperparameters (including the group-level step) and re-bind the FQN-keyed state to the current parameters, mirroring the Megatron-FSDP branch of

Meth:

DistributedOptimizer.load_state_dict.

sharded_state_dict(
model_sharded_state_dict: core.dist_checkpointing.mapping.ShardedStateDict,
is_loading: bool = False,
metadata: Optional[dict] = None,
) → core.dist_checkpointing.mapping.ShardedStateDict#

Build the fsdp_dtensor sharded optimizer state dict.

The layout mirrors :meth:DistributedOptimizer.sharded_param_state_fsdp_dtensor, so v1 and v2 write the same on-disk format::

{"state": {fqn: {"exp_avg": DTensor, "exp_avg_sq": DTensor}},
 "param_to_group_meta": {fqn: {...group hyperparameters...}}}

That is also why the FQNs come from get_global_unique_param_name instead of torch’s :func:~torch.distributed.checkpoint.state_dict.get_optimizer_state_dict, whose names are the plain named_parameters ones rather than MCore’s PP/EP-unique checkpoint names.

Rank consistency is the load-bearing invariant here. Under TE < 2.18 the empty-shard filter (see :meth:_trainable_parameters) leaves each rank with a different set of parameters, while preprocess_state_dict_for_uneven_dtensor runs one all_gather_object per DTensor in sorted key order, so a rank-dependent DTensor keyspace deadlocks. Every trainable parameter is therefore emitted: the parameters this rank’s optimizer holds contribute their real state, and the rest contribute empty-local placeholders. On a newer TE no parameter is ever missing, so the loop below takes the real-state branch throughout and the result is the same.

Parameters:

  • model_sharded_state_dict – Accepted for interface parity; the optimizer state is read from the wrapped optimizer directly.

  • is_loading – Whether the state dict will be filled by a load. If so, the optimizer state is materialized first so the load has DTensors to write into.

  • metadata – Accepted for interface parity; the fsdp_dtensor format takes no sharding options.

Returns:

The optimizer state dict, in the fsdp_dtensor format described above.

get_grad_norm()#

Compute the global gradient L2 norm from each gradient’s own DTensor layout.

MFSDP v2 gradients are DTensors that record how they are distributed, and the dense and expert gradients do not share a device mesh: with EP=2 over eight ranks the dense gradients live on all eight while the expert gradients live on the four-rank expert-DP stripe. Reading the layout off each gradient keeps the norm correct without assuming a single mesh for all of them.

Each rank contributes ||local||^2 divided by the product of its replicated mesh-axis sizes. A sharded axis holds disjoint pieces that must all be added; a replicated axis holds identical copies that must be counted once. Summing that over the grad-stats group is then exact, because every shard is held by exactly one rank in that group.

get_grad_norm_fp32 cannot do this: get_main_grads_for_grad_norm replaces each DTensor with grad._local_tensor before it runs, so get_data_parallel_group_if_dtensor always sees plain tensors, returns None, and the layout is gone by the time the norm is taken.

count_zeros() → float#

Count zero gradient entries from each gradient’s own DTensor layout.

count_zeros_fp32 has the same single-mesh assumption as the grad-norm path, and additionally rejects the combination of a Megatron-FSDP parameter with a DTensor-derived data-parallel group. Counting here keeps MFSDP v2 off that path, and matches how get_grad_norm reduces: each rank contributes its own shard, divided by the size of any replicated mesh axis, summed over the grad-stats group.

zero_grad(set_to_none: bool = True) → None#

Clear optimizer-visible sharded grads.

_copy_model_grads_to_main_grads() → None#

Install optimizer-compatible gradients for non-precision-aware optimizers.

step_with_ready_grads() → bool#

Step the optimizer and restore MFSDP gradient dtypes.

_copy_main_params_to_model_params() → None#

Refresh MFSDP V2 compute weights after updating optimizer weights.

_copy_model_params_to_main_params(state_dict=None) → None#

No-op: model loads already write into MFSDP v2’s main weights.