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

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 whose type is used to construct the data-parallel mesh. Defaults to CUDA.
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#: MFSDP v2 gradient reduction is complete when backward returns.

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#: MFSDP v2 communication is complete when backward returns.

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#