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.

MFSDP v2 optimizer checkpointing needs an FSDP-native DTensor state contract. Keep this intentionally unsupported for the prototype instead of falling back to DDP-buffer assumptions.

abstractmethod load_state_dict(state_dict)#

Load optimizer state.

abstractmethod 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 a sharded optimizer state dict.

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.