[DO NOT MERGE] Draft mfsdp_refactor to dev branch#5665
Draft
shjwudp wants to merge 36 commits into
Draft
Conversation
…Tongliu) - Add FSDP v2 mixed precision policy with MXFP8 support - Add param group for mixed precision parameter management - Add FSDP v2 hooks for mixed precision workflow - Add bucket allocator for mixed precision buffer management
… core (by Jack Chang) - Add FSDP v2 checkpoint module with save/load and online format conversion - Add DP buffer, FSDP module, fully_shard API, and utils for v2 - Add NVFP4 design doc and checkpoint design doc - Add FSDP v2 hooks in mcore_fsdp_adapter for v2 integration - Add use_megatron_fsdp_v2 flag in DDP config and arguments - Wire FSDP v2 into distrib_optimizer state_dict/load paths - Update checkpointing.py with v2 import, loading, and format support - Add v2 unit tests (checkpoint, fully_shard, nd_parallel, param_group, allocator) - Update uneven_dtensor, megatron_fsdp, and fsdp_dtensor_checkpoint for v2
…upport - README: fix architecture diagram (add missing files), replace stale unit test references, remove false limitations (state_dict_for_save_checkpoint and stop_communication are already wired), add FullyShardMixedPrecisionPolicy docs, add NVFP4 non-distributed path limitation, fix installation and toy example commands - tp_support_design.md: draft TP support plan referencing v1 patterns (two-step DTensor construction, 2D DeviceMesh, gradient decoupling) - fsdp_toy.py: add --activation-checkpoint flag and ToyBlock checkpointing - test_fully_shard.py: add 6 activation checkpointing tests (forward/backward, multi-step stability, overlap compatibility, nested FSDP, numerical equivalence, per-layer mixed checkpointing)
- fsdp_toy.py: use Megatron get_state_dict in AppState to attach uneven DTensor chunk metadata during checkpoint save; remove dead sys.modules hack; add warnings import fallback - checkpoint.py: rename load_torch_dist_into_fsdp_v2 to _load_torch_dist_into_megatron_fsdp_v2 (private convention) - checkpointing.py: add import of _load_torch_dist_into_megatron_fsdp_v2 - test_param_group.py: fix dimension comments and simplify assertion
…rename - Rename FullyShardMixedPrecisionPolicy to MixedPrecisionPolicy across all imports, exports, docs, and tests - Add ZeRO-1 (optim) and ZeRO-2 (optim_grads) sharding strategy support alongside existing ZeRO-3 (optim_grads_params); reject no_shard in v2 - Refactor DataParallelBuffer.unshard() to support replicated buffer refresh via _dirty flag for ZeRO-1/2 weight consistency after optimizer - Refactor DataParallelBuffer.reduce_grad() to handle all three ZeRO strategies (reduce-scatter vs all-reduce vs delayed reduce-scatter) - Add FSDPModule.finish_grad_sync() for delayed ZeRO-1 reduce-scatter at iteration boundary; delegate finish_grad_sync through adapter - Handle empty trace pools gracefully in allocator.plan() - Add strict iter-equivalence tests for ZeRO strategies with parameter snapshot comparison - Add MXFP8 ZeRO-1/2 smoke tests for replicated quantized buffer refresh - Update design.md with ZeRO-1/2 workflow docs and replicated weight refresh design
…en dtensor split, improve checkpoint inspector - Add _sync_module_states_after_load to FSDPModule and call it after loading Megatron FSDP v2 checkpoints to ensure model weights are consistent - Add get_root_module to _FSDPRootContext and FSDPModule - Harden split_dtensor: assert __create_chunk_list__ exists instead of falling back to gather_and_compute_chunk_metadata - Rewrite compare_two_checkpoint in checkpoint_inspector with batched, memory-aware loading and DeviceMesh support for large distributed ckpts - Improve print_tensor output with L2 norm and local shard data display
Restructure step() to fire all grad gathers up front (so later gathers overlap earlier params' Newton-Schulz) instead of the prior depth-2 prefetch pipeline; gather Work handles renamed handles->reqs, the inner ParameterGroup local renamed pg->param_group. Trim the step() / class docstrings to a concise step-by-step description and drop the stale module docstring that still described the removed reduce-scatter-stream prefetch design. Add an assert that no closure is passed (Muon needs none). No behavior change: 4-rank GB200 e2e (hidden=8192, optim_grads_params) trains 10/10 iters, finite decreasing loss. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Restructure test_dist_muon.py around a single five-step matches-reference test and parametrize it over shape scale (small/medium/large, each sized so matrices straddle a shard boundary and exercise the gather/scatter), nesterov, and sharding strategy. Fold the 1D-bias skip check into the same param set so the bit-exact comparison also asserts Muon never touches non-2D params. Make the reference mirror step()'s exact ops so torch.equal holds bit for bit: the nesterov look-ahead is a single alpha-add (g + m*buf), not g + (m*buf). Verified 18/18 on a 4-rank GB200 run, stable across two consecutive runs. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…olicy - Mark optim (ZeRO-1) and optim_grads (ZeRO-2) as Supported, note no param-gather overlap - Clarify FullyShardFP8Policy / FullyShardNVFP4Policy are recipe dataclasses embedded in MixedPrecisionPolicy, not standalone policies - Add NVFP4 usage example - Update Known Limitations to reflect actual support status
…age2 Add support of muon optimizer
Steps so far (step() not yet rewired — still uses the back-ref path): - FullyShardV2Muon now takes an explicit `grads` arg (DTensor per param, replacing the unused `pg_collection`); asserts params and grads are all DTensors and aligned 1:1. - Init reconstructs every DP rank's (flat_offset, size) shard range for each param purely from the grad DTensor sharding + one all_gather per dp_group (_gather_shard_ranges), with a temporary cross-check against the old ParameterGroup _grad_gather_plans to prove equivalence. - master(fp32) -> model(bf16) cast moved out of the inner step() into the FullyShardV2MuonOptimizer wrapper, calling the v2 FSDPModule _copy_main_weights_to_model_weights() (not the v1 param_and_grad_buffer shim). - Factory builds ONE Muon over all model chunks (global root balance / per-layer packaging / a2a pipelining see every layer together); the wrapper now pulls the 2D-matrix dist_params + grad DTensors straight from the FSDP v2 _fsdp_param_groups in layer order, so the factory no longer news the inner optimizer or uses back-refs. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Rewrite FullyShardV2Muon's gather/scatter to per-package batch_isend_irecv that reads/writes the grad, full-grad and orth tensors directly, removing all pack/unpack buffers and copies: - gather/scatter plans are static (group, send, recv), built once; the P2Ps irecv straight into the preallocated full grad (split roots) and isend straight from the post-momentum shard. - split-root full grads are preallocated once; single-holder params view their grad shard as the full grad (no buffer, no P2P). - Phase 1 writes the post-momentum shard into its final location (full-grad slice / grad shard / momentum buffer), so no copy. - Phase 3b reads each root's own orth segment straight from the orth it computed; non-root holders read the scatter-delivered grad. - gather(i+1) is issued before NS(i) so the P2Ps overlap the NS. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
- fully_shard_v2_muon.py: trim duplicated/over-explaining comments (no logic change). - param_group.py / dp_buffer.py: remove the dead single-root Muon plumbing the copy-free DTensor optimizer no longer uses — ParameterGroup.unshard_grad_to_root_async / scatter_grad_from_root_async, the per-init _grad_gather_plans build, the dist_param._fsdp_param_group / _fsdp_orig_param back-refs, and DataParallelBuffer.item_root_ranks / _build_item_root_ranks. Confirmed a closed dead cluster (only referenced each other); the Adam/FSDP path never touched them. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Feat/fsdp v2 muon refactor
CUDA Graph: - Add FSDPCudaGraphRunner for per-module CUDA graph capture/replay - Support enable_cuda_graph on leaf FSDP modules only - Integrate with TracePoolAllocator for stable buffer addresses - Mark as experimental feature with FutureWarning and docs - Add cuda_graph_design.md TracePoolAllocator: - Implement trace/plan/optimized lifecycle with interval coloring - Add allocator_design.md with CUDA graph integration lifecycle - Stable slot-based allocation for fixed addresses across micro-batches Checkpoint: - Add chunk metadata source tags (init/preprocess/split/propagate:*) for provenance tracing when debugging numel mismatches - Add _verify_chunk_metadata with numel consistency check and detailed failure diagnostics (key, shapes, chunk_list, source, device_mesh) - Fix _maybe_wrap_as_uneven_dtensor dead isinstance(DTensor) check - Add warning for handle_experts_in_state_dict identity transform on non-local experts (local_expert_start=0) - Make model an explicit parameter to _preprocess_and_verify_v2_state_dict - Guard Phase 4 slicing for _extra_state tensors without chunk metadata - Update mcore_fsdp_checkpoint_design.md: clarify v1/v2 split, fix feature matrix, annotate unchecked PP/multi-optimizer items DP Buffer: - Move _dirty flag from buffer attribute to tensor attribute - Refactor is_unsharded() and unshard() for non-distributed buffers - Change _get_item_global_range to return (start, end) for consistency with _get_item_self_range and _get_item_local_range Hooks: - Merge _register_fine_grained_forward_pre_hooks into _register_forward_pre_hook with fine_grained parameter - Handle CUDA graph activation correctly for root vs child modules Distributed Muon: - Refactor toward DTensor-driven, ParameterGroup-free design - Make copy-free via batched isend/irecv - Clean up comments and dead single-root plumbing Docs: - Update v2/README.md with experimental CUDA graph note, fix cuda_graph_active replay description, add known limitations - Add experimental note to examples/megatron_fsdp/README.md - Update checkpoint design doc across multiple sections
…2-muon-refactor # Conflicts: # megatron/core/distributed/fsdp/src/megatron_fsdp/v2/README.md # megatron/core/distributed/fsdp/src/megatron_fsdp/v2/mixed_precision.py
Support no-shard in FSDP v2
Feat/fsdp v2 muon refactor
Skip no-op unshard for clean replicated buffers
Initialize clean replicated FSDP weight buffers
…/resume - Lazy gradient buffer: _init_dist_grads() defers main_grad_buffer.data allocation to first backward; _maybe_free_grad_data() frees between steps; _grad_buffer_is_fresh flag ensures overwrite (not accumulate) on first reduce-scatter after alloc or zero_grad. torch.empty + overwrite avoids zero-init cost. Saves ~35 GB per rank for 70B/bf16/dp=8 during forward and between training steps. - CPU offload infrastructure: _is_on_cpu / _ensure_data_on_gpu / _move_data_to primitives on DataParallelBuffer; _ensure_buffers_on_gpu auto-reload at every FSDP entry point; _rebuild_dist_views updates _local_tensor in-place after device moves. - TracePoolAllocator lifecycle: release() frees slot tensors while preserving plan metadata; _auto_resume() re-allocates on next alloc/free; resume() for explicit restore. FSDPModule.release_memory_pool() tears down CUDA graphs, clears sentinels, and releases allocator slots. - get_state_dict zero-collective upgrade: copy chunk metadata from model dist_params instead of all_gather_object. - CLI: --fsdp-trace-pool flag; adapter wires both fsdp_double_buffer and fsdp_trace_pool to enable_trace_pool. fsdp_toy.py --release-memory-pool. - CI: unblock checkpoint unit tests in test_fully_shard.py; add fsdp_trace_pool coverage to test_mcore_nd_parallel and test_mcore_checkpoint. - Docs: lazy_grad_buffer_design.md with memory lifecycle, edge cases, and torch.empty bug analysis. - _is_torchdynamo_compiling() helper guards free_storage/alloc_storage.
…, CPU offload, and design docs ## Overview 21 files, +1,473 / -336. Integrates M-FSDP v2 into the 1F1B pipeline schedule, adds lazy gradient buffer allocation, CPU offload/reload infrastructure, and reorganizes design documentation. ## Key Changes ### 1F1B + FSDP v2 Integration - Refactored hooks from closures to standalone functions (pre/post-forward, pre/post-backward) with `skip_final_callback` parameter for EP overlap schedules to suppress auto-enqueue and trigger manually. - Unified fine-grained hook registration replaces the old EP-overlap system; `fine_grained_hooks`, `skip_backward_callback`, `skip_final_backward_callback` flags decompose the monolithic `enable_ep_overlap`. - combined_1f1b.py: dispatch pre/post-backward setup, reshard hooks, and post-backward final callback conditionally for v1 vs v2 FSDP. - Root detection now raises RuntimeError when v2 modules exist but none is marked `_is_root` (previously silently disabled FSDP). ### Lazy Gradient Buffers - `_init_dist_grads()` allocates gradient buffers on first use instead of eagerly at init, using `torch.empty` for zero-cost allocation. - `_grad_buffer_is_fresh` flag replaces `is_zero_grad`; correctly resolves stale-data FIXME when `set_to_none=True` is used for zero_grad. ### CPU Offload / Reload - `offload_to_cpu()`: moves buffers to CPU sorted largest-first, with optional `max_cpu_bytes` budget and `pin_memory` support. - `reload_to_gpu()`: explicit pre-warm to hide first-touch latency. - `release_memory_pool()`: tears down CUDA graphs and releases allocator. - `_ensure_buffers_on_gpu()` auto-reloads on any GPU access path. ### Design Docs - Moved all design docs to `v2/design/` directory. - New: `1f1b_ep_overlap_fsdp_design.md` (544 lines), `hooks_api.md` (136). - Existing: `lazy_grad_buffer_design.md`, `allocator_design.md`, `cuda_graph_design.md`, `nvfp4_design.md`, `tp_support_design.md`, `mcore_fsdp_checkpoint_design.md` relocated. ### Removals - `enable_ep_overlap` flag and `_register_ep_overlap_hooks` deleted. - `_ep_submodule_fwd_total/bwd_total/bwd_done` counters removed from `_FSDPState`. - No callers remain in the repo for the removed API.
…ooks` for MFSDP v2 (#21) Replace the custom FSDPCudaGraphRunner with a batch CUDA graph capture approach built on te-graph-runtime (https://github.com/buptzyb/te-graph-runtime) vendored at v2/te_graph_runtime/. Core changes ------------ - cuda_graph_runner.py: rewritten — CudaGraphRunner records sample args per module, batch-captures all graphs in a single make_graphed_callables. - capture_time_hooks run FSDP unshard/reshard outside CUDA graph capture. - hooks.py: @torch.compiler.disable on all 5 FSDP hook functions; CG recording logic replaces per-module capture; capture trigger in post_backward_final_callback guarded by TracePoolAllocator asserts. - te_graph_runtime/: vendored with local modifications (None-safe guards, positional arg replay, empty_cache between warmup/capture, optional capture_stream, separate warmup stream for torch.compile). - fully_shard.py: _fsdp_class_cache avoids torch.compile recompilation from dynamically-created class id changes. Examples & docs --------------- - test_qwenimage.py: end-to-end training test for QwenImage (FA2/FA3, per-block compile, CG+trace pool, memory tracking, OOM dump, benchmarks). - README.md / design/cuda_graph_design.md: updated for te-graph-runtime architecture; TracePoolAllocator highlighted as CG enabler. Benchmarks (QwenImageTransformer2DModel, bs=4, bf16, torch.compile, FA2) ----------------------------------------------------------------------- | Backend | 8×H100 | 4×GB200 | |------------|----------------|----------------| | fsdp1 | 729ms / 60.2GB | 679ms / 75.4GB | | mfsdpv2 | 769ms / 59.3GB | 647ms / 74.7GB | | mfsdpv2+cg | 674ms / 68.3GB | 364ms / 88.7GB | Files: 18 changed, +4199 / -721 " From github.com:shjwudp/Megatron-LM * branch mfsdp_refactor -> FETCH_HEAD [mfsdp_refactor_diffuser 2234409e3] CUDA graph capture via vendored te-graph-runtime with capture_time_hooks for MFSDP v2 18 files changed, 4199 insertions(+), 721 deletions(-) create mode 100644 examples/qwenimage_mfsdp/README.md create mode 100644 examples/qwenimage_mfsdp/test_qwenimage.py create mode 100644 megatron/core/distributed/fsdp/src/megatron_fsdp/v2/cuda_graph_memory_analysis.md create mode 100644 megatron/core/distributed/fsdp/src/megatron_fsdp/v2/te_graph_runtime/README.md create mode 100644 megatron/core/distributed/fsdp/src/megatron_fsdp/v2/te_graph_runtime/__init__.py create mode 100644 megatron/core/distributed/fsdp/src/megatron_fsdp/v2/te_graph_runtime/graph.py …
Squash of mfsdp_refactor_cg_test: - fsdp_toy.py: add --use-real-data (teacher-student regression) with per-step loss + final convergence assert; default CUDA graph off to align Megatron-FSDP and PyTorch FSDP2 paths; scientific-notation loss. - test_qwenimage.py: add --real-data (fixed flow-matching overfit) with cross-rank mean loss and convergence assert.
…tream # Conflicts: # megatron/core/distributed/distributed_data_parallel_config.py # megatron/core/distributed/fsdp/src/megatron_fsdp/distributed_data_parallel_config.py # megatron/core/optimizer/clip_grads.py # megatron/core/pipeline_parallel/combined_1f1b.py # megatron/core/transformer/moe/router.py
6 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do?
Issue tracking
For PRs from open-source community contributors:
Linked issue:
Contribution process
Pre-checks
Code review
Feel free to message or comment @NVIDIA/mcore-oncall to help accelerate your merge into main. The less complex your PR is, the faster it will be approved and merged!
All PRs start as draft. If you open a non-draft PR, it will be automatically converted to draft.
Step 1: Mark PR as "Ready for Review"
.github/CODEOWNERS.Final Review might get declined if these requirements are not fulfilled.
Step 2: Final Review
For PRs that change
megatron/core, once all expert reviewers have approved, theFinal Reviewlabel is applied automatically and final reviewers are assigned.For PRs outside
megatron/core, this step is skipped.Step 3: Approved
Once all required reviewers have approved, the
Approvedlabel is applied automatically.Merge
Any member of mcore-engineers will be able to merge your PR.