Is this page helpful?

core.distributed.fsdp.mcore_fsdp_adapter#

Public factory and re-exports for the MCore FSDP wrappers.

Module Contents#

Functions#

FullyShardedDataParallel

Construct the configured Megatron-FSDP implementation.

Data#

__all__

API#

core.distributed.fsdp.mcore_fsdp_adapter.__all__#

[‘FullyShardedDataParallel’, ‘FullyShardedDataParallelV1’, ‘FullyShardedDataParallelV2’]

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.