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

Functions#

`FullyShardedDataParallel`	Construct the configured Megatron-FSDP implementation.
`_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`

API#

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

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

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.

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

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.

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#