`core.distributed.fsdp.mcore_fsdp_adapter`#

Module Contents#

Classes#

`FullyShardedDataParallelV1`	Fully Sharded Data Parallel (FSDP) wrapper for the Megatron model.
`FullyShardedDataParallelV2`	MFSDP v2 wrapper for the Megatron model.
`_AxisPlacements`	How one mesh axis places each of MFSDP’s three buffers.

Functions#

`_materialize_meta_module`	Materialize and initialize one module’s direct meta parameters.
`_materialize_owned_meta_modules`	Materialize meta parameters reachable from one FSDP unit.
`FullyShardedDataParallel`	Construct the configured Megatron-FSDP implementation.
`_build_expert_mesh_and_placements`	Build the expert-DP mesh and placements, or return neither when EP is disabled.
`_build_hybrid_dp_mesh`	Build the (“dp_outer”, “dp_shard”) mesh for a hybrid data-parallel domain.
`_get_hsdp_tp_mesh`
`_get_dp_tp_mesh`
`_check_mesh_ranks_and_group_ranks_are_consistent`
`_get_rng_state_dict`
`_load_rng_state_dict`

Data#

`__all__`
`logger`
`_DATA_PARALLEL_PLACEMENTS`

API#

core.distributed.fsdp.mcore_fsdp_adapter.__all__#: [‘FullyShardedDataParallel’]

core.distributed.fsdp.mcore_fsdp_adapter.logger#: ‘getLogger(…)’

core.distributed.fsdp.mcore_fsdp_adapter._materialize_meta_module( module: torch.nn.Module, device: torch.device | None, ) → None#: Materialize and initialize one module’s direct meta parameters.

core.distributed.fsdp.mcore_fsdp_adapter._materialize_owned_meta_modules( module: torch.nn.Module, device: torch.device | None, ) → None#

Materialize meta parameters reachable from one FSDP unit.

PyTorch and Transformer Engine differ in reset_parameters(): PyTorch typically initializes existing Parameters in place, while TE may replace a Parameter via setattr(). Run resets before fully_shard() so FSDP sees the final Parameter objects.

class core.distributed.fsdp.mcore_fsdp_adapter.FullyShardedDataParallelV1( config: megatron.core.transformer.transformer_config.TransformerConfig, ddp_config: megatron.core.distributed.distributed_data_parallel_config.DistributedDataParallelConfig, module: torch.nn.Module, fsdp_unit_modules: Optional[List[Type[torch.nn.Module]]] = None, disable_bucketing: bool = False, device: Optional[torch.device] = None, pg_collection: Optional[megatron.core.process_groups_config.ProcessGroupCollection] = None, )#

Bases: megatron.core.distributed.data_parallel_base._BaseDataParallel

Fully Sharded Data Parallel (FSDP) wrapper for the Megatron model.

Initialization

_MODULE_TYPE_REGISTRY: Dict[str, set]#: None

static _fine_grained_recurse_module_types( config: megatron.core.transformer.transformer_config.TransformerConfig, ddp_config: megatron.core.distributed.distributed_data_parallel_config.DistributedDataParallelConfig, ) → Tuple[Type[torch.nn.Module], ...]#: Module classes needing parameters(recurse=True) for fine-grained hooks.

load_state_dict(state_dict, strict=True)#: Load the state dictionary into the module.

_detect_parallelism_type( param_name: str, module: torch.nn.Module, ) → Optional[str]#

Infer tensor-parallelism type for a parameter under a given module (forked from Megatron-Bridge).

Returns:: “column”, “row”, or “replicated” if a type can be inferred, else None.

_annotate_tensor_parallelism(root_module: torch.nn.Module) → None#

Annotate parameters under root_module with inferred tensor-parallel metadata.

Each parameter that can be classified will get a _tensor_parallel_mode attribute set to one of: “column”, “row”, or “replicated”.

_init_dist_index(pg_collection)#: Initialize the distributed index for the module.

stop_communication()#: Stop communication for the module.

sync_rng_states_across_tp_group()#: Synchronize the tensor parallel random number generator states.

class core.distributed.fsdp.mcore_fsdp_adapter.FullyShardedDataParallelV2( config: megatron.core.transformer.transformer_config.TransformerConfig, ddp_config: megatron.core.distributed.distributed_data_parallel_config.DistributedDataParallelConfig, module: torch.nn.Module, fsdp_unit_modules: Optional[List[Type[torch.nn.Module]]] = None, disable_bucketing: bool = False, device: Optional[torch.device] = None, pg_collection: Optional[megatron.core.process_groups_config.ProcessGroupCollection] = None, )#

Bases: megatron.core.distributed.data_parallel_base._BaseDataParallel

MFSDP v2 wrapper for the Megatron model.

Supports MXFP8 compute with either high-precision or MXFP8 parameter gathers. For MXFP8 gathers, enable both config.fp8_param and ddp_config.fp8_param_gather; MCore’s model initialization preserves the high-precision values used for main weights.

Initialization

Initialize the MFSDP v2 wrapper.

Parameters:

config – Transformer configuration for the model.
ddp_config – Data-parallel and sharding configuration.
module – Root model module to shard.
fsdp_unit_modules – Module types to shard as child FSDP units. If unspecified, transformer, MoE transformer, and Mamba layers are used.
disable_bucketing – Compatibility argument that must remain False for MFSDP v2.
device – Device used to construct the data-parallel mesh and, when this wrapper opens its own fully_shard_context, that context’s device. Defaults to CUDA. When the caller has already opened an ambient context (multi-chunk wrapping), this wrapper joins it and the ambient context’s device and use_symmetric_memory settings apply instead.
pg_collection – Explicit process groups. The dp_cp group defines the data-parallel mesh.

Raises:

ImportError – If the Megatron FSDP implementation is unavailable.
ValueError – If required process groups are missing or the configuration requests a feature unsupported by MFSDP v2.

static _validate_config( config: megatron.core.transformer.transformer_config.TransformerConfig, ddp_config: megatron.core.distributed.distributed_data_parallel_config.DistributedDataParallelConfig, module: torch.nn.Module, pg_collection: megatron.core.process_groups_config.ProcessGroupCollection, disable_bucketing: bool, ) → None#

Validate that the model and configuration are supported by MFSDP v2.

Parameters:

config – Transformer configuration describing the requested model topology.
ddp_config – Data-parallel and sharding configuration to validate.
module – Model whose parameters are checked for expert parallelism.
pg_collection – Materialized process groups whose topology must match the supported MFSDP v2 topology.
disable_bucketing – Whether parameter bucketing is disabled.

Raises:

ValueError – If a required process group is missing or the model, topology, or data-parallel configuration uses an unsupported feature.

no_sync()#

Suppress gradient finalization for a non-final microbatch.

HSDP/HFSDP leave the DP-outer axis Partial across microbatches and reduce it on the last backward of a step, so MFSDP has to be told which backward that is. Without it every backward finalizes that axis and marks the accumulation buffer stale, so the next microbatch zeroes it and only the last microbatch’s gradient reaches the optimizer.

MCore’s schedules wrap every microbatch but the last in no_sync_func.

start_param_sync(*unused, **unused_kwargs) → None#: No-op: MFSDP v2 gathers parameters from its forward pre-hooks.

start_grad_sync(*unused, **unused_kwargs) → None#: MFSDP v2 reduces gradients during backward.

finish_grad_sync(*unused, **unused_kwargs) → None#: Wait on all submitted reductions, including delayed weight gradients.

synchronize_param_gather(*unused, **unused_kwargs) → None#: MFSDP v2 parameter gathers complete inside module hooks.

abstractmethod broadcast_params() → None#: Reject parameter broadcast, which is unsupported by MFSDP v2.

stop_communication() → None#: Wait for any pending gradient reductions.

core.distributed.fsdp.mcore_fsdp_adapter.FullyShardedDataParallel( config: megatron.core.transformer.transformer_config.TransformerConfig, ddp_config: megatron.core.distributed.distributed_data_parallel_config.DistributedDataParallelConfig, module: torch.nn.Module, fsdp_unit_modules: Optional[List[Type[torch.nn.Module]]] = None, disable_bucketing: bool = False, device: Optional[torch.device] = None, pg_collection: Optional[megatron.core.process_groups_config.ProcessGroupCollection] = None, ) → megatron.core.distributed.data_parallel_base._BaseDataParallel#

Construct the configured Megatron-FSDP implementation.

This is a factory function, not a wrapper type. Use the explicit V1 or V2 implementation classes for type checks.

class core.distributed.fsdp.mcore_fsdp_adapter._AxisPlacements#

Bases: typing.NamedTuple

How one mesh axis places each of MFSDP’s three buffers.

parameter: torch.distributed.tensor.placement_types.Placement#: None

gradient: torch.distributed.tensor.placement_types.Placement#: None

optimizer: torch.distributed.tensor.placement_types.Placement#: None

core.distributed.fsdp.mcore_fsdp_adapter._DATA_PARALLEL_PLACEMENTS#: None

core.distributed.fsdp.mcore_fsdp_adapter._build_expert_mesh_and_placements( config: megatron.core.transformer.transformer_config.TransformerConfig, ddp_config: megatron.core.distributed.distributed_data_parallel_config.DistributedDataParallelConfig, pg_collection: megatron.core.process_groups_config.ProcessGroupCollection, device_type: str, ) → Tuple[torch.distributed.DeviceMesh | None, megatron.core.distributed.fsdp.src.megatron_fsdp.experimental.Placements | None]#: Build the expert-DP mesh and placements, or return neither when EP is disabled.

core.distributed.fsdp.mcore_fsdp_adapter._build_hybrid_dp_mesh(outer_group, inner_group, device_type)#: Build the (“dp_outer”, “dp_shard”) mesh for a hybrid data-parallel domain.

core.distributed.fsdp.mcore_fsdp_adapter._get_hsdp_tp_mesh( outer_fsdp_dp_group, dp_cp_group, tp_group, ep_size=1, )#

core.distributed.fsdp.mcore_fsdp_adapter._get_dp_tp_mesh(dp_cp_group, tp_group, ep_size=1)#

core.distributed.fsdp.mcore_fsdp_adapter._check_mesh_ranks_and_group_ranks_are_consistent( mesh_ranks, group_ranks, )#

core.distributed.fsdp.mcore_fsdp_adapter._get_rng_state_dict()#

core.distributed.fsdp.mcore_fsdp_adapter._load_rng_state_dict(rng_state_dict)#

core.distributed.fsdp.mcore_fsdp_adapter#