diff --git a/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py b/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py index a8393fcf147..bf5dc19b58a 100644 --- a/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py +++ b/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py @@ -564,6 +564,11 @@ def __init__( placements = Placements( dp_axes=[0], parameter=[Flat()], gradient=[Flat()], optimizer=[Flat()] ) + # NCCL symmetric memory requires UB. MFSDP v2 intentionally does not support UB + # without symmetric memory: it uses ncclCommRegister rather than the more performant + # ncclCommWindowRegister: + # https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/usage/bufferreg.html#window-registration + use_symm_mem = ddp_config.nccl_ub with fully_shard_context(device=device): for submodule in reversed(list(module.modules())): if submodule is module: @@ -576,9 +581,14 @@ def __init__( mesh=mesh, placements=placements, mixed_precision_policy=self.mp_policy, + use_symm_mem=use_symm_mem, ) fully_shard( - module, mesh=mesh, placements=placements, mixed_precision_policy=self.mp_policy + module, + mesh=mesh, + placements=placements, + mixed_precision_policy=self.mp_policy, + use_symm_mem=use_symm_mem, ) super().__init__(config=config, module=module) @@ -642,12 +652,8 @@ def _validate_config( raise ValueError( "MFSDP v2 requires data_parallel_sharding_strategy='optim_grads_params'." ) - if ddp_config.num_distributed_optimizer_instances != 1: - raise ValueError("MFSDP v2 does not currently support HSDP.") if ddp_config.outer_dp_sharding_strategy != "no_shard": raise ValueError("MFSDP v2 does not currently support outer DP sharding.") - if ddp_config.overlap_grad_reduce or ddp_config.overlap_param_gather: - raise ValueError("MFSDP v2 does not currently support communication overlap modes.") if config.gradient_accumulation_fusion: raise ValueError("MFSDP v2 does not currently support gradient accumulation fusion.") if config.calculate_per_token_loss: @@ -657,16 +663,13 @@ def _validate_config( if config.cuda_graph_impl != "none" or ddp_config.megatron_fsdp_cuda_graph_mode: raise ValueError("MFSDP v2 does not currently support CUDA graphs.") - if ddp_config.fsdp_double_buffer: - raise ValueError("MFSDP v2 does not support fsdp_double_buffer.") if ddp_config.fsdp_db_use_persist_buf_on_alloc_fail: - raise ValueError("MFSDP v2 does not support fsdp_db_use_persist_buf_on_alloc_fail.") - if ddp_config.fsdp_all_gather_in_start_param_sync: - raise ValueError("MFSDP v2 does not support fsdp_all_gather_in_start_param_sync.") - if ddp_config.nccl_ub: - raise ValueError("MFSDP v2 does not support nccl_ub.") - if ddp_config.disable_symmetric_registration: - raise ValueError("MFSDP v2 does not support disable_symmetric_registration.") + raise ValueError( + "MFSDP v2 does not support fsdp_db_use_persist_buf_on_alloc_fail: " + "it allocates communication buffers from PyTorch memory pools." + ) + if ddp_config.nccl_ub and ddp_config.disable_symmetric_registration: + raise ValueError("MFSDP v2 requires symmetric registration when nccl_ub is enabled.") if ddp_config.fsdp_manual_registration: raise ValueError("MFSDP v2 does not support fsdp_manual_registration.") if ddp_config.delay_wgrad_compute: diff --git a/tests/unit_tests/distributed/mfsdp_v2/test_mcore_adapter.py b/tests/unit_tests/distributed/mfsdp_v2/test_mcore_adapter.py index a132ff88139..c99fd52fc43 100644 --- a/tests/unit_tests/distributed/mfsdp_v2/test_mcore_adapter.py +++ b/tests/unit_tests/distributed/mfsdp_v2/test_mcore_adapter.py @@ -7,6 +7,7 @@ import pytest import torch +import megatron.core.distributed.fsdp.mcore_fsdp_adapter as mcore_fsdp_adapter from megatron.core.distributed import DistributedDataParallelConfig from megatron.core.distributed.fsdp.mcore_fsdp_adapter import FullyShardedDataParallel from megatron.core.distributed.fsdp.src.megatron_fsdp.experimental.module import FsdpModule @@ -73,7 +74,6 @@ def test_wraps_fsdp_unit_modules_before_root(self): data_parallel_sharding_strategy="optim_grads_params", megatron_fsdp_main_params_dtype=torch.float32, megatron_fsdp_main_grads_dtype=torch.float32, - fsdp_all_gather_in_start_param_sync=False, ), module=model, fsdp_unit_modules=[TransformerLayer], @@ -100,6 +100,40 @@ def test_wraps_fsdp_unit_modules_before_root(self): assert child_parameter_names assert root_parameter_names == {"1.weight", "1.bias"} + def test_nccl_ub_enables_symmetric_memory(self, monkeypatch): + config = TransformerConfig( + num_layers=1, + hidden_size=16, + num_attention_heads=4, + ffn_hidden_size=32, + bf16=True, + params_dtype=torch.bfloat16, + ) + model = torch.nn.Linear(config.hidden_size, config.hidden_size).to( + device="cuda", dtype=config.params_dtype + ) + fully_shard_calls = [] + original_fully_shard = mcore_fsdp_adapter.fully_shard + + def record_fully_shard(*args, **kwargs): + fully_shard_calls.append(kwargs["use_symm_mem"]) + return original_fully_shard(*args, **kwargs) + + monkeypatch.setattr(mcore_fsdp_adapter, "fully_shard", record_fully_shard) + FullyShardedDataParallel( + config=config, + ddp_config=DistributedDataParallelConfig( + use_megatron_fsdp=True, + megatron_fsdp_version=2, + data_parallel_sharding_strategy="optim_grads_params", + nccl_ub=True, + ), + module=model, + pg_collection=self.pg_collection, + ) + + assert fully_shard_calls == [True] + def test_build_train_and_step(self): config = TransformerConfig( num_layers=2, @@ -128,7 +162,6 @@ def test_build_train_and_step(self): data_parallel_sharding_strategy="optim_grads_params", megatron_fsdp_main_params_dtype=torch.float32, megatron_fsdp_main_grads_dtype=torch.bfloat16, - fsdp_all_gather_in_start_param_sync=False, ), module=model, pg_collection=self.pg_collection,