Draft [2/N] Megatron FSDP: Introduce Megatron-FSDP2 with per-module fully_shard() API - #4849
Draft
shjwudp wants to merge 68 commits into
Draft
Draft [2/N] Megatron FSDP: Introduce Megatron-FSDP2 with per-module fully_shard() API#4849shjwudp wants to merge 68 commits into
shjwudp wants to merge 68 commits into
Conversation
2. init fully_shard v2 api
Preserve the Megatron-FSDP fully_shard optimizer contract by wiring main-gradient access into the DTensor-backed path and copying optimizer-updated main weights back into model-weight buffers after optimizer.step(). This squash also folds in the review fixes needed for the debug branch changes: - allocate gradient reduce buffers on demand during reduce-scatter - align ParameterGroup.reduce_grad with the fully_shard caller - propagate NaN-check flags to every FSDP module - validate unsharded parameters across all parameter groups - drop leftover post-backward debug code and guard optional grad checks Documentation for this merge is captured in this commit message per request. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
- Move fully_shard_v2.py, param_group.py, dp_buffer.py, allocator.py to megatron/core/distributed/fsdp/src/megatron_fsdp/fully_shard_rewrite/ - Reference uneven_dtensor from parent module instead of copying - Fix test imports in test_param_group.py and test_allocator.py - Add __init__.py with exports - Add README.md documentation
- Rename fully_shard_v2 to fully_shard_rewrite in fsdp_toy.py and adapter - Add extensive docstrings to fully_shard.py and param_group.py - Fix ParameterGroup call: pass mesh=mesh instead of params[0]._dtype
…_rewrite - Add _FSDPRootContext with dedicated CUDA streams for all-gather and reduce-scatter to enable computation-communication overlap - Implement unshard prefetch that pipelines parameter loading to the next module in forward/backward execution order - Enable async reduce-scatter overlapped with backward pass - Plumb mixed precision policy (main_params_dtype, main_grads_dtype) from adapter layer through to ParameterGroup gradient buffers - Simplify dp_buffer.reduce_grad to always accumulate reduced shards synchronously into persistent buffer for gradient accumulation - Fix no_sync property (nullcontext() -> nullcontext) - Add --use-fsdp-fully-shard-api training argument
Fix multiple correctness bugs in the unshard-prefetch and reduce-scatter overlap implementation for fully_shard_rewrite: - Fix `_init_fsdp_state` to pass `enable_unshard_prefetch`/`enable_async_reduce_grad` to child FSDP modules (was calling without required args, causing TypeError). - Fix `reduce_grad_buckets` from positional list to `Dict[id(module)]` mapping, fixing incorrect bucket index lookup that used `forward_order.index(self)`. - Fix async reduce_grad path: remove premature `event.wait()` + `release_grad_buffer()` after recording event, replacing with deferred sliding drain (2-module lag). - Fix `_post_backward_final_callback` to skip modules already handled via `post_backward_issued` flag and drain remaining buckets correctly. - Fix `_scale_gradients` to operate on `dist_grad._local_tensor` instead of `main_grad_buffer.data`. - Fix `_copy_main_weights_to_model_weights` to zero main grads afterward to avoid stale gradients. - Fix `param_group.py`: include `"optim"` in `shard_grads` condition, fix gradient buffer dtype selection, fix `_init_dist_params` for `no_shard` and missing buffers, fix `release_grad_buffer` to drop `main_grad` views (prevent memory leak with TE gradient-accumulation-fusion). - Fix `allocator.py`: change `param_group_id` from `int` to `ParamGroupIdx`, fix `free` to delete bucket after freeing storage, add `_free_storage` safety helper, remove `_resize_` in `allocate` that incorrectly truncated reused buckets. - Add `grad_added_to_main_grad` flag support for TE gradient-accumulation-fusion. - Wire `enable_unshard_prefetch`/`enable_async_reduce_grad` config flags through `mcore_fsdp_adapter.py`. - Extract `ParamGroupIdx`, `RegisterFSDPBackwardFunction`, `_replace_module_parameter` into new `utils.py`. - Add comprehensive `design_overlap.md` documenting the full overlap architecture.
…lap path Bug fixes: - Add stream.wait_stream(current_stream) in unshard() before launching async all-gather on ag_stream, ensuring main-stream writes to parameter data complete before the NCCL all-gather reads them. This was the root cause of convergence divergence in the computation-communication overlap path where stale or partially-written parameter shards were gathered. - Fix shard_grads in _init_buffers() to exclude "optim" strategy: ZeRO-1 should shard optimizer states only, not distribute the gradient buffer. New: - Implement stop_communication() for the fully_shard path (was raising NotImplementedError): waits on ag_stream and rs_stream to bring all communication into the main CUDA stream before the optimizer step. - Add NVTX ranges (MFSDP unshard/reshard/reduce_grad) for profiling. Memory: - Free original full-parameter storage via _free_storage() after copying data into weight buffers during _init_buffers(), reducing peak memory. Docs: - Rename design_overlap.md → design.md and expand to cover stream barriers, NVTX profiling, memory optimization, stop_communication(). - Fix README.md: correct file tree, sharding strategy table (optim row was incorrectly marked as sharding weights/grads), add ZeRO analogues. - Remove stale FIXME comment, polish code comments and docstrings.
This patchset resolves multiple convergence-critical correctness bugs discovered during Llama3-8B convergence testing with the rewrite (`fully_shard_rewrite`) Megatron-FSDP path, and adds diagnostic tooling to prevent future regressions. ## Correctness Fixes (Convergence Critical) ### dist_param attribute propagation (mcore_fsdp_adapter.py) Copy `is_embedding_parameter`, `is_embedding_or_output_parameter`, `sequence_parallel`, `partition_dim`, `partition_stride`, `_tensor_parallel_mode`, and other metadata from original parameters to the FSDP DTensor dist_params. The optimizer param-group builder relies on these attributes; missing them causes parameters to be assigned to wrong optimizer groups (wrong weight-decay / LR multipliers), leading directly to convergence divergence. ### Zero-numel DTensor grads cause silent FusedAdam corruption (param_group.py) When a parameter's local shard has numel=0 on a DP rank, creating a DTensor with an empty local tensor and passing it to fused multi-tensor optimizers (e.g. TE FusedAdam) silently corrupts updates for neighboring non-empty parameters in the same group. The optimizer runs without error — only convergence divergence reveals the bug. Fix: skip DTensor creation when `grad_data.numel() == 0`, record `None` instead. Also guard `_scale_gradients` with `dist_grad is None`. Documented in design.md § Pitfall and README § Gotchas. ### FusedAdam master_weights always disabled (optimizer/__init__.py) Remove the `use_precision_aware_optimizer` gate; FusedAdam's internal master_weights are redundant with Megatron-FSDP's main-weight buffers and cause correctness issues with 1D DTensor parameters. Always disable them. ### Fragment bin-packing in dp_buffer (dp_buffer.py) Replace ad-hoc leftover-fragment placement with a proper bin-packing algorithm that fills chunk_size_factor-sized alignment grids. This fixes buffer layout edge cases where misaligned fragments caused overlaps or out-of-bounds accesses. ### Reduce-scatter output aliasing (dp_buffer.py) Use a separate output tensor in `reduce_scatter_tensor` instead of slicing the full input buffer directly. Aliasing the input can cause the collective to overwrite its own input, silently corrupting gradients. ### Inverted async reduce wait condition (fully_shard.py) `_wait_for_previous_async_reduce_grad` returned early when `enable_async_reduce_grad=True`, opposite of the intended behavior (the legacy code waited only in the async path). This prevented waiting for in-flight reduce events, causing premature grad-buffer release. Fix: invert condition to `if not ctx.enable_async_reduce_grad`. ### Meta-device init DP-rank parameter sync (fully_shard.py) After materializing meta parameters with `reset_parameters()`, each DP rank may have different random values (due to divergent RNG states). Broadcast full parameters from DP rank 0 within the DP-CP mesh *before* FSDP param groups create DTensors, so every rank's shard is a correct slice of the same full parameter. Also fix `broadcast_params` in adapter: changed from `not_implemented_op` to `noop` (the training loop calls it under `--data-parallel-random-init`; the init-time broadcast above covers the sync). ### Meta-device init RNG tracker forking (fully_shard.py) The legacy path wraps `reset_parameters()` with `ResetParametersContext` which forks the model-parallel RNG tracker for non-TE modules when TE >= 0.9.0 is present. Without this fork, TP ranks consume different RNG sequences during init, producing inconsistent values for TP-duplicated parameters (LayerNorm weights, biases). Add equivalent forking in `_materialize_meta_module`. ### post_backward_issued guard (fully_shard.py) Use `getattr(module, "post_backward_issued", False)` to avoid AttributeError when the attribute is missing. ## New Features ### Gradient scaling factor & ReduceOp (fully_shard.py, dp_buffer.py, mcore_fsdp_adapter.py) Support `calculate_per_token_loss`, `average_in_collective`, and default (1/dp_world_size) scaling. Use `_make_nccl_premul_sum` for FP32/FP16 and manual pre-scaling for BF16. ### Stream synchronization (mcore_fsdp_adapter.py) Implement `finish_grad_sync` and `synchronize_param_gather` for the rewrite path (formerly no-ops), synchronizing `rs_stream` and `ag_stream` with the main CUDA stream. ### grad_added_to_main_grad handling (fully_shard.py) When TE gradient-accumulation-fusion writes directly to `main_grad`, discard the dummy `.grad` tensor instead of zeroing/overwriting `main_grad`. ## Debug / Diagnostic Tooling ### per-param norm logging (fully_shard.py, megatron_fsdp.py, training.py, config) New `--log-per-param-norm` config flag logs per-parameter L2 norms for both params and grads, globally reduced across DP ranks. ### Parameter group diagnostics (fully_shard.py) `_log_parameter_groups()` prints compact buffer-layout summaries with memory metrics. `check_all_fsdp_buffers()` validates no local-slice overlaps in any FSDP buffer at runtime. ### Logging callbacks wired to adapter (mcore_fsdp_adapter.py) `log_per_param_norms`, `compute_per_param_norms`, `log_parameter_groups` are now accessible from the `FullyShardedDataParallel` wrapper. ## Documentation - design.md: new "Pitfall: Zero-Numel Gradient Shards" section - README.md: new "Gotchas / Pitfalls" section
Refactor fully shard runtime modules
refactor(mfsdp): simplify mixed precision policy defaults
…ergence fixes Squashed from mfsdp_refactor_main_ep (12 commits up to 00adbef): - Activation recompute / gradient checkpointing: Add backward-phase coordination via _FSDPRootContext with derived backward_module tracking and persistent unshard_done_events. Prevents redundant all-gathers and premature resharding during recompute. (930843c, ca847c1) - TracePoolAllocator: Three-phase bucket allocator (trace/plan/optimized) using greedy left-edge interval coloring. Eliminates per-call torch.empty overhead by pre-allocating a static pool with cursor-based replay across micro-batches. (5410f39, a21041b, e8d434e) - Bucket allocators for weight/grad buffers: StorageFreeingBucketAllocator and TemporaryBucketAllocator manage pooled memory for unsharded parameter and gradient buffers. (5032b6c) - Convergence fixes: Fix zero-numel gradient shard handling to prevent fused optimizer corruption. Add loss detach for E2E correctness. Fix ND parallel setting bug in unit tests. (8cfffdb, adaf543, 4ffa3d6, a91e3da) - Docs: design.md, gap analysis, and allocator documentation. (00adbef, dee8bd4)
- Add Apache 2.0 Copyright headers to 10 new files, update year to 2026 - Replace print() with logging: logger.warning for validation, logger.info for param/grad norms, logger.debug for trace dumps - Add missing docstrings: Bucket, phase, ParamGroupIdx, make_uneven_dtensor, get_state_dict - Remove unused Optional import from allocator.py - Remove misleading # TODO from fully_shard.py mp_policy param - Revert megatron/core/optimizer/__init__.py - isort clean
fix(mfsdp): align v2 token-loss scaling
## Overview Adds full checkpoint save/load support for Megatron FSDP v2 (use_megatron_fsdp_v2, formerly use_fully_shard_api), including model + optimizer state, cross-setting resharding, and online conversion from MFSDP v1 and ND-parallel (torch_dist) formats. ## Core changes ### checkpoint.py — Path A save/load and post-processing - _apply_mcore_postprocess: unified post-processing (FP8 cleanup, SwiGLU/GDN split, expert key remapping) on model state dict; optimizer keys kept canonical for DistributedOptimizer compat - _build_dtensor_optim_sd: wraps FusedAdam plain-tensor states as DTensors using model param mesh/placements (required for DCP) - _propagate_chunk_metadata_to_state_dict: copies chunk metadata from model params to state dict DTensors (zero collectives) - _split_dtensor_v2: unified split supporting both DTensor and plain tensor inputs; uses chunk metadata for uneven sharding - _get_fsdp_slice_from_dtensor: derives FSDP slice from __create_chunk_list__ metadata (handles uneven sharding) - MegatronFSDPStateful: handles optimizer=None gracefully ### uneven_dtensor.py — chunk metadata and split utilities - split_dtensor: general DTensor split with locally-derived chunk metadata (one collective upfront, zero per split) - compute_split_offsets_and_sizes, copy_chunk_metadata, get_fsdp_slice_from_uneven_dtensor: moved from checkpoint.py - make_uneven_dtensor: accepts chunk_metadata/copy_chunk_meta_from params; handles 0-numel edge case ### checkpointing.py — training loop integration - _is_megatron_fsdp_v2: detects v2 models (FSDPModule or FullyShardedDataParallel wrapper) - preprocess_fsdp_dtensor_state_dict: v2 path delegates to _apply_mcore_postprocess; v1 path preserved for baseline - load_checkpoint: auto-detect checkpoint format for fsdp_dtensor/torch_dist before building skeleton - _load_torch_dist_into_megatron_fsdp_v2: loads torch_dist checkpoint (model weights) into v2 skeleton via key mapping - optimizer state wrapping deferred to caller (save_checkpoint / _load_base_checkpoint after raw_optimizer_state_dict capture) - opt_param_scheduler preserved in checkpoint metadata ### distrib_optimizer.py — expert key handling - _param_name: skips handle_experts_in_state_dict for v2 (keeps global expert indices; applied later in post-process) ### arguments.py — rename and auto-set - --use-fully-shard-api renamed to --use-megatron-fsdp-v2 - use_megatron_fsdp_v2=True auto-sets use_megatron_fsdp=True ### Config dataclasses - use_fully_shard_api -> use_megatron_fsdp_v2 in two config files ## Test suite (test_mcore_checkpoint.py) - Single parametrized test_checkpoint_online_convert with (source_type, source_configs, target_configs) covering: - V2 round-trip: 2 cases (optim_grads_params, optim_grads) - V2 cross-setting: 2 cases (strategy A -> strategy B and reverse) - V1 -> V2: 3 cases (different sharding strategies) - ND-parallel -> V2: 2 cases (dp_reshardable, fully_reshardable) - All use train_step() + save_checkpoint() (native MCore APIs) - Model weights verified with exact match; optimizer verified for FSDP source types (V1/V2) via gather-to-full-tensor ## Design doc (mcore_fsdp_checkpoint_design.md) - Updated function table, feature support matrix, dual-dict pattern (Section 5.11), and save/load flow diagrams
shjwudp
changed the title
**load_torch_dist_into_fsdp_v2 — 5-phase pipeline:**
- Phase 1: _preprocess_and_verify_v2_state_dict — shadow dict for DCP load
(DTensors share storage with original plain tensors), verifies chunk metadata
- Phase 2: _build_torch_dist_to_v2_map — key matching with debug logging
- Phase 3: DCP load of mapped state dict
- Phase 4: _load_expert_params_from_torch_dist — expert flatten→individual split
- Phase 5: strictness verification (raises or warns via --strict flag)
**Code consolidation:**
- _split_fused_params_v2 merges handle_swiglu/gdn_in_state_dict_v2 (~130→75 lines)
- _maybe_wrap_as_uneven_dtensor shared by two wrapping functions
- Merged double iteration in preprocessing, simplified apply_mcore_postprocess
**Bug fixes:**
- _match_gdn_key missing dtensor arg (TypeError in GDN optimizer path)
- _load_expert_params_from_torch_dist: tuple unpack mismatch, multi-chunk
copy overwrite, module. prefix mismatch in optim_matched, td_flat_key
construction with removeprefix('module.'), debug print cleanup
**Dead code removal:**
- save_checkpoint/load_checkpoint from __all__, _get_tp_world_size,
_intersection, _offset_slice, _expert_param_local_key, unused imports
**checkpointing.py:** strict param threaded through call chain
**Design doc:** updated function table, section references, normalization table
… develop/mfsdp-refactor-main-stage2 # Conflicts: # megatron/core/distributed/fsdp/src/megatron_fsdp/v2/param_group.py
Introduce NVFP4 primary-weights support in Megatron FSDP v2. ### Core: NVFP4 storage handling - **BufferIndex.compact(factor, compact_shapes)**: Proportionally scale all indices for packed storage (factor=0.5 for NVFP4 2:1 packing). All buffers build with logical shapes + shared chunk_size_factor; only model_weight and transpose_weight buffers are compacted. Preserves proportional item-offset mapping across buffers. - **get_param_storage_shapes()**: Return packed shapes for NVFP4 (last dim halved), logical shapes otherwise. Drives compact and DTensor creation. - **FullyShardNVFP4Policy**: NVFP4 policy with enabled/recipe, wired through FullyShardMixedPrecisionPolicy and mcore_fsdp_adapter. - **fp4_param_gather**: Added to DistributedDataParallelConfig, mirrors fp8_param_gather. ### Mixed precision NVFP4 path - quantize_main_weights_to_nvfp4: Convert fp32 main shards to packed NVFP4 model-weight shards using TE quantize_master_weights. - post_unshard: Calls post_all_gather_processing for NVFP4 params. - is_nvfp4_param, get_nvfp4_raw_data, model_init_context NVFP4 support. ### Checkpoint refactor - Generic fused-layer matching (_match_fused_key) replaces hardcoded expert-key patterns. Supports regular fused layers, GroupedMLP, and SequentialMLP experts. - Phase 4 slicing in load_torch_dist_into_fsdp_v2 uses DTensor shards to avoid OOM on large fused tensors. - _find_param_in_map: While-loop strips multiple module. prefixes, fixing optimizer state dict lookups for double-wrapped models. - Metadata propagation: _propagate_chunk_metadata_to_state_dict at save time; _apply_mcore_postprocess validates __create_chunk_list__ on all DTensors. - Removed tautological assert in _match_fused_key debug logging. ### API renames - _get_item_offset → _get_item_global_range - _get_item_slice_in_shard → _get_item_self_range - Merged _get_item_local_index/_get_item_local_shard_index → _get_item_local_range with as_shard kwarg - fetch_unsharded_buffer → fetch_buffer - only_shard → as_shard kwarg on get_item ### Documentation - Added nvfp4_design.md: Full NVFP4 design including storage shape distinction, compact pattern, chunk_size_factor rationale, checkpoint flow, and BufferIndex API reference (summon_full_params marked planned). - Updated mcore_fsdp_checkpoint_design.md for Phase 4 generic slicing. - Updated README.md sharding strategy status (only optim_grads_params). ### Tests - NVFP4 checkpoint round-trip test (v2_rt_nvfp4_optim_grads_params). - NVFP4 E2E test in test_mcore_nd_parallel. - Skipped no_shard/optim/optim_grads tests (not yet supported). - Reduced layer dimensions in test_param_group for faster UT runs. ### Fixes - _match_fused_key: Delete tautological assert after add_to_set. - checkpointing.py: Recursive v2 detection, metadata propagation for fsdp_dtensor saves, reject dp_reshardable for v2. - dp_buffer.py: init_buffers remove dead model_weight_shapes variable. - fsdp_module.py: Gradient dtype cast replaces failing assert. - utils.py: Fix torch._C._get_accelerator crash on CPU.
fix(mfsdp): support fsdp v2 zero-1 and zero-2
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 ?
Part-1 link: #4435
Summary
Introduce Megatron-FSDP2 — a refactored per-module sharding implementation (
fully_shard_v2) with FSDP2-compatible API, communication-compute overlap, activation recompute support, and a pooled memory allocator for Megatron Core.Key Features
fully_shard()API: FSDP2-compatible wrapping interface that converts modules toFSDPModuledynamically, groups parameters intoParameterGroups with flatDataParallelBuffers, and installs forward/backward hooks for the unshard → forward → reshard → backward → reduce lifecycle.ag_stream): all-gathers parameters for the next module while the current module computes forward.rs_stream): reduce-scatters gradients while later modules compute backward. Sliding drain rule keeps ≤ 2 gradient buffers live at any time.backward_moduletracking + persistentunshard_done_eventsprevent redundant all-gathers and premature resharding during gradient checkpointing recompute passes.TracePoolAllocator: Three-phase (trace → plan → optimized) pooled memory allocator using greedy left-edge interval coloring. Eliminates per-calltorch.emptyoverhead with a static pool replayed across micro-batches.no_shard,optim(ZeRO-1),optim_grads(ZeRO-2),optim_grads_params(ZeRO-3), with uneven DTensor and distributed checkpoint support.Files Changed (27 files, +5,014 / −9)
v2/package__init__.py,fully_shard.py,fsdp_module.py,hooks.py,param_group.py,dp_buffer.py,allocator.py,utils.py,mixed_precision.py,design.md,README.mdmcore_fsdp_adapter.py,src/megatron_fsdp/__init__.py,megatron_fsdp.py,distributed_data_parallel_config.pyexamples/megatron_fsdp/fsdp_toy.pytests/.../v2/test_allocator.py,test_param_group.py,test_mcore_fully_shard_api.py,test_checkpoint_online_convert.pyExperimental Results — Per-Module Sharding Rewrite
Verification was performed comparing the refactored implementation against the baseline.
W&B link: https://wandb.ai/nvidia/megatron-fsdp/reports/M-FSDP-Rewrite-Convergence-Test-llama3-8b---VmlldzoxNjgzNTg3OQ
Issue tracking
For PRs from open-source community contributors:
Linked issue:
Contribution process
Pre-checks
Code review
Feel free to message or comment the @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.
For MRs into `dev` branch
The proposed review process for `dev` branch is under active discussion.MRs are mergable after one approval by either
eharper@nvidia.comorzijiey@nvidia.com.