Skip to content

Refine MFSDP v2 configuration validation - #6333

Open
wujingyue wants to merge 2 commits into
NVIDIA:mainfrom
wujingyue:agent/mfsdp-v2-ignore-partial-distopt
Open

Refine MFSDP v2 configuration validation#6333
wujingyue wants to merge 2 commits into
NVIDIA:mainfrom
wujingyue:agent/mfsdp-v2-ignore-partial-distopt

Conversation

@wujingyue

@wujingyue wujingyue commented Aug 6, 2026

Copy link
Copy Markdown
Contributor

Summary

  • Refine MFSDP v2 configuration validation: retain restrictions that are genuinely unsupported while allowing v2-native overlap, HSDP, and double-buffer behavior.
  • Enable NCCL user buffers through fully_shard_context(use_symmetric_memory=ddp_config.nccl_ub). Reject only the incompatible combination of NCCL user buffers with disabled symmetric registration.
  • Accept the default fsdp_all_gather_in_start_param_sync=True as a v2 no-op. MFSDP v2 gathers parameters from forward pre-hooks, so start_param_sync() does not issue an all-gather.
  • Add adapter coverage for construction-scoped symmetric memory.

Validation

  • Black check (locked version)
  • isort --check-only and ruff check on the changed adapter and test files
  • python -m torch.distributed.run --nproc-per-node 2 -m pytest -q tests/unit_tests/distributed/mfsdp_v2/test_mcore_adapter.py (3 passed)
  • git diff --check

@copy-pr-bot

copy-pr-bot Bot commented Aug 6, 2026

Copy link
Copy Markdown

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

@wujingyue
wujingyue force-pushed the agent/mfsdp-v2-ignore-partial-distopt branch 2 times, most recently from a713f67 to 107da0b Compare August 6, 2026 23:55
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:

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.

@wujingyue
wujingyue force-pushed the agent/mfsdp-v2-ignore-partial-distopt branch from 107da0b to c60c1de Compare August 7, 2026 00:49
Comment thread megatron/core/distributed/fsdp/mcore_fsdp_adapter.py
@wujingyue
wujingyue marked this pull request as ready for review August 7, 2026 00:51
@wujingyue
wujingyue requested review from a team as code owners August 7, 2026 00:51
@wujingyue
wujingyue requested a review from skyw August 7, 2026 00:51
@wujingyue

Copy link
Copy Markdown
Contributor Author

/ok to test c60c1de

@wujingyue

Copy link
Copy Markdown
Contributor Author

/ok to test 6b81672

Signed-off-by: Jingyue Wu <jingyuew@nvidia.com>
Signed-off-by: Jingyue Wu <jingyuew@nvidia.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Approved All necessary approvals have been made complexity: low Run functional tests

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Integrate MFSDP v2 with Megatron Core training

6 participants