Skip to content

Refine MFSDP v2 configuration validation #6333

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wujingyue wants to merge 1 commit into
base: main
Choose a base branch
from
+52 −16
Open

Refine MFSDP v2 configuration validation #6333

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
31 changes: 17 additions & 14 deletions megatron/core/distributed/fsdp/mcore_fsdp_adapter.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

@skyw skyw Aug 7, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Consider comment why symmetric memory depends on UB use here.

Q: Any reason not using long name "use_symmetric_memory"? the other PR has gone length to rename it that way.

@wujingyue wujingyue Aug 7, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Q: Any reason not using long name "use_symmetric_memory"? the other PR has gone length to rename it that way.

Yes. I'll do that after the other PR (#6127) is merged.

@wujingyue wujingyue Aug 7, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Also added a code comment as you requested.

with fully_shard_context(device=device):
for submodule in reversed(list(module.modules())):
if submodule is module:
Expand All @@ -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)

Expand Down Expand Up @@ -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:

@wujingyue wujingyue Aug 6, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

MFSDP v2 unconditionally overlaps reduce-scatter and parameter all-gathers, and always uses double buffering. We intentionally accept and ignore the legacy overlap_grad_reduce, overlap_param_gather, and fsdp_double_buffer flags for compatibility: their defaults are False, so rejecting those values would require every existing v2 caller to override them just to preserve current behavior. The tradeoff is that setting them to False does not disable v2 overlap or double buffering.

Expand All @@ -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:

@wujingyue wujingyue Aug 7, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This flag doesn't make sense for MFSDP v2 because it doesn't rely on this flag to prefetch the first bucket.

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:
Expand Down
37 changes: 35 additions & 2 deletions tests/unit_tests/distributed/mfsdp_v2/test_mcore_adapter.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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],
Expand All @@ -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,
Expand Down Expand Up @@ -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,
Expand Down
Loading