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.

_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.

_add_empty_shard_state(state_by_fqn: dict[str, Any]) → None#

Add checkpoint placeholders for empty shards omitted by TE < 2.18.

Uneven-DTensor preprocessing needs matching state keys on every rank of a parameter’s mesh. Gather the keys within each dense or expert-DP mesh.

_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.

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.