Is this page helpful?

core.distributed.fsdp.src.megatron_fsdp.experimental.muon#

Muon with sharded momentum and owner-compute updates for MFSDP v2.

Module Contents#

Classes#

_FsdpGroupAssignment

An FSDP group and its cached optimizer-group and owner-rank assignments.

FsdpMuon

Muon for 2-D parameters with sharded optimizer buffers on a one-dimensional DP mesh.

API#

class core.distributed.fsdp.src.megatron_fsdp.experimental.muon._FsdpGroupAssignment#

An FSDP group and its cached optimizer-group and owner-rank assignments.

fsdp_group: core.distributed.fsdp.src.megatron_fsdp.experimental.parameter_group.FsdpParameterGroup#

None

tensor_to_optimizer_group_index: dict[int, int]#

None

tensor_to_owner_rank: dict[int, int]#

None

class core.distributed.fsdp.src.megatron_fsdp.experimental.muon.FsdpMuon(*args: Any, **kwargs: Any)#

Bases: emerging_optimizers.orthogonalized_optimizers.muon.Muon

Muon for 2-D parameters with sharded optimizer buffers on a one-dimensional DP mesh.

Momentum remains sharded. Owner assignments balance orthogonalization across a fixed set of buffers and are computed during initialization. Each step gathers, orthogonalizes, scatters, and applies one buffer’s updates at a time on the caller’s CUDA stream. Full matrices exist only on their assigned owners. The mesh is inferred from the FSDP groups, which must all share the same mesh. Parameters and optimizer groups must have the same order on every DP rank, and ranks must agree on which parameters have gradients. Parameter membership and optimizer-group order must remain fixed after construction.

This optimizer refreshes its FSDP model weights itself, so it does not need fully_shard_optimizer. Momentum is initialized lazily as DTensors in the standard optimizer state.

Limitations: - Non-2D parameters are skipped. Use a separate optimizer for them and for matrices you exclude from Muon, such as embeddings and output projections. - Tensor parallelism is unsupported.

Parameters:

  • *args – Positional arguments passed unchanged to EO’s Muon.

  • **kwargs – Keyword arguments passed unchanged to EO’s Muon.

Initialization

_fsdp_group_assignments: list[core.distributed.fsdp.src.megatron_fsdp.experimental.muon._FsdpGroupAssignment]#

None

_validate_parameter_groups() → None#

Validate optimizer parameters and the FSDP buffers that contain them.

step(
closure: collections.abc.Callable[[], float] | None = None,
) → None#

Update sharded parameters and refresh their model weights.

_step_group(
fsdp_group: core.distributed.fsdp.src.megatron_fsdp.experimental.parameter_group.FsdpParameterGroup,
tensor_to_optimizer_group_index: dict[int, int],
tensor_to_owner: dict[int, int],
) → None#

Gather momentum, orthogonalize on owners, and scatter updates.

Parameters:

  • fsdp_group – FSDP buffers and parameters to update.

  • tensor_to_optimizer_group_index – Active indices into fsdp_group.fsdp_parameters mapped to indices into self.param_groups for their optimizer settings.

  • tensor_to_owner – The same active tensor indices mapped to global owner ranks.

_compute_momentum_inputs(
fsdp_group: core.distributed.fsdp.src.megatron_fsdp.experimental.parameter_group.FsdpParameterGroup,
tensor_to_optimizer_group_index: dict[int, int],
) → core.distributed.fsdp.src.megatron_fsdp.experimental.dbuffer.DBuffer#

Apply local weight decay, update momentum, and build inputs for orthogonalization.

_orthogonalize_updates(
fsdp_group: core.distributed.fsdp.src.megatron_fsdp.experimental.parameter_group.FsdpParameterGroup,
tensor_to_optimizer_group_index: dict[int, int],
full_updates: dict[int, torch.Tensor],
) → None#

Replace this rank’s gathered inputs with full orthogonalized updates.

_apply_local_updates(
fsdp_group: core.distributed.fsdp.src.megatron_fsdp.experimental.parameter_group.FsdpParameterGroup,
tensor_to_optimizer_group_index: dict[int, int],
full_updates: dict[int, torch.Tensor],
local_updates: core.distributed.fsdp.src.megatron_fsdp.experimental.dbuffer.DBuffer,
) → None#

Apply owned result slices and received updates to local parameter shards.