From 31b968bbf50484b73e3ebd6bce57a974d61e6bae Mon Sep 17 00:00:00 2001 From: Tong Liu Date: Thu, 28 May 2026 07:30:42 +0800 Subject: [PATCH 01/28] feat(mfsdp-v2): mixed-precision infrastructure and MXFP8 support (by 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 --- .../fsdp/src/megatron_fsdp/v2/allocator.py | 564 ++++++++++++++ .../fsdp/src/megatron_fsdp/v2/hooks.py | 256 ++++++ .../src/megatron_fsdp/v2/mixed_precision.py | 727 ++++++++++++++++++ .../fsdp/src/megatron_fsdp/v2/param_group.py | 357 +++++++++ 4 files changed, 1904 insertions(+) create mode 100644 megatron/core/distributed/fsdp/src/megatron_fsdp/v2/allocator.py create mode 100644 megatron/core/distributed/fsdp/src/megatron_fsdp/v2/hooks.py create mode 100644 megatron/core/distributed/fsdp/src/megatron_fsdp/v2/mixed_precision.py create mode 100644 megatron/core/distributed/fsdp/src/megatron_fsdp/v2/param_group.py diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/allocator.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/allocator.py new file mode 100644 index 00000000000..80ba21e750e --- /dev/null +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/allocator.py @@ -0,0 +1,564 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import dataclasses +from typing import Dict, Hashable, List, Optional, Tuple + +import torch + +AllocatorKey = Hashable + + +def _resolve_key(key: Optional[AllocatorKey], param_group_id: Optional[AllocatorKey]): + if key is not None: + return key + # Backward-compatible alias for older callers that keyed only by param group. + assert param_group_id is not None, "allocator key is required" + return param_group_id + + +@dataclasses.dataclass +class Bucket: + """Lightweight container for a temporary allocated tensor buffer.""" + + data: torch.Tensor + + +class BucketAllocator: + """Interface for allocating and freeing temporary buckets.""" + + def allocate( + self, + key: Optional[AllocatorKey] = None, + size: int = 0, + dtype: Optional[torch.dtype] = None, + device: Optional[torch.device] = None, + *, + param_group_id: Optional[AllocatorKey] = None, + ) -> Bucket: + """Allocate a bucket for the given key.""" + raise NotImplementedError + + def free( + self, key: Optional[AllocatorKey] = None, *, param_group_id: Optional[AllocatorKey] = None + ) -> None: + """Free the bucket associated with the given key.""" + raise NotImplementedError + + +class TemporaryBucketAllocator(BucketAllocator): + """Manages temporary flat buffers keyed by a caller-provided key. + + Used by DataParallelBuffer for unshard (all-gather) and gradient + reduction (reduce-scatter) operations. + """ + + def __init__(self): + super().__init__() + self.buckets = {} + + def allocate( + self, + key: Optional[AllocatorKey] = None, + size: int = 0, + dtype: Optional[torch.dtype] = None, + device: Optional[torch.device] = None, + *, + param_group_id: Optional[AllocatorKey] = None, + ) -> Bucket: + key = _resolve_key(key, param_group_id) + assert dtype is not None and device is not None + if key not in self.buckets: + self.buckets[key] = Bucket(data=torch.empty(size, dtype=dtype, device=device)) + return self.buckets[key] + + def free( + self, key: Optional[AllocatorKey] = None, *, param_group_id: Optional[AllocatorKey] = None + ) -> None: + key = _resolve_key(key, param_group_id) + if key in self.buckets: + _free_storage(self.buckets[key].data) + del self.buckets[key] + + +class StorageFreeingBucketAllocator(BucketAllocator): + """Manages temporary flat buffers keyed by caller-provided allocation key. + + Freeing releases the underlying storage without deleting the bucket entry, + so the same tensor object can be reused on the next allocation. + """ + + def __init__(self): + super().__init__() + self.buckets = {} + + def allocate( + self, + key: Optional[AllocatorKey] = None, + size: int = 0, + dtype: Optional[torch.dtype] = None, + device: Optional[torch.device] = None, + *, + param_group_id: Optional[AllocatorKey] = None, + ) -> Bucket: + key = _resolve_key(key, param_group_id) + assert dtype is not None and device is not None + if key not in self.buckets: + self.buckets[key] = Bucket(data=torch.empty(size, dtype=dtype, device=device)) + return self.buckets[key] + _alloc_storage(self.buckets[key].data, torch.Size([size])) + return self.buckets[key] + + def free( + self, key: Optional[AllocatorKey] = None, *, param_group_id: Optional[AllocatorKey] = None + ) -> None: + key = _resolve_key(key, param_group_id) + if key in self.buckets: + _free_storage(self.buckets[key].data) + + +class TracePoolAllocator(BucketAllocator): + """Two-phase bucket allocator that eliminates per-call ``torch.empty`` overhead. + + **Design** + + The FSDP framework allocates and frees temporary flat buffers (for + all-gather input/output and gradient accumulation) in a deterministic, + repeatable order across micro-batches. ``TracePoolAllocator`` exploits + this by profiling one pass and then serving all subsequent passes from + a pre-allocated pool. + + **Phase 1 — Trace** (``plan()`` not yet called) + + Behaves like ``TemporaryBucketAllocator``: ``allocate`` creates a + ``torch.empty`` bucket on first use, ``free`` releases its storage. + Additionally, every alloc/free call is recorded as a ``_TraceEvent`` + with a monotonic ``seq`` number, and metadata ``(size, dtype, device)`` + is stored per allocation key for later planning. + + **Phase 2 — Plan** (``plan()``) + + The trace is replayed to extract *intervals*: for each alloc/free pair + an ``_Interval(alloc_seq, free_seq, size)`` is built. Intervals are + grouped by ``(dtype, device)`` and then colored with a greedy + left-edge algorithm: + + 1. Sort intervals by ``alloc_seq``. + 2. For each interval, try to reuse a *slot* whose previous occupant + freed before this interval starts (``slot_free_seq < alloc_seq``). + 3. If no slot is free, allocate a new one. + 4. Grow the slot's capacity to ``max(size, current)``. + 5. Record the assignment: append the slot index to the per-key list + in ``_slot_map``. + + After coloring, slots are laid out contiguously and a single + ``torch.empty`` per ``(dtype, device)`` group is allocated. + + **Phase 3 — Optimized** (after ``plan()``) + + ``allocate`` returns a ``Bucket`` with a slice-view into the pool; + ``free`` marks the slot as unused but never releases storage. Because + the same allocation key can appear in multiple intervals (e.g., + forward unshard → free → backward unshard → free), ``_slot_map`` maps + each key to a **list** of slot indices in alloc order. A per-key + ``_slot_cursors`` counter tracks which index to consume next. Call + ``reset_cursor()`` at the start of each micro-batch to rewind all + cursors to 0. + + The trace pattern must be **repeatable** — the same alloc/free call + sequence is expected every micro-batch. + """ + + # -- Inner types ---------------------------------------------------- # + + class _Slot: + """A contiguous slice of the pool tensor assigned to one or more + non-overlapping intervals.""" + + __slots__ = ("offset", "size", "dtype", "device", "in_use") + + def __init__(self, offset: int, size: int, dtype: torch.dtype, device: torch.device): + self.offset = offset + self.size = size + self.dtype = dtype + self.device = device + self.in_use = False + + @dataclasses.dataclass + class _TraceEvent: + """A single alloc or free recorded during the trace phase.""" + + seq: int + op: str # "alloc" | "free" + key: AllocatorKey + + @dataclasses.dataclass + class _Interval: + """An allocation's lifetime: from alloc_seq to free_seq with a given size.""" + + key: AllocatorKey + size: int + alloc_seq: int + free_seq: int + + # -- Init ----------------------------------------------------------- # + + def __init__(self) -> None: + super().__init__() + # Phase bookkeeping + self._phase: str = "trace" # "trace" | "optimized" + self._seq: int = 0 # monotonic alloc/free counter + + # Trace state + self._trace: List["TracePoolAllocator._TraceEvent"] = [] + self._trace_meta: Dict[AllocatorKey, Tuple[int, torch.dtype, torch.device]] = {} + self._buckets: Dict[AllocatorKey, Bucket] = {} # only used in trace phase + + # Pool state — populated by plan(), used in optimized phase + self._pools: Dict[Tuple[torch.dtype, torch.device], torch.Tensor] = {} + self._slot_map: Dict[AllocatorKey, List[int]] = {} # key -> [slot indices] + self._slot_cursors: Dict[AllocatorKey, int] = {} # key -> next index to use + self._slots: List["TracePoolAllocator._Slot"] = [] + + # -- Phase 1: trace -------------------------------------------------- # + + def allocate( + self, + key: Optional[AllocatorKey] = None, + size: int = 0, + dtype: Optional[torch.dtype] = None, + device: Optional[torch.device] = None, + *, + param_group_id: Optional[AllocatorKey] = None, + ) -> Bucket: + """Dispatch to trace or pool path depending on phase.""" + key = _resolve_key(key, param_group_id) + assert dtype is not None and device is not None + if self._phase != "optimized": + return self._trace_allocate(key, size, dtype, device) + return self._pool_allocate(key, size, dtype, device) + + def free( + self, key: Optional[AllocatorKey] = None, *, param_group_id: Optional[AllocatorKey] = None + ) -> None: + """Dispatch to trace or pool path depending on phase.""" + key = _resolve_key(key, param_group_id) + if self._phase != "optimized": + self._trace_free(key) + else: + self._pool_free(key) + + def _trace_allocate( + self, key: AllocatorKey, size: int, dtype: torch.dtype, device: torch.device + ) -> Bucket: + """Trace-phase allocate: record the event and create a bucket on first use. + + Duplicate allocs (without an intervening free) do NOT generate + new trace events — they are no-ops that return the existing bucket. + """ + if key not in self._buckets: + self._trace.append(self._TraceEvent(seq=self._seq, op="alloc", key=key)) + self._seq += 1 + self._trace_meta[key] = (size, dtype, device) + self._buckets[key] = Bucket(data=torch.empty(size, dtype=dtype, device=device)) + return self._buckets[key] + + def _trace_free(self, key: AllocatorKey) -> None: + """Trace-phase free: record the event and release the bucket storage.""" + self._trace.append(self._TraceEvent(seq=self._seq, op="free", key=key)) + self._seq += 1 + if key in self._buckets: + _free_storage(self._buckets[key].data) + del self._buckets[key] + + # -- Phase 2: plan --------------------------------------------------- # + + def plan(self) -> int: + """Build the static pool from the recorded trace. + + 1. Replay the trace to pair alloc/free events into ``_Interval`` objects. + 2. Group intervals by ``(dtype, device)``. + 3. Color each group with the greedy left-edge algorithm. + 4. Allocate one flat pool tensor per group. + + Returns: + Total pool size in **elements** (sum across all groups). + Multiply by ``element_size(dtype)`` for bytes. + """ + assert self._phase == "trace", "plan() can only be called in trace phase" + assert len(self._trace) > 0, "empty trace — nothing to plan" + + # ---- step 1: build intervals from alloc/free pairs ---- + alloc_stack: Dict[AllocatorKey, List[int]] = {} # key -> [alloc_seq, ...] + intervals: List["TracePoolAllocator._Interval"] = [] + + for ev in self._trace: + if ev.op == "alloc": + alloc_stack.setdefault(ev.key, []).append(ev.seq) + else: # "free" + if ev.key in alloc_stack and alloc_stack[ev.key]: + alloc_seq = alloc_stack[ev.key].pop(0) + meta = self._trace_meta.get(ev.key) + if meta is not None: + size, dtype, device = meta + intervals.append( + self._Interval( + key=ev.key, size=size, alloc_seq=alloc_seq, free_seq=ev.seq + ) + ) + + assert len(intervals) > 0, "no paired alloc/free intervals found in trace" + + # ---- step 2 & 3: color and allocate ---- + return self._assign_pool(intervals) + + def _assign_pool(self, intervals: List["TracePoolAllocator._Interval"]) -> int: + """Group intervals by (dtype, device), color each group, sum sizes.""" + groups: Dict[Tuple[torch.dtype, torch.device], List["TracePoolAllocator._Interval"]] = {} + for iv in intervals: + meta = self._trace_meta[iv.key] + dtype_device = (meta[1], meta[2]) + groups.setdefault(dtype_device, []).append(iv) + + # Clear any previous plan state before rebuilding + self._slot_map.clear() + self._slot_cursors.clear() + self._slots.clear() + self._pools.clear() + + total_elems = 0 + for (dtype, device), group in groups.items(): + total_elems += self._color_group(group, dtype, device) + + self._phase = "optimized" + return total_elems + + def _color_group( + self, + intervals: List["TracePoolAllocator._Interval"], + dtype: torch.dtype, + device: torch.device, + ) -> int: + """Greedy left-edge interval coloring for one (dtype, device) group. + + Algorithm: + + 1. Sort intervals by ``alloc_seq``. + 2. Maintain a *free-list* of ``(slot_index, free_seq)`` — slots that + become free at ``free_seq``. + 3. For each interval: + a. Scan the free-list for a slot whose ``free_seq < alloc_seq``. + If found, reuse it (grow its size if needed) and update its + free time to this interval's ``free_seq``. + b. If no slot is free, create a new one. + c. Append the assigned slot index to ``_slot_map[key]``. + + After all intervals are colored, slots are laid out contiguously + and a single ``torch.empty`` is issued for the group. + + Complexity: O(n²) worst-case (linear scan per interval), negligible + for the typical tens-to-hundreds of param groups per dtype/device. + """ + intervals = sorted(intervals, key=lambda iv: iv.alloc_seq) + + # free_slots: list of (local_slot_index, free_seq) + free_slots: List[Tuple[int, int]] = [] + group_slots: List["TracePoolAllocator._Slot"] = [] # slots for this group + local_to_global: Dict[int, int] = {} # local index -> global index + + for iv in intervals: + assigned = False + # Try to reuse an existing slot whose previous occupant has already freed + for i, (slot_idx, slot_free_seq) in enumerate(free_slots): + if slot_free_seq < iv.alloc_seq: + slot = group_slots[slot_idx] + if iv.size > slot.size: + slot.size = iv.size # grow if needed + free_slots[i] = (slot_idx, iv.free_seq) # update free time + self._slot_map.setdefault(iv.key, []).append(local_to_global[slot_idx]) + assigned = True + break + + if not assigned: + # No reusable slot — allocate a new one + local_idx = len(group_slots) + global_idx = len(self._slots) + local_to_global[local_idx] = global_idx + slot = self._Slot(offset=0, size=iv.size, dtype=dtype, device=device) + group_slots.append(slot) + self._slots.append(slot) + free_slots.append((local_idx, iv.free_seq)) + self._slot_map.setdefault(iv.key, []).append(global_idx) + + # Lay out slots contiguously within the group pool + offset = 0 + for slot in group_slots: + slot.offset = offset + offset += slot.size + + if offset > 0: + self._pools[(dtype, device)] = torch.empty(offset, dtype=dtype, device=device) + return offset + + # -- Phase 3: optimized runtime ------------------------------------- # + # + # Each micro-batch replays the same alloc/free sequence. Because a + # single allocation key may appear multiple times (e.g., forward + # unshard → free → backward unshard → free), ``_slot_map[key]`` is + # a **list** of slot indices in alloc order. A per-key cursor in + # ``_slot_cursors`` tracks which index to consume next. Between + # micro-batches, ``reset_cursor()`` rewinds all cursors to 0. + + def _pool_allocate( + self, key: AllocatorKey, size: int, dtype: torch.dtype, device: torch.device + ) -> Bucket: + """Return a ``Bucket`` whose data is a slice of the pre-allocated pool. + + Advances the per-key slot cursor after consuming the slot. + """ + slot_list = self._slot_map[key] + cursor = self._slot_cursors.get(key, 0) + assert cursor < len(slot_list), ( + f"no slot available for key={key} " f"(cursor={cursor}, slots={slot_list})" + ) + slot_idx = slot_list[cursor] + self._slot_cursors[key] = cursor + 1 + + slot = self._slots[slot_idx] + assert not slot.in_use, ( + f"slot {slot_idx} already in use (key={key}, " + f"cursor={cursor}, slot_list={slot_list})" + ) + assert size <= slot.size, f"requested {size} > slot capacity {slot.size} (key={key})" + pool = self._pools[(slot.dtype, slot.device)] + slot.in_use = True + self._seq += 1 + return Bucket(data=pool[slot.offset : slot.offset + size]) + + def _pool_free(self, key: AllocatorKey) -> None: + """Mark the most recently allocated slot for this key as free. + + Double-frees are silently ignored (idempotent). + """ + # The last allocated slot for this key is at cursor - 1. + slot_idx = self._slot_map[key][self._slot_cursors.get(key, 1) - 1] + slot = self._slots[slot_idx] + self._seq += 1 + if not slot.in_use: + return + slot.in_use = False + + # -- Debug ---------------------------------------------------------- # + + def dump_trace(self) -> str: + """Return a human-readable dump of the trace and pool plan. + + Useful for debugging slot-conflict errors (e.g., "slot already in use"). + """ + lines = [] + lines.append(f"=== TracePoolAllocator (phase={self._phase}, seq={self._seq}) ===") + lines.append(f"trace events: {len(self._trace)}") + for ev in self._trace: + meta = self._trace_meta.get(ev.key) + size_str = f"size={meta[0]}" if meta else "size=?" + dtype_str = f"dtype={meta[1]}" if meta else "dtype=?" + device_str = f"device={meta[2]}" if meta else "device=?" + lines.append( + f" seq={ev.seq:>4} {ev.op:>5} key={ev.key} " + f"{size_str} {dtype_str} {device_str}" + ) + + if self._phase == "optimized": + lines.append(f"\nslots: {len(self._slots)}") + for i, slot in enumerate(self._slots): + lines.append( + f" slot[{i}]: offset={slot.offset} size={slot.size} " + f"dtype={slot.dtype} device={slot.device}" + ) + lines.append("\nslot_map (key -> [slot indices]):") + for key, slot_list in self._slot_map.items(): + cursor = self._slot_cursors.get(key, 0) + lines.append(f" {key} -> {slot_list} cursor={cursor}") + + return "\n".join(lines) + + # -- Lifecycle ------------------------------------------------------- # + + def reset_cursor(self) -> None: + """Reset all slot cursors to 0 for the next iteration / micro-batch. + + Must be called between micro-batches in the optimized phase so that + the alloc/free sequence replays from the start of the slot lists. + """ + self._slot_cursors.clear() + self._seq = 0 + + def reset(self) -> None: + """Reset to trace phase, discarding the pool and all recorded state.""" + self._phase = "trace" + self._seq = 0 + self._trace.clear() + self._trace_meta.clear() + self._buckets.clear() + self._pools.clear() + self._slot_map.clear() + self._slot_cursors.clear() + self._slots.clear() + + @property + def phase(self) -> str: + """Current allocator phase: 'trace', 'plan', or 'optimized'.""" + return self._phase + + @property + def total_pool_bytes(self) -> int: + """Total pool size in bytes across all dtype/device groups.""" + total = 0 + for (dtype, _), pool in self._pools.items(): + total += pool.numel() * pool.element_size() + return total + + +def _free_storage(tensor: torch.Tensor) -> None: + """Free the underlying storage of ``tensor`` by resizing it to 0.""" + with torch.no_grad(): + if not torch.distributed._functional_collectives.is_torchdynamo_compiling(): + already_freed = tensor._typed_storage()._size() == 0 + if not already_freed: + assert tensor.storage_offset() == 0, ( + "Freeing a tensor's storage is unsafe when it is not the sole occupant\n" + f"storage offset: {tensor.storage_offset()}\n" + f"storage size: {tensor._typed_storage()._size()}\n" + f"tensor shape: {tensor.shape}" + ) + tensor._typed_storage()._resize_(0) + + +def _alloc_storage(tensor: torch.Tensor, size: torch.Size) -> None: + """Re-allocate storage for ``tensor`` to the given ``size``. + + Requires that the tensor's storage has been freed (resized to 0) + before calling. The caller must ensure ``size`` matches the tensor's + existing shape. + """ + with torch.no_grad(): + if not torch.distributed._functional_collectives.is_torchdynamo_compiling(): + already_allocated = tensor._typed_storage()._size() == size.numel() + if not already_allocated: + tensor_storage_size = tensor._typed_storage()._size() + assert tensor_storage_size == 0, ( + "Tensor storage should have been resized to 0 but got " + f"{tensor_storage_size} (shape={tensor.shape})" + ) + tensor._typed_storage()._resize_(size.numel()) diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/hooks.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/hooks.py new file mode 100644 index 00000000000..30a269e9437 --- /dev/null +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/hooks.py @@ -0,0 +1,256 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Forward and backward hook registration for Megatron-FSDP2.""" + +import functools +import logging +from typing import Any, Callable, Dict, List, Tuple + +import torch +import torch.nn as nn +from torch.autograd import Variable +from torch.utils._pytree import tree_flatten, tree_map, tree_unflatten + +from .allocator import TracePoolAllocator +from .fsdp_module import FSDPModule, _FSDPState +from .utils import RegisterFSDPBackwardFunction + +logger = logging.getLogger(__name__) + + +def _register_forward_pre_hook(fsdp_module: FSDPModule, hook_module: nn.Module | None = None): + """Register a pre-forward hook to unshard parameters for this FSDP unit.""" + + def unshard_param_groups(_hook_module, *unused): + ctx = fsdp_module._fsdp_root_context + if ctx.backward_phase: + fsdp_module.unshard(async_op=ctx.enable_unshard_prefetch, bwd_pass=True) + fsdp_module.unshard(async_op=ctx.enable_unshard_prefetch, bwd_pass=False) + + # In the normal path, the hook is installed on the FSDP module itself. + # Fine-grained recompute can replay a child forward without entering the + # FSDP module's forward, so the same owner hook may need to be installed + # on a child module while still unsharding the owning FSDP module above. + module_to_hook = hook_module if hook_module is not None else fsdp_module + return module_to_hook.register_forward_pre_hook(unshard_param_groups, prepend=True) + + +def _register_fine_grained_forward_pre_hooks(module: FSDPModule): + """Register pre-forward hooks on this FSDP unit and its owned submodules.""" + + nested_fsdp_modules = set() + for submodule in module.modules(): + if submodule is not module and isinstance(submodule, FSDPModule): + nested_fsdp_modules.update(submodule.modules()) + + for submodule in module.modules(): + if submodule in nested_fsdp_modules: + continue + _register_forward_pre_hook(module, hook_module=submodule) + + +def _register_forward_hook(module: FSDPModule): + """Register post-forward hook to reshard parameters.""" + + def reshard_param_groups(module, *unused): + ctx = module._fsdp_root_context + if ctx.backward_phase and id(module) == ctx.backward_module: + return + module.reshard() + + module._mfsdp_forward_hook = module.register_forward_hook(reshard_param_groups) + + +def _register_backward_pre_hook(module: FSDPModule): + """ + Register backward pre-hook to handle gradient computation. + + This uses a custom backward hook that attaches to output tensors + to trigger unshard at the right time during backward pass. + """ + + def create_custom_backward_hook(module: FSDPModule, custom_backward_handler: Callable): + """Create a custom backward hook attached to output tensors.""" + + @torch.compiler.disable + def forward_hook(_module, inputs, output): + # View-as to avoid output being the same tensor object + output = tree_map(lambda t: t.view_as(t) if torch.is_tensor(t) else t, output) + + # Collect tensor outputs + output_list = [] + if isinstance(output, torch.Tensor): + output_list = [output] + elif isinstance(output, (tuple, list)): + output_list = [t for t in output if isinstance(t, torch.Tensor)] + + # Register pre-backward hook on output tensors. + # This triggers when gradients are computed. + torch.autograd.graph.register_multi_grad_hook( + output_list, lambda grads: custom_backward_handler(_module, grads), mode="any" + ) + return output + + return module.register_forward_hook(forward_hook) + + def pre_backward_hook(module: FSDPModule, grads): + """Hook called before backward pass for this module.""" + ctx = module._fsdp_root_context + if module._fsdp_state._is_root: + ctx.backward_done_modules.clear() + ctx.backward_phase = True + ctx._advance_backward_module() + setattr(module, "post_backward_issued", False) + for param_group in module._fsdp_param_groups: + for param in param_group.params: + setattr(param, "grad_added_to_main_grad", False) + # When gradient_accumulation_fusion is active, the TE backward + # kernel writes weight gradients directly into param.main_grad. + # By default TE accumulates (adds) into main_grad, which silently + # doubles gradients when the buffer isn't zeroed between + # micro-batches. Setting overwrite_main_grad tells TE to + # overwrite instead of accumulate, preventing NaN from corrupted + # gradient buffers. + if param_group.sharding_strategy in ["optim_grads_params", "optim_grads"]: + setattr(param, "overwrite_main_grad", True) + if module._fsdp_state._is_root and not module._fsdp_state._post_backward_callback_queued: + _register_post_backward_final_callback(module._fsdp_state, module) + module.unshard(async_op=ctx.enable_unshard_prefetch, bwd_pass=True) + + module._mfsdp_backward_pre_hook = create_custom_backward_hook( + module, custom_backward_handler=pre_backward_hook + ) + + +def _register_backward_hook(module: FSDPModule): + """ + Register backward hook using autograd Function. + + This inserts a RegisterFSDPBackwardFunction in the backward pass + that triggers reshard and reduce_grad after gradients are computed. + """ + + def post_backward(module: FSDPModule): + """Hook called after backward pass for this module.""" + ctx = module._fsdp_root_context + ctx.backward_done_modules.add(id(module)) + ctx._advance_backward_module() + module.reshard() + module.reduce_grad(async_op=ctx.enable_async_reduce_grad) + module.post_backward_issued = True + + @torch.compiler.disable + def _register_post_backward_hook( + post_backward_hook: Callable, + module: nn.Module, + args: Tuple[Any, ...], + kwargs: Dict[str, Any], + ): + """ + Register a post-backward hook by inserting an autograd Function. + + This approach works by registering a pre-forward hook that wraps + input tensors in an autograd Function. The Function's backward + calls the post_backward_hook after gradients are computed. + """ + if not torch.is_grad_enabled(): + return args, kwargs + + # Flatten args and kwargs + args_list, args_spec = tree_flatten(args) + kwargs_list, kwargs_spec = tree_flatten(kwargs) + args_kwargs_list = list(args_list) + list(kwargs_list) + + # Filter to tensors with gradients + inp_tensor_indices: List[int] = [] + inp_tensors: List[torch.Tensor] = [] + for i, obj in enumerate(args_kwargs_list): + if torch.is_tensor(obj) and obj.requires_grad: + inp_tensor_indices.append(i) + inp_tensors.append(obj) + + if len(inp_tensors) == 0: + return args, kwargs + + # Wrap inputs in autograd Function. + # The Function's backward will call post_backward_hook. + inp_tensors = RegisterFSDPBackwardFunction.apply( + functools.partial(post_backward_hook, module), *inp_tensors + ) + + # Restore args and kwargs + for inp_tensor_idx, inp_tensor in zip(inp_tensor_indices, inp_tensors): + args_kwargs_list[inp_tensor_idx] = inp_tensor + args_list = args_kwargs_list[: len(args_list)] + kwargs_list = args_kwargs_list[len(args_list) :] + args = tree_unflatten(args_list, args_spec) + kwargs = tree_unflatten(kwargs_list, kwargs_spec) + + return args, kwargs + + module._mfsdp_backward_hook = module.register_forward_pre_hook( + functools.partial(_register_post_backward_hook, post_backward), with_kwargs=True + ) + + +def _register_post_backward_final_callback(state: _FSDPState, module: nn.Module) -> None: + """ + Register the final callback that runs after all backward passes complete. + + This is only registered by the root FSDP module to avoid duplicate + callbacks. It reshards all modules and reduces gradients at the end + of the backward pass. + """ + assert state._is_root, "Only root FSDP should register post-backward callback" + if state._post_backward_callback_queued: + return + + def _post_backward_final_callback(root_state: _FSDPState, root_module: nn.Module): + """Final callback: reshard all modules and reduce gradients.""" + ctx = root_module._fsdp_root_context + stream = ctx.rs_stream + for module in reversed(ctx.forward_order): + if getattr(module, "post_backward_issued", False): + continue + module.reshard() + module.reduce_grad(async_op=ctx.enable_async_reduce_grad) + for buckets in ctx.reduce_grad_buckets.values(): + while len(buckets) > 0: + event, param_group = buckets.pop() + event.wait() + param_group.release_grad_buffer() + torch.cuda.current_stream().wait_stream(stream) + root_state._post_backward_callback_queued = False + ctx.backward_phase = False + ctx.backward_module = None + ctx.backward_done_modules.clear() + + # After the first backward pass, we have the full trace of bucket allocations + # and releases. We can now plan the memory pool based on this trace. + if isinstance(ctx.bucket_allocator, TracePoolAllocator): + bucket_alloc = ctx.bucket_allocator + if bucket_alloc.phase == "trace": + if torch.distributed.get_rank() == 0: + logger.debug(bucket_alloc.dump_trace()) + bucket_alloc.plan() + elif bucket_alloc.phase == "optimized": + bucket_alloc.reset_cursor() + else: + raise ValueError(f"Unexpected bucket allocator phase: {bucket_alloc.phase}") + + state._post_backward_callback_queued = True + Variable._execution_engine.queue_callback( + functools.partial(_post_backward_final_callback, state, module) + ) diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/mixed_precision.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/mixed_precision.py new file mode 100644 index 00000000000..ed0d44cab83 --- /dev/null +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/mixed_precision.py @@ -0,0 +1,727 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Mixed precision policy helpers for Megatron-FSDP2. + +This module owns the v2 policy data model. Translation from Megatron/MCore +config objects belongs in the adapter layer. +""" + +import inspect +from contextlib import ExitStack, contextmanager, nullcontext +from dataclasses import dataclass, field +from importlib.metadata import version +from typing import List, Optional, Tuple + +import torch +from packaging.version import Version as PkgVersion + +try: + import transformer_engine as te + + if hasattr(te, "__version__"): + TE_VERSION = PkgVersion(str(te.__version__)) + else: + TE_VERSION = PkgVersion(version("transformer-engine")) +except Exception: + TE_VERSION = None + +try: + from transformer_engine.pytorch.tensor import QuantizedTensor + + FP8_TENSOR_CLASS = QuantizedTensor + HAVE_TE_FP8 = True +except Exception: + try: + from transformer_engine.pytorch.float8_tensor import Float8Tensor + + FP8_TENSOR_CLASS = Float8Tensor + HAVE_TE_FP8 = True + except Exception: + FP8_TENSOR_CLASS = None + HAVE_TE_FP8 = False + +try: + from transformer_engine.pytorch.tensor.float8_tensor import Float8Tensor + + HAVE_TE_FLOAT8TENSOR = True +except Exception: + try: + from transformer_engine.pytorch.float8_tensor import Float8Tensor + + HAVE_TE_FLOAT8TENSOR = True + except Exception: + Float8Tensor = None + HAVE_TE_FLOAT8TENSOR = False + +try: + from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Tensor + + HAVE_TE_MXFP8 = True +except Exception: + MXFP8Tensor = None + HAVE_TE_MXFP8 = False + +try: + from transformer_engine.pytorch.tensor.nvfp4_tensor import NVFP4Tensor as NVFP4_TENSOR_CLASS + + HAVE_TE_NVFP4 = True +except Exception: + NVFP4_TENSOR_CLASS = None + HAVE_TE_NVFP4 = False + +try: + from transformer_engine.common.recipe import NVFP4BlockScaling + + HAVE_TE_NVFP4_RECIPE = True +except Exception: + NVFP4BlockScaling = None + HAVE_TE_NVFP4_RECIPE = False + +try: + from transformer_engine.pytorch.tensor.utils import quantize_master_weights + + HAVE_TE_QUANTIZE_MASTER_WEIGHTS = True +except Exception: + HAVE_TE_QUANTIZE_MASTER_WEIGHTS = False + +try: + from transformer_engine.pytorch.tensor.utils import cast_master_weights_to_fp8 + + HAVE_TE_CAST_MASTER_WEIGHTS_TO_FP8 = True +except Exception: + HAVE_TE_CAST_MASTER_WEIGHTS_TO_FP8 = False + +try: + from transformer_engine.pytorch.tensor.utils import post_all_gather_processing + + HAVE_TE_POST_ALL_GATHER_PROCESSING = True +except Exception: + HAVE_TE_POST_ALL_GATHER_PROCESSING = False + +try: + from transformer_engine.pytorch import quantized_model_init + + QUANTIZED_MODEL_INIT_CLASS = quantized_model_init +except Exception: + try: + from transformer_engine.pytorch import fp8_model_init + + QUANTIZED_MODEL_INIT_CLASS = fp8_model_init + except Exception: + QUANTIZED_MODEL_INIT_CLASS = nullcontext + +try: + from transformer_engine.pytorch.module.base import TransformerEngineBaseModule +except Exception: + TransformerEngineBaseModule = None + +try: + from megatron.core.tensor_parallel import get_cuda_rng_tracker +except Exception: + from ..utils import get_cuda_rng_tracker + +from ..utils import _MODEL_PARALLEL_RNG_TRACKER_NAME + +if not HAVE_TE_CAST_MASTER_WEIGHTS_TO_FP8: + try: + from transformer_engine.pytorch.optimizers import multi_tensor_applier, multi_tensor_scale + + multi_tensor_scale_impl = multi_tensor_scale + except ImportError: + try: + import amp_C + from apex.multi_tensor_apply import multi_tensor_applier + + multi_tensor_scale_impl = amp_C.multi_tensor_scale + except ImportError: + + def local_multi_tensor_applier(op, noop_flag_buffer, tensor_lists, *args): + return op(2048 * 32, noop_flag_buffer, tensor_lists, *args) + + def local_multi_tensor_scale(chunk_size, noop_flag, tensor_lists, scale): + for src, dst in zip(tensor_lists[0], tensor_lists[1]): + dst.copy_(src * scale) + + multi_tensor_applier = local_multi_tensor_applier + multi_tensor_scale_impl = local_multi_tensor_scale + + def _multi_tensor_copy_this_to_that( + this: List[torch.Tensor], + that: List[torch.Tensor], + overflow_buf: Optional[torch.Tensor] = None, + ): + if overflow_buf is not None: + overflow_buf.fill_(0) + multi_tensor_applier(multi_tensor_scale_impl, overflow_buf, [this, that], 1.0) + else: + for this_, that_ in zip(this, that): + that_.copy_(this_) + + +@dataclass(frozen=True) +class FullyShardFP8Policy: + """FP8 recipe settings owned by the v2 ``fully_shard`` path.""" + + enabled: bool = False + recipe: Optional[str] = None + keep_transpose_cache: bool = False + + +@dataclass(frozen=True) +class FullyShardNVFP4Policy: + """NVFP4 recipe settings owned by the v2 ``fully_shard`` path.""" + + enabled: bool = False + recipe: Optional[str] = None + + +@dataclass(frozen=True) +class FullyShardMixedPrecisionPolicy: + """Mixed precision policy owned by the v2 ``fully_shard`` path.""" + + main_params_dtype: Optional[torch.dtype] = None + main_grads_dtype: Optional[torch.dtype] = None + grad_comm_dtype: Optional[torch.dtype] = None + fp8_param_gather: bool = False + fp8_recipe: Optional[str] = None + keep_fp8_transpose_cache: bool = False + use_decoupled_grad: bool = False + fp4_param_gather: bool = False + fp4_recipe: Optional[str] = None + fp8: Optional[FullyShardFP8Policy] = field(default=None, repr=False) + nvfp4: Optional[FullyShardNVFP4Policy] = field(default=None, repr=False) + + def __post_init__(self) -> None: + if self.fp8 is None: + fp8 = FullyShardFP8Policy( + enabled=self.fp8_param_gather, + recipe=self.fp8_recipe, + keep_transpose_cache=self.keep_fp8_transpose_cache, + ) + object.__setattr__(self, "fp8", fp8) + else: + object.__setattr__(self, "fp8_param_gather", self.fp8.enabled) + object.__setattr__(self, "fp8_recipe", self.fp8.recipe) + object.__setattr__(self, "keep_fp8_transpose_cache", self.fp8.keep_transpose_cache) + + if self.nvfp4 is None: + nvfp4 = FullyShardNVFP4Policy(enabled=self.fp4_param_gather, recipe=self.fp4_recipe) + object.__setattr__(self, "nvfp4", nvfp4) + else: + object.__setattr__(self, "fp4_param_gather", self.nvfp4.enabled) + object.__setattr__(self, "fp4_recipe", self.nvfp4.recipe) + + @contextmanager + def model_init_context(self, module: Optional[torch.nn.Module] = None): + """Return the model-init context for mixed precision parameter creation.""" + with ExitStack() as stack: + if self.fp8.enabled: + # TE initializes FP8 parameters while preserving a high-precision value + # that can seed the optimizer main-weight buffer. + assert ( + QUANTIZED_MODEL_INIT_CLASS is not nullcontext + ), "Transformer Engine is required for FP8 parameter initialization" + args = {"enabled": True} + if ( + "preserve_high_precision_init_val" + in inspect.signature(QUANTIZED_MODEL_INIT_CLASS).parameters + ): + args["preserve_high_precision_init_val"] = True + stack.enter_context(QUANTIZED_MODEL_INIT_CLASS(**args)) + + if self.nvfp4.enabled: + assert not self.fp8.enabled, "NVFP4 and FP8 are mutually exclusive" + assert ( + QUANTIZED_MODEL_INIT_CLASS is not nullcontext + ), "Transformer Engine is required for NVFP4 parameter initialization" + assert ( + HAVE_TE_NVFP4_RECIPE + ), "NVFP4BlockScaling recipe requires Transformer Engine >= 2.7.0.dev0" + args = {"enabled": True, "recipe": NVFP4BlockScaling()} + if ( + "preserve_high_precision_init_val" + in inspect.signature(QUANTIZED_MODEL_INIT_CLASS).parameters + ): + args["preserve_high_precision_init_val"] = True + stack.enter_context(QUANTIZED_MODEL_INIT_CLASS(**args)) + + if ( + module is not None + and TransformerEngineBaseModule is not None + and TE_VERSION is not None + and TE_VERSION >= PkgVersion("0.9.0") + and not isinstance(module, TransformerEngineBaseModule) + ): + cuda_rng_tracker = get_cuda_rng_tracker() + if _MODEL_PARALLEL_RNG_TRACKER_NAME in cuda_rng_tracker.states_: + stack.enter_context(cuda_rng_tracker.fork()) + + yield + + def group_key_dtype(self, tensor: torch.Tensor): + """Return the parameter grouping dtype key.""" + if self.is_fp8_param(tensor): + return ("quantized", type(tensor).__name__, self.fp8.recipe) + if self.is_nvfp4_param(tensor): + return ("quantized", type(tensor).__name__, self.nvfp4.recipe) + return tensor.dtype + + def validate_param_group(self, params: List[torch.Tensor]) -> None: + """Validate that one ParameterGroup has one mixed-precision storage kind.""" + if len(params) == 0: + return + group_key = self.group_key_dtype(params[0]) + assert all(self.group_key_dtype(param) == group_key for param in params), ( + "Parameters with different mixed-precision storage kinds must not " + "share a ParameterGroup" + ) + + def is_fp8_param(self, tensor: torch.Tensor) -> bool: + """Return whether ``tensor`` is managed as an FP8 parameter.""" + return is_fp8_param(tensor) + + def is_nvfp4_param(self, tensor: torch.Tensor) -> bool: + """Return whether ``tensor`` is managed as an NVFP4 parameter.""" + return is_nvfp4_param(tensor) + + def fine_grained_forward_hooks_required(self, param_groups) -> bool: + """Return whether submodule forward hooks are needed for these parameter groups.""" + for param_group in param_groups: + for param in param_group.params: + if self.is_fp8_param(param): + return True + return False + + @staticmethod + def model_weight_buffer_dtype(tensor: torch.Tensor) -> torch.dtype: + """Return the model-weight buffer dtype for ``tensor``.""" + if is_fp8_param(tensor) or is_nvfp4_param(tensor): + return torch.uint8 + return tensor.dtype + + def get_param_storage_shapes(self, params: List[torch.Tensor]) -> Optional[List[torch.Size]]: + """Return real parameter shapes for ``params``. + + For NVFP4 params the storage is packed (2 values per byte), so the + last dimension is halved. Returns ``None`` when no shape transform + is required, in which case the caller falls back to ``param.shape``. + """ + if not HAVE_TE_NVFP4 or not any(self.is_nvfp4_param(p) for p in params): + return [p.shape for p in params] + shapes = [] + for p in params: + if self.is_nvfp4_param(p): + packed = list(p.shape) + packed[-1] = packed[-1] // 2 + shapes.append(torch.Size(packed)) + else: + shapes.append(p.shape) + return shapes + + def get_param_data(self, tensor: torch.Tensor, transpose: bool = False) -> torch.Tensor: + """Return the parameter data view to store in the model-weight buffer.""" + if self.is_fp8_param(tensor): + return get_fp8_raw_data(tensor, transpose=transpose) + if self.is_nvfp4_param(tensor): + return get_nvfp4_raw_data(tensor) + return tensor.detach() + + def bind_unsharded_param( + self, tensor: torch.Tensor, data: torch.Tensor, buffer_role: str + ) -> None: + """Bind a parameter to an unsharded model-weight buffer view.""" + if self.is_fp8_param(tensor): + if buffer_role == "transpose_weight": + assert self.needs_transpose_weight_buffer( + tensor + ), f"Type {type(tensor)} has no transpose payload" + attr = "_columnwise_data" + else: + attr = "_rowwise_data" if hasattr(tensor, "_rowwise_data") else "_data" + + # Rebind TE's raw uint8 payload to the all-gathered FSDP buffer view. + old_data = getattr(tensor, attr) + if old_data is not None: + assert ( + old_data.dtype == data.dtype + ), f"FP8 raw dtype mismatch: {old_data.dtype} vs {data.dtype}" + assert ( + old_data.shape == data.shape + ), f"FP8 raw shape mismatch: {old_data.shape} vs {data.shape}" + setattr(tensor, attr, data) + return + if self.is_nvfp4_param(tensor): + attr = "_rowwise_data" + old_data = getattr(tensor, attr) + if old_data is not None: + assert ( + old_data.dtype == data.dtype + ), f"NVFP4 raw dtype mismatch: {old_data.dtype} vs {data.dtype}" + assert ( + old_data.shape == data.shape + ), f"NVFP4 raw shape mismatch: {old_data.shape} vs {data.shape}" + setattr(tensor, attr, data) + return + tensor.data = data + + def storage_tensors_to_free( + self, tensor: torch.Tensor, model_weight_buffer, main_weight_buffer + ) -> List[torch.Tensor]: + """Return original parameter storages that FSDP buffers have replaced.""" + # The buffers are ownership signals, not data sources: non-FP8/F4P params can + # be backed by either model or main weight storage, while FP8/NVFP4 params are + # only safe to free after their raw payload is copied to the model buffer. + if model_weight_buffer is None and main_weight_buffer is None: + return [] + + if self.is_nvfp4_param(tensor): + if model_weight_buffer is None: + return [] + return [get_nvfp4_raw_data(tensor)] + + if not self.is_fp8_param(tensor): + return [tensor.data] + + if model_weight_buffer is None: + return [] + + tensors = [get_fp8_raw_data(tensor)] + if self.needs_transpose_weight_buffer(tensor): + tensors.append(get_fp8_raw_data(tensor, transpose=True)) + return tensors + + def weight_buffers_for_unshard( + self, model_weight_buffer, transpose_weight_buffer=None, *, bwd_pass: bool = False + ) -> List: + """Return the weight buffer needed for forward or backward compute.""" + if bwd_pass and transpose_weight_buffer is not None: + return [transpose_weight_buffer] + return [model_weight_buffer] + + def needs_transpose_weight_buffer(self, tensor: torch.Tensor) -> bool: + """Return whether ``tensor`` needs an extra transpose/columnwise buffer.""" + return HAVE_TE_MXFP8 and isinstance(tensor, MXFP8Tensor) + + def main_params_dtype_for_param(self, tensor: torch.Tensor) -> Optional[torch.dtype]: + """Return the main-parameter dtype for a parameter group.""" + if self.is_fp8_param(tensor) and self.main_params_dtype is None: + return torch.float32 + if self.is_nvfp4_param(tensor) and self.main_params_dtype is None: + return torch.float32 + return self.main_params_dtype + + def main_grads_dtype_for_param(self, tensor: torch.Tensor) -> torch.dtype: + """Return the main-gradient dtype for a parameter group. + + Defaults to float32 for stable gradient accumulation (avoids + precision loss in reduce-scatter and optimizer steps). + """ + if self.main_grads_dtype is not None: + return self.main_grads_dtype + return torch.float32 + + def get_high_precision_value(self, tensor: torch.Tensor) -> torch.Tensor: + """Return a high-precision value for initializing optimizer main weights.""" + if not self.is_fp8_param(tensor) and not self.is_nvfp4_param(tensor): + return tensor.detach() + if hasattr(tensor, "get_high_precision_init_val"): + item = tensor.get_high_precision_init_val() + tensor.clear_high_precision_init_val() + return item + + # If TE did not preserve the FP32 init value, recover a main-weight + # value from the quantized tensor before the original full parameter + # is freed. + return tensor.dequantize() + + def post_unshard(self, params: List[torch.Tensor], bwd_pass: bool = False) -> None: + """Run post-unshard mixed precision processing for a parameter group.""" + fp8_params = [param for param in params if self.is_fp8_param(param)] + nvfp4_params = [param for param in params if self.is_nvfp4_param(param)] + + if len(fp8_params) > 0: + if self.needs_transpose_weight_buffer(fp8_params[0]): + # Match v1: forward only rebinds rowwise raw data. Do not mark the + # MXFP8 columnwise payload unavailable since TE may request the + # backward workspace from inside the forward call stack. + if bwd_pass: + if HAVE_TE_POST_ALL_GATHER_PROCESSING: + post_all_gather_processing(fp8_params) + else: + for param in fp8_params: + if hasattr(param, "_create_transpose"): + param._create_transpose() + else: + param._create_columnwise() + return + + # TE rebuilds recipe-specific state after FSDP all-gather for recipes + # where columnwise data is derived from the all-gathered rowwise data. + if HAVE_TE_POST_ALL_GATHER_PROCESSING: + post_all_gather_processing(fp8_params) + else: + for param in fp8_params: + if hasattr(param, "_create_transpose"): + param._create_transpose() + else: + param._create_columnwise() + for param in fp8_params: + if hasattr(param, "update_usage"): + param.update_usage(rowwise_usage=not bwd_pass, columnwise_usage=True) + + if len(nvfp4_params) > 0 and bwd_pass: + # TE rebuilds recipe-specific state after FSDP all-gather for NVFP4. + # Only call post_all_gather_processing if the params have valid + # internal state; skip if _rowwise_data or _rowwise_scale_inv is + # None (e.g. buffer already unsharded from prior pass and state + # hasn't been rebound). + valid_nvfp4 = [ + p + for p in nvfp4_params + if getattr(p, "_rowwise_data", None) is not None + and getattr(p, "_rowwise_scale_inv", None) is not None + ] + if valid_nvfp4: + if HAVE_TE_POST_ALL_GATHER_PROCESSING: + post_all_gather_processing(valid_nvfp4) + + def post_reshard(self, params: List[torch.Tensor]) -> None: + """Run post-reshard mixed precision processing for a parameter group.""" + if self.fp8.keep_transpose_cache: + return + for param in params: + if self.is_nvfp4_param(param): + continue + if not self.is_fp8_param(param): + continue + # Drop full-weight transpose/cache views after FSDP reshard releases + # the unsharded payload. + if hasattr(param, "_transpose_invalid"): + param._transpose_invalid = True + param._transpose = None + elif not self.needs_transpose_weight_buffer(param): + param.update_usage(rowwise_usage=True, columnwise_usage=False) + + def copy_main_weights_to_model_weights( + self, + params: List[torch.Tensor], + param_idx: dict, + data_parallel_group: torch.distributed.ProcessGroup, + model_weight_buffer, + main_weight_buffer, + transpose_weight_buffer=None, + ) -> None: + """Install optimized main weights into model compute weights.""" + if main_weight_buffer is None: + return + + assert model_weight_buffer is not None, "main weights require a model-weight buffer" + + if self.is_nvfp4_param(params[0]): + quantize_main_weights_to_nvfp4( + params, param_idx, data_parallel_group, model_weight_buffer, main_weight_buffer + ) + return + + if not self.is_fp8_param(params[0]): + model_weight_buffer.data.copy_(main_weight_buffer.data) + return + + fp8_params = [] + main_params = [] + start_offsets = [] + model_param_shards = [] + for param in params: + item_id = param_idx[param] + model_shard = model_weight_buffer.get_item(item_id, as_shard=True) + if model_shard.numel() == 0: + fp8_params.append(param) + main_params.append(None) + start_offsets.append(None) + model_param_shards.append((None, None)) + continue + + transpose_shard = None + if transpose_weight_buffer is not None: + transpose_shard = transpose_weight_buffer.get_item(item_id, as_shard=True) + main_weight = main_weight_buffer.get_item(item_id, as_shard=True) + start_offset, _ = model_weight_buffer.buffer_index._get_item_self_range(item_id) + fp8_params.append(param) + main_params.append(main_weight) + start_offsets.append(start_offset) + model_param_shards.append((model_shard, transpose_shard)) + + quantize_main_weights_to_fp8( + fp8_params, main_params, start_offsets, data_parallel_group, model_param_shards + ) + + +def is_fp8_param(tensor: torch.Tensor) -> bool: + """Return True if the parameter is backed by a Transformer Engine FP8 tensor. + + Excludes NVFP4 tensors since they share the QuantizedTensor base class. + """ + return HAVE_TE_FP8 and isinstance(tensor, FP8_TENSOR_CLASS) and not is_nvfp4_param(tensor) + + +def get_fp8_raw_data(tensor: torch.Tensor, transpose: bool = False) -> torch.Tensor: + """Return the raw uint8 storage owned by a Transformer Engine FP8 tensor.""" + assert is_fp8_param(tensor), f"Type {type(tensor)} is not a Transformer Engine FP8 tensor" + if transpose: + assert HAVE_TE_MXFP8 and isinstance( + tensor, MXFP8Tensor + ), f"Type {type(tensor)} has no transpose payload" + return getattr(tensor, "_columnwise_data") + attr = "_rowwise_data" if hasattr(tensor, "_rowwise_data") else "_data" + return getattr(tensor, attr) + + +def quantize_main_weights_to_fp8( + model_params: List[torch.Tensor], + main_params: List[Optional[torch.Tensor]], + start_offsets: List[Optional[int]], + data_parallel_group: torch.distributed.ProcessGroup, + fsdp_shard_model_params: List[Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]], +) -> None: + """Quantize FP32 main-weight shards into FP8/MXFP8 model-weight shards.""" + if len(model_params) == 0: + return + + fsdp_shard_model_params = [x[0] if x[1] is None else x for x in fsdp_shard_model_params] + if HAVE_TE_CAST_MASTER_WEIGHTS_TO_FP8: + args = [ + model_params, + main_params, + start_offsets, + data_parallel_group, + fsdp_shard_model_params, + ] + kwargs = {} + if HAVE_TE_POST_ALL_GATHER_PROCESSING: + kwargs["manual_post_all_gather_processing"] = True + cast_master_weights_to_fp8(*args, **kwargs) + return + + assert HAVE_TE_FLOAT8TENSOR, "Transformer Engine Float8Tensor is required for FP8 fallback" + for model_param, main_param, start_offset, shard_model_param in zip( + model_params, main_params, start_offsets, fsdp_shard_model_params + ): + if main_param is None: + continue + if shard_model_param is None: + shard_model_param = get_fp8_raw_data(model_param).view(-1)[ + start_offset : start_offset + main_param.numel() + ] + + quantizer = model_param._quantizer + main_param = main_param.to(model_param.dtype) + out = Float8Tensor( + shape=main_param.size(), + dtype=model_param.dtype, + requires_grad=False, + data=shard_model_param, + fp8_scale_inv=model_param._scale_inv, + fp8_dtype=model_param._fp8_dtype, + quantizer=quantizer, + ) + quantizer.update_quantized(main_param, out) + + amaxes = [] + scales = [] + scale_invs = [] + for model_param in model_params: + quantizer = model_param._quantizer + amaxes.append(quantizer.amax.view(1)) + scales.append(quantizer.scale.view(1)) + scale_invs.append(model_param._scale_inv.view(1)) + model_param._reset_caches() + + dummy_overflow_buf = torch.tensor([0], dtype=torch.int, device="cuda") + packed_scales = torch.empty(len(scales), dtype=torch.float32, device=scales[0].device) + packed_scale_views = [packed_scales[i].view(1) for i in range(len(scales))] + _multi_tensor_copy_this_to_that(scales, packed_scale_views, dummy_overflow_buf) + torch.reciprocal(packed_scales, out=packed_scales) + _multi_tensor_copy_this_to_that(packed_scale_views, scale_invs, dummy_overflow_buf) + + packed_amaxes = torch.empty(len(amaxes), dtype=torch.float32, device=amaxes[0].device) + packed_amax_views = [packed_amaxes[i].view(1) for i in range(len(amaxes))] + _multi_tensor_copy_this_to_that(amaxes, packed_amax_views, dummy_overflow_buf) + torch.distributed.all_reduce( + packed_amaxes, op=torch.distributed.ReduceOp.MAX, group=data_parallel_group + ) + _multi_tensor_copy_this_to_that(packed_amax_views, amaxes, dummy_overflow_buf) + + +def is_nvfp4_param(tensor: torch.Tensor) -> bool: + """Return True if the parameter is backed by a Transformer Engine NVFP4 tensor.""" + return HAVE_TE_NVFP4 and isinstance(tensor, NVFP4_TENSOR_CLASS) + + +def get_nvfp4_raw_data(tensor: torch.Tensor) -> torch.Tensor: + """Return the raw uint8 storage owned by a Transformer Engine NVFP4 tensor.""" + assert is_nvfp4_param(tensor), f"Type {type(tensor)} is not a Transformer Engine NVFP4 tensor" + return getattr(tensor, "_rowwise_data") + + +def quantize_main_weights_to_nvfp4( + model_params: List[torch.Tensor], + param_idx: dict, + data_parallel_group: torch.distributed.ProcessGroup, + model_weight_buffer, + main_weight_buffer, +) -> None: + """Quantize FP32 main-weight shards into NVFP4 model-weight shards.""" + if not HAVE_TE_QUANTIZE_MASTER_WEIGHTS: + raise RuntimeError("quantize_master_weights requires Transformer Engine >= 2.7.0.dev0") + + if len(model_params) == 0: + return + + te_model_params = [] + te_main_params = [] + te_start_offsets = [] + + wbuf = model_weight_buffer + if not wbuf.is_distributed: + raise RuntimeError("FIXME: implement non-distributed NVFP4 quantization path") + + full_weight_buffer = wbuf.fetch_buffer() + wbuf._bind_buffer_to_params(full_weight_buffer) + + for param in model_params: + item_id = param_idx[param] + main_weight_shard = main_weight_buffer.get_item(item_id, as_shard=True) + if main_weight_shard.numel() == 0: + main_weight_shard = None + + # Compute the start offset in LOGICAL element space using the main + # weight buffer index (full shapes), not the model weight buffer + # index (packed shapes for NVFP4). + # + # WARNING: Do NOT use wbuf.buffer_index._get_item_self_range() here. + # The model weight buffer for NVFP4 uses packed shapes (last dim + # halved, 2 values per uint8 byte). Its self_range returns offsets + # in *packed-byte* space, but TE's quantize_master_weights expects + # offsets in *logical-element* space. Using the wrong offset on + # non-zero DP ranks silently corrupts the model weight buffer because + # TE writes to the wrong byte position. Always derive this offset + # from the main_weight_buffer index, which uses full logical shapes. + shard_offset, _ = main_weight_buffer.buffer_index._get_item_self_range( + item_id, as_shard=True + ) + te_model_params.append(param) + te_main_params.append(main_weight_shard) + te_start_offsets.append(shard_offset) + + kwargs = {} + if HAVE_TE_POST_ALL_GATHER_PROCESSING: + kwargs["manual_post_all_gather_processing"] = True + + quantize_master_weights( + te_model_params, te_main_params, te_start_offsets, data_parallel_group, **kwargs + ) + + wbuf.data.copy_(wbuf.fetch_buffer(as_shard=True)) + + # Don't forget to reshard the model weight buffer after directly writing into its payload + wbuf.reshard() diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/param_group.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/param_group.py new file mode 100644 index 00000000000..2cb0648853e --- /dev/null +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/param_group.py @@ -0,0 +1,357 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Parameter Group for FSDP + +Groups parameters that share the same (device, dtype, requires_grad) and +manages their buffers collectively. This enables efficient memory management +and collective operations across parameters. +""" + +import math +from typing import Dict, List, Optional + +import torch +from torch.distributed.tensor import DeviceMesh +from torch.distributed.tensor.placement_types import Replicate, Shard + +from ..uneven_dtensor import make_uneven_dtensor, update_uneven_dtensor_chunk_metadata +from .allocator import BucketAllocator, TemporaryBucketAllocator, _free_storage +from .dp_buffer import DataParallelBuffer +from .mixed_precision import FullyShardMixedPrecisionPolicy +from .utils import ParamGroupIdx + + +class ParameterGroup: + """ + Groups parameters sharing same properties for collective buffer management. + + All parameters in a group have the same: + - device (cuda device) + - dtype (data type) + - requires_grad (whether gradients are needed) + + The group manages: + - model_weight_buffer: stores sharded model weights + - main_weight_buffer: optional high-precision copy for mixed precision + - main_grad_buffer: accumulates gradients before reduction + - dist_params: DTensor views into the buffer + - dist_grads: DTensor gradient views + """ + + def __init__( + self, + params: List[torch.nn.Parameter], + param_group_id: ParamGroupIdx, + *, + mp_policy: FullyShardMixedPrecisionPolicy, + mesh: Optional[DeviceMesh] = None, + sharding_strategy: str = "optim_grads_params", + gradient_scaling_factor: Optional[float] = None, + allocator: Optional[BucketAllocator] = None, + ): + self.params = params + self.param_idx: Dict[torch.nn.Parameter, int] = {p: i for i, p in enumerate(params)} + + # Assume all params have same device/dtype/require_grad + # TODO: validate all params have same properties + self.device = params[0].device + self.dtype = params[0].dtype + self.requires_grad = params[0].requires_grad + self.mp_policy = mp_policy + self.mp_policy.validate_param_group(params) + + # Setup device mesh and derived process group + if mesh is None: + mesh = DeviceMesh( + self.device.type, + list(range(torch.distributed.get_world_size(torch.distributed.group.WORLD))), + ) + assert mesh.ndim == 1, "Only 1D mesh is supported" + self.mesh = mesh + self.dp_group = mesh.get_group() + + # FIXME: no_shard, optim, and optim_grads sharding strategies are not yet supported in v2. + # Currently only optim_grads_params is fully implemented and tested. + # See README.md § "Sharding Strategies" for details. + # We will add support for these strategies in a follow-up change. + if sharding_strategy not in ("optim_grads_params",): + raise NotImplementedError( + f"Sharding strategy '{sharding_strategy}' is not yet supported in FSDP v2. " + f"Currently only 'optim_grads_params' is implemented. " + f"We will add support for 'no_shard', 'optim', and 'optim_grads' in a follow-up change." + ) + self.sharding_strategy = sharding_strategy + self.param_group_id = param_group_id + + # Compute chunk size factor for alignment + # LCM ensures params align to common boundary for efficient sharding + if len(params) > 0 and any(p.shape[1:].numel() > 0 for p in params): + self.chunk_size_factor = max(1, math.lcm(*[p.shape[1:].numel() for p in params])) + else: + self.chunk_size_factor = 1 + + self.gradient_scaling_factor = gradient_scaling_factor + self.allocator = allocator if allocator is not None else TemporaryBucketAllocator() + + # Buffer references (initialized in _init_buffers) + self.model_weight_buffer: Optional[DataParallelBuffer] = None + self.transpose_weight_buffer: Optional[DataParallelBuffer] = None + self.main_weight_buffer: Optional[DataParallelBuffer] = None + self.main_grad_buffer: Optional[DataParallelBuffer] = None + self.hsdp_wbuf: Optional[DataParallelBuffer] = None + self.hsdp_gbuf: Optional[DataParallelBuffer] = None + self.hsdp_comm_gbuf: Optional[DataParallelBuffer] = None + + # Initialize buffers and distributed parameters + self._init_buffers() + + def set_allocator(self, allocator: BucketAllocator) -> None: + self.allocator = allocator + for buffer in ( + self.model_weight_buffer, + self.transpose_weight_buffer, + self.main_weight_buffer, + self.main_grad_buffer, + ): + if buffer is not None: + buffer.allocator = allocator + + def _create_buffer( + self, dtype: torch.dtype, is_distributed: bool, role: str + ) -> DataParallelBuffer: + """Create a buffer and namespace its temporary bucket by role.""" + return DataParallelBuffer( + params=self.params, + param_idx=self.param_idx, + dtype=dtype, + device=self.device, + dp_group=self.dp_group, + allocator=self.allocator, + buffer_role=role, + is_distributed=is_distributed, + param_group_id=self.param_group_id, + gradient_scaling_factor=self.gradient_scaling_factor, + chunk_size_factor=self.chunk_size_factor, + sharding_strategy=self.sharding_strategy, + mp_policy=self.mp_policy, + ) + + # FIXME: The branching below currently only handles optim_grads_params since + # no_shard, optim, and optim_grads are gated by a NotImplementedError at init. + # When support for those strategies is added, the logic below must be validated. + def _init_buffers(self) -> None: + """ + Initialize all buffers based on sharding strategy. + + Buffer creation logic: + - model_weight_buffer: always created unless "no_shard" + - main_weight_buffer: created if mp_policy.main_params_dtype is specified + - main_grad_buffer: created if requires_grad + """ + s = self.sharding_strategy + shard_weights = s == "optim_grads_params" + shard_main_weights = s != "no_shard" + shard_grads = s in ("optim_grads", "optim_grads_params") + + # Create model weight buffers. The policy owns dtype-sensitive storage + # choices and exposes the tensor view that should be packed. + if s != "no_shard": + model_weight_dtype = self.mp_policy.model_weight_buffer_dtype(self.params[0]) + wbuf = self._create_buffer(model_weight_dtype, shard_weights, "model_weight") + wbuf.init_data(torch.empty(wbuf.data_size, dtype=wbuf.dtype, device=self.device)) + for i, p in enumerate(self.params): + wbuf.set_item(i, self.mp_policy.get_param_data(p)) + self.model_weight_buffer = wbuf + + if self.mp_policy.needs_transpose_weight_buffer(self.params[0]): + tbuf = self._create_buffer(torch.uint8, shard_weights, "transpose_weight") + tbuf.init_data(torch.empty(tbuf.data_size, dtype=tbuf.dtype, device=self.device)) + for i, p in enumerate(self.params): + tbuf.set_item(i, self.mp_policy.get_param_data(p, transpose=True)) + self.transpose_weight_buffer = tbuf + + # Create main weight buffer for mixed precision + main_params_dtype = self.mp_policy.main_params_dtype_for_param(self.params[0]) + if main_params_dtype is not None: + mbuf = self._create_buffer(main_params_dtype, shard_main_weights, "main_weight") + mbuf.init_data(torch.empty(mbuf.data_size, dtype=mbuf.dtype, device=self.device)) + for i, p in enumerate(self.params): + item = self.mp_policy.get_high_precision_value(p) + mbuf.set_item(i, item.detach().to(main_params_dtype)) + self.main_weight_buffer = mbuf + + # Free the original full parameter tensors now that their data has been + # copied into the weight buffers. The module holds DTensor shard views and + # unshard() rebinds .data to the all-gathered buffer, so the original + # storage is never accessed again. + for p in self.params: + # Pass the replacement buffers so the policy can tell whether this + # parameter's original storage has been copied into FSDP-owned storage. + for tensor in self.mp_policy.storage_tensors_to_free( + p, self.model_weight_buffer, self.main_weight_buffer + ): + _free_storage(tensor) + + # Create gradient buffer + if self.requires_grad: + main_grads_dtype = self.mp_policy.main_grads_dtype_for_param(self.params[0]) + gbuf = self._create_buffer(main_grads_dtype, shard_grads, "main_grad") + gbuf.init_data(torch.zeros(gbuf.data_size, dtype=gbuf.dtype, device=self.device)) + self.main_grad_buffer = gbuf + + # Create distributed parameter views + self._init_dist_params() + + def unshard(self, async_op: bool = False, bwd_pass: bool = False): + """ + Unshard model weights by all-gathering from sharded buffer. + + After unshard, parameters point to full unsharded storage. FP8 + parameters rebind their TE raw payload instead of ``param.data``. + """ + work = None + for weight_buffer in self.mp_policy.weight_buffers_for_unshard( + self.model_weight_buffer, self.transpose_weight_buffer, bwd_pass=bwd_pass + ): + if weight_buffer is None: + continue + if weight_buffer.is_distributed and weight_buffer.is_unsharded(): + continue + _, weight_work = weight_buffer.unshard(async_op=async_op) + if work is None: + work = weight_work + + self.mp_policy.post_unshard(self.params, bwd_pass=bwd_pass) + return work + + def has_unsharded_weight_buffers(self, bwd_pass: bool = False) -> bool: + """Return whether all weight buffers needed for this forward/backward phase are unsharded.""" + for weight_buffer in self.mp_policy.weight_buffers_for_unshard( + self.model_weight_buffer, self.transpose_weight_buffer, bwd_pass=bwd_pass + ): + if weight_buffer is None: + continue + if not weight_buffer.is_unsharded(): + return False + return True + + def reshard(self): + """Reshard model weights by releasing unsharded buffer.""" + self.model_weight_buffer.reshard() + if self.transpose_weight_buffer is not None: + self.transpose_weight_buffer.reshard() + self.mp_policy.post_reshard(self.params) + + @torch.no_grad() + def copy_main_weights_to_model_weights(self): + """Install optimized main weights into model compute weights.""" + self.mp_policy.copy_main_weights_to_model_weights( + self.params, + self.param_idx, + self.dp_group, + self.model_weight_buffer, + self.main_weight_buffer, + self.transpose_weight_buffer, + ) + + def reduce_grad(self): + """ + Reduce gradients across DP ranks. + + For distributed buffers: reduce-scatter the full gradient + For non-distributed buffers: all-reduce in-place + """ + self.main_grad_buffer.reduce_grad(grad_comm_dtype=self.mp_policy.grad_comm_dtype) + + def release_grad_buffer(self): + """Release the main gradient buffer to free memory.""" + if self.main_grad_buffer is not None: + # Drop weight.main_grad views that layers.py stores during gradient-accumulation-fusion + # backward. Those views keep _unsharded_buffer alive even after reshard() sets the + # internal reference to None, causing the grad buffer to leak until the next backward. + for param in self.params: + if hasattr(param, 'main_grad'): + del param.main_grad + self.main_grad_buffer.reshard() + + def _init_dist_params(self): + """ + Initialize distributed parameter views (DTensors) into the buffers. + + Creates DTensor views of model weights and gradients based on sharding strategy: + - "optim_grads_params": weights and grads sharded, full ZeRO-3 + - "optim_grads": grads sharded, weights replicated (ZeRO-2) + - "optim": optimizer state sharding only, weights and grads replicated (ZeRO-1) + - "no_shard": replicated, no sharding (DDP-equivalent) + """ + self.dist_params = [] + self.dist_grads = [] + s = self.sharding_strategy + + # Determine placement based on sharding strategy + is_param_shard = s in ("optim", "optim_grads", "optim_grads_params") + placements = [Shard(dim=0)] if is_param_shard else [Replicate()] + + # Create parameter DTensor views + for param in self.params: + if self.main_weight_buffer is not None: + mbuf = self.main_weight_buffer + data = mbuf.get_item(self.param_idx[param], as_shard=is_param_shard) + param_shape = param.shape + elif self.model_weight_buffer is not None: + wbuf = self.model_weight_buffer + data = wbuf.get_item(self.param_idx[param], as_shard=is_param_shard) + param_shape = self.mp_policy.get_param_storage_shapes([param])[0] + else: + data = param.data.detach() + param_shape = param.shape + + dist_param = torch.nn.Parameter( + make_uneven_dtensor( + data, param_shape, self.mesh, placements, post_process_uneven=True + ), + requires_grad=param.requires_grad, + ) + # Mark as FSDP parameter for special handling + setattr(param, "__fsdp_param__", True) + setattr(dist_param, "__fsdp_param__", True) + self.dist_params.append(dist_param) + + # Update dist_param chunk metadata for checkpointing and debugging. + for dist_param in self.dist_params: + update_uneven_dtensor_chunk_metadata(dist_param) + + # Create gradient DTensor views. Some groups, e.g. uint8 FP8 model + # payloads, do not require grads and therefore have no grad buffer. + if self.main_grad_buffer is None: + self.dist_grads = [None for _ in self.params] + return + + is_grad_shard = is_param_shard + for p in self.params: + gbuf = self.main_grad_buffer + grad_data = gbuf.get_item(self.param_idx[p], as_shard=is_grad_shard) + # NOTE: Do not remove the grad_data.numel() > 0 check. + # Empty local grad shards are semantically no-ops, but materializing + # them as DTensor grads can pass zero-numel tensors into fused + # multi-tensor optimizers such as TE FusedAdam. That can break + # updates for neighboring non-empty shards. + if p.requires_grad and grad_data.numel() > 0: + self.dist_grads.append( + make_uneven_dtensor(grad_data, p.shape, self.mesh, placements) + ) + else: + self.dist_grads.append(None) From 7c841b12ed4752fe73bca2ed25771a15bd659a2d Mon Sep 17 00:00:00 2001 From: jianbinc Date: Thu, 28 May 2026 07:39:48 +0800 Subject: [PATCH 02/28] feat(mfsdp-v2): checkpoint infrastructure, NVFP4 support, and FSDP v2 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 --- examples/megatron_fsdp/fsdp_toy.py | 260 ++++ .../distributed_data_parallel_config.py | 4 + megatron/core/distributed/fsdp/checkpoint.py | 1372 +++++++++++++++++ .../distributed/fsdp/mcore_fsdp_adapter.py | 316 +++- .../fsdp/src/megatron_fsdp/__init__.py | 1 + .../distributed_data_parallel_config.py | 10 +- .../fsdp/src/megatron_fsdp/megatron_fsdp.py | 111 +- .../fsdp/src/megatron_fsdp/uneven_dtensor.py | 348 ++++- .../fsdp/src/megatron_fsdp/v2/README.md | 133 ++ .../fsdp/src/megatron_fsdp/v2/__init__.py | 61 + .../fsdp/src/megatron_fsdp/v2/design.md | 780 ++++++++++ .../fsdp/src/megatron_fsdp/v2/dp_buffer.py | 710 +++++++++ .../fsdp/src/megatron_fsdp/v2/fsdp_module.py | 830 ++++++++++ .../fsdp/src/megatron_fsdp/v2/fully_shard.py | 113 ++ .../v2/mcore_fsdp_checkpoint_design.md | 1009 ++++++++++++ .../fsdp/src/megatron_fsdp/v2/nvfp4_design.md | 366 +++++ .../fsdp/src/megatron_fsdp/v2/utils.py | 83 + megatron/core/optimizer/distrib_optimizer.py | 33 +- megatron/core/optimizer/optimizer_config.py | 3 + megatron/training/arguments.py | 10 + megatron/training/checkpointing.py | 39 +- megatron/training/config/training_config.py | 3 + megatron/training/training.py | 8 + .../distributed/megatron_fsdp/v2/conftest.py | 15 + .../megatron_fsdp/v2/test_allocator.py | 99 ++ .../megatron_fsdp/v2/test_fully_shard.py | 762 +++++++++ .../megatron_fsdp/v2/test_mcore_checkpoint.py | 567 +++++++ .../v2/test_mcore_nd_parallel.py | 312 ++++ .../megatron_fsdp/v2/test_param_group.py | 313 ++++ 29 files changed, 8540 insertions(+), 131 deletions(-) create mode 100644 examples/megatron_fsdp/fsdp_toy.py create mode 100644 megatron/core/distributed/fsdp/checkpoint.py create mode 100644 megatron/core/distributed/fsdp/src/megatron_fsdp/v2/README.md create mode 100644 megatron/core/distributed/fsdp/src/megatron_fsdp/v2/__init__.py create mode 100644 megatron/core/distributed/fsdp/src/megatron_fsdp/v2/design.md create mode 100644 megatron/core/distributed/fsdp/src/megatron_fsdp/v2/dp_buffer.py create mode 100644 megatron/core/distributed/fsdp/src/megatron_fsdp/v2/fsdp_module.py create mode 100644 megatron/core/distributed/fsdp/src/megatron_fsdp/v2/fully_shard.py create mode 100644 megatron/core/distributed/fsdp/src/megatron_fsdp/v2/mcore_fsdp_checkpoint_design.md create mode 100644 megatron/core/distributed/fsdp/src/megatron_fsdp/v2/nvfp4_design.md create mode 100644 megatron/core/distributed/fsdp/src/megatron_fsdp/v2/utils.py create mode 100644 tests/unit_tests/distributed/megatron_fsdp/v2/conftest.py create mode 100644 tests/unit_tests/distributed/megatron_fsdp/v2/test_allocator.py create mode 100644 tests/unit_tests/distributed/megatron_fsdp/v2/test_fully_shard.py create mode 100644 tests/unit_tests/distributed/megatron_fsdp/v2/test_mcore_checkpoint.py create mode 100644 tests/unit_tests/distributed/megatron_fsdp/v2/test_mcore_nd_parallel.py create mode 100644 tests/unit_tests/distributed/megatron_fsdp/v2/test_param_group.py diff --git a/examples/megatron_fsdp/fsdp_toy.py b/examples/megatron_fsdp/fsdp_toy.py new file mode 100644 index 00000000000..6e7299f4411 --- /dev/null +++ b/examples/megatron_fsdp/fsdp_toy.py @@ -0,0 +1,260 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import argparse +import os +import sys +from pathlib import Path +from typing import Tuple + +import torch +import torch.distributed as dist +import torch.distributed.checkpoint as dcp +import torch.nn as nn +from torch.distributed.checkpoint.state_dict import get_state_dict, set_state_dict +from torch.distributed.checkpoint.stateful import Stateful +from torch.distributed.device_mesh import init_device_mesh +from torch.distributed.tensor import DTensor + +# ----------------------- +# Model definitions +# ----------------------- + +class ToyBlock(nn.Module): + def __init__(self, dim: int): + super().__init__() + self.linear1 = nn.Linear(dim, dim) + self.linear2 = nn.Linear(dim, dim) + + def forward(self, x): + return self.linear2(torch.relu(self.linear1(x))) + + +class ToyModel(nn.Module): + def __init__(self, dim: int, n_layers: int): + super().__init__() + self.layers = nn.ModuleList(ToyBlock(dim) for _ in range(n_layers)) + self.out = nn.Linear(dim, dim) + + def forward(self, x): + for layer in self.layers: + x = layer(x) + return self.out(x) + + +# ----------------------- +# Distributed init / mesh +# ----------------------- + +def init_distributed() -> torch.distributed.device_mesh.DeviceMesh: + """Initialize process group and device mesh.""" + if not dist.is_initialized(): + dist.init_process_group("nccl") + world_size = dist.get_world_size() + rank = dist.get_rank() + torch.cuda.set_device(rank) + mesh = init_device_mesh("cuda", mesh_shape=(world_size,)) + return mesh + + +def build_fsdp_model( + dim: int, + n_layers: int, + use_megatron_fsdp: bool, +) -> Tuple["FSDPModule", torch.distributed.device_mesh.DeviceMesh]: + if use_megatron_fsdp: + from megatron.core.distributed.fsdp.src.megatron_fsdp.uneven_dtensor import get_state_dict + from megatron.core.distributed.fsdp.src.megatron_fsdp.v2 import FSDPModule, fully_shard + sys.modules["get_state_dict"] = get_state_dict + else: + from torch.distributed.fsdp import FSDPModule, fully_shard + + mesh = init_distributed() + model = ToyModel(dim=dim, n_layers=n_layers).to("cuda") + + # Example: per-layer sharding + for layer in model.layers: + fully_shard(layer, mesh=mesh) + + # Optionally shard the root as well + fully_shard(model, mesh=mesh) + + assert isinstance(model, ToyModel) + assert isinstance(model, FSDPModule) + + p = next(model.parameters()) + assert isinstance(p, DTensor) + return model, mesh + + +# ----------------------- +# Checkpoint helpers +# ----------------------- + +class AppState(Stateful): + """This is a useful wrapper for checkpointing the Application State. Since this object is compliant + with the Stateful protocol, DCP will automatically call state_dict/load_stat_dict as needed in the + dcp.save/load APIs. + + Note: We take advantage of this wrapper to hande calling distributed state dict methods on the model + and optimizer. + """ + + def __init__(self, model, optimizer=None): + self.model = model + self.optimizer = optimizer + + def state_dict(self): + # this line automatically manages FSDP FQN's, as well as sets the default state dict type to FSDP.SHARDED_STATE_DICT + model_state_dict, optimizer_state_dict = get_state_dict(self.model, self.optimizer) + return { + "model": model_state_dict, + "optim": optimizer_state_dict + } + + def load_state_dict(self, state_dict): + # sets our state dicts on the model and optimizer, now that we've loaded + set_state_dict( + self.model, + self.optimizer, + model_state_dict=state_dict["model"], + optim_state_dict=state_dict["optim"] + ) + + +def save_checkpoint( + model: nn.Module, + optimizer: torch.optim.Optimizer, + step: int, + ckpt_dir: str, +) -> None: + rank = dist.get_rank() + Path(ckpt_dir).mkdir(parents=True, exist_ok=True) + + ckpt_step_dir = os.path.join(ckpt_dir, f"step_{step:06d}") + + state = {"app": AppState(model, optimizer), "step": step} + dcp.save(state_dict=state, checkpoint_id=ckpt_step_dir) + + if rank == 0: + print(f"[rank0] Saved checkpoint to {ckpt_dir}") + + +def load_checkpoint_if_available( + model: nn.Module, + optimizer: torch.optim.Optimizer, + ckpt_dir: str, +) -> int: + """ + Load the latest checkpoint if present. + Returns starting step (step+1) for training. + """ + if not os.path.exists(ckpt_dir): + return 0 + + all_ckpts = sorted( + [f for f in os.listdir(ckpt_dir) if f.startswith("step_")] + ) + last_ckpt = os.path.join(ckpt_dir, all_ckpts[-1]) if all_ckpts else None + + if last_ckpt is None: + return 0 + + step = torch.zeros([1]) + state = {"app": AppState(model, optimizer), "step": step} + dcp.load(state_dict=state, checkpoint_id=last_ckpt) + + return int(step.item()) + 1 + + +# ----------------------- +# Training loop +# ----------------------- + +def train( + args: argparse.Namespace, + model: nn.Module, + optimizer: torch.optim.Optimizer, + start_step: int = 0, +) -> None: + rank = dist.get_rank() + world_size = dist.get_world_size() + + model.train() + step = start_step + + for epoch in range(args.epochs): + for _ in range(args.steps_per_epoch): + # Dummy data + x = torch.randn(args.batch_size, args.model_dim, device="cuda") + y = model(x) + loss = y.sum() / (world_size * args.batch_size) + loss.backward() + optimizer.step() + optimizer.zero_grad(set_to_none=True) + + if step % args.log_interval == 0 and rank == 0: + print(f"[rank0] epoch={epoch} step={step} loss={loss.item():.4f}") + + if args.ckpt_dir and step % args.ckpt_interval == 0 and step > 0: + save_checkpoint(model, optimizer, step, args.ckpt_dir) + + step += 1 + + # Final checkpoint + if args.ckpt_dir: + save_checkpoint(model, optimizer, step, args.ckpt_dir) + + +# ----------------------- +# __main__ entry +# ----------------------- + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Toy FSDP2 training example") + parser.add_argument("--model-dim", type=int, default=1024) + parser.add_argument("--n-layers", type=int, default=3) + parser.add_argument("--batch-size", type=int, default=8) + parser.add_argument("--epochs", type=int, default=2) + parser.add_argument("--steps-per-epoch", type=int, default=10) + parser.add_argument("--lr", type=float, default=1e-3) + parser.add_argument("--ckpt-dir", type=str, default="checkpoints") + parser.add_argument("--ckpt-interval", type=int, default=20) + parser.add_argument("--log-interval", type=int, default=5) + parser.add_argument("--use-megatron-fsdp", action="store_true", help="Use Megatron-FSDP instead of PyTorch FSDP2") + return parser.parse_args() + + +def main() -> None: + args = parse_args() + + model, _ = build_fsdp_model( + dim=args.model_dim, + n_layers=args.n_layers, + use_megatron_fsdp=args.use_megatron_fsdp, + ) + optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr) + + start_step = 0 + if args.ckpt_dir: + start_step = load_checkpoint_if_available(model, optimizer, args.ckpt_dir) + + train(args, model, optimizer, start_step=start_step) + + dist.barrier() + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/megatron/core/distributed/distributed_data_parallel_config.py b/megatron/core/distributed/distributed_data_parallel_config.py index 50878e149de..f4dbcddab2c 100644 --- a/megatron/core/distributed/distributed_data_parallel_config.py +++ b/megatron/core/distributed/distributed_data_parallel_config.py @@ -206,6 +206,10 @@ class DistributedDataParallelConfig: main gradients to parameter dtype for `.grad`. """ + use_megatron_fsdp_v2: bool = False + """If true, use the `fully_shard` API for FSDP sharding the model. + """ + def __post_init__(self): import os diff --git a/megatron/core/distributed/fsdp/checkpoint.py b/megatron/core/distributed/fsdp/checkpoint.py new file mode 100644 index 00000000000..0c753a23a38 --- /dev/null +++ b/megatron/core/distributed/fsdp/checkpoint.py @@ -0,0 +1,1372 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Checkpoint save/load helpers for Megatron FSDP v2. + +Provides: + +- ``MegatronFSDPStateful`` — ``Stateful`` wrapper that integrates with + PyTorch DCP, handles uneven DTensor chunk metadata via ``get_state_dict``, + and applies MCore post-processing. +- Post-processing functions for Megatron FSDP v2 state dicts: + ``handle_swiglu_in_state_dict_v2``, ``handle_gdn_in_state_dict_v2``, + ``handle_experts_in_state_dict``, ``handle_fp8_extra_state_case``. +""" + +import logging +import re +from collections.abc import Mapping +from typing import Optional, Tuple + +import torch +import torch.nn as nn +from torch.distributed.checkpoint.state_dict import set_state_dict as _set_state_dict +from torch.distributed.checkpoint.stateful import Stateful +from torch.distributed.tensor import DTensor, Shard + +import megatron.core.parallel_state as mpu +from megatron.core import dist_checkpointing +from megatron.core.distributed.fsdp.src.megatron_fsdp.uneven_dtensor import copy_chunk_metadata +from megatron.core.distributed.fsdp.src.megatron_fsdp.uneven_dtensor import ( + get_state_dict as _get_state_dict, +) +from megatron.core.distributed.fsdp.src.megatron_fsdp.uneven_dtensor import ( + make_uneven_dtensor, + preprocess_state_dict_for_uneven_dtensor, + redistribute_uneven_dtensor_to_replicated, + split_dtensor, +) +from megatron.core.transformer.transformer_layer import TransformerLayer + +logger = logging.getLogger(__name__) + +__all__ = [ + "MegatronFSDPStateful", + "add_module_prefix", + "strip_module_prefix", + "get_model_state_dict", + "get_optimizer_state_dict", + "_build_dtensor_optim_sd", + "handle_fp8_extra_state_case", + "handle_experts_in_state_dict", + "handle_swiglu_in_state_dict_v2", + "handle_gdn_in_state_dict_v2", +] + +_MODULE_PREFIX = "module." + + +# ------------------------------------------------------------------ +# Stateful wrapper +# ------------------------------------------------------------------ + + +class MegatronFSDPStateful(Stateful): + """Stateful wrapper for Megatron FSDP v2 model + optimizer checkpointing. + + Implements DCP's ``Stateful`` protocol. ``state_dict()`` uses + ``get_state_dict`` from ``uneven_dtensor`` (which preprocesses DTensors + for uneven sharding and produces FQN-keyed optimizer states), then applies + MCore post-processing (SwiGLU, GDN, FP8, expert remapping — see individual + ``handle_*_v2`` functions). ``load_state_dict()`` uses PyTorch's + ``set_state_dict``. + + Parameters: + model: FSDP v2 wrapped model. + optimizer: Optional optimizer to checkpoint alongside the model. + args: Optional MCore args namespace for post-processing config + (``swiglu``, ``num_experts``, ``gdn``). + """ + + def __init__( + self, + model: nn.Module, + optimizer: Optional[torch.optim.Optimizer] = None, + args: Optional[object] = None, + ): + self.model = model + self.optimizer = optimizer + self.args = args + + def state_dict(self): + if self.optimizer is not None: + model_sd, optim_sd = _get_state_dict(self.model, self.optimizer) + else: + model_sd = self.model.state_dict() + preprocess_state_dict_for_uneven_dtensor(model_sd) + optim_sd = None + + state_dict = {"model": model_sd} + if optim_sd is not None: + state_dict["optimizer"] = optim_sd + + if self.args is not None: + state_dict = _apply_mcore_postprocess(state_dict, self.args, self.model) + return state_dict + + def load_state_dict(self, state_dict): + _set_state_dict( + self.model, + self.optimizer, + model_state_dict=state_dict.get("model"), + optim_state_dict=state_dict.get("optimizer"), + ) + + +# ------------------------------------------------------------------ +# FP8 artifact cleanup (format-agnostic) +# ------------------------------------------------------------------ + + +def handle_fp8_extra_state_case(model_state_dict: dict) -> None: + """Remove ``._extra_state`` keys (artifact of FP8 training).""" + keys_to_remove = [k for k in model_state_dict if k.endswith("._extra_state")] + for k in keys_to_remove: + del model_state_dict[k] + + +# ------------------------------------------------------------------ +# Expert key remapping (format-agnostic) +# ------------------------------------------------------------------ + + +def handle_experts_in_state_dict(model_state_dict: dict, num_experts: int) -> dict: + """Rename expert keys to reflect expert-parallel sharding. + + Standalone implementation (no dependency on ``fsdp_dtensor_checkpoint``). + """ + ep_size = mpu.get_expert_model_parallel_world_size() + ep_rank = mpu.get_expert_model_parallel_rank() + local_expert_start = ep_rank * (num_experts // ep_size) if num_experts else 0 + local_expert_end = local_expert_start + (num_experts // ep_size) if num_experts else 0 + + def _should_keep(expert_index): + if expert_index is None: + return True + return local_expert_start <= expert_index < local_expert_end + + def _replace(key, expert_index, sd): + new_idx = expert_index + local_expert_start + # GroupedMLP: 'mlp.experts.linear_fc1.weight0', 'mlp.experts.linear_fc2.weight0' + if 'mlp.experts.linear_fc1.weight' in key or 'mlp.experts.linear_fc2.weight' in key: + if key.endswith('_w') or key.endswith('_v'): + new_key = key.replace( + f'weight{expert_index}{key[-2:]}', f'weight{new_idx}{key[-2:]}' + ) + else: + new_key = key.replace(f'weight{expert_index}', f'weight{new_idx}') + # SequentialMLP: index is between 'local_experts.' and next '.' + elif 'mlp.experts.local_experts' in key: + new_key = key.replace(f'local_experts.{expert_index}.', f'local_experts.{new_idx}.') + else: + raise ValueError(f"Unexpected expert key format: {key}") + sd[new_key] = sd.pop(key) + + model_state_dict = model_state_dict.copy() + for key in list(model_state_dict.keys()): + ei = _get_expert_index_from_key(key) + if not _should_keep(ei): + _replace(key, ei, model_state_dict) + + return model_state_dict + + +def _get_expert_index_from_key(key: str): + """Extract expert index from key (``weight{N}`` or ``local_experts.N.``).""" + m = re.search(r'(?:weight(\d+)$)|(?:local_experts\.(\d+)\.)', key) + if m: + return int(m.group(1) or m.group(2)) + return None + + +# ------------------------------------------------------------------ +# Shared helpers +# ------------------------------------------------------------------ + + +def _strip_wrappers(path: str) -> str: + """Strip DDP/FSDP wrapper prefixes (``module.``, ``model.``) from a path.""" + parts = path.split('.') + while parts and parts[0] in ('module', 'model'): + parts = parts[1:] + return '.'.join(parts) + + +def _get_dist_param(model: nn.Module, key: str) -> nn.Parameter: + """Get a model parameter by state dict key, handling ``module.`` prefix. + + Tries multiple key variants to handle both wrapped (``FullyShardedDataParallel`` + with ``self.module``) and unwrapped (raw ``FSDPModule``) models: + 1. Key as-is. + 2. Key with ``module.`` prefix added. + 3. Key with ``module.`` prefix stripped (if present). + """ + try: + return model.get_parameter(key) + except AttributeError: + pass + try: + return model.get_parameter(f"module.{key}") + except AttributeError: + pass + if key.startswith("module."): + try: + return model.get_parameter(key[len("module.") :]) + except AttributeError: + pass + raise AttributeError( + f"Parameter '{key}' not found in model " + f"(tried as-is, with 'module.' prefix, and without 'module.' prefix)" + ) + + +def _detect_glu_layers(model: nn.Module) -> dict: + """Return ``{layer_path: uses_gated_linear_unit}`` for TransformerLayers.""" + _layer_glu = {} + for name, module in model.named_modules(): + if isinstance(module, TransformerLayer): + _layer_glu[_strip_wrappers(name)] = getattr(module.config, 'gated_linear_unit', False) + return _layer_glu + + +def _key_in_glu_layer(key: str, layer_glu: dict) -> bool: + """Return True if *key* belongs to a TransformerLayer with GLU enabled.""" + norm_key = _strip_wrappers(key) + best_glu, best_len = None, -1 + for layer_path, uses_glu in layer_glu.items(): + if norm_key.startswith(layer_path + '.') and len(layer_path) > best_len: + best_glu, best_len = uses_glu, len(layer_path) + if best_glu is None: + return True # no TransformerLayer found — assume GLU for backward compat + return best_glu + + +# ------------------------------------------------------------------ +# Fused-parameter splitting — shared skeleton (Megatron FSDP v2) +# ------------------------------------------------------------------ +# +# SwiGLU and GDN both require splitting fused parameter tensors into +# per-component pieces (SwiGLU: fc1 → _w/_v; GDN: linear_qkv → query/key/value). +# This function handles both. +# +# The caller provides a *detector* that returns ``(sizes, names, dim)`` for +# parameters that need splitting, or ``None`` otherwise, and a *key formatter* +# to build sub-parameter names. + +_DEFAULT_STATE_KEYS = ("exp_avg", "exp_avg_sq") + + +def _split_fused_params_v2( + model: nn.Module, + model_state_dict: dict, + optimizer_state_dict: Optional[dict], + detector, + key_fmt, + tag: str, + state_keys: Tuple[str, ...] = _DEFAULT_STATE_KEYS, +) -> Tuple[dict, Optional[dict]]: + """Split fused DTensor parameters in model and optimizer state dicts. + + Args: + model: The model (for ``_get_dist_param``). + model_state_dict: Model state dict (a copy is returned). + optimizer_state_dict: Optional optimizer state dict (a copy is returned). + detector(key, dtensor, model) -> ``(sizes, names, dim)`` or ``None``. + key_fmt(key, sub_name) -> new_key. + tag: Log tag (e.g. ``"SwiGLU v2"``). + state_keys: Optimizer state keys to split (default: ``exp_avg``, ``exp_avg_sq``). + + Returns: + ``(model_state_dict, optimizer_state_dict)`` — modified copies. + """ + # ---- Model state dict ---- + model_state_dict = model_state_dict.copy() + split_count = 0 + for key in list(model_state_dict.keys()): + value = model_state_dict[key] + match = detector(key, value, model) + if match is None: + continue + sizes, names, dim = match + + dist_param = _get_dist_param(model, key) + assert isinstance( + dist_param, DTensor + ), f"Expected DTensor for {key}, got {type(dist_param).__name__}" + sub_tensors = split_dtensor(value, sizes, dim) + for sub_name, tensor in zip(names, sub_tensors): + model_state_dict[key_fmt(key, sub_name)] = tensor + del model_state_dict[key] + split_count += 1 + + if split_count > 0: + logger.info(f"[{tag}] Split {split_count} fused keys.") + + # ---- Optimizer state dict ---- + if optimizer_state_dict is not None: + optimizer_state_dict = optimizer_state_dict.copy() + if len(optimizer_state_dict.get("state", {})) != 0: + opt_state = optimizer_state_dict["state"] + new_opt_state = {} + for key in list(opt_state.keys()): + param_key = key + if param_key.startswith("module."): + param_key = param_key[len("module.") :] + dist_param = _get_dist_param(model, param_key) + + match = detector(key, dist_param, model) + if match is None: + new_opt_state[key] = opt_state[key] + continue + sizes, names, dim = match + + for sub_name in names: + new_opt_state[key_fmt(key, sub_name)] = opt_state[key].copy() + for sk in state_keys: + sub_tensors = split_dtensor(opt_state[key][sk], sizes, dim) + for sub_name, tensor in zip(names, sub_tensors): + new_opt_state[key_fmt(key, sub_name)][sk] = tensor + optimizer_state_dict["state"] = new_opt_state + + return model_state_dict, optimizer_state_dict + + +# ------------------------------------------------------------------ +# SwiGLU key patterns and helpers +# ------------------------------------------------------------------ + +_SWIGLU_KEY_PATTERNS = [ + r"(.*)\.mlp\.linear_fc1\.weight$", + r"(.*)\.mlp\.linear_fc1\.bias$", + r"(.*)\.mlp\.experts\.linear_fc1\.weight(\d+)$", + r"(.*)\.mlp\.experts\.linear_fc1\.bias(\d+)$", + r"(.*)\.mlp\.experts\.local_experts\.(\d+)\.linear_fc1\.weight$", + r"(.*)\.mlp\.experts\.local_experts\.(\d+)\.linear_fc1\.bias$", + r"(.*)\.mlp\.shared_experts\.linear_fc1\.weight$", + r"(.*)\.mlp\.shared_experts\.linear_fc1\.bias$", +] + + +def _is_swiglu_key(key: str) -> bool: + return any(re.search(pat, key) for pat in _SWIGLU_KEY_PATTERNS) + + +def _swiglu_detector(key, dtensor, model, layer_glu): + if not _is_swiglu_key(key): + return None + if not _key_in_glu_layer(key, layer_glu): + return None + dim = 0 + assert ( + dtensor.shape[dim] % 2 == 0 + ), f"Expected SwiGLU fc1 weight size divisible by 2, got {dtensor.shape[dim]}" + half = dtensor.shape[dim] // 2 + return ([half, half], ["w", "v"], dim) + + +def handle_swiglu_in_state_dict_v2( + model: nn.Module, model_state_dict: dict, optimizer_state_dict: Optional[dict] +) -> Tuple[dict, Optional[dict]]: + """Split SwiGLU fc1 parameters in model and optimizer state dicts. + + Megatron FSDP v2 version — only processes layers with + ``gated_linear_unit=True``. Delegates to :func:`_split_fused_params_v2`. + """ + layer_glu = _detect_glu_layers(model) + + def detector(key, dtensor, _model): + return _swiglu_detector(key, dtensor, _model, layer_glu) + + return _split_fused_params_v2( + model, + model_state_dict, + optimizer_state_dict, + detector, + lambda k, s: f"{k}_{s}", + "SwiGLU v2", + ) + + +# ------------------------------------------------------------------ +# GDN key patterns and helpers +# ------------------------------------------------------------------ + +GDN_CONV1D_NAMES = ["query", "key", "value"] + +_GDN_KEY_PATTERNS = [ + r"(.*)\.self_attention\.linear_proj\.weight$", + r"(.*)\.self_attention\.linear_qkv\.weight$", + r"(.*)\.self_attention\.linear_qkv\.bias$", +] + + +def _match_gdn_key(key: str, dtensor: DTensor): + for pat in _GDN_KEY_PATTERNS: + m = re.match(pat, key) + if m: + dim = 0 + size = dtensor[dim] + assert ( + size % 3 == 0 + ), f"Expected GDN projection size divisible by 3, got {size} for key {key}" + qkv_size = size // 3 + return [qkv_size, qkv_size, qkv_size], GDN_CONV1D_NAMES, dim + return None + + +def handle_gdn_in_state_dict_v2( + model: nn.Module, model_state_dict: dict, optimizer_state_dict: Optional[dict] +) -> Tuple[dict, Optional[dict]]: + """Split fused GDN projection parameters into per-component DTensors. + + Megatron FSDP v2 version. Delegates to :func:`_split_fused_params_v2`. + """ + return _split_fused_params_v2( + model, + model_state_dict, + optimizer_state_dict, + _match_gdn_key, + lambda k, s: f"{k}.{s}", + "GDN v2", + ) + + +# ------------------------------------------------------------------ +# Unified post-processing +# ------------------------------------------------------------------ + + +def _find_param_in_map(key: str, param_map: dict) -> Optional[DTensor]: + """Look up *key* in *param_map*, trying ``module.`` prefix variants.""" + param = param_map.get(key) + if param is not None: + return param + stripped = key + while stripped.startswith(_MODULE_PREFIX): + stripped = stripped[len(_MODULE_PREFIX) :] + param = param_map.get(stripped) + if param is not None: + return param + return param_map.get(f"{_MODULE_PREFIX}{key}") + + +def _propagate_chunk_metadata_to_state_dict(model: nn.Module, state_dict: dict) -> None: + """Copy chunk metadata from model parameters to state dict DTensors. + + ``model.state_dict()`` returns fresh DTensor objects that lack + ``__create_chunk_list__`` / ``__create_write_items__``. The model + parameters (from ``named_parameters()``) already have them. This + function copies the closures locally — no collectives. + """ + param_map = {} + for name, param in model.named_parameters(): + if isinstance(param, DTensor) and hasattr(param._local_tensor, "__create_chunk_list__"): + param_map[name] = param + + for key, value in state_dict.items(): + if not isinstance(value, DTensor): + continue + param = _find_param_in_map(key, param_map) + if param is not None: + copy_chunk_metadata(param, value) + + +def _apply_mcore_postprocess(raw_state_dict, args, model): + """Apply MCore-specific state dict post-processing. + + Copies *raw_state_dict*, wraps optimizer states as DTensors, then + applies FP8 cleanup, SwiGLU/GDN split, and expert key remapping. + The original *raw_state_dict* is not mutated. + """ + state_dict = raw_state_dict.copy() + _propagate_chunk_metadata_to_state_dict(model, state_dict["model"]) + + if "optimizer" in state_dict: + state_dict["optimizer"] = _build_dtensor_optim_sd(state_dict["optimizer"], model) + handle_fp8_extra_state_case(state_dict["model"]) + + if getattr(args, "swiglu", False): + opt_sd = state_dict.get("optimizer") + model_sd, new_opt = handle_swiglu_in_state_dict_v2(model, state_dict["model"], opt_sd) + state_dict["model"] = model_sd + if new_opt is not None: + state_dict["optimizer"] = new_opt + + if getattr(args, "gdn", False): + opt_sd = state_dict.get("optimizer") + model_sd, new_opt = handle_gdn_in_state_dict_v2(model, state_dict["model"], opt_sd) + state_dict["model"] = model_sd + if new_opt is not None: + state_dict["optimizer"] = new_opt + + num_experts = getattr(args, "num_experts", None) + if num_experts: + state_dict["model"] = handle_experts_in_state_dict(state_dict["model"], num_experts) + if "optimizer" in state_dict: + optim_sd = state_dict["optimizer"] + optim_sd["state"] = handle_experts_in_state_dict(optim_sd["state"], num_experts) + optim_sd["param_to_group_meta"] = handle_experts_in_state_dict( + optim_sd["param_to_group_meta"], num_experts + ) + + flattened_sd = flatten_state_dict_keys(state_dict) + for key, value in flattened_sd.items(): + if isinstance(value, DTensor): + assert hasattr(value._local_tensor, "__create_chunk_list__"), ( + f"DTensor for key '{key}' is missing chunk metadata. This may cause issues" + f" with checkpointing. Ensure that the model parameters have chunk metadata" + f" and that it is properly propagated." + ) + + return state_dict + + +def flatten_state_dict_keys(state_dict, parent_key="", sep="."): + """ + Recursively flatten nested mappings inside a state_dict. + + Returns a dict: { "flat.key": value } + """ + items = {} + for k, v in state_dict.items(): + new_key = parent_key + sep + k if parent_key else k + if isinstance(v, Mapping): + # Recurse into nested dicts (e.g., when you saved extra metadata) + items.update(flatten_state_dict_keys(v, new_key, sep=sep)) + else: + items[new_key] = v + return items + + +# ------------------------------------------------------------------ +# Model state dict prefix alignment +# ------------------------------------------------------------------ + + +def add_module_prefix(state_dict: dict) -> dict: + """Add ``module.`` prefix to all keys in a state dict. + + Megatron FSDP v2 applies ``fully_shard`` directly to the model without + a ``MegatronFSDP`` wrapper, so ``model.state_dict()`` produces keys + without the ``module.`` prefix (e.g., ``embedding.word_embeddings.weight``). + Megatron's checkpoint format expects the prefix to be present. This + function adds it for alignment. + """ + return {f"{_MODULE_PREFIX}{k}": v for k, v in state_dict.items()} + + +def strip_module_prefix(state_dict: dict) -> dict: + """Remove ``module.`` prefix from all keys in a state dict. + + Inverse of :func:`add_module_prefix`. Used when loading a Megatron + checkpoint (which stores keys with ``module.`` prefix) back into a + Megatron FSDP v2 model that does not use the prefix. + """ + return { + k[len(_MODULE_PREFIX) :] if k.startswith(_MODULE_PREFIX) else k: v + for k, v in state_dict.items() + } + + +def _model_has_module_prefix(model: nn.Module) -> bool: + """Detect whether the model's state dict keys already carry ``module.`` prefix.""" + for name, _ in model.named_parameters(): + return name.startswith(_MODULE_PREFIX) + return False + + +# ------------------------------------------------------------------ +# Model state dict (Megatron FSDP v2) +# ------------------------------------------------------------------ + + +def get_model_state_dict(model: nn.Module) -> dict: + """Get model state dict with ``module.`` prefix for Megatron FSDP v2. + + If the model's parameters already carry the ``module.`` prefix (e.g., + when the model is accessed through the ``FullyShardedDataParallel`` + adapter), the state dict is returned as-is. Otherwise the prefix is + added to align with Megatron's checkpoint key convention. + + Returns: + Model state dict with ``module.``-prefixed keys containing DTensors. + """ + model_sd = model.state_dict() + if not _model_has_module_prefix(model): + model_sd = add_module_prefix(model_sd) + return model_sd + + +# ------------------------------------------------------------------ +# Optimizer state dict — Path A (sharded_param_state_fsdp_dtensor) +# ------------------------------------------------------------------ + + +def get_optimizer_state_dict(optimizer, is_loading: bool = False) -> Optional[dict]: + """Get optimizer state dict following Path A (``sharded_param_state_fsdp_dtensor``). + + Delegates to ``optimizer.sharded_state_dict()`` with the ``fsdp_dtensor`` + sharding type, which internally calls ``sharded_param_state_fsdp_dtensor``. + + For loading, this also initializes optimizer states with dummy values so + that the returned dict has correctly-sized tensors for DCP in-place load. + + Args: + optimizer: A ``DistributedOptimizer`` (or compatible) instance. + is_loading: If True, pre-allocates optimizer state tensors. + + Returns: + Optimizer state dict with structure:: + + { + "state": {: {"exp_avg": DTensor, "exp_avg_sq": DTensor, ...}}, + "param_to_group_meta": {: {"lr": ..., "weight_decay": ...}}, + } + + Returns None if ``optimizer`` is None or is a stub optimizer. + """ + if optimizer is None: + return None + if getattr(optimizer, "is_stub_optimizer", False): + return None + return optimizer.sharded_state_dict( + model_sharded_state_dict={}, + is_loading=is_loading, + metadata={"distrib_optim_sharding_type": "fsdp_dtensor"}, + ) + + +# ------------------------------------------------------------------ +# Optimizer state DTensor wrapping +# ------------------------------------------------------------------ + + +def _maybe_wrap_as_uneven_dtensor(tensor, dist_param: DTensor): + """Wrap a plain tensor as an uneven DTensor if it matches the param's local shard. + + FusedAdam stores ``exp_avg`` / ``exp_avg_sq`` as plain tensors; DCP needs + DTensors. Returns the original tensor unchanged if it is already a DTensor + or its shape does not match the parameter's local shard. + """ + if isinstance(tensor, DTensor): + dt = tensor + else: + if isinstance(tensor, DTensor) or not isinstance(tensor, torch.Tensor): + return tensor + if tensor.shape != dist_param._local_tensor.shape: + return tensor + dt = make_uneven_dtensor( + tensor, + shape=dist_param.size(), + dp_mesh=dist_param.device_mesh, + placements=dist_param.placements, + ) + copy_chunk_metadata(dist_param, dt) + return dt + + +def _build_dtensor_optim_sd(raw_opt_state_dict: dict, model: nn.Module) -> dict: + """Return a copy of *raw_opt_state_dict* with optimizer state tensors converted to DTensors. + + FusedAdam stores ``exp_avg`` / ``exp_avg_sq`` as plain tensors matching + the parameter's local DTensor shard. DCP requires all sharded data to + be DTensors. This returns a new dict where those plain tensors are + converted to uneven DTensors using the model parameter's mesh/placements. + *raw_opt_state_dict* is **not** mutated. + """ + opt_state_dict = raw_opt_state_dict.copy() + if "state" not in opt_state_dict: + return opt_state_dict + + param_map = {} + for name, param in model.named_parameters(): + if isinstance(param, DTensor): + param_map[name] = param + + if not param_map: + return opt_state_dict + + old_state = opt_state_dict["state"] + new_state = {} + for key, param_states in old_state.items(): + if not isinstance(param_states, dict): + new_state[key] = param_states + continue + new_param_states = dict(param_states) + dist_param = _find_param_in_map(key, param_map) + if dist_param is not None: + for state_key, state_val in new_param_states.items(): + new_param_states[state_key] = _maybe_wrap_as_uneven_dtensor(state_val, dist_param) + new_state[key] = new_param_states + + opt_state_dict["state"] = new_state + return opt_state_dict + + +def _preprocess_and_verify_v2_state_dict(v2_state_dict): + """Preprocess and verify the Megatron FSDP v2 state dict before DCP loading. + + DCP requires DTensors for distributed load, but ``optimizer.load_state_dict`` + only accepts plain tensors. This function builds a *shadow* optimizer state + dict (``v2_optim_state``) where plain tensors are wrapped as uneven DTensors + sharing storage with the originals. DCP loads into this shadow dict, and the + data lands in the original plain tensors via shared storage. The original + ``v2_state_dict`` is **not** mutated. + + Also verifies that every model and shadow optimizer DTensor has + ``__create_chunk_list__`` and ``__create_write_items__`` metadata. + + Returns: + v2_by_canonical: ``{canonical_name: model_DTensor}`` mapping. + v2_optim_state: ``{canonical_name: {state_key: DTensor}}`` shadow dict. + """ + v2_model = v2_state_dict["model"] + v2_optim_state_raw = v2_state_dict.get("optimizer", {}).get("state", {}) + + # ---- Propagate chunk metadata from model to state dict DTensors ---- + model = v2_state_dict.get("_model") + if model is not None: + if isinstance(model, (list, tuple)): + model = model[0] + _propagate_chunk_metadata_to_state_dict(model, v2_model) + + # ---- Strip ``module.`` prefix from optimizer state keys ---- + # Copy inner dicts so DTensor wrapping below does not mutate + # the original state dict (optimizer.load_state_dict needs plain tensors). + v2_optim_state = {} + for k, v in v2_optim_state_raw.items(): + canonical = k + while canonical.startswith("module."): + canonical = canonical[len("module.") :] + v2_optim_state[canonical] = dict(v) + + # ---- Build canonical model param map + param lookup ---- + v2_by_canonical = {} + param_map = {} + for v2_key, v2_val in v2_model.items(): + canonical = v2_key + while canonical.startswith("module."): + canonical = canonical[len("module.") :] + v2_by_canonical[canonical] = v2_val + if isinstance(v2_val, DTensor): + param_map[canonical] = v2_val + + # ---- Wrap plain tensors as uneven DTensors ---- + for param_name, states in dict(v2_optim_state).items(): + dist_param = param_map.get(param_name) + if dist_param is None: + continue + for sk, sv in dict(states).items(): + v2_optim_state[param_name][sk] = _maybe_wrap_as_uneven_dtensor(sv, dist_param) + + # ---- Verify model DTensors have chunk metadata ---- + for canonical, v2_val in v2_by_canonical.items(): + if hasattr(v2_val, "_local_tensor"): + lt = v2_val._local_tensor + assert hasattr(lt, "__create_chunk_list__"), ( + f"Expected v2 model DTensor '{canonical}' to have " + f"__create_chunk_list__ metadata" + ) + assert hasattr(lt, "__create_write_items__"), ( + f"Expected v2 model DTensor '{canonical}' to have " + f"__create_write_items__ metadata" + ) + + # ---- Verify optimizer state DTensors have chunk metadata ---- + # Also propagate metadata for states that are already DTensors + # (_maybe_wrap_as_uneven_dtensor only copies metadata when wrapping + # plain tensors, not for pre-existing DTensors). + for param_name, states in v2_optim_state.items(): + dist_param = param_map.get(param_name) + for sk, sv in states.items(): + if ( + isinstance(sv, DTensor) + and dist_param is not None + and hasattr(dist_param._local_tensor, "__create_chunk_list__") + ): + copy_chunk_metadata(dist_param, sv) + if hasattr(sv, "_local_tensor"): + lt = sv._local_tensor + assert hasattr(lt, "__create_chunk_list__"), ( + f"Expected optimizer state DTensor '{param_name}.{sk}' " + f"to have __create_chunk_list__ metadata" + ) + assert hasattr(lt, "__create_write_items__"), ( + f"Expected optimizer state DTensor '{param_name}.{sk}' " + f"to have __create_write_items__ metadata" + ) + + return v2_by_canonical, v2_optim_state + + +# ------------------------------------------------------------------ +# Torch_dist → FSDP v2 checkpoint key normalization +# ------------------------------------------------------------------ + + +def normalize_torch_dist_key(key: str) -> str: + """Normalize a torch_dist checkpoint key to v2 canonical form. + + Maps structural naming differences: + - ``transformer_layer`` → ``mtp_model_layer`` + """ + if ".transformer_layer." in key: + key = key.replace(".transformer_layer.", ".mtp_model_layer.") + return key + + +def reverse_normalize_torch_dist_key(key: str) -> str: + """Reverse the v2 canonical key back to torch_dist naming. + + Inverse of :func:`normalize_torch_dist_key`: + - ``mtp_model_layer`` → ``transformer_layer`` + """ + key = key.replace(".mtp_model_layer", ".transformer_layer") + return key + + +# ------------------------------------------------------------------ +# Torch_dist → FSDP v2 name mapping +# ------------------------------------------------------------------ + + +# Regex patterns for detecting fused vs per-layer/per-expert keys. +# Torch_dist fuses across layers: .layers.{field} (no layer index) +# v2 has per-layer: .layers.{N}.{field} +# Torch_dist fuses experts: .layers.{field}.mlp.experts.experts.{fc}.weight +# v2 has per-expert: .layers.{N}.{field}.mlp.experts.{fc}.weight{M} +_LAYERED_KEY_RE = re.compile(r'^(.+\.)?layers\.(\d+)\.(.+)$') +_EXPERT_V2_KEY_RE = re.compile(r'(.+\.)?mlp\.experts\.linear_fc([12])\.weight(\d+)$') +_SEQ_EXPERT_FIELD_RE = re.compile( + r'(?:.+\.)?mlp\.experts\.local_experts\.(\d+)\.(linear_fc[12])\.weight$' +) + + +def _strip_module_prefix(key: str) -> str: + """Strip leading 'module.' from a key string.""" + return key[len("module.") :] if key.startswith("module.") else key + + +def _build_torch_dist_to_v2_map(metadata, v2_by_canonical, v2_optim_state): + """Build maps from torch_dist checkpoint keys to Megatron FSDP v2 DTensors. + + Automatically handles fused layer tensors by matching + ``decoder.layers.{N}.{field}`` (v2) against + ``decoder.layers.{field}`` (torch_dist) — no hardcoded field paths. + + Returns: + regular_model: ``{td_key: v2_DTensor}`` for 1:1 params. + fused_layer_groups: ``{td_key: [(v2_key, v2_val, layer_idx), ...]}`` + for fused layer params to be loaded/sliced in Phase 4. + hi_prec_model: ``{td_param_name: v2_DTensor}`` for high-precision copies. + optim_keys: ``{td_key: v2_state_DTensor}`` for 1:1 optimizer states. + optim_matched: ``set`` of canonical param names with matched optimizer states. + """ + regular_model = {} + fused_layer_groups = {} + hi_prec_td_keys = set() + hi_prec_model = {} + optim_keys = {} + optim_matched = set() + + # ------------------------------------------------------------------ + # Helper: given a v2 canonical key with layer index, derive the + # torch_dist fused key (if the fused key exists in metadata). + # ------------------------------------------------------------------ + def _match_fused_key(v2_canonical, value=None): + """Return match info for a v2 key whose fused counterpart exists in metadata. + + Returns: + ``(td_key, layer_idx)`` for regular fused-layer keys and GroupedMLP keys. + ``(td_key, local_expert_idx, "seq_mlp")`` for SequentialMLP expert keys. + ``None`` if no fused counterpart was found. + """ + m = _LAYERED_KEY_RE.match(v2_canonical) + if not m: + return None + prefix = m.group(1) or "" # e.g. "decoder." or "" + layer_idx = int(m.group(2)) + field = m.group(3) + + # ---- Check SequentialMLP expert format first ---- + # v2: ``{sub_layer}mlp.experts.local_experts.{N}.linear_fc{1|2}.weight`` + # (sub_layer may be empty or e.g. ``mtp_model_layer.`` for MTP) + # td: ``{prefix}layers.{layer_idx}.{sub_layer}mlp.experts.experts.linear_fc{1|2}.weight`` + seq_exp_m = _SEQ_EXPERT_FIELD_RE.match(field) + if seq_exp_m: + local_expert_idx = int(seq_exp_m.group(1)) + fc_type = seq_exp_m.group(2) + # Extract sub-layer prefix (if any) before ``mlp.experts.local_experts`` + sub_layer_prefix = "" + if not field.startswith("mlp.experts.local_experts."): + idx = field.find("mlp.experts.local_experts.") + if idx > 0: + sub_layer_prefix = field[:idx] # e.g. "mtp_model_layer." + td_base = ( + f"{prefix}layers.{layer_idx}." + f"{sub_layer_prefix}mlp.experts.experts.{fc_type}.weight" + ) + td_base = reverse_normalize_torch_dist_key(td_base) + candidates = [td_base] + candidates.append(f"model.{td_base}") + candidates.append(f"module.{td_base}") + candidates.append(f"model.module.{td_base}") + for c in candidates: + if c in metadata: + return c, local_expert_idx, "seq_mlp" + return None + + # ---- Build fused key (strip layer index) ---- + # v2: ``{prefix}layers.{N}.{field}`` + # td: ``{prefix}layers.{field}`` (no layer index) + fused_base = f"{prefix}layers.{field}" + + candidates = [fused_base] + # Torch_dist metadata may have "model." prefix + candidates.append(f"model.{fused_base}") + candidates.append(f"module.{fused_base}") + candidates.append(f"model.module.{fused_base}") + + # For GroupedMLP expert params: v2 uses "experts.linear_fc1.weight{N}" + # torch_dist uses "experts.experts.linear_fc1.weight" + exp_m = _EXPERT_V2_KEY_RE.search(field) + if exp_m: + fc_type = exp_m.group(2) + field_exp = ( + _EXPERT_V2_KEY_RE.sub(rf"mlp.experts.experts.linear_fc\2.weight", field).rstrip( + ".weight" + ) + + ".weight" + ) + fused_exp = f"{prefix}layers.{field_exp}" + candidates.append(fused_exp) + candidates.append(f"model.{fused_exp}") + + for c in candidates: + if c in metadata: + return c, layer_idx + + # Debug: log first failure per unique field + if not hasattr(_match_fused_key, '_logged'): + _match_fused_key._logged = set() + if field not in _match_fused_key._logged: + _match_fused_key._logged.add(field) + if not v2_canonical.endswith("._extra_state"): + logger.warning( + "Fused key not found in metadata for v2_key='%s' (field='%s'). " + "Tried candidates: %s", + v2_canonical, + field, + candidates, + ) + return None + + # ------------------------------------------------------------------ + # Phase A: Model weights + # ------------------------------------------------------------------ + # Debug: log sample metadata keys on first call + _sample_keys = sorted(k for k in metadata if not k.startswith("optimizer.")) + logger.debug( + "DCP metadata has %d non-optimizer keys. Sample: %s", len(_sample_keys), _sample_keys[:10] + ) + + for v2_key, v2_val in v2_by_canonical.items(): + result = _match_fused_key(v2_key, v2_val) + if result: + if len(result) == 3 and result[2] == "seq_mlp": + # SequentialMLP: (td_key, local_expert_idx, "seq_mlp") + td_key, local_expert_idx, _ = result + fused_layer_groups.setdefault(td_key, []).append( + (v2_key, v2_val, local_expert_idx, "seq_mlp") + ) + else: + # Fused layer or GroupedMLP: (td_key, layer_idx) + td_key, layer_idx = result + exp_m = _EXPERT_V2_KEY_RE.match(v2_key) + if exp_m: + expert_idx = int(exp_m.group(3)) + fused_layer_groups.setdefault(td_key, []).append( + (v2_key, v2_val, layer_idx, expert_idx) + ) + else: + fused_layer_groups.setdefault(td_key, []).append((v2_key, v2_val, layer_idx)) + continue + # Not fused — check 1:1 + if v2_key in metadata: + regular_model[v2_key] = v2_val + continue + m_key = f"module.{v2_key}" + if m_key in metadata: + regular_model[m_key] = v2_val + continue + + # ------------------------------------------------------------------ + # Phase B: Optimizer states (exp_avg / exp_avg_sq) & hi-prec (param) + # ------------------------------------------------------------------ + + # Phase B1: Build fused_layer_groups for optimizer state params + # that have fused counterparts in the torch_dist checkpoint. + # Handles three formats: regular fused layer, GroupedMLP, SequentialMLP. + for param_name, states in v2_optim_state.items(): + m = _LAYERED_KEY_RE.match(param_name) + if not m: + continue + prefix = m.group(1) or "" + layer_idx = int(m.group(2)) + field = m.group(3) + + # Check SequentialMLP expert format first + seq_exp_m = _SEQ_EXPERT_FIELD_RE.match(field) + if seq_exp_m: + local_expert_idx = int(seq_exp_m.group(1)) + fc_type = seq_exp_m.group(2) + # Extract sub-layer prefix (if any) before ``mlp.experts.local_experts`` + sub_layer_prefix = "" + if not field.startswith("mlp.experts.local_experts."): + idx = field.find("mlp.experts.local_experts.") + if idx > 0: + sub_layer_prefix = field[:idx] # e.g. "mtp_model_layer." + for sk, sv in states.items(): + if sk not in ("exp_avg", "exp_avg_sq"): + continue + fused_base = ( + f"optimizer.state.{sk}." + f"{prefix}layers.{layer_idx}." + f"{sub_layer_prefix}mlp.experts.experts.{fc_type}.weight" + ) + fused_base = reverse_normalize_torch_dist_key(fused_base) + td_opt_key = None + for candidate in [fused_base, f"model.{fused_base}"]: + if candidate in metadata: + td_opt_key = candidate + break + if td_opt_key is None: + continue + fused_layer_groups.setdefault(td_opt_key, []).append( + (param_name, sv, local_expert_idx, "seq_mlp") + ) + optim_matched.add(param_name) + continue + + # Check GroupedMLP expert format + fused_base = f"{prefix}layers.{field}" + exp_m = _EXPERT_V2_KEY_RE.search(field) + if exp_m: + fc_type = exp_m.group(2) + field_exp = ( + _EXPERT_V2_KEY_RE.sub(rf"mlp.experts.experts.linear_fc\2.weight", field).rstrip( + ".weight" + ) + + ".weight" + ) + fused_base = f"{prefix}layers.{field_exp}" + expert_idx = int(exp_m.group(3)) + else: + expert_idx = None + + for sk, sv in states.items(): + if sk not in ("exp_avg", "exp_avg_sq"): + continue + td_opt_key = None + for candidate in [ + f"optimizer.state.{sk}.{fused_base}", + f"optimizer.state.{sk}.model.{fused_base}", + ]: + if candidate in metadata: + td_opt_key = candidate + break + if td_opt_key is None: + continue + if expert_idx is not None: + fused_layer_groups.setdefault(td_opt_key, []).append( + (param_name, sv, layer_idx, expert_idx) + ) + else: + fused_layer_groups.setdefault(td_opt_key, []).append((param_name, sv, layer_idx)) + optim_matched.add(param_name) + + # Phase B2: Handle remaining metadata entries (1:1 optimizer keys + hi-prec) + for td_key in metadata: + if not td_key.startswith("optimizer.state."): + continue + rest = td_key[len("optimizer.state.") :] + parts = rest.split(".", 1) + + if rest.startswith("param."): + param_name_td = rest[len("param.") :] + param_canonical = _canonicalize_td_key(param_name_td, strip_model_prefix=False) + param_canonical = _strip_module_prefix(param_canonical) + if param_canonical in v2_by_canonical: + hi_prec_model[param_name_td] = v2_by_canonical[param_canonical] + hi_prec_td_keys.add(param_canonical) + continue + + if len(parts) != 2: + continue + state_key, param_name_td = parts + param_canonical = _canonicalize_td_key(param_name_td, strip_model_prefix=False) + param_canonical = _strip_module_prefix(param_canonical) + + # Try 1:1 match. Fused-layer matching is handled by Phase B1 above. + if param_canonical in v2_optim_state and state_key in v2_optim_state[param_canonical]: + optim_keys[td_key] = v2_optim_state[param_canonical][state_key] + optim_matched.add(param_canonical) + + if logger.isEnabledFor(logging.DEBUG): + total_fused = sum(len(v) for v in fused_layer_groups.values()) + logger.debug( + "torch_dist → v2: %d metadata keys → %d regular + %d fused (%d groups) " + "+ %d hi-prec + %d optim matched", + len(metadata), + len(regular_model), + total_fused, + len(fused_layer_groups), + len(hi_prec_model), + len(optim_matched), + ) + return regular_model, fused_layer_groups, hi_prec_model, optim_keys, optim_matched + + +# ------------------------------------------------------------------ +# Torch_dist → FSDP v2 online checkpoint conversion +# ------------------------------------------------------------------ + + +def _assert_dcp_keys_in_metadata(mapped_sd: dict, metadata: dict) -> None: + """Verify every DCP load key exists in the torch_dist metadata.""" + missing = [] + for top_key, subtree in mapped_sd.items(): + _collect_missing_keys(f"{top_key}.", subtree, metadata, missing) + if missing: + raise RuntimeError( + f"{len(missing)} DCP load keys not found in torch_dist metadata. " + f"Missing: {sorted(missing)}" + ) + + +def _collect_missing_keys(prefix: str, subtree, metadata: dict, missing: list) -> None: + """Recursively check that every leaf key in *subtree* exists in *metadata*.""" + if isinstance(subtree, dict): + for k, v in subtree.items(): + _collect_missing_keys(f"{prefix}{k}.", v, metadata, missing) + elif isinstance(subtree, torch.Tensor): + key = prefix.rstrip(".") + if key not in metadata: + missing.append(key) + + +def _canonicalize_td_key(td_key, *, strip_model_prefix=True): + """Strip the top-level DCP prefix, shard suffix, and normalize naming.""" + if strip_model_prefix and td_key.startswith("model."): + key = td_key[len("model.") :] + else: + key = td_key + key = normalize_torch_dist_key(key) + return key + + +def load_torch_dist_into_fsdp_v2(args, checkpoint_name, v2_state_dict, strict=True): + """Load a torch_dist checkpoint into a Megatron FSDP v2 skeleton via DCP. + + This is the entry point for online checkpoint conversion from + legacy ``torch_dist`` format (ND-parallel) to ``fsdp_dtensor`` + (Megatron FSDP v2). + + The conversion proceeds in five phases: + + 1. **Preprocess & verify** — wrap optimizer states as uneven DTensors + and verify ``__create_chunk_list__`` / ``__create_write_items__`` metadata. + 2. **Build name mapping** — match torch_dist metadata keys to v2 DTensors, + handling fused layers, GroupedMLP, and SequentialMLP expert formats. + 3. **DCP load 1:1 entries** — load regular model weights and optimizer + states that have a direct 1:1 correspondence in the torch_dist checkpoint. + 4. **Load fused entries** — load fused (multi-layer and/or multi-expert) + tensors and slice them into per-layer/per-expert v2 DTensors. + 5. **Verify** — when ``strict=True``, ensure every v2 model param and + optimizer state was loaded from the torch_dist checkpoint. When + ``strict=False``, unmatched entries are logged as warnings instead + of raising errors. + """ + import torch.distributed.checkpoint as dcp + from torch.distributed.checkpoint import FileSystemReader + from torch.distributed.checkpoint.default_planner import DefaultLoadPlanner + + # ---- Phase 1: Preprocess & verify v2 state dict ---- + v2_by_canonical, v2_optim_state = _preprocess_and_verify_v2_state_dict(v2_state_dict) + + # ---- Phase 2: Read torch_dist metadata & build name mapping ---- + reader = FileSystemReader(checkpoint_name) + metadata = reader.read_metadata().state_dict_metadata + regular_model, fused_layer_groups, hi_prec_model, optim_keys, optim_matched = ( + _build_torch_dist_to_v2_map(metadata, v2_by_canonical, v2_optim_state) + ) + + # ---- Phase 3: Build & load mapped state dict via DCP ---- + mapped_sd = {} + if regular_model: + mapped_sd.update(regular_model) + + opt_state = {} + for td_key, v2_val in optim_keys.items(): + rest = td_key[len("optimizer.state.") :] + state_key, param_name_td = rest.split(".", 1) + opt_state.setdefault(state_key, {})[param_name_td] = v2_val + if hi_prec_model: + opt_state["param"] = hi_prec_model + if opt_state: + mapped_sd["optimizer"] = {"state": opt_state} + + _assert_dcp_keys_in_metadata(mapped_sd, metadata) + if mapped_sd: + dcp.load( + state_dict=mapped_sd, + checkpoint_id=checkpoint_name, + planner=DefaultLoadPlanner(allow_partial_load=True), + ) + + # Merge hi-precision model params back so the model subtree is complete. + if hi_prec_model: + mapped_sd.update({k: v for k, v in hi_prec_model.items() if k not in mapped_sd}) + + # ---- Phase 4: Load fused layer / expert params by slicing ---- + # Handles three tensor formats: + # Regular fused: shape (num_layers, ...) → flat[layer_idx] + # GroupedMLP: shape (num_layers, num_experts, ...) → flat[layer_idx, expert_idx] + # SequentialMLP: shape (num_global_experts, ...) → flat[global_expert_idx] + # + # Entry types by tag: + # (v2_key, v2_val, layer_idx) → regular fused layer + # (v2_key, v2_val, layer_idx, expert_idx) → GroupedMLP expert + # (v2_key_or_param_name, v2_val, local_expert_idx, "seq_mlp") → SequentialMLP expert + for td_key, entries in fused_layer_groups.items(): + td_meta = metadata.get(td_key) + assert td_meta is not None, f"Missing metadata for fused key '{td_key}'" + first_entry = entries[0] + is_seq_mlp = len(first_entry) >= 4 and first_entry[-1] == "seq_mlp" + is_grouped_mlp = len(first_entry) >= 4 and not is_seq_mlp + + entries.sort(key=lambda x: x[2]) # sort by layer_idx / local_expert_idx + example_val = entries[0][1] # v2_val (model) or optimizer state tensor + + # Build the full fused tensor shape and compute slicing + fused_shape = list(example_val.shape) + if is_seq_mlp: + num_total_experts = td_meta.size[0] + fused_shape.insert(0, num_total_experts) + num_local = len(entries) + ep_rank = mpu.get_expert_model_parallel_rank() + elif is_grouped_mlp: + num_layers = td_meta.size[0] + num_experts = td_meta.size[1] + fused_shape.insert(0, num_experts) + fused_shape.insert(0, num_layers) + else: + num_layers = td_meta.size[0] + fused_shape.insert(0, num_layers) + + # Use a Shard(0) DTensor so the large fused tensor is distributed + # across DP ranks during DCP load, avoiding OOM on a single rank + # (e.g. for GroupedMLP with many layers × experts). + device_mesh = example_val.device_mesh if isinstance(example_val, DTensor) else None + if device_mesh is not None: + flat = torch.distributed.tensor.empty( + fused_shape, dtype=example_val.dtype, device_mesh=device_mesh, placements=[Shard(0)] + ) + else: + flat = torch.empty(fused_shape, dtype=example_val.dtype, device=example_val.device) + + dcp.load( + state_dict={td_key: flat}, + checkpoint_id=checkpoint_name, + planner=DefaultLoadPlanner(allow_partial_load=True), + ) + + # Gather shards back into a full local tensor for per-layer slicing. + if device_mesh is not None and isinstance(flat, DTensor): + flat = redistribute_uneven_dtensor_to_replicated(flat).to_local() + + for i, entry in enumerate(entries): + if is_seq_mlp: + v2_key_or_param, v2_val, local_expert_idx = entry[:3] + chunk = flat[ep_rank * num_local + local_expert_idx] + elif is_grouped_mlp: + v2_key_or_param, v2_val, layer_idx, expert_idx = entry[:4] + chunk = flat[layer_idx, expert_idx] + else: + v2_key_or_param, v2_val, layer_idx = entry[:3] + chunk = flat[layer_idx] + if isinstance(v2_val, DTensor): + lt = v2_val._local_tensor + if lt.numel() != 0: + local_off = 0 + for c in lt.__create_chunk_list__(): + off = c.offsets + sz = c.sizes + src_idx = tuple(slice(o, o + s) for o, s in zip(off, sz)) + src = chunk[src_idx] + dst_idx = (slice(local_off, local_off + sz[0]),) + tuple( + slice(0, s) for s in sz[1:] + ) + lt[dst_idx].copy_(src) + local_off += sz[0] + else: + v2_val.copy_(chunk) + mapped_sd[v2_key_or_param] = v2_val + + # ---- Phase 5: Verify all v2 params & optimizer states were loaded ---- + loaded = set(_canonicalize_td_key(k) for k in mapped_sd) + all_v2 = set(v2_by_canonical.keys()) + unloaded = all_v2 - loaded - {k for k in all_v2 if k.endswith("._extra_state")} + + v2_optim_unmatched = set(v2_optim_state.keys()) - optim_matched + + if strict: + if unloaded: + raise RuntimeError( + f"{len(unloaded)} v2 model parameters were not loaded from the " + f"torch_dist checkpoint. Unloaded params: {sorted(unloaded)}" + f" loaded: {sorted(loaded)}" + ) + if v2_optim_unmatched: + raise RuntimeError( + f"{len(v2_optim_unmatched)} v2 optimizer state entries were " + f"not matched to torch_dist data: {sorted(v2_optim_unmatched)}" + ) + else: + if unloaded: + logger.warning( + "%d v2 model parameters were not loaded from the torch_dist " + "checkpoint (strict=False): %s", + len(unloaded), + sorted(unloaded), + ) + if v2_optim_unmatched: + logger.warning( + "%d v2 optimizer state entries were not matched to " + "torch_dist data (strict=False): %s", + len(v2_optim_unmatched), + sorted(v2_optim_unmatched), + ) + + # ---- Load common state dict ---- + common_info = dist_checkpointing.load_common_state_dict(checkpoint_name) + if "optimizer" in common_info: + opt_common = common_info.pop("optimizer") + if isinstance(opt_common, dict): + if "optimizer" in opt_common: + opt_common = opt_common["optimizer"] + v2_state_dict.update(common_info) + + return v2_state_dict diff --git a/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py b/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py index 2b02b1a9f57..d78efb61c56 100644 --- a/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py +++ b/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py @@ -14,6 +14,7 @@ import logging import random +from contextlib import nullcontext from typing import Dict, List, Optional try: @@ -50,12 +51,22 @@ MegatronFSDP, MixedPrecisionPolicy, ) + from megatron.core.distributed.fsdp.src.megatron_fsdp.v2 import FSDPModule + from megatron.core.distributed.fsdp.src.megatron_fsdp.v2.mixed_precision import ( + FullyShardFP8Policy, + FullyShardMixedPrecisionPolicy, + FullyShardNVFP4Policy, + ) HAVE_MEGATRON_FSDP = True except ImportError as import_megatron_fsdp_error: IMPORT_MEGATRON_FSDP_ERROR = import_megatron_fsdp_error HAVE_MEGATRON_FSDP = False +from megatron.core.transformer.moe.experts import SequentialMLP, TEGroupedMLP + +from .checkpoint import _propagate_chunk_metadata_to_state_dict + logger = logging.getLogger(__name__) @@ -68,7 +79,6 @@ class FullyShardedDataParallel(_BaseDataParallel): _MODULE_TYPE_REGISTRY: Dict[str, set] = { "column": { "ColumnParallelLinear", - "LinearCrossEntropyModule", "TEColumnParallelLinear", "TELayerNormColumnParallelLinear", "TEColumnParallelGroupedLinear", @@ -104,6 +114,19 @@ def __init__( if not HAVE_MEGATRON_FSDP: raise IMPORT_MEGATRON_FSDP_ERROR + if ddp_config.use_megatron_fsdp_v2: + self._init_with_fully_shard( + config, + ddp_config, + module, + fsdp_unit_modules, + disable_bucketing, + device, + pg_collection, + ) + self._set_dist_param_unique_name() + return + if has_config_logger_enabled(config): log_config_to_disk(config, locals(), prefix=type(self).__name__) @@ -158,26 +181,6 @@ def __init__( self._annotate_tensor_parallelism(module) - if config.overlap_moe_expert_parallel_comm: - assert not ddp_config.fsdp_double_buffer, ( - "1F1B overlap with FSDP does not support double buffer. " - "Please set fsdp_double_buffer=False in the ddp config." - ) - assert config.cuda_graph_impl in ("none", "full_iteration"), ( - "1F1B overlap with FSDP does not support per-layer CUDA graphs " - f"(cuda_graph_impl={config.cuda_graph_impl!r}). " - "Use cuda_graph_impl='full_iteration' or disable CUDA graphs " - "(cuda_graph_impl='none')." - ) - - if ( - config.overlap_moe_expert_parallel_comm - and ddp_config.data_parallel_sharding_strategy == "optim_grads_params" - ): - assert self.fsdp_unit_modules == [TransformerLayer], ( - "EP overlap with FSDP currently requires fsdp_unit_modules " - f"to be [TransformerLayer], got {self.fsdp_unit_modules}." - ) super().__init__( config=config, module=MegatronFSDP( @@ -190,20 +193,8 @@ def __init__( dist_index=self.megatron_fsdp_dist_index, calculate_per_token_loss=config.calculate_per_token_loss, init_model_with_meta_device=config.init_model_with_meta_device, - # EP overlap schedule calls sub-modules directly instead of - # TransformerLayer.forward(), so fine-grained hooks are needed - # to manage _training_state and all-gather each sub-module's - # parameters individually. This applies to all sharding - # strategies (not only optim_grads_params) because the hooks - # also maintain per-module training-state bookkeeping that the - # gradient-reduction pipeline relies on. enable_fine_grained_param_gather_hook=( - (config.fp8_recipe == "mxfp8" and ddp_config.fp8_param_gather) - or config.overlap_moe_expert_parallel_comm - ), - enable_fine_grained_param_gather_backward_hook=( - config.overlap_moe_expert_parallel_comm - and ddp_config.data_parallel_sharding_strategy == "optim_grads_params" + config.fp8_recipe == "mxfp8" and ddp_config.fp8_param_gather ), ), ) @@ -214,18 +205,214 @@ def __init__( self.finish_grad_sync = self.module.finish_grad_sync self.scale_gradients = self.module.scale_gradients self.zero_grad_buffer = self.module.zero_grad_buffer + self.log_per_param_norms = self.module._log_per_param_norms + self.compute_per_param_norms = self.module._compute_per_param_norms self.broadcast_params = self.module.broadcast_params self.synchronize_param_gather = self.module.synchronize_param_gather self.module.state_dict_for_save_checkpoint = self.module.state_dict self.state_dict_for_save_checkpoint = self.state_dict self.module.config = config - self.sync_rng_states_across_tp_group() + self._set_dist_param_unique_name() + + def _init_with_fully_shard( + self, + config: TransformerConfig, + ddp_config: DistributedDataParallelConfig, + module: torch.nn.Module, + fsdp_unit_modules: Optional[List[torch.nn.Module]] = None, + disable_bucketing: bool = False, + device: Optional[torch.device] = None, + pg_collection: Optional[ProcessGroupCollection] = None, + ): + if ddp_config.use_megatron_fsdp: + from megatron.core.distributed.fsdp.src.megatron_fsdp.v2 import fully_shard + else: + from torch.distributed.fsdp import fully_shard + + if ( + fsdp_unit_modules is None + and ddp_config.data_parallel_sharding_strategy == "optim_grads_params" + ): + fsdp_unit_modules = [TransformerLayer] + + if pg_collection is None: + pg_collection = ProcessGroupCollection.use_mpu_process_groups() + + edp_mesh = _init_dp_mesh(pg_collection, edp=True) + dp_mesh = _init_dp_mesh(pg_collection, edp=False) + + fully_shard_mp_policy = FullyShardMixedPrecisionPolicy( + main_params_dtype=ddp_config.megatron_fsdp_main_params_dtype, + main_grads_dtype=( + torch.float32 + if ddp_config.grad_reduce_in_fp32 + else ddp_config.megatron_fsdp_main_grads_dtype + ), + grad_comm_dtype=( + torch.float32 + if ddp_config.grad_reduce_in_fp32 + else ddp_config.megatron_fsdp_grad_comm_dtype + ), + use_decoupled_grad=ddp_config.megatron_fsdp_use_decoupled_grad, + fp8=FullyShardFP8Policy( + enabled=ddp_config.fp8_param_gather, + recipe=config.fp8_recipe, + keep_transpose_cache=ddp_config.keep_fp8_transpose_cache, + ), + nvfp4=FullyShardNVFP4Policy( + enabled=ddp_config.fp4_param_gather, recipe=config.fp4_recipe + ), + ) + kwargs = { + "mp_policy": fully_shard_mp_policy, + "enable_unshard_prefetch": ddp_config.overlap_param_gather, + "enable_async_reduce_grad": ddp_config.overlap_grad_reduce, + "enable_trace_pool": ddp_config.fsdp_double_buffer, + } + if config.calculate_per_token_loss: + gradient_scaling_factor = None + expert_gradient_scaling_factor = None + elif ddp_config.average_in_collective: + dp_world_size = pg_collection.dp.size() + expt_dp_world_size = pg_collection.expt_dp.size() + gradient_scaling_factor = 1.0 + expert_gradient_scaling_factor = expt_dp_world_size / dp_world_size + else: + dp_world_size = pg_collection.dp.size() + gradient_scaling_factor = 1.0 / dp_world_size + expert_gradient_scaling_factor = 1.0 / dp_world_size + + for m in module.modules(): + if isinstance(m, (TEGroupedMLP, SequentialMLP)): + fully_shard( + m, + mesh=edp_mesh, + gradient_scaling_factor=expert_gradient_scaling_factor, + **kwargs, + ) + if fsdp_unit_modules is not None: + for m in module.modules(): + if isinstance(m, tuple(fsdp_unit_modules)): + fully_shard( + m, mesh=dp_mesh, gradient_scaling_factor=gradient_scaling_factor, **kwargs + ) + fully_shard(module, mesh=dp_mesh, gradient_scaling_factor=gradient_scaling_factor, **kwargs) + + # Propagate relevant attributes from original parameters to the new + # distributed parameters created by FSDP. This is REQUIRED for + # correctness: the optimizer's param group builder + # (_get_param_groups) relies on attributes like + # is_embedding_parameter and is_embedding_or_output_parameter to + # classify parameters into groups. If these attributes are missing + # on the DTensor dist_params, the optimizer may assign a param to + # the wrong group, causing skipped weight decay on embeddings or + # incorrect learning-rate multipliers, which leads to convergence + # divergence. + for child in module.modules(): + if not isinstance(child, FSDPModule): + continue + for param_group in child._fsdp_param_groups: + for param, dist_param in zip(param_group.params, param_group.dist_params): + for attr_name in [ + # allreduce: expert params have allreduce=False set by + # te layers. Missing this causes is_expert_parallel + # misclassification in _get_param_groups, which can + # produce NaN gradients when wgrad fusion writes to + # the wrong buffer or when expert gradient scaling is + # incorrect. + "allreduce", + "requires_grad", + "sequence_parallel", + "shared", + "tensor_model_parallel", + "partition_dim", + "partition_stride", + "is_embedding_or_output_parameter", + "is_embedding_parameter", + "_tensor_parallel_mode", + ]: + if hasattr(param, attr_name): + setattr(dist_param, attr_name, getattr(param, attr_name)) + + # Per-module NaN checking is disabled by default on the fully_shard + # path to avoid the per-parameter synchronization overhead on every + # unshard. Enable via a manual call to module._set_nan_check(True). + # if ddp_config.check_for_nan_in_grad: + # module._set_nan_check(True) + + super().__init__(config=config, module=module) + + noop = lambda *args, **kwargs: None + + def not_implemented_op(): + raise NotImplementedError( + "This operation is not implemented for the fully_shard API path. " + ) + + from unittest.mock import Mock + + self.param_and_grad_buffer = Mock() + self.param_and_grad_buffer.parameter_groups = [] + self.param_and_grad_buffer.copy_main_weights_to_model_weights = ( + self.module._copy_main_weights_to_model_weights + ) + self.ddp_config = ddp_config + self.no_sync = nullcontext + self.start_param_sync = noop + self.start_grad_sync = noop + + def finish_grad_sync(force_all_reduce: Optional[bool] = False): + ctx = self.module._fsdp_root_context + torch.cuda.current_stream().wait_stream(ctx.rs_stream) + + def synchronize_param_gather(): + ctx = self.module._fsdp_root_context + torch.cuda.current_stream().wait_stream(ctx.ag_stream) + + self.finish_grad_sync = finish_grad_sync + self.scale_gradients = self.module._scale_gradients + self.zero_grad_buffer = self.module._zero_grad_buffer + self.log_per_param_norms = self.module._log_per_param_norms + self.compute_per_param_norms = self.module._compute_per_param_norms + self.log_parameter_groups = self.module._log_parameter_groups + # Parameter broadcast is handled during _materialize_meta_module + # for the fully_shard path (params are synced across DP ranks + # immediately after meta-device init, before DTensor wrapping). + # For non-meta init, all ranks share the same seed so no sync is + # needed. This is intentionally a no-op instead of + # not_implemented_op to avoid crashing when the training loop + # calls broadcast_params() with --data-parallel-random-init. + self.broadcast_params = noop + self.synchronize_param_gather = synchronize_param_gather + self.module.state_dict_for_save_checkpoint = module.state_dict + self.state_dict_for_save_checkpoint = lambda *args, **kwargs: self.state_dict() + + # Register state dict post-hook on the root module to propagate + # chunk metadata (``__create_chunk_list__``) from model params to + # the DTensors in the state dict. The hook receives the complete + # state dict so all parameter names can be matched in one pass. + def _post_hook(module, state_dict, prefix, local_metadata): + _propagate_chunk_metadata_to_state_dict(module, state_dict) + + self.module.register_state_dict_post_hook(_post_hook) + + if torch.distributed.get_rank() == 0: + self.module._log_parameter_groups() + + def _set_dist_param_unique_name(self): + for name, param in self.module.named_parameters(): + if not hasattr(param, "_unique_name"): + setattr(param, "_unique_name", name) def load_state_dict(self, state_dict, strict=True): """ Load the state dictionary into the module. """ + if self.ddp_config.use_megatron_fsdp_v2: + super().load_state_dict(state_dict, strict=strict) + return + custom_state_dict = {} for key, value in state_dict.items(): if self.config.fp8 and key.endswith('._extra_state'): @@ -447,8 +634,20 @@ def _init_dist_index(self, pg_collection): def stop_communication(self): """ - Stop communication for the module. + Cease all pending FSDP communication and synchronize main stream. + + For the fully_shard path: waits on both ag_stream (async all-gather) + and rs_stream (async reduce-scatter) to bring their work into the + main CUDA stream. + For the Megatron-FSDP path: calls synchronize_gradient_reduce and + synchronize_param_gather. """ + if self.ddp_config.use_megatron_fsdp_v2: + ctx = self.module._fsdp_root_context + torch.cuda.current_stream().wait_stream(ctx.ag_stream) + torch.cuda.current_stream().wait_stream(ctx.rs_stream) + return + self.module.synchronize_gradient_reduce() self.module.synchronize_param_gather() @@ -467,6 +666,51 @@ def sync_rng_states_across_tp_group(self): _load_rng_state_dict(broadcast_list[0]) +def _reset_parameters(module): + """ + Recursively reset parameters for the module and its submodules. + This is used to ensure that all ranks start with the same initial parameters + before sharding, which is important for correctness when using FSDP. + """ + parent_fsdp_module_map = {} + for m in module.modules(): + if isinstance(m, FSDPModule): + for child in m.module.modules(): + parent_fsdp_module_map[child] = m + + for m in module.modules(): + if hasattr(m, "reset_parameters"): + parent_fsdp_module_map[m].unshard() + m.reset_parameters() + parent_fsdp_module_map[m].reshard() + + +def _init_dp_mesh(pg_collection, edp=False): + assert HAVE_DTENSOR, ( + "DTensor support is required to initialize the device mesh. " + "Please install a compatible version of PyTorch." + ) + + if pg_collection is None: + pg_collection = ProcessGroupCollection.use_mpu_process_groups() + if edp: + mesh = DeviceMesh.from_group( + device_type="cuda", + group=pg_collection.expt_dp, + mesh=dist.get_process_group_ranks(pg_collection.expt_dp), + mesh_dim_names=("edp",), + ) + else: + mesh = DeviceMesh.from_group( + device_type="cuda", + group=pg_collection.dp_cp, + mesh=dist.get_process_group_ranks(pg_collection.dp_cp), + mesh_dim_names=("dp",), + ) + + return mesh + + def _get_hsdp_tp_mesh(outer_fsdp_dp_group, dp_cp_group, tp_group, ep_size=1): assert HAVE_EINOPS, "einops is not installed. Please install it with `pip install einops`." world_size = dist.get_world_size() diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/__init__.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/__init__.py index a1ae54e622d..f1982fa6842 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/__init__.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/__init__.py @@ -30,6 +30,7 @@ __version__, ) from .utils import FSDPDistributedIndex +from .v2 import FSDPModule, fully_shard __all__ = [ "DistributedDataParallelConfig", diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/distributed_data_parallel_config.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/distributed_data_parallel_config.py index a2feb99cb23..6c360b92fbf 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/distributed_data_parallel_config.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/distributed_data_parallel_config.py @@ -36,7 +36,11 @@ class DistributedDataParallelConfig: fp8_param_gather: bool = False """If true, keep the compute param in fp8 (do not use any other intermediate dtype) and - perform the param all-gather in fp8.""" + perform the param all-gather in fp8.""" + + fp4_param_gather: bool = False + """If true, keep the compute param in fp4 (do not use any other intermediate dtype) and + perform the param all-gather in fp4.""" data_parallel_sharding_strategy: str = 'no_shard' """Sharding strategy for FSDP. Valid values are 'no_shard', 'optim', @@ -151,6 +155,10 @@ class DistributedDataParallelConfig: main gradients to parameter dtype for `.grad`. """ + use_megatron_fsdp_v2: bool = False + """If true, use the `fully_shard` API for FSDP sharding the model. + """ + def __post_init__(self): import os diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py index 1c71cd33c74..b589a5eece3 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py @@ -17,7 +17,6 @@ import logging from contextlib import contextmanager from enum import Enum, auto -from functools import partial from typing import Any, Dict, List, Optional, Tuple import torch @@ -47,12 +46,11 @@ try: # Default to Megatron-LM FW. + logger.info("Detected Megatron Core, using Megatron-FSDP with Megatron.") from megatron.core.distributed.distributed_data_parallel_config import ( DistributedDataParallelConfig, ) from megatron.core.utils import is_submodule - - logger.info("Detected Megatron Core, using Megatron-FSDP with Megatron.") except ImportError: # Megatron-LM is not installed, use Megatron-FSDP as a standalone module. logger.info("Megatron Core is not installed, Megatron-FSDP will run without Megatron Core.") @@ -216,7 +214,6 @@ def __init__( fsdp_db_use_persist_buf_on_alloc_fail: bool = False, disable_symmetric_registration: bool = False, enable_fine_grained_param_gather_hook: bool = False, - enable_fine_grained_param_gather_backward_hook: bool = False, report_nan_in_param_grad: bool = False, ): super().__init__() @@ -269,9 +266,6 @@ def __init__( self.calculate_per_token_loss = calculate_per_token_loss self.init_model_with_meta_device = init_model_with_meta_device self.enable_fine_grained_param_gather_hook = enable_fine_grained_param_gather_hook - self.enable_fine_grained_param_gather_backward_hook = ( - enable_fine_grained_param_gather_backward_hook - ) self.report_nan_in_param_grad = report_nan_in_param_grad # FSDPDistributedIndex stores the process groups and meshes used by Megatron-FSDP. @@ -699,12 +693,18 @@ def _process_post_backward_gradients(param_list): # Filter out shared parameters whose gradients are handled by the root hook. param_list = [p for p in param_list if not getattr(p, "_is_shared", False)] - # Make sure for delayed wgrad params, the grad_acc_hooks are registered. - for p in param_list: - if getattr(p, 'skip_backward_post_hook', False): - assert hasattr( - p, 'post_wgrad_grad_acc_hook' - ), "Missing grad accumulation hook for delayed_wgrad_compute param." + # Filter out parameters whose gradient processing is deferred to a delayed + # wgrad accumulation hook (post_wgrad_grad_acc_hook). If skip_backward_post_hook + # is set but the delayed hook was never installed, process the parameter + # immediately as a safety fallback to avoid silently dropping gradients. + param_list = [ + p + for p in param_list + if not ( + getattr(p, 'skip_backward_post_hook', False) + and hasattr(p, 'post_wgrad_grad_acc_hook') + ) + ] if not param_list: return @@ -884,7 +884,7 @@ def _pre_backward_param_unshard(module: nn.Module, *unused): self._root_pre_backward_hook_issued = False - def _root_pre_backward(module: nn.Module, *unused, skip_backward_hook: bool = False): + def _root_pre_backward(module: nn.Module, *unused): """Marks the module's training state as PRE_BACKWARD before the backprop, this function is registered on the root module. @@ -920,8 +920,6 @@ def _root_pre_backward(module: nn.Module, *unused, skip_backward_hook: bool = Fa param.grad_added_to_main_grad = False # Queue the root post-backward hook to reduce leftover gradients after # the backward pass. - if skip_backward_hook: - return torch.autograd.Variable._execution_engine.queue_callback(_root_post_backward) @torch.compiler.disable @@ -1009,24 +1007,12 @@ def _register_pre_backward_param_unshard_hook(module): create_custom_backward_hook(module, _pre_backward_param_unshard) ) - # These hooks need to be exposed for manual management by 1F1B Overlapping - # and triggered by 1F1B Overlapped execution pipeline, except for - # `param_unshard` hook that needs to be installed at param level, - # such that non-overlapped params like embedding layer are also correctly - # unsharded. - self.post_forward_release_module = partial(_post_forward, input=None, output=None) - self.post_backward_release_module = _post_backward_release_module - self.pre_backward = partial(_root_pre_backward, module=None, skip_backward_hook=True) - self.post_backward = _root_post_backward - fsdp_modules = [] for name, module in root_module.named_modules(): # Set post backward hook for TE grouped gemm if enabled comm overlap setup_delayed_wgrad_acc_hook(module, _process_post_backward_gradients) if self.enable_fine_grained_param_gather_hook: _register_pre_forward_param_unshard_hook(module) - if self.enable_fine_grained_param_gather_backward_hook: - _register_pre_backward_param_unshard_hook(module) # Skip if the module is already registered in fsdp_modules. if any(is_submodule(module, fsdp_module) for fsdp_module in fsdp_modules): @@ -1081,11 +1067,7 @@ def _register_pre_backward_param_unshard_hook(module): continue self.grad_acc_hooks[f"grad_acc and reduce for {self.param_to_name[param]}"] = ( param.register_post_accumulate_grad_hook( - lambda p: ( - None - if getattr(p, 'skip_backward_post_hook', False) - else _process_post_backward_gradients([p]) - ) + lambda p: _process_post_backward_gradients([p]) ) ) @@ -1242,9 +1224,6 @@ def start_param_sync(self, *unused, force_sync: bool = False, force_dispatch: bo """ self._replace_param_with_raw_if_needed() - if self.data_parallel_sharding_strategy == "no_shard": - return - if not force_sync and self.ddp_config.overlap_param_gather: # All-gather the first bucket before the forward pass. if self.ddp_config.fsdp_all_gather_in_start_param_sync: @@ -1432,6 +1411,66 @@ def install_optimized_model_weights(self): """ self.param_and_grad_buffer.copy_main_weights_to_model_weights() + def _compute_per_param_norms(self): + """ + Compute per-parameter L2 norms for params and grads. + + Returns a dict {param_name: {"param_norm": float, "grad_norm": float}}. + The norms are globally reduced across DP ranks for direct comparison. + """ + results = {} + pg_buffer = self.param_and_grad_buffer + param_to_group = pg_buffer.param_to_param_group + dp_group = self.dist_index.get_dp_group(is_expert_parallel=False) + + for full_name, param in self.module.named_parameters(): + if not param.requires_grad: + continue + + results[full_name] = {"param_norm": 0.0, "grad_norm": 0.0} + if hasattr(param, '_local_tensor'): + local_param = param._local_tensor + else: + local_param = param + if local_param.numel() > 0: + results[full_name]["param_norm"] = local_param.float().norm(p=2).item() ** 2 + + if param in param_to_group: + group = pg_buffer.parameter_groups[param_to_group[param]] + if group.requires_grad: + gbuf = ( + group.hfsdp_helper_gbuf + if group.hfsdp_helper_gbuf + else group.main_grad_buffer + ) + if gbuf is not None and hasattr(gbuf, 'data') and gbuf.data is not None: + item_id = group.param_idx[param] + local_grad = gbuf.get_item(item_id, only_shard=True) + if local_grad.numel() > 0: + results[full_name]["grad_norm"] = ( + local_grad.float().norm(p=2).item() ** 2 + ) + + for param_name in results: + for key in ("param_norm", "grad_norm"): + t = torch.tensor([results[param_name][key]], device="cuda") + torch.distributed.all_reduce(t, group=dp_group) + results[param_name][key] = t.sqrt().item() + return results + + def _log_per_param_norms(self, iteration: int, prefix: str = ""): + """Log per-parameter param and gradient L2 norms via print (rank 0 only).""" + norms = self._compute_per_param_norms() + if torch.distributed.get_rank() != 0: + return + for param_name in sorted(norms.keys()): + pn = norms[param_name]["param_norm"] + gn = norms[param_name]["grad_norm"] + logger.info( + f"{prefix} iter={iteration} param={param_name} " + f"param_norm={pn:.6f} grad_norm={gn:.6f}" + ) + def broadcast_params(self): """ Syncs parameters across all DP ranks. diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/uneven_dtensor.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/uneven_dtensor.py index cffdf09362f..d42c10494e7 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/uneven_dtensor.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/uneven_dtensor.py @@ -12,10 +12,11 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Iterable, List, Union +from typing import Iterable, List, Optional, Union import torch import torch.distributed as dist +import torch.nn as nn from torch.distributed._tensor import DTensor from torch.distributed.checkpoint.metadata import ( ChunkStorageMetadata, @@ -23,7 +24,9 @@ TensorProperties, ) from torch.distributed.checkpoint.planner import TensorWriteData, WriteItem, WriteItemType -from torch.distributed.tensor.placement_types import Replicate, Shard, _StridedShard +from torch.distributed.checkpoint.state_dict import get_state_dict as _get_state_dict +from torch.distributed.tensor import DeviceMesh +from torch.distributed.tensor.placement_types import Placement, Replicate, Shard, _StridedShard def gather_and_compute_chunk_metadata(dtensor: DTensor) -> ChunkStorageMetadata: @@ -99,30 +102,6 @@ def update_uneven_dtensor_chunk_metadata(dtensor: DTensor) -> dict: This function modifies the DTensor in-place to include chunk metadata and write items closures for saving and loading. """ - - def _chunk_list_closure(chunk_meta): - return lambda: chunk_meta - - def _write_items_closure(uneven_chunk_meta): - def _write_items(fqn: str, tensor: DTensor) -> List[WriteItem]: - if tensor.to_local().numel() == 0: - # If the tensor is empty, return an empty list - return [] - - return [ - WriteItem( - type=WriteItemType.SHARD, - index=MetadataIndex(fqn, uneven_chunk_meta.offsets), - tensor_data=TensorWriteData( - chunk=uneven_chunk_meta, - properties=TensorProperties.create_from_tensor(tensor.to_local()), - size=tensor.size(), - ), - ) - ] - - return _write_items - # Get uneven chunk metadata for the DTensor # TODO: Optimize gather_and_compute_chunk_metadata synchronization: # 1. Add pre-check validation to verify tensor shape consistency @@ -132,8 +111,7 @@ def _write_items(fqn: str, tensor: DTensor) -> List[WriteItem]: uneven_chunk_meta = gather_and_compute_chunk_metadata(dtensor) # Set the chunk list and write items closure for the DTensor - dtensor._local_tensor.__create_chunk_list__ = _chunk_list_closure([uneven_chunk_meta]) - dtensor._local_tensor.__create_write_items__ = _write_items_closure(uneven_chunk_meta) + _set_chunk_metadata(dtensor, uneven_chunk_meta.offsets, uneven_chunk_meta.sizes) def validate_uneven_dtensor(dtensor: DTensor) -> None: @@ -414,26 +392,92 @@ def split_dtensor( update_uneven_dtensor_chunk_meta: bool = False, ) -> Iterable[DTensor]: """ - Splits a DTensor into smaller DTensors along a specified dimension. + Split a DTensor into smaller DTensors along a specified dimension. + + This function handles uneven sharding correctly by computing chunk metadata + for every output DTensor without extra collective operations. A single call + to ``gather_and_compute_chunk_metadata`` is made upfront; all subsequent + per-split metadata is derived locally from that result using pure integer + arithmetic. + + Unlike the native PyTorch ``DTensor.split``, this function does **not** + redistribute ``Replicate`` placements, which avoids OOM issues when the + full tensor is large. + + Chunk metadata assignment strategy + ------------------------------------ + For each split window ``[split_start, split_end)`` along *dim*: + + * The rank's local chunk covers the global interval + ``[local_start, local_end)`` where:: - This function manages uneven sharding by accurately assigning chunk metadata - for each split. Unlike the native PyTorch DTensor split functionality, - it does not redistribute `Replicate` placements, which helps avoid Out-Of-Memory (OOM) issues. + local_start = chunk_meta.offsets[dim] + local_end = local_start + chunk_meta.sizes[dim] + + * The overlap with the split window is:: + + overlap_start = max(local_start, split_start) + overlap_end = min(local_end, split_end) + + * If ``overlap_start < overlap_end`` the rank owns part of this split: + + - ``new_offsets[dim] = overlap_start - split_start`` + (offset is relative to the split's own global origin) + - ``new_sizes[dim] = overlap_end - overlap_start`` + + * Otherwise the rank owns nothing in this split: + + - ``new_offsets[dim] = 0``, ``new_sizes[dim] = 0`` + + All other dimensions are copied unchanged from the parent's chunk metadata. Args: - dtensor (DTensor): The DTensor to split. - split_size_or_sections (int or list of int): If int, defines the size of each chunk. - If a list, specifies the sizes of each chunk in order. - dim (int, optional): The axis along which to split. Default is 0. - update_uneven_dtensor_chunk_meta (bool, optional): Whether to update chunk - metadata for each resulting DTensor. Default is False. + dtensor (DTensor): The DTensor to split. Must be compatible with + ``gather_and_compute_chunk_metadata`` (placements are ``Shard``, + ``_StridedShard``, or ``Replicate``). + split_size_or_sections (int | list[int]): If an ``int``, each chunk + has this size (the last chunk may be smaller). If a ``list``, + each element is the exact size of the corresponding chunk; the + sizes must sum to ``dtensor.shape[dim]``. + dim (int, optional): Dimension along which to split. Default: ``0``. + update_uneven_dtensor_chunk_meta (bool, optional): If ``True``, call + ``update_uneven_dtensor_chunk_metadata`` for every output DTensor + instead of using the locally-derived metadata. This triggers one + collective per split chunk and is only needed when you require a + full round-trip validation or the parent's chunk metadata cannot + be trusted. Default: ``False``. Yields: - DTensor: Sub-DTensor resulting from the split, maintaining correct metadata. + DTensor: Split sub-DTensors in order. Each yielded DTensor: - Example: - >>> for chunk in split_dtensor(dt, 3, dim=1): - ... print(chunk) + * Shares the same ``placements`` and ``device_mesh`` as *dtensor*. + * Has shape ``dtensor.shape`` with ``shape[dim]`` replaced by the + split-window size. + * Has ``__create_chunk_list__`` / ``__create_write_items__`` set on + its local tensor (via ``_set_chunk_metadata``), so checkpoint + writers can consume it directly without further collectives. + + Raises: + ValueError: If the locally-computed split slice is internally + inconsistent (negative start after offset correction). + + Note: + ``gather_and_compute_chunk_metadata`` is called exactly **once**, + regardless of the number of splits. When + ``update_uneven_dtensor_chunk_meta=False`` (the default) no further + collective operations are performed. + + Example:: + + # Split a [1024, 512] DTensor sharded on dim-0 across 4 ranks + # into four [256, 512] chunks — one collective, zero extra collectives + # for metadata. + for chunk in split_dtensor(dt, split_size_or_sections=256, dim=0): + process(chunk) + + # Variable-size split, chunk metadata derived locally + for chunk in split_dtensor(dt, split_size_or_sections=[300, 300, 424], dim=0): + process(chunk) """ tensor_size = dtensor.shape[dim] @@ -446,25 +490,35 @@ def split_dtensor( for size in split_size_or_sections: split_points.append(split_points[-1] + size) - chunk_meta = gather_and_compute_chunk_metadata(dtensor) + # One collective call — result reused for all splits below. + if not hasattr(dtensor._local_tensor, "__create_chunk_list__"): + chunk_meta = gather_and_compute_chunk_metadata(dtensor) + else: + chunk_meta = dtensor.__create_chunk_list__()[0] chunk_slice = slice(chunk_meta.offsets[dim], chunk_meta.offsets[dim] + chunk_meta.sizes[dim]) local_offset = chunk_meta.offsets[dim] local_tensor = dtensor.to_local() - # Create chunks using manual slicing for i in range(len(split_points) - 1): split_slice = slice(split_points[i], split_points[i + 1]) + + # Compute the intersection of this rank's local chunk with the split + # window. _intersection returns slice(0, 0) when there is no overlap, + # which produces an empty tensor via narrow() — that is the correct + # behaviour for ranks that own nothing in this split window. s = _intersection(split_slice, chunk_slice) if s.start < s.stop: + # Translate to local-tensor coordinates. s = _offset_slice(s, -local_offset) - if s.start < 0 or s.stop < s.start and torch.distributed.get_rank() == 0: + if s.start < 0: raise ValueError( f"Invalid split slice {s} for DTensor with shape {dtensor.shape} " - f"and local offset {local_offset} on dimension {dim}." + f"and local offset {local_offset} on dimension {dim}. " + "This is a bug — please report it." ) - # Slice the local tensor + # Slice the local tensor along `dim`. sliced_tensor = local_tensor.narrow(dim, s.start, s.stop - s.start) out_shape = list(dtensor.shape) out_shape[dim] = split_slice.stop - split_slice.start @@ -478,6 +532,206 @@ def split_dtensor( ) if update_uneven_dtensor_chunk_meta: + # Triggers one collective per split — use only when a validated + # round-trip is required. update_uneven_dtensor_chunk_metadata(new_dtensor) + else: + # Derive chunk metadata locally from the already-computed + # chunk_meta — zero additional collectives. + # + # Compute the intersection in *global* coordinates so that the + # resulting offset is expressed relative to the split's own origin. + global_local_start = chunk_meta.offsets[dim] + global_local_end = global_local_start + chunk_meta.sizes[dim] + global_split_start = split_points[i] + global_split_end = split_points[i + 1] + + overlap_start = max(global_local_start, global_split_start) + overlap_end = min(global_local_end, global_split_end) + + new_offsets = list(chunk_meta.offsets) + new_sizes = list(chunk_meta.sizes) + + if overlap_start < overlap_end: + # This rank owns [overlap_start, overlap_end) of the global + # tensor; the offset within the split chunk is the distance + # from the split's start. + new_offsets[dim] = overlap_start - global_split_start + new_sizes[dim] = overlap_end - overlap_start + else: + # This rank owns nothing in this split window. + new_offsets[dim] = 0 + new_sizes[dim] = 0 + + _set_chunk_metadata(new_dtensor, tuple(new_offsets), tuple(new_sizes)) yield new_dtensor + + +def make_uneven_dtensor( + local_tensor: torch.Tensor, + shape: torch.Size, + dp_mesh: DeviceMesh, + placements: List[Placement], + *, + post_process_uneven: bool = False, + copy_chunk_meta_from: Optional[DTensor] = None, + chunk_metadata: Optional[tuple] = None, +): + """Create a DTensor from a possibly uneven local shard with known global shape. + + Args: + local_tensor: Local shard tensor. + shape: Global shape of the full DTensor. + dp_mesh: 1D device mesh. + placements: DTensor placements (e.g., [Shard(0)]). + post_process_uneven: If True, call ``update_uneven_dtensor_chunk_metadata``. + copy_chunk_meta_from: If set, copy ``__create_chunk_list__`` / + ``__create_write_items__`` from this DTensor. + chunk_metadata: ``(offsets, sizes)`` tuple where *offsets* and *sizes* + are tuples of ints (one per dimension). Sets chunk metadata + closures without collectives. + """ + assert dp_mesh.ndim == 1, "Only 1D mesh is supported for now" + if local_tensor.numel() == 0: + local_shape = (0,) + tuple(shape[1:]) if len(shape) > 1 else (0,) + local_tensor = local_tensor.reshape(local_shape) + else: + local_tensor = local_tensor.view(-1, *shape[1:]) + dtensor = DTensor.from_local( + local_tensor=local_tensor, + device_mesh=dp_mesh, + placements=placements, + run_check=False, + shape=shape, + stride=torch.empty(shape, device="meta").stride(), + ) + if post_process_uneven: + update_uneven_dtensor_chunk_metadata(dtensor) + elif copy_chunk_meta_from is not None: + # This branch is used for the case where we are creating a new DTensor that has the same + # sharding as an existing uneven DTensor, so we can copy the chunk metadata from the + # existing uneven DTensor instead of recomputing it. + copy_chunk_metadata(copy_chunk_meta_from, dtensor) + elif chunk_metadata is not None: + _set_chunk_metadata(dtensor, *chunk_metadata) + return dtensor + + +def get_state_dict( + model: nn.Module, + optimizers: Union[torch.optim.Optimizer, Iterable[torch.optim.Optimizer]], + *, + submodules: Optional[set[nn.Module]] = None, + options: Optional["StateDictOptions"] = None, +) -> tuple[dict[str, "ValueType"], "OptimizerStateType"]: + """Produce model and optimizer state dicts with uneven DTensor preprocessing.""" + for param in model.parameters(): + assert isinstance(param, DTensor), "Expected all parameters to be DTensors" + + model_state_dict, optimizer_state_dict = _get_state_dict( + model=model, optimizers=optimizers, submodules=submodules, options=options + ) + preprocess_state_dict_for_uneven_dtensor(model_state_dict) + preprocess_state_dict_for_uneven_dtensor(optimizer_state_dict) + return model_state_dict, optimizer_state_dict + + +# ------------------------------------------------------------------ +# Chunk metadata helpers (zero-collective) +# ------------------------------------------------------------------ + + +def copy_chunk_metadata(src: DTensor, dst: DTensor) -> None: + """Copy ``__create_chunk_list__`` / ``__create_write_items__`` from *src* to *dst*.""" + dst._local_tensor.__create_chunk_list__ = src._local_tensor.__create_chunk_list__ + dst._local_tensor.__create_write_items__ = src._local_tensor.__create_write_items__ + + +def compute_split_offsets_and_sizes(dist_param, split_dim, comp_idx, total_split, comp_data): + """Compute chunk offsets/sizes for a split component, derived from *dist_param*'s metadata. + + Pure local computation — no collectives. + """ + chunk_list = dist_param._local_tensor.__create_chunk_list__() + orig = chunk_list[0] + global_shape = list(dist_param.size()) + + comp_size = global_shape[split_dim] // total_split + comp_start = comp_idx * comp_size + comp_end = comp_start + comp_size + + offsets = list(orig.offsets) + sizes = list(comp_data.shape) + + o = offsets[split_dim] + s = orig.sizes[split_dim] + + if o < comp_end and o + s > comp_start: + offsets[split_dim] = max(o, comp_start) - comp_start + sizes[split_dim] = min(o + s, comp_end) - max(o, comp_start) + else: + offsets[split_dim] = 0 + sizes[split_dim] = 0 + + return tuple(offsets), tuple(sizes) + + +def get_fsdp_slice_from_uneven_dtensor(dist_param: DTensor) -> slice: + """Compute the FSDP slice (flattened range) from a v2 DTensor. + + Uses the uneven chunk metadata (``__create_chunk_list__``) attached by + ``update_uneven_dtensor_chunk_metadata`` to correctly handle uneven + sharding where ranks may own different-sized slices. + + The DTensor must have ``__create_chunk_list__`` set on its local tensor + (via ``preprocess_state_dict_for_uneven_dtensor`` or + ``update_uneven_dtensor_chunk_metadata``) before calling this function. + """ + local_numel = dist_param._local_tensor.numel() + if local_numel == 0: + return slice(0, 0) + + assert hasattr(dist_param._local_tensor, "__create_chunk_list__"), ( + "get_fsdp_slice_from_uneven_dtensor requires the DTensor to have " + "__create_chunk_list__ metadata. Call update_uneven_dtensor_chunk_metadata " + "or preprocess_state_dict_for_uneven_dtensor first." + ) + + chunk_list = dist_param._local_tensor.__create_chunk_list__() + assert len(chunk_list) == 1, f"Expected exactly one chunk per rank, got {len(chunk_list)}" + chunk_meta = chunk_list[0] + offsets = chunk_meta.offsets + sizes = chunk_meta.sizes + + global_shape = dist_param.size() + strides = torch.empty(global_shape, device="meta").stride() + + start = sum(o * s for o, s in zip(offsets, strides)) + return slice(start, start + local_numel) + + +def _set_chunk_metadata(dtensor: DTensor, offsets: tuple, sizes: tuple) -> None: + """Set ``__create_chunk_list__`` / ``__create_write_items__`` closures on *dtensor*. + + No collective ops — *offsets* and *sizes* are computed locally. + """ + chunk_meta = ChunkStorageMetadata(offsets=tuple(offsets), sizes=tuple(sizes)) + + def _write_items(fqn: str, tensor: DTensor) -> list: + if tensor.to_local().numel() == 0: + return [] + return [ + WriteItem( + type=WriteItemType.SHARD, + index=MetadataIndex(fqn, chunk_meta.offsets), + tensor_data=TensorWriteData( + chunk=chunk_meta, + properties=TensorProperties.create_from_tensor(tensor.to_local()), + size=tensor.size(), + ), + ) + ] + + dtensor._local_tensor.__create_chunk_list__ = lambda: [chunk_meta] + dtensor._local_tensor.__create_write_items__ = _write_items diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/README.md b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/README.md new file mode 100644 index 00000000000..f3a9e90e05d --- /dev/null +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/README.md @@ -0,0 +1,133 @@ +# Megatron FSDP2 (fully_shard_v2) + +This directory contains the **fully_shard_v2** implementation — a PyTorch FSDP2-compatible API layer for Megatron Core. + +## Architecture + +``` +v2/ +├── README.md # This file +├── __init__.py # Public exports +├── fully_shard.py # Public fully_shard() API +├── fsdp_module.py # FSDPModule runtime state and methods +├── hooks.py # Forward/backward hook registration +├── param_group.py # ParameterGroup — groups params with shared buffers +├── dp_buffer.py # DataParallelBuffer — flat buffer management +├── allocator.py # BucketAllocator implementations +├── utils.py # Internal utility functions +└── design.md # Detailed design documentation (overlap, memory, sync) +``` + +For the broader `megatron_fsdp` package layout, see the parent directory. + +## Key Concepts + +### fully_shard(module, mesh=...) + +Wraps a module with FSDP sharding semantics: + +1. Converts the module class to `FSDPModule` dynamically +2. Groups parameters by (device, dtype, requires_grad) +3. Creates `ParameterGroup` for each group with dedicated buffers +4. Registers forward/backward hooks for unshard/reshard/reduce +5. Replaces module parameters with `DTensor` representations + +### FSDPModule + +Mixin class added to wrapped modules. Methods: + +| Method | When Called | Purpose | +|--------|-------------|---------| +| `unshard()` | Pre-forward | All-gather params from sharded buffer | +| `reshard()` | Post-forward, post-backward | Release unsharded buffer | +| `reduce_grad()` | Post-backward | All-reduce or reduce-scatter gradients | + +### DataParallelBuffer + +Flat buffer managing (a shard of) parameter/gradient data: + +- `unshard()` — all-gather to full tensor +- `reshard()` — free temporary buffer +- `reduce_grad()` — all-reduce or reduce-scatter gradients +- Uses `BufferIndex` to track parameter layout within the buffer + +### ParameterGroup + +Groups parameters sharing the same (device, dtype, requires_grad): + +- `model_weight_buffer` — stores sharded model weights +- `main_weight_buffer` — optional high-precision copy +- `main_grad_buffer` — accumulates gradients before reduce +- `dist_params` — DTensor views into the buffer + +### Uneven DTensor Handling + +See the parent directory `..` for `uneven_dtensor.py` which provides: + +- `gather_and_compute_chunk_metadata()` — computes global offsets/sizes for each shard +- `update_uneven_dtensor_chunk_metadata()` — attaches chunk metadata to DTensor for checkpointing +- `preprocess_state_dict_for_uneven_dtensor()` — enables distributed checkpoint of uneven shards +- `split_dtensor()` — splits DTensor with proper chunk metadata preservation +- `get_state_dict()` — PyTorch DCP-compatible state dict with uneven shard support + +## Sharding Strategies + +| Strategy | Shard Weights | Shard Gradients | Status | Notes | +|----------|---------------|-----------------|--------|-------| +| `optim_grads_params` | Yes | Yes | **Supported** | Like ZeRO-3: full parameter/gradient/optimizer sharding | +| `no_shard` | No | No | **Not yet supported** | Like DDP: no sharding | +| `optim` | No | No | **Not yet supported** | Like ZeRO-1: shard optimizer states only | +| `optim_grads` | No | Yes | **Not yet supported** | Like ZeRO-2: shard optimizer states + gradients | + +> **FIXME:** `no_shard`, `optim`, and `optim_grads` sharding strategies are not yet supported in v2. +> Currently only `optim_grads_params` is fully implemented and tested. +> These strategies will be added in a follow-up change. + +## Integration with Megatron + +`FullyShardedDataParallel` in `mcore_fsdp_adapter.py` supports two code paths: + +1. **Legacy path** (`use_fully_shard_api=False`): Uses `MegatronFSDP` from the original megatron_fsdp module +2. **fully_shard path** (`use_fully_shard_api=True`): Uses PyTorch or Megatron's `fully_shard_v2` + +### Using fully_shard_v2 + +```python +from megatron.core.distributed.fsdp.mcore_fsdp_adapter import FullyShardedDataParallel + +# In your config: +ddp_config.data_parallel_sharding_strategy = "optim_grads_params" +ddp_config.use_fully_shard_api = True + +# During model setup: +model = FullyShardedDataParallel( + config=transformer_config, + ddp_config=ddp_config, + module=model, +) +``` + +## Toy Example + +See `examples/megatron_fsdp/fsdp_toy.py` for a standalone example showing: + +- Basic model wrapping with `fully_shard()` +- Training loop with gradient accumulation +- Distributed checkpointing with `torch.distributed.checkpoint` + +```bash +torchrun --nproc_per_node=2 examples/megatron_fsdp/fsdp_toy.py \ + --model-dim 512 --n-layers 2 --batch-size 4 +``` + +## Gotchas / Pitfalls + +- **Zero-numel gradient shards and fused optimizers.** When a parameter's local shard is empty on some DP ranks (e.g., small biases on high DP counts), creating a `DTensor` gradient with `numel() == 0` and passing it to fused multi-tensor optimizers (TE `FusedAdam`) can silently corrupt updates for neighboring non-empty parameters. This manifests only as convergence divergence with no error — see [design.md § Pitfall](design.md) for details and the fix in `param_group.py`. +- **Temporary communication buckets must record their CUDA stream.** `DataParallelBuffer.reshard()` and `release_grad_buffer()` can return temporary all-gather / reduce-scatter bucket storage to the allocator while side-stream CUDA work is still queued. A CUDA event wait orders streams, but it does not tell PyTorch's caching allocator that the freed storage is still used by `ag_stream` or `rs_stream`. If the block is recycled on another stream, the communication payload can be silently corrupted. Call `record_stream()` on any temporary CUDA tensor that participates in side-stream communication before it may be freed; a full synchronization only masks the bug and hurts overlap. + +## Unit Tests + +```bash +# Run FSDP2 API tests +pytest -xvs tests/unit_tests/distributed/megatron_fsdp/test_mcore_fsdp_fully_shard_v2_api.py +``` diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/__init__.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/__init__.py new file mode 100644 index 00000000000..21879c4717a --- /dev/null +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/__init__.py @@ -0,0 +1,61 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# FIXME: The following sharding strategies are not yet supported in FSDP v2: +# - no_shard +# - optim +# - optim_grads +# Currently only optim_grads_params is fully implemented and tested. +# We will add support for these strategies in a follow-up change. +# When using ParameterGroup directly or via fully_shard(), passing an unsupported +# sharding_strategy will raise NotImplementedError. See README.md for details. + +from ..uneven_dtensor import ( + gather_and_compute_chunk_metadata, + get_state_dict, + make_uneven_dtensor, + preprocess_state_dict_for_uneven_dtensor, + redistribute_uneven_dtensor_to_replicated, + split_dtensor, + uneven_dtensor_to_full_tensor, +) +from .allocator import Bucket, TemporaryBucketAllocator +from .dp_buffer import BufferIndex, DataParallelBuffer +from .fully_shard import FSDPModule, fully_shard +from .mixed_precision import ( + FullyShardFP8Policy, + FullyShardMixedPrecisionPolicy, + FullyShardNVFP4Policy, +) +from .param_group import ParameterGroup + +__all__ = [ + "FSDPModule", + "fully_shard", + "FullyShardFP8Policy", + "FullyShardMixedPrecisionPolicy", + "FullyShardNVFP4Policy", + "ParameterGroup", + "BufferIndex", + "DataParallelBuffer", + "Bucket", + "TemporaryBucketAllocator", + "make_uneven_dtensor", + "get_state_dict", + "preprocess_state_dict_for_uneven_dtensor", + "gather_and_compute_chunk_metadata", + "split_dtensor", + "uneven_dtensor_to_full_tensor", + "redistribute_uneven_dtensor_to_replicated", +] diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/design.md b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/design.md new file mode 100644 index 00000000000..0a6c53dd3bc --- /dev/null +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/design.md @@ -0,0 +1,780 @@ +# Design: Megatron-FSDP2 Implementation + +--- + +## File Map + +| File | Role in overlap | +|---|---| +| `fully_shard.py` | Public `fully_shard()` API and allocator selection | +| `fsdp_module.py` | `FSDPModule`, `_FSDPRootContext`, `_FSDPState`, `unshard()`, `reshard()`, `reduce_grad()` | +| `hooks.py` | Forward/backward hook registration and final callback | +| `param_group.py` | `ParameterGroup.unshard(async_op)`, `reduce_grad()`, `release_grad_buffer()`, `_init_buffers()` (memory optimization) | +| `dp_buffer.py` | `DataParallelBuffer.unshard(async_op)` (all-gather + `p.data` rebind), `reduce_grad()` (reduce-scatter + shard accumulation) | +| `allocator.py` | `BucketAllocator` hierarchy: `TemporaryBucketAllocator`, `StorageFreeingBucketAllocator`, `TracePoolAllocator` — pooled memory for unsharded parameter and gradient buffers | +| `mcore_fsdp_adapter.py` | `FullyShardedDataParallel.stop_communication()` — synchronizes ag_stream and rs_stream into main stream | + +No changes to `utils.py` are needed for either overlap feature. + +--- + +## `_FSDPRootContext` — Shared Coordination Object + +One instance is created by the root `FSDPModule` at `_init_fsdp_state()` time and stored as +`_fsdp_root_context` on **every** FSDP module (root and all children). + +```python +@dataclass +class _FSDPRootContext: + # --- CUDA streams --- + ag_stream: torch.cuda.Stream # all-gather (unshard) side stream + rs_stream: torch.cuda.Stream # reduce-scatter side stream + # When the corresponding feature flag is False, these are set to + # torch.cuda.current_stream() so stream-context switches become no-ops. + + # --- Temporary bucket allocator --- + bucket_allocator: BucketAllocator + # One allocator handles all temporary buckets. Allocation keys include + # both parameter-group identity and buffer role. + + # --- Static execution order (set at init, never mutated) --- + forward_order: List[FSDPModule] + # Populated as: [m for m in root.modules() if isinstance(m, FSDPModule)] + + # --- Unshard prefetch tracking --- + unshard_done_events: Dict[int, Optional[torch.cuda.Event]] + # module_id -> Event: signals when that module's all-gather is complete. + # None means "not yet launched" or "already consumed by a wait". + + # --- Reduce-grad overlap tracking --- + reduce_grad_buckets: Dict[int, List[Tuple[torch.cuda.Event, ParameterGroup]]] + # module_id -> [(event, param_group), ...] + # Each entry: event signals RS complete; param_group holds the grad buffer. + + # --- Feature flags --- + enable_unshard_prefetch: bool + enable_async_reduce_grad: bool + + # --- Activation recompute support --- + backward_phase: bool = False + # True from the root backward pre-hook until the final callback. + + backward_module: Optional[int] = None + # ``id(module)`` of the FSDP module whose backward is pending next. + # Derived from ``_reversed_order`` and ``backward_done_modules`` — NOT + # set by any hook directly. Updated by ``_advance_backward_module()``. + + backward_done_modules: set = field(default_factory=set) + # Set of ``id(module)`` for FSDP modules whose backward has completed. + # Populated in ``post_backward``, cleared in the root backward pre-hook. + + _reversed_order: List[FSDPModule] = field(default_factory=list) + # ``list(reversed(forward_order))`` — precomputed backward processing order. + + def _advance_backward_module(self) -> None: + """Set ``backward_module`` to the first module in ``_reversed_order`` + that is NOT in ``backward_done_modules``.""" + for m in self._reversed_order: + if id(m) not in self.backward_done_modules: + self.backward_module = id(m) + return + self.backward_module = None +``` + +### Initialization in `_init_fsdp_state()` + +```python +bucket_allocator = TracePoolAllocator() if enable_trace_pool else StorageFreeingBucketAllocator() +module._init_named_param_groups(..., bucket_allocator=bucket_allocator) + +forward_order = [child for child in self.modules() if isinstance(child, FSDPModule)] +root_context = _FSDPRootContext( + ag_stream=torch.cuda.Stream() if enable_unshard_prefetch else torch.cuda.current_stream(), + rs_stream=torch.cuda.Stream() if enable_async_reduce_grad else torch.cuda.current_stream(), + bucket_allocator=bucket_allocator, + forward_order=forward_order, + reduce_grad_buckets={id(m): [] for m in forward_order}, + unshard_done_events={id(m): None for m in forward_order}, + enable_unshard_prefetch=enable_unshard_prefetch, + enable_async_reduce_grad=enable_async_reduce_grad, +) +# Root and children share one context and one bucket allocator: +for module in forward_order: + for param_group in module._fsdp_param_groups: + param_group.set_allocator(root_context.bucket_allocator) +for child in self.modules(): + if child is not self and isinstance(child, FSDPModule): + child._fsdp_state._is_root = False + setattr(child, "_fsdp_root_context", root_context) +``` + +`forward_order` is **static** (module tree topology, computed once). There is no first-pass +dynamic recording phase. + +**Safety constraint.** `_init_fsdp_state()` must be called **before** any forward/backward pass +runs. The method includes a runtime guard that rejects re-initialization if any child +FSDPModule is still unsharded (`unshard_done_events` live) or has pending reduce-scatter +operations (`reduce_grad_buckets` non-empty). Violating this constraint would overwrite a +running module's `_fsdp_root_context` while its hooks are still firing, causing undefined +behavior. + +--- + +## Feature 1: Unshard Prefetch + +### Hook entry points + +```python +# _register_forward_pre_hook: +module.unshard(async_op=ctx.enable_unshard_prefetch, bwd_pass=False) + +# _register_backward_pre_hook (called inside register_multi_grad_hook): +module.unshard(async_op=ctx.enable_unshard_prefetch, bwd_pass=True) +``` + +### `FSDPModule.unshard(async_op, bwd_pass)` + +```python +stream = ctx.ag_stream if async_op else torch.cuda.current_stream() + +# *** Critical: synchronize ag_stream with current_stream before launching AG *** +# This ensures main-stream writes to parameter data (e.g. reshard after forward, +# or tensor-parallel slice writes) are visible before the all-gather reads them. +# Without this barrier, stale or partially-written parameter shards may be +# gathered, causing convergence divergence. +if async_op: + stream.wait_stream(torch.cuda.current_stream()) + +# Build the work list: self + (optionally) next module to prefetch +if async_op: + prefetch = _get_prefetch_next_modules(bwd_pass) # returns [next_module] or [] +else: + prefetch = [] + +for module in [self] + prefetch: + if all(pg.has_unsharded_weight_buffers(bwd_pass=bwd_pass) for pg in module._fsdp_param_groups): + continue # Required buffers are already unsharded — skip + + with torch.cuda.stream(stream): + for _, param_group in module._named_param_groups: + param_group.unshard() + # → DataParallelBuffer.unshard(): + # allocate unsharded bucket, launch all_gather_into_tensor, + # rebind p.data → unsharded buffer slice (even before AG done!) + # NOTE: async_op is NOT passed; the stream context handles dispatch. + + if async_op: + event = stream.record_event() + ctx.unshard_done_events[id(module)] = event # store completion signal + +# Synchronize self: block main stream until this module's AG is done. +# The event is NOT cleared here; it is only cleared by reshard(). Buffer +# readiness, not event presence, determines whether a later unshard can skip. +if ctx.unshard_done_events[id(self)] is not None: + ctx.unshard_done_events[id(self)].wait() # main stream waits on ag_stream event + +# Install full parameter tensors into the nn.Module (safe after event.wait) +for param_names, param_group in self._named_param_groups: + for name, param in zip(param_names, param_group.params): + _replace_module_parameter(self, name, param) # swaps nn.Parameter object +``` + +**Important: `p.data` rebind race.** +`DataParallelBuffer.unshard()` rebinds `p.data` to the unsharded buffer slice +**inside the `with torch.cuda.stream(stream)` block**, before the all-gather completes (when +`async_op=True`). The memory is already allocated and the slice indices are correct; only the +NCCL fill is in-flight. The outer `unshard()` guards correctness by calling `event.wait()` +before calling `_replace_module_parameter`, so the module's parameters are safe to read by +the time the forward kernel uses them. + +**Stream ordering barrier.** When `async_op=True`, `unshard()` inserts a +`stream.wait_stream(torch.cuda.current_stream())` on `ag_stream` before launching any +all-gather. This ensures that writes performed on the main stream (e.g., reshard after a +previous forward, or tensor-parallel slice updates) are fully visible to the all-gather +kernel. Without this barrier, stale or partially-written parameter shards may be read by +the NCCL collective, causing convergence divergence. + +**NVTX profiling.** `unshard()`, `reshard()`, and `reduce_grad()` each push/pop a +`torch.cuda.nvtx` range (`"MFSDP unshard"`, `"MFSDP reshard"`, `"MFSDP reduce_grad"`) +for profiling visibility in tools like Nsight Systems. + +Prefetched modules' data also becomes valid when their own pre-hook later calls `event.wait()` +for them. If a module's pre-hook arrives and its event is already set (prefetch was launched +by the previous module), it just waits on the event and skips re-launching the AG. + +### `_get_prefetch_next_modules(bwd_pass)` + +```python +order = list(reversed(ctx.forward_order)) if bwd_pass else ctx.forward_order +i = order.index(self) +return [order[i + 1]] if i + 1 < len(order) else [] +``` + +Exactly one module is prefetched per step. Multi-module lookahead is a future extension. + +### `FSDPModule.reshard()` + +```python +for param_names, param_group in self._named_param_groups: + param_group.reshard() # → DataParallelBuffer.reshard() + # frees TemporaryBucketAllocator bucket + # sets _unsharded_buffer = None + for name, dist_param in zip(param_names, param_group.dist_params): + _replace_module_parameter(self, name, dist_param) # reinstall sharded DTensor +ctx.unshard_done_events[id(self)] = None # reset so next iteration can prefetch again +``` + +--- + +## Feature 2: Reduce-Grad Overlap + +### Hook entry point + +Inside the `post_backward` closure registered by `_register_backward_hook`: + +```python +module.reshard() +module.reduce_grad(async_op=ctx.enable_async_reduce_grad) +module.post_backward_issued = True +``` + +### `FSDPModule.reduce_grad(async_op)` + +```python +stream = ctx.rs_stream if async_op else torch.cuda.current_stream() + +# --- Step 1: Sliding drain — free grad buffers 2 positions back in backward order --- +if async_op: + backward_order = list(reversed(ctx.forward_order)) + for i, module in enumerate(backward_order): + if i - 2 >= 0: + for event, param_group in drain(ctx.reduce_grad_buckets[id(backward_order[i-2])]): + event.wait() + param_group.release_grad_buffer() + # → deletes param.main_grad views (prevents TE grad-accum-fusion leak) + # → DataParallelBuffer.reshard() (frees unsharded grad bucket) + if module is self: break + +# --- Step 2: Copy .grad → main_grad_buffer (on main stream, fast memcpy) --- +for param_names, param_group in self._named_param_groups: + if not param_group.requires_grad: continue + + for name, param in zip(param_names, param_group.params): + main_grad = param.get_main_grad() + if param.grad is None: + if not getattr(param, 'grad_added_to_main_grad', False): + main_grad.zero_() # no TE fusion: zero the slot + else: + main_grad.copy_(param.grad.detach()) # normal backward: copy .grad + del param.grad + + # --- Step 3: Reduce-scatter on rs_stream --- + if async_op: + stream.wait_stream(torch.cuda.current_stream()) # ensure .grad copy is visible to rs_stream + with torch.cuda.stream(stream): + param_group.reduce_grad() + # → DataParallelBuffer.reduce_grad() (synchronous within this stream): + # fetch_unsharded_buffer() allocates full grad buffer + # reduce_scatter_tensor(output=grad_shard, input=full_grad) + # self.data[local_idx:...] += grad_shard + event = stream.record_event() + ctx.reduce_grad_buckets[id(self)].append((event, param_group)) + # param_group.release_grad_buffer() is NOT called here; deferred until drain/final CB + else: + param_group.reduce_grad() + param_group.release_grad_buffer() + + # --- Step 4: Install dist_grad on dist_param (runs in stream context) --- + for name, param, dist_param, dist_grad in zip( + param_names, param_group.params, param_group.dist_params, param_group.dist_grads + ): + if param.requires_grad and dist_grad is not None: + with torch.cuda.stream(stream): + dist_grad = dist_grad.to(dist_param.dtype) # dtype cast on rs_stream + setattr(dist_param, "grad", dist_grad) # Python ref, no GPU dependency +``` + +**Key design point — `DataParallelBuffer.reduce_grad()` has no `async_op` parameter.** +The operation is inherently synchronous *within whatever stream is current* when called. The +"async" behavior is achieved entirely by the caller dispatching into `rs_stream` via +`with torch.cuda.stream(stream)`. This avoids any API changes to `DataParallelBuffer`. + +**`grad_added_to_main_grad` and `overwrite_main_grad` flags:** +When TransformerEngine's `gradient_accumulation_fusion` is active, the backward kernel writes +directly into `param.main_grad` (bypassing `.grad`). Two flags coordinate this: + +- **`grad_added_to_main_grad`**: Set to `False` in `pre_backward_hook` before each backward + pass; the kernel sets it to `True` after writing. In `reduce_grad`, the `zero_()` call is + skipped when `True` to preserve the fused-gradient value. + +- **`overwrite_main_grad`**: Set to `True` in `pre_backward_hook` for sharded parameters + (`optim_grads_params` / `optim_grads`). By default TE **accumulates** (adds) into + `main_grad` — useful for micro-batch gradient accumulation in non-FSDP settings. With FSDP + the gradient buffer is re-used across micro-batches; accumulation would silently **double** + gradients and produce NaN after the first step. This flag tells TE to **overwrite** instead + of accumulate. + +### Sliding Drain: The `i-2` Rule + +The drain loop ensures at most **2 modules' gradient buffers** are live at any time: + +``` +Backward processing order (reversed forward): + layer[N] ← current (i=0): i-2=-2 → no drain + layer[N-1] ← current (i=1): i-2=-1 → no drain + layer[N-2] ← current (i=2): i-2=0 → drain layer[N] (i-2=0) + layer[N-3] ← current (i=3): i-2=1 → drain layer[N-1] (i-2=1) + ... +``` + +By the time RS for `layer[N-2]` starts, `layer[N]`'s RS event is expected to be done +(two backward steps of compute have elapsed). `event.wait()` makes this explicit and safe +even if the timing estimate is wrong. + +### `_post_backward_final_callback` + +Registered on the root by `_register_post_backward_final_callback()` via +`Variable._execution_engine.queue_callback`. Fires after all autograd ops complete. + +```python +def _post_backward_final_callback(root_state, root_module): + ctx = root_module._fsdp_root_context + stream = ctx.rs_stream + + # Handle modules whose post_backward hook was never triggered + # (e.g. modules with no grad-requiring inputs on this micro-batch) + for module in reversed(ctx.forward_order): + if module.post_backward_issued: + continue + module.reshard() + module.reduce_grad(async_op=ctx.enable_async_reduce_grad) + + # Drain ALL remaining buckets (anything not drained by the sliding rule above) + for buckets in ctx.reduce_grad_buckets.values(): + while buckets: + event, param_group = buckets.pop() + event.wait() + param_group.release_grad_buffer() + + # Ensure main stream sees all rs_stream work before optimizer step + torch.cuda.current_stream().wait_stream(stream) + + root_state._post_backward_callback_queued = False +``` + +--- + +## Feature 3: Activation Recomputation (Gradient Checkpointing) + +### Problem + +When activation checkpointing re-runs a forward pass during backward, the FSDP +forward hooks fire again. Without mitigation this causes two problems: + +1. **Redundant all-gather**: `forward_pre_hook` → `unshard()` launches a second + all-gather even though parameters are already unsharded. +2. **Premature reshard**: `forward_hook` → `reshard()` releases the unsharded + parameter buffer before backward gradient computation has consumed it. + +The baseline Megatron-FSDP addresses this by switching submodules into a +pre-backward mode before backprop (`megatron_fsdp.py:900-938`). + +### Solution Overview + +Two mechanisms: + +| Mechanism | Effect | +|---|---| +| **Derived `backward_module`** | `_advance_backward_module()` scans `_reversed_order` for the first module **not** in `backward_done_modules`. This identifies the pending module even when activation recompute fires **before** any layer's `pre_backward_hook` (which is always the case — the checkpoint wrapper triggers recompute, then backward flows through the recomputed graph). | +| **Buffer readiness check** | `has_unsharded_weight_buffers(bwd_pass=...)` skips redundant all-gathers for the buffers required by the current forward/backward phase. | + +The `backward_phase` flag gates the forward post-hook check; `backward_done_modules` +drives both the derived pointer and the prefetch guard. + +### Hook Entry Points + +```python +# _register_forward_hook → reshard_param_groups: +if ctx.backward_phase and id(module) == ctx.backward_module: + return # skip reshard — this is the pending module + +# _register_backward_pre_hook → pre_backward_hook (root only): +ctx.backward_done_modules.clear() +ctx.backward_phase = True +ctx._advance_backward_module() # picks first non-done in _reversed_order + +# _register_backward_hook → post_backward: +ctx.backward_done_modules.add(id(module)) +ctx._advance_backward_module() # advances to next pending module +module.reshard() + +# _register_post_backward_final_callback: +ctx.backward_phase = False +ctx.backward_module = None +ctx.backward_done_modules.clear() +``` + +### Prefetch Constraint + +During backward, `unshard(bwd_pass=True)` prefetches the next module in +`_reversed_order`. An extra guard skips modules whose backward is already done: + +```python +# fsdp_module.py — unshard() +if bwd_pass and id(module) in ctx.backward_done_modules: + continue # backward already done — skip prefetch +``` + +### Timeline + +Consider two FSDP-wrapped layers L1, L2 checkpointed together. +`forward_order = [root, L1, L2]`, `_reversed_order = [L2, L1, root]`. + +``` +----- FORWARD (normal) ---------------------------------- +L1: pre → unshard(L1) → forward → reshard(L1) +L2: pre → unshard(L2) → forward → reshard(L2) + (checkpoint drops intermediates) + +----- BACKWARD (root enters phase) ---------------------- +root pre_backward: + clear done_modules, backward_phase = True + _advance → backward_module = L2 (first not done) + unshard(root) + +----- ACTIVATION RECOMPUTE (L1→L2, inside checkpoint backward) -- +L1 pre → unshard(L1) → forward +L1 post: L1 ≠ backward_module(L2) → reshard(L1) +L2 pre → unshard(L2) (event[L2] set, persistent) +L2 post: L2 == backward_module → skip reshard + +----- L2 BACKWARD ---------------------------------------- +L2 pre_backward → unshard(L2) (event set → skip) +L2 backward compute +L2 post_backward: + done_modules.add(L2), _advance → L1, reshard + +----- L1 BACKWARD ---------------------------------------- +L1 pre_backward → unshard(L1) (re-allocates, all-gathers) +L1 backward (gradients already computed → copies .grad) +L1 post_backward: + done_modules.add(L1), _advance → root, reshard + +----- FINAL CALLBACK -------------------------------------- +backward_phase = False +backward_module = None +done_modules.clear() +``` + +### Key Design Decisions + +1. **`backward_module` is derived, not set by hooks.** Activation recompute + always fires before any layer's `pre_backward_hook`. Deriving from the done + set + `_reversed_order` correctly identifies the pending module regardless + of timing. + +2. **`_advance_backward_module()` is called at exactly two points:** root + `pre_backward_hook` (after clearing the done set) and `post_backward` + (after adding a done module). These are the only mutations to `backward_done_modules`. + +3. **`backward_done_modules` serves dual purpose:** drives the derived pointer + AND gates the prefetch guard in `unshard()`. + +4. **Event persists between `unshard()` and `reshard()`.** `unshard()` no + longer clears its own event. Prevents redundant all-gathers. + +### Edge Cases + +- **Sync mode (`enable_unshard_prefetch=False`):** No event is recorded, + so the persistent-event mechanism does not apply. `backward_module` still + prevents premature resharding. +- **Module not reached by backward:** The final callback runs `reshard()` + for untouched modules. +- **Multiple micro-batches:** All state is reset in the final callback. + +--- + +## Complete Timeline + +``` +FORWARD PASS (enable_unshard_prefetch=True) +--------------------------------------------------------- +main stream: |← compute L[0] →|← compute L[1] →|← compute L[2] →| +ag_stream: |AG(L[0]) AG(L[1])| AG(L[2])| | + +pre-hook L[0]: async unshard L[0] + prefetch L[1] on ag_stream + event[L[0]].wait() → main stream unblocks + _replace_module_parameter(L[0]) + +pre-hook L[1]: event[L[1]] already set → wait (likely done) + _replace_module_parameter(L[1]) + async prefetch L[2] on ag_stream + +pre-hook L[2]: event[L[2]].wait() → main stream unblocks + _replace_module_parameter(L[2]) + +BACKWARD PASS (enable_async_reduce_grad=True) +--------------------------------------------------------- +main stream: |bwd L[2]|copy grad[2]|bwd L[1]|copy grad[1]|bwd L[0]|copy grad[0]| +ag_stream: |AG(L[1]) prefetch |AG(L[0]) prefetch | | +rs_stream: | RS(L[2]) ------| RS(L[1]) ------| RS(L[0])---| + +post-bwd L[2]: reshard, copy grad[2]→main_grad, rs_stream.wait(main), RS(L[2]), event[2] +post-bwd L[1]: drain event[2-2]? (i=1, no drain), copy grad[1], RS(L[1]), event[1] +post-bwd L[0]: drain event[L[2]] (i=2, drain backward_order[0]=L[2]), RS(L[0]), event[0] + +final_callback: + drain event[L[1]], event[L[0]] + main_stream.wait_stream(rs_stream) + ← optimizer step safe → +``` + +--- + +## `DataParallelBuffer.reduce_grad()` — Implementation Note + +No `async_op` parameter is needed. The method is purely synchronous within the calling stream: + +```python +def reduce_grad(self): + if not self.is_distributed: + torch.distributed.all_reduce(self.data, group=self.dp_group) + return + full_grad = self.fetch_unsharded_buffer() # allocates bucket if not present + sm = self.buffer_index.shard_meta + grad_shard = full_grad[sm.bucket_data_index : sm.bucket_data_index + sm.size] + torch.distributed.reduce_scatter_tensor( + output=grad_shard, input=full_grad, group=self.dp_group + ) + # Accumulate into persistent shard — supports multi-micro-batch grad accumulation + self.data[sm.local_data_index : sm.local_data_index + sm.size] += grad_shard +``` + +The caller (`FSDPModule.reduce_grad`) provides the stream context; `DataParallelBuffer` +just does the collective. This clean separation means `DataParallelBuffer` requires no +modifications for the overlap feature. + +--- + +## Pitfall: Zero-Numel Gradient Shards and Fused Optimizers + +**Problem.** When a parameter is sharded across DP ranks, its local shard on a given rank +may contain **zero elements** (e.g., a small bias or embedding table on a high-DP-count setup). +Materializing a `DTensor` gradient for such a shard creates a tensor with `numel() == 0`. + +Fused multi-tensor optimizers (e.g., TransformerEngine `FusedAdam`) operate on **all** +gradients in a parameter group in a single fused kernel launch. Passing a zero-numel +tensor into these fused ops can silently corrupt the weight updates for **neighboring +non-empty parameters** in the same group. The optimizer does not crash or raise an error +— it produces numerically incorrect steps that manifest only as **convergence divergence**, +making this extremely difficult to attribute and debug. + +**Symptoms (hard to diagnose):** +- Training loss diverges or fails to converge despite correct hyperparameters +- No NaN or Inf in gradients — the corruption is a numerical perturbation +- Occurs only at certain DP-world-size / model-size combinations where sharding produces empty local slices +- Bisecting the codebase is unhelpful because the optimizer runs without error + +**Fix in `param_group.py`:** +```python +# DO NOT REMOVE THIS CHECK: +if p.requires_grad and grad_data.numel() > 0: + self.dist_grads.append(make_uneven_dtensor(...)) +else: + self.dist_grads.append(None) # zero-numel shard → no DTensor grad +``` + +By recording `None` for zero-numel shards instead of a DTensor with an empty local tensor, +the fused optimizer never receives the empty tensor and cannot corrupt neighboring updates. +The optimizer already handles `None` grads correctly (parameters without a grad are +simply skipped during the fused update). + +**Additional safeguard in `_scale_gradients`:** +```python +for dist_grad in param_group.dist_grads: + if dist_grad is None: + continue # skip zero-numel shards + dist_grad._local_tensor.mul_(scaling_factor) +``` + +--- + +## Pitfall: Attribute Propagation from Original Params to DTensor Dist Params + +**Problem.** `_init_buffers()` in `ParameterGroup` creates DTensor views (`dist_params`) into +sharded buffers and `_replace_module_parameter` registers these DTensors on the module. +However, critical metadata set on the **original** `nn.Parameter` objects by upstream layers +(e.g. TE linear layers from `transformer_engine.py`) is **not** automatically transferred to +the new DTensor wrappers. + +The adapter (`mcore_fsdp_adapter.py:310-330`) copies a fixed list of attributes from original +params to dist_params. If an attribute is missing from this list, downstream consumers that +inspect the registered module parameters will see the wrong metadata. + +**Affected attributes and their consumers:** + +| Attribute | Set by | Consumer | Failure mode | +|-----------|--------|----------|-------------| +| `allreduce` | `transformer_engine.py:841` — set to `False` on expert MLP weights | `_get_param_groups` (`optimizer/__init__.py:348`) — classifies `is_expert_parallel` | Expert params misclassified as non-expert, causing wrong gradient scaling, clipping group assignment, and optimizer partition placement | +| `is_embedding_parameter` | Various embedding layers | `_get_param_groups` — controls weight decay exclusion | Embeddings incorrectly decayed → convergence divergence | +| `is_embedding_or_output_parameter` | Embedding / output layers | Same as above | Same | +| `sequence_parallel` | TE layers | `parallel_state` / loss computation | Incorrect SP semantics | +| `tensor_model_parallel`, `partition_dim`, `partition_stride` | TE layers | Distributed checkpointing / state dict | Incorrect checkpoint sharding | +| `requires_grad` | All layers | Optimizer | Frozen params may receive updates | + +**Fix.** When adding a new metadata attribute to TE layers or custom modules that are +consumed by downstream code (optimizer, checkpointing, mixed precision), add its name to +the `attr_name` list in `mcore_fsdp_adapter.py` to ensure it propagates to the DTensor +dist_params. + +**Debugging.** Misattributed params can be detected by dumping +`model._log_parameter_groups()` output and verifying that expert params appear in the +`is_expert_parallel` group. NaN after a single step with `gradient_accumulation_fusion` +is a strong indicator of missing `allreduce` propagation. + +--- + +## Memory Optimization: Freeing Original Parameter Storage + +After `ParameterGroup._init_buffers()` copies parameter data into the internal weight buffers +(`model_weight_buffer` and optionally `main_weight_buffer`), the original full parameter tensors +are freed via `_free_storage(p.data)`. The module holds DTensor shard views and `unshard()` +rebinds `.data` to the all-gathered buffer, so the original storage is dead and freeing it +reduces peak memory during model construction. + +--- + +## Configuration + +```python +fully_shard( + module, + mesh=mesh, + enable_unshard_prefetch=True, # pipeline AG on ag_stream while current module computes + enable_async_reduce_grad=True, # pipeline RS on rs_stream while later modules compute bwd +) +``` + +Setting either flag to `False` assigns `torch.cuda.current_stream()` to the corresponding +stream variable, making all `with torch.cuda.stream(stream)` blocks no-ops — zero overhead, +identical to baseline. + +--- + +## `stop_communication()` — Main Stream Synchronization + +The `FullyShardedDataParallel.stop_communication()` method (in `mcore_fsdp_adapter.py`) ensures +all pending FSDP communication is complete and visible to the main CUDA stream. This is called +before the optimizer step to guarantee that gradient reductions and parameter updates are +synchronized. + +For the `fully_shard` path, the implementation was previously `NotImplementedError`. It now +calls: + +```python +torch.cuda.current_stream().wait_stream(ctx.ag_stream) # finish all-gather work +torch.cuda.current_stream().wait_stream(ctx.rs_stream) # finish reduce-scatter work +``` + +This brings both communication streams into the main stream, ensuring the optimizer sees +fully-synchronized parameters and gradients. + +--- + +## Known Gaps / Recommended Follow-ups + +1. **Single-module prefetch only.** `_get_prefetch_next_modules` returns at most one module. + For networks with many small modules, a size-aware multi-step lookahead + (analogous to `suggested_AG_prefetch_size` in the old `AllGatherPipeline`) would + yield better overlap. + +2. **Outer-DP / HSDP.** `_FSDPRootContext` does not carry an outer-DP stream for the second + all-gather needed in hybrid-sharding (outer-DP × inner-FSDP) setups. This mirrors the + `outer_fsdp_group_param_gather_stream` in the old `AllGatherPipeline`. + +--- + +## Bucket Allocator Hierarchy + +`allocator.py` provides a polymorphic allocator family via the `BucketAllocator` +interface, letting callers swap allocation strategies without changing +`DataParallelBuffer` or `ParameterGroup`. + +``` +BucketAllocator (interface) +|-- TemporaryBucketAllocator — legacy: allocates per key, frees + deletes +|-- StorageFreeingBucketAllocator — allocates per key, frees storage but keeps bucket +| (same tensor object reused on next allocation) +\-- TracePoolAllocator — two-phase: trace → plan → static pool +``` + +### `TracePoolAllocator` + +**Purpose.** During parameter unshard and gradient reduction the FSDP +framework allocates and frees temporary flat buffers (all-gather input/output, +gradient accumulation) in a deterministic, repeatable order. `TracePoolAllocator` +replaces per-call `torch.empty` + `_free_storage` with a one-time planned pool +that eliminates allocation overhead and fragmentation. + +**Design — three phases.** + +| Phase | Behaviour | +|---|---| +| **Trace** (``plan()`` not yet called) | Records every ``allocate`` / ``free`` call as a ``(seq, op, key)`` event. Also stores ``(size, dtype, device)`` metadata per key. Buckets are allocated via ``torch.empty`` as usual. Duplicate allocs (without an intervening free) do not generate new trace events. | +| **Plan** (``plan()``) | Replays the trace to build intervals ``(alloc_seq, free_seq, size)`` for each matched alloc/free pair, groups them by ``(dtype, device)``, and runs a greedy left-edge interval-coloring algorithm per group. Each color is a **slot** in a contiguous flat pool tensor. Because the same key may appear in multiple intervals, ``_slot_map[key]`` is a **list** of slot indices in alloc order. | +| **Optimized** (after ``plan()``) | ``allocate`` returns a ``Bucket`` with a slice-view into the pool, advancing a per-key cursor through the slot list. ``free`` marks the most recently allocated slot as unused (idempotent). ``reset_cursor()`` rewinds all cursors between micro-batches so the same sequence replays. | + +**Slot lists and cursors.** A single allocation key can appear in multiple +intervals — e.g., forward unshard → free → backward unshard → free. The plan +may assign these intervals to *different* slots (if they overlap) or *reuse* +the same slot (if they don't). ``_slot_map`` therefore maps each key to a +**list** of slot indices in the exact alloc order. During optimized-phase +runtime a per-key **cursor** tracks which list entry to consume next; between +micro-batches ``reset_cursor()`` rewinds all cursors to 0. + +**Greedy left-edge coloring.** For each ``(dtype, device)`` group, intervals are +sorted by ``alloc_seq``. For each interval the algorithm tries to reuse a slot +whose previous occupant has already freed (``slot_free_seq < alloc_seq``). If no +slot is free a new one is allocated. The slot is sized to the maximum bucket +assigned to it. After coloring, slots are laid out contiguously and a single +``torch.empty`` is issued per group. + +**Properties.** + +- **Optimal slot count:** left-edge produces the minimum number of slots for + interval graphs — it is impossible to use fewer without causing a conflict. +- **Repeatable trace required:** the same allocate/free call sequence must + repeat across micro-batches. Call ``reset_cursor()`` between micro-batches; + call ``reset()`` to re-profile if the pattern changes. +- **Double-free safe:** ``free`` is idempotent (silently returns if the slot + is already free), matching ``TemporaryBucketAllocator``'s behavior. + +**API.** + +```python +allocator = TracePoolAllocator() +# … run one iteration (trace phase) … +pool_elems = allocator.plan() # returns total element count +# … subsequent micro-batches use the pool … +allocator.reset_cursor() # between micro-batches +print(allocator.total_pool_bytes) # bytes across all groups +allocator.reset() # back to trace phase +``` + +**Lifecycle diagram for one allocation key across two micro-batches.** + +``` +Trace phase Optimized phase +----------- --------------- +allocate(key) → torch.empty --. allocate(key) → pool slot 0 (cursor 0→1) +free(key) → _free_storage | plan free(key) → slot free +allocate(key) → torch.empty --' allocate(key) → pool slot 1 (cursor 1→2) +free(key) → _free_storage free(key) → slot free + -- reset_cursor() -- + allocate(key) → pool slot 0 (cursor 0→1) + free(key) → slot free + ... +``` + +No ``torch.empty`` or storage resizing occurs in the optimized phase — the pool +owns all memory, and buckets are lightweight views. diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/dp_buffer.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/dp_buffer.py new file mode 100644 index 00000000000..313fbcc3e02 --- /dev/null +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/dp_buffer.py @@ -0,0 +1,710 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import logging +import math +from collections import namedtuple +from typing import Dict, List, Optional, Tuple + +import torch + +from .allocator import BucketAllocator, TemporaryBucketAllocator +from .mixed_precision import FullyShardMixedPrecisionPolicy +from .utils import ParamGroupIdx + +logger = logging.getLogger(__name__) + + +class BufferIndex: + """Describes how params are laid out in a flat buffer, including global layout + and per-rank shard information. + + Each DataParallelBuffer owns its own independent BufferIndex instance. + """ + + ItemIndex = namedtuple("ItemIndex", ["global_data_index", "size", "item_id", "shape"]) + BucketMeta = namedtuple("BucketMeta", ["global_data_index", "size", "items"]) + ShardMeta = namedtuple( + "ShardMeta", ["global_data_index", "local_data_index", "bucket_data_index", "size"] + ) + + def __init__( + self, + param_shapes: List[torch.Size], + dp_rank: int, + dp_world_size: int, + is_distributed: bool, + param_group_id: ParamGroupIdx, + chunk_size_factor: int = 1, + sharding_strategy: str = "no_shard", + ): + self.param_group_id = param_group_id + self.is_distributed = is_distributed + self.dp_rank = dp_rank + self.dp_world_size = dp_world_size + self.chunk_size_factor = chunk_size_factor + self.sharding_strategy = sharding_strategy + self.item_index_map, self.bucket_meta = self._build_layout( + param_shapes, dp_world_size, chunk_size_factor, sharding_strategy + ) + self.shard_meta = self._build_shard_meta( + self.bucket_meta, is_distributed, dp_world_size, dp_rank + ) + + # ------------------------------------------------------------------ # + # Layout construction (global, rank-independent) + # ------------------------------------------------------------------ # + + @classmethod + def _build_layout( + cls, + param_shapes: List[torch.Size], + dp_world_size: int, + chunk_size_factor: int, + sharding_strategy: str, + ) -> Tuple[Dict[int, "BufferIndex.ItemIndex"], "BufferIndex.BucketMeta"]: + """ + Compute global buffer layout for a list of parameter shapes. + + Regular parameters (numel >= chunk_size_factor) are placed first. + When a regular parameter has a remainder modulo chunk_size_factor, + its trailing grid is shared with either another regular parameter + or filled with "fragment" parameters. Leftover fragments are + bin-packed into chunk_size_factor-sized grids so that every grid + is aligned and can be split evenly across DP ranks. + + Returns: + item_index_map: Maps item_id -> ItemIndex (global position). + bucket_meta: Describes the full padded buffer. + """ + + def _pad(n: int, divisor: int) -> int: + return int(math.ceil(n / divisor) * divisor) + + def _pad_if_needed(data_index: int) -> int: + if sharding_strategy != "no_shard": + return _pad(data_index, dp_world_size * chunk_size_factor) + return data_index + + def add_item(item_id, shape, offset, index_map): + index_map[item_id] = cls.ItemIndex( + global_data_index=offset, size=shape.numel(), item_id=item_id, shape=shape + ) + + # Separate regular and fragment parameters. + regular_items: List[Tuple[int, torch.Size]] = [] + fragment_items: List[Tuple[int, torch.Size]] = [] + for item_id, shape in enumerate(param_shapes): + if shape.numel() < chunk_size_factor: + fragment_items.append((item_id, shape)) + else: + regular_items.append((item_id, shape)) + + # Sort fragments largest-first for best gap-filling. + fragment_items.sort(key=lambda x: -x[1].numel()) + + item_index_map: Dict[int, cls.ItemIndex] = {} + data_index = 0 + + # ---- First pass: place regular parameters, fill gaps with fragments ---- + while len(regular_items) > 0: + item_id, shape = regular_items.pop(0) + add_item(item_id, shape, data_index, item_index_map) + + if shape.numel() % chunk_size_factor == 0: + data_index += shape.numel() + continue + + gap_offset = data_index + shape.numel() + data_index += (shape.numel() // chunk_size_factor + 1) * chunk_size_factor + remain = shape.numel() % chunk_size_factor + space = chunk_size_factor - remain + + # Try to pair another regular param whose remainder fits. + rhs_found_id = None + rhs_found_shape = None + for id_rhs in regular_items[:]: + rhs_id, rhs = id_rhs + if rhs.numel() % chunk_size_factor == 0: + continue + rhs_remain = rhs.numel() % chunk_size_factor + if remain + rhs_remain <= chunk_size_factor: + rhs_found_id, rhs_found_shape = rhs_id, rhs + regular_items.remove(id_rhs) + break + + if rhs_found_id is not None: + rhs_remain = rhs_found_shape.numel() % chunk_size_factor + # Place the paired param so its LAST ``rhs_remain`` elements + # land in the same alignment grid as the current param's remainder. + # The bulk of the param extends backward from the grid boundary. + add_item(rhs_found_id, rhs_found_shape, data_index - rhs_remain, item_index_map) + space -= rhs_remain + # Advance past the aligned portion of the paired param. + data_index += (rhs_found_shape.numel() // chunk_size_factor) * chunk_size_factor + + # Fill remaining space in this grid with fragments. + for id_frag in fragment_items[:]: + frag_id, frag = id_frag + if frag.numel() > space: + continue + add_item(frag_id, frag, gap_offset, item_index_map) + space -= frag.numel() + gap_offset += frag.numel() + fragment_items.remove(id_frag) + + # ---- helper: bin-pack leftover fragments into chunk_size_factor grids ---- + def pack_fragments( + fragments: List[Tuple[int, torch.Size]], capacity: int + ) -> List[List[Tuple[int, torch.Size]]]: + """ + Bin-pack fragments into fixed-capacity slots (grids). + + Each slot has size *capacity* (== chunk_size_factor). Returns a + list of slots, each containing one or more (param_id, shape) pairs. + """ + sorted_frags = sorted(fragments, key=lambda p: -p[1].numel()) + slots: List[List[Tuple[int, torch.Size]]] = [] + for pid, pshape in sorted_frags: + psize = pshape.numel() + placed = False + for slot in slots: + used = sum(p[1].numel() for p in slot) + if used + psize <= capacity: + slot.append((pid, pshape)) + placed = True + break + if not placed: + slots.append([(pid, pshape)]) + return slots + + # ---- Second pass: bin-pack any remaining fragments into aligned grids ---- + if fragment_items: + fragment_slots = pack_fragments(fragment_items, chunk_size_factor) + for slot in fragment_slots: + offset_within_grid = 0 + for pid, pshape in slot: + add_item(pid, pshape, data_index + offset_within_grid, item_index_map) + offset_within_grid += pshape.numel() + data_index += chunk_size_factor + + bucket_meta = cls.BucketMeta( + global_data_index=0, + size=_pad_if_needed(data_index), + items=list(item_index_map.values()), + ) + return item_index_map, bucket_meta + + # ------------------------------------------------------------------ # + # Shard meta construction (per-rank) + # ------------------------------------------------------------------ # + + @classmethod + def _build_shard_meta( + cls, + bucket_meta: "BufferIndex.BucketMeta", + is_distributed: bool, + dp_world_size: int, + dp_rank: int, + ) -> "BufferIndex.ShardMeta": + shard_size = bucket_meta.size // dp_world_size + bucket_data_index = shard_size * dp_rank + global_data_index = bucket_meta.global_data_index + bucket_data_index + + if is_distributed: + return cls.ShardMeta( + global_data_index=global_data_index, + local_data_index=0, + bucket_data_index=bucket_data_index, + size=shard_size, + ) + else: + return cls.ShardMeta( + global_data_index=global_data_index, + # For non-distributed buffers, each rank has the full buffer, so + # the local index is the same as the global index. + local_data_index=global_data_index, + bucket_data_index=bucket_data_index, + size=shard_size, + ) + + # ------------------------------------------------------------------ # + # Compaction — scale indices proportionally for packed storage + # ------------------------------------------------------------------ # + + def compact(self, factor: float, compact_shapes: List[torch.Size]) -> None: + """Scale all indices proportionally for packed storage. + + Args: + factor: Scale factor (0.5 for NVFP4 2-values-per-byte packing). + compact_shapes: Per-item shapes for the packed layout (same + length and order as the original param_shapes). + """ + new_map: Dict[int, "BufferIndex.ItemIndex"] = {} + for item_id, item in self.item_index_map.items(): + new_map[item_id] = self.ItemIndex( + global_data_index=int(item.global_data_index * factor), + size=int(item.size * factor), + item_id=item.item_id, + shape=compact_shapes[item_id], + ) + self.item_index_map = new_map + + self.bucket_meta = self.BucketMeta( + global_data_index=0, + size=int(self.bucket_meta.size * factor), + items=list(new_map.values()), + ) + self.shard_meta = self._build_shard_meta( + self.bucket_meta, self.is_distributed, self.dp_world_size, self.dp_rank + ) + + # ------------------------------------------------------------------ # + # Internal index query methods — three coordinate domains: + # + # _get_item_self_range → (start, end) relative to the item's own + # start. Tells what portion of this item + # falls within the current rank's shard. + # _get_item_local_range → (start, end) within self.data (the local + # GPU buffer). Where to read/write bytes. + # _get_item_global_range → (global_data_index, size) in the full + # logical (unsharded) buffer, same on all + # ranks. + # ------------------------------------------------------------------ # + + def _get_item_global_range(self, item_id: int) -> Tuple[int, int]: + """Return (global_data_index, size) for the given item.""" + idx = self.item_index_map[item_id] + return (idx.global_data_index, idx.size) + + def _get_item_self_range(self, item_id: int, *, as_shard: bool = True) -> Tuple[int, int]: + """Return coordinates relative to the item's own start. + + When ``as_shard=True`` (default), returns the portion of the item + that falls within this rank's shard — the slice ``(start, end)`` + within the item. When ``as_shard=False``, returns ``(0, size)`` + representing the full item. + """ + idx = self.item_index_map[item_id] + if not as_shard: + return (0, idx.size) + + item_start = idx.global_data_index + item_end = item_start + idx.size + shard_start = self.shard_meta.global_data_index + shard_end = shard_start + self.shard_meta.size + + if item_start > shard_end or item_end < shard_start: + return (0, 0) + + start = max(item_start, shard_start) - item_start + end = min(item_end, shard_end) - item_start + return (start, end) + + def _get_item_local_range(self, item_id: int, *, as_shard: bool = False) -> Tuple[int, int]: + """Return coordinates within self.data for the item. + + Parameters + ---------- + as_shard : bool + If True, compute the shard intersection even when the buffer + is not distributed. Default (False) returns the full item + range for non-distributed buffers. + """ + if not self.is_distributed and not as_shard: + idx = self.item_index_map[item_id] + return (idx.global_data_index, idx.global_data_index + idx.size) + + slice_start, slice_end = self._get_item_self_range(item_id) + if slice_start == slice_end: + return (0, 0) + + idx = self.item_index_map[item_id] + offset = ( + idx.global_data_index + - self.shard_meta.global_data_index + + self.shard_meta.local_data_index + ) + return (offset + slice_start, offset + slice_end) + + +class DataParallelBuffer: + """Manages a flat buffer that stores (a shard of) a group of parameters. + + On construction it builds its own BufferIndex describing the layout and + shard ownership. External callers interact via init_data / set_item / + get_item only. + """ + + def __init__( + self, + params: List[torch.nn.Parameter], + param_idx: Dict[torch.nn.Parameter, int], + dtype: torch.dtype, + device: torch.device, + dp_group: torch.distributed.ProcessGroup, + param_group_id: ParamGroupIdx, + mp_policy: FullyShardMixedPrecisionPolicy, + *, + allocator: Optional[BucketAllocator] = None, + buffer_role: str = "model_weight", + is_distributed: bool = False, + gradient_scaling_factor: Optional[float] = None, + chunk_size_factor: int = 1, + sharding_strategy: str = "no_shard", + ): + assert mp_policy is not None, "DataParallelBuffer requires a mixed-precision policy" + self.params = params + self.param_idx = param_idx + self.dtype = dtype + self.device = device + self.dp_group = dp_group + self.allocator = allocator if allocator is not None else TemporaryBucketAllocator() + self.buffer_role = buffer_role + self.alloc_key = (param_group_id, buffer_role) + self.mp_policy = mp_policy + self.is_distributed = is_distributed + self.gradient_scaling_factor = gradient_scaling_factor + + dp_rank = torch.distributed.get_rank(dp_group) + dp_world_size = torch.distributed.get_world_size(dp_group) + + # Always build layout with logical shapes and shared chunk_size_factor + # so that all buffers share the same proportional item-offset mapping. + _logical_shapes = [p.shape for p in params] + self.buffer_index = BufferIndex( + param_shapes=_logical_shapes, + dp_rank=dp_rank, + dp_world_size=dp_world_size, + is_distributed=is_distributed, + chunk_size_factor=chunk_size_factor, + sharding_strategy=sharding_strategy, + param_group_id=param_group_id, + ) + + # Compact NVFP4 weight buffers: scale all indices proportionally so + # the buffer holds only the packed data without fragment-binning waste. + if buffer_role in ("model_weight", "transpose_weight") and any( + mp_policy.is_nvfp4_param(p) for p in params + ): + compact_shapes = mp_policy.get_param_storage_shapes(params) + self.buffer_index.compact(0.5, compact_shapes) + + if is_distributed: + self.data_size = self.buffer_index.shard_meta.size + else: + self.data_size = self.buffer_index.bucket_meta.size + + self.data: Optional[torch.Tensor] = None + self._unsharded_buffer: Optional[torch.Tensor] = None + + # ------------------------------------------------------------------ # + # Public API + # ------------------------------------------------------------------ # + + def init_data(self, data: torch.Tensor) -> None: + """Bind an externally allocated tensor as the persistent storage.""" + assert data.dtype == self.dtype, f"dtype mismatch: {data.dtype} vs {self.dtype}" + assert data.numel() == self.data_size, f"size mismatch: {data.numel()} vs {self.data_size}" + self.data = data + + def check_no_local_overlap(self, label: str = "") -> bool: + """ + Runtime check: verify no two items' local slices overlap within ``self.data``. + + Returns True if layout is valid (no overlaps, all slices in bounds). + Returns False and prints diagnostic info if any overlap or bound violation is found. + """ + if self.data is None: + return True + + items = self.buffer_index.item_index_map + n_items = len(items) + if n_items == 0: + return True + + label_prefix = f"[{label}] " if label else "" + + # Collect (local_start, local_end, item_id, global_start, size) for each item + slices = [] + for item_id in range(n_items): + local_start, local_end = self.buffer_index._get_item_local_range(item_id) + idx = self.buffer_index.item_index_map[item_id] + slices.append((local_start, local_end, item_id, idx.global_data_index, idx.size)) + + # Sort by local_start + slices.sort(key=lambda x: x[0]) + + valid = True + data_nel = self.data.numel() + + for i in range(len(slices)): + s_start, s_end, s_id, g_start, size = slices[i] + shape = items[s_id].shape + + # Bounds check: end must not exceed data size + if s_end > data_nel: + logger.warning( + f"{label_prefix}OVERFLOW: item {s_id} shape={list(shape)} " + f"local=[{s_start}, {s_end}) but data.numel()={data_nel} " + f"(global=[{g_start}, {g_start + size}))" + ) + valid = False + + # Overlap check with next item + if i + 1 < len(slices): + n_start, n_end, n_id, n_gstart, n_size = slices[i + 1] + if s_end > n_start: + overlap = s_end - n_start + logger.warning( + f"{label_prefix}OVERLAP: item {s_id} shape={list(shape)} " + f"local=[{s_start}, {s_end}) overlaps item {n_id} " + f"local=[{n_start}, {n_end}) by {overlap} elements " + f"(global_{s_id}=[{g_start}, {g_start + size}), " + f"global_{n_id}=[{n_gstart}, {n_gstart + n_size}))" + ) + valid = False + + return valid + + def check_no_global_overlap(self, label: str = "") -> bool: + """ + Runtime check: verify no two items' **global** slices overlap. + + This checks the logical layout (should never fail if _build_layout is correct). + """ + items = self.buffer_index.item_index_map + n_items = len(items) + if n_items == 0: + return True + + label_prefix = f"[{label}] " if label else "" + + ranges = [] + for item_id in range(n_items): + idx = items[item_id] + ranges.append( + (idx.global_data_index, idx.global_data_index + idx.size, item_id, idx.shape) + ) + + ranges.sort(key=lambda x: x[0]) + + valid = True + for i in range(len(ranges) - 1): + a_start, a_end, a_id, a_shape = ranges[i] + b_start, b_end, b_id, b_shape = ranges[i + 1] + if a_end > b_start: + logger.warning( + f"{label_prefix}GLOBAL OVERLAP: item {a_id} shape={list(a_shape)} " + f"[{a_start}, {a_end}) vs item {b_id} shape={list(b_shape)} " + f"[{b_start}, {b_end}) overlap={a_end - b_start}" + ) + valid = False + + if valid: + pass # silent on success + return valid + + def set_item(self, item_id: int, item_data: torch.Tensor) -> None: + """Write a parameter tensor into the corresponding region of the buffer.""" + if self.is_distributed: + slice_start, slice_end = self.buffer_index._get_item_self_range(item_id) + item_data = item_data.flatten()[slice_start:slice_end] + + local_start, local_end = self.buffer_index._get_item_local_range(item_id) + shard = self.data[local_start:local_end] + if shard.numel() > 0: + shard.data.copy_(item_data.flatten()) + + def get_item(self, item_id: int, *, as_shard: bool = False) -> torch.Tensor: + """Read a parameter tensor (or its shard) from the buffer.""" + start, end = self.buffer_index._get_item_local_range(item_id, as_shard=as_shard) + return self.data[start:end] + + def is_unsharded(self) -> bool: + """Return whether this buffer currently has a full unsharded view.""" + return not self.is_distributed or self._unsharded_buffer is not None + + @torch.no_grad() + def unshard( + self, async_op: bool = True + ) -> Tuple[torch.Tensor, Optional[torch.distributed.Work]]: + """All-gather the full buffer from all shards and bind parameter storage. + + For non-distributed buffers self.data is already full, so + (self.data, None) is returned directly. + """ + work = None + if self.is_distributed: + bucket = self.allocator.allocate( + key=self.alloc_key, + size=self.buffer_index.bucket_meta.size, + dtype=self.dtype, + device=self.device, + ) + self._unsharded_buffer = bucket.data + + sm = self.buffer_index.shard_meta + shard = self.data[sm.local_data_index : sm.local_data_index + sm.size] + + work = torch.distributed.all_gather_into_tensor( + output_tensor=self._unsharded_buffer, + input_tensor=shard, + group=self.dp_group, + async_op=async_op, + ) + if self._unsharded_buffer.is_cuda: + # The temporary bucket may be released from another stream before this + # collective finishes; record the producer stream for allocator safety. + self._unsharded_buffer.record_stream(torch.cuda.current_stream()) + full_buffer = self._unsharded_buffer + else: + full_buffer = self.data + + # Bind the full buffer to the params according to the layout. + self._bind_buffer_to_params(full_buffer) + + return (full_buffer, work) + + def _bind_buffer_to_params(self, buffer: torch.Tensor) -> None: + """Bind the given buffer to the params according to the layout.""" + assert buffer.numel() == self.buffer_index.bucket_meta.size, ( + f"Buffer size {buffer.numel()} does not match expected size " + f"{self.buffer_index.bucket_meta.size}" + ) + for p in self.params: + item_id = self.param_idx[p] + offset, size = self.buffer_index._get_item_global_range(item_id) + idx_shape = self.buffer_index.item_index_map[item_id].shape + param_data = buffer[offset : offset + size].view(idx_shape) + self.mp_policy.bind_unsharded_param(p, param_data, self.buffer_role) + + @torch.no_grad() + def reshard(self) -> None: + """Release the temporary unsharded buffer allocated by unshard().""" + if not self.is_distributed: + return + self.allocator.free(self.alloc_key) + self._unsharded_buffer = None + + def fetch_buffer(self, *, as_shard: bool = False) -> torch.Tensor: + """Return the buffer, allocating the full unsharded view if needed. + + Parameters + ---------- + as_shard : bool + If True, return only this rank's shard slice of the full buffer. + Default (False) returns the full unsharded buffer. + """ + if self.is_distributed: + if self._unsharded_buffer is None: + bucket = self.allocator.allocate( + key=self.alloc_key, + size=self.buffer_index.bucket_meta.size, + dtype=self.dtype, + device=self.device, + ) + self._unsharded_buffer = bucket.data + full = self._unsharded_buffer + else: + full = self.data + + if as_shard: + sm = self.buffer_index.shard_meta + return full[sm.bucket_data_index : sm.bucket_data_index + sm.size] + return full + + @torch.no_grad() + def reduce_grad(self, grad_comm_dtype: Optional[torch.dtype] = None): + """Reduce gradients across the data-parallel group. + + For distributed buffers: reduce-scatter the full gradient into each + rank's shard, then accumulate into self.data. + For non-distributed buffers: all-reduce in-place. + If grad_comm_dtype differs from self.dtype, communicate with a temporary + casted tensor and cast the reduced result back before accumulation. + """ + grad_comm_dtype = grad_comm_dtype or self.dtype + + if self.gradient_scaling_factor in (None, 1.0): + op = torch.distributed.ReduceOp.SUM + prescale = False + elif grad_comm_dtype != torch.bfloat16: + op = torch.distributed._make_nccl_premul_sum(self.gradient_scaling_factor) + prescale = False + else: + op = torch.distributed.ReduceOp.SUM + prescale = True + + if not self.is_distributed: + comm_data = ( + self.data if grad_comm_dtype == self.dtype else self.data.to(grad_comm_dtype) + ) + if prescale: + comm_data.mul_(self.gradient_scaling_factor) + torch.distributed.all_reduce(comm_data, group=self.dp_group, op=op) + if comm_data is not self.data: + self.data.copy_(comm_data.to(self.dtype)) + return + + full_grad = self.fetch_buffer() + comm_input = full_grad if grad_comm_dtype == self.dtype else full_grad.to(grad_comm_dtype) + if full_grad.is_cuda: + # Keep temporary reduce-scatter buffers tied to the stream that uses them. + full_grad.record_stream(torch.cuda.current_stream()) + if prescale: + comm_input.mul_(self.gradient_scaling_factor) + + sm = self.buffer_index.shard_meta + local_grad_shard = self.data[sm.local_data_index : sm.local_data_index + sm.size] + reduced_grad_shard = comm_input[sm.bucket_data_index : sm.bucket_data_index + sm.size] + + torch.distributed.reduce_scatter_tensor( + output=reduced_grad_shard, input=comm_input, group=self.dp_group, op=op + ) + + local_grad_shard += reduced_grad_shard + + +def check_all_fsdp_buffers(module) -> bool: + """ + Scan every FSDPModule in *module* and verify no local slice overlaps + in any buffer (model_weight, main_weight, main_grad). + + Call this at any point after FSDP initialization to catch runtime + corruption. Returns True if all buffers are clean. + """ + import torch.distributed as dist + + from .fsdp_module import FSDPModule + + rank = dist.get_rank() if dist.is_initialized() else -1 + all_ok = True + + for name, child in module.named_modules(): + if not isinstance(child, FSDPModule): + continue + for param_names, param_group in child._named_param_groups: + gid = f"mod={name} pg={param_group.param_group_id} rank={rank}" + if param_group.model_weight_buffer is not None: + ok = param_group.model_weight_buffer.check_no_local_overlap(gid + " wbuf") + all_ok = all_ok and ok + if param_group.main_weight_buffer is not None: + ok = param_group.main_weight_buffer.check_no_local_overlap(gid + " mbuf") + all_ok = all_ok and ok + if param_group.main_grad_buffer is not None: + ok = param_group.main_grad_buffer.check_no_local_overlap(gid + " gbuf") + all_ok = all_ok and ok + + return all_ok diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/fsdp_module.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/fsdp_module.py new file mode 100644 index 00000000000..ad831a56dff --- /dev/null +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/fsdp_module.py @@ -0,0 +1,830 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""FSDPModule implementation for Megatron-FSDP2.""" + +import logging +from contextlib import nullcontext +from dataclasses import dataclass, field +from typing import Dict, List, Optional, Tuple + +import torch +import torch.nn as nn +from torch.distributed.device_mesh import DeviceMesh +from torch.distributed.tensor import DTensor + +from .allocator import BucketAllocator +from .mixed_precision import FullyShardMixedPrecisionPolicy +from .param_group import ParameterGroup +from .utils import ParamGroupIdx, _replace_module_parameter + +logger = logging.getLogger(__name__) + + +class _FSDPState: + """ + Internal state for FSDP module tracking. + + Attributes: + _is_root: Whether this is the root FSDP module (handles final callback). + _post_backward_callback_queued: Whether callback is queued for execution. + """ + + def __init__(self): + self._is_root = True + self._post_backward_callback_queued = False + + +@dataclass +class _FSDPRootContext: + """ + Runtime context shared across all FSDP modules within a single root. + + This object coordinates CUDA streams, execution ordering, and async + communication overlap (all-gather / reduce-scatter) during forward + and backward passes. + """ + + # ------------------------------------------------------------------ + # CUDA streams (communication overlap) + # ------------------------------------------------------------------ + ag_stream: torch.cuda.Stream # all-gather / unshard stream + rs_stream: torch.cuda.Stream # reduce-scatter stream + + # ------------------------------------------------------------------ + # Bucket allocator + # ------------------------------------------------------------------ + bucket_allocator: BucketAllocator + """ + Bucket allocator for temporary all-gather and reduce-scatter buffers. + + Buffer roles are part of each allocation key, so one allocator can safely + manage all DataParallelBuffer temporary buckets without requiring separate + weight/gradient allocator instances. + """ + + # ------------------------------------------------------------------ + # Forward execution ordering + # ------------------------------------------------------------------ + forward_order: List["FSDPModule"] = field(default_factory=list) + """ + FSDP modules in actual forward execution order. + + This ordering is used to: + - Schedule prefetching of parameters (unshard) + - Ensure correct overlap between compute and communication + """ + + # ------------------------------------------------------------------ + # Unshard (all-gather) tracking + # ------------------------------------------------------------------ + unshard_done_events: Dict[int, Optional[torch.cuda.Event]] = field(default_factory=dict) + """ + Maps module_id -> CUDA event signaling completion of parameter unshard. + + Used to enforce correct dependency between all-gather and compute. + """ + + enable_unshard_prefetch: bool = True + """Whether to prefetch (pipeline) parameter unshard for upcoming modules.""" + + # ------------------------------------------------------------------ + # Reduce-scatter (gradient sync) tracking + # ------------------------------------------------------------------ + reduce_grad_buckets: Dict[int, List[Tuple[torch.cuda.Event, "ParameterGroup"]]] = field( + default_factory=dict + ) + """ + Maps module_id -> list of (event, parameter_group) tuples. + + Each entry corresponds to a module and contains a list of: + (event, parameter_group) + + - event: signals gradient readiness + - parameter_group: gradients to be reduced + + This structure enables ordered overlap of backward compute and + gradient synchronization. + """ + + enable_async_reduce_grad: bool = True + """Whether to overlap gradient reduction with backward computation.""" + + # ------------------------------------------------------------------ + # Activation recompute / gradient checkpointing support + # ------------------------------------------------------------------ + backward_phase: bool = False + """True from the root backward pre-hook until the final callback.""" + + backward_module: Optional[int] = None + """``id(module)`` of the FSDP module whose backward is pending next. + Derived from ``_reversed_order`` and ``backward_done_modules`` — NOT + set by any hook directly. Updated by ``_advance_backward_module()``.""" + + backward_done_modules: set = field(default_factory=set) + """Set of ``id(module)`` for FSDP modules whose backward has completed. + Populated in ``post_backward``, cleared in the root backward pre-hook.""" + + _reversed_order: List["FSDPModule"] = field(default_factory=list) + """``list(reversed(forward_order))`` — precomputed backward processing order.""" + + def _advance_backward_module(self) -> None: + """Set ``backward_module`` to the first module in ``_reversed_order`` + that is NOT in ``backward_done_modules``.""" + for m in self._reversed_order: + if id(m) not in self.backward_done_modules: + self.backward_module = id(m) + return + self.backward_module = None + + def get_prefetch_next_modules( + self, module: "FSDPModule", bwd_pass: bool = False + ) -> List["FSDPModule"]: + """Return the next FSDP module to prefetch in forward or backward order.""" + module_order = list(reversed(self.forward_order)) if bwd_pass else self.forward_order + + for module_index, candidate_module in enumerate(module_order): + if candidate_module is module: + if module_index + 1 >= len(module_order): + return [] + return [module_order[module_index + 1]] + + raise AssertionError("Current module not found in forward module order") + + +class FSDPModule(nn.Module): + """ + Mixin class for FSDP-wrapped modules. + + This class is dynamically added to wrapped modules and provides + methods for managing parameter sharding state: + - unshard(): All-gather parameters before forward + - reshard(): Release unsharded buffer after forward + - reduce_grad(): Reduce-scatter gradients after backward + """ + + def _init_named_param_groups( + self, + mesh: Optional[DeviceMesh], + ignored_params: Optional[set], + mp_policy: FullyShardMixedPrecisionPolicy, + bucket_allocator: BucketAllocator, + gradient_scaling_factor: Optional[float] = None, + ): + """ + Initialize parameter groups and build param name mapping. + + This method: + 1. Collects ignored modules (nested FSDP modules) + 2. Materializes meta modules to actual devices + 3. Groups parameters by (device, dtype, requires_grad) + 4. Builds parameter name to parameter mapping + """ + ignored_params = ignored_params or set() + ignored_modules = set() + + # Collect nested FSDP modules as ignored + for _, child in self.named_modules(): + if child is not self and isinstance(child, FSDPModule): + ignored_params.update(child.parameters()) + for child_submodule in child.modules(): + ignored_modules.add(child_submodule) + + # Materialize meta parameters to actual device + self._materialize_meta_module(ignored_modules, mesh=mesh, mp_policy=mp_policy) + + # Create parameter groups + fsdp_param_groups = _get_module_fsdp_param_groups( + self, + mp_policy=mp_policy, + mesh=mesh, + ignored_params=ignored_params, + allocator=bucket_allocator, + gradient_scaling_factor=gradient_scaling_factor, + ) + setattr(self, "_fsdp_param_groups", fsdp_param_groups) + + # Build param name to param mapping for later lookup + param_to_name = {p: n for n, p in self.named_parameters()} + self._named_param_groups = [] + + for fsdp_param_group in fsdp_param_groups: + param_names = [] + for param in fsdp_param_group.params: + param_name = param_to_name[param] + param_names.append(param_name) + self._named_param_groups.append((param_names, fsdp_param_group)) + + def _init_param_main_grad_func(self): + """ + Initialize main gradient getter function for each parameter. + + This creates a closure that fetches the gradient from the + gradient buffer when accessed. It handles both sharded and + unsharded gradient buffers. + """ + + def main_grad_getter(p): + """Get main gradient from buffer with proper offset/size.""" + gbuf = p._gbuf + item_id = p._item_id + + gbuf_data = gbuf.fetch_buffer() + assert gbuf_data is not None + assert gbuf_data.numel() > 0 + + # Get offset and size from buffer index + offset, size = gbuf.buffer_index._get_item_global_range(item_id) + param_shape = gbuf.buffer_index.item_index_map[item_id].shape + grad_data = gbuf_data[offset : offset + size].view(param_shape) + + return grad_data + + # Attach getter to each parameter + for param_group in self._fsdp_param_groups: + for param in param_group.params: + setattr(param, "_gbuf", param_group.main_grad_buffer) + setattr(param, "_item_id", param_group.param_idx[param]) + param.get_main_grad = main_grad_getter.__get__(param) + + def _materialize_meta_module( + self, + ignored_modules: set, + mesh: Optional[DeviceMesh] = None, + mp_policy: Optional[FullyShardMixedPrecisionPolicy] = None, + ): + """ + Materialize meta parameters to actual device and initialize. + + This is needed for large models that cannot fit in a single GPU. + Meta parameters are moved to the current device and reset. + After materialization, full parameters are broadcast from DP rank 0 + before DTensor wrapping so every rank shards the same initialized value. + """ + materialization_device = torch.cuda.current_device() + for name, m in self.named_modules(): + if m in ignored_modules: + continue + # Skip modules that don't have meta parameters + if all(not p.is_meta for p in m.parameters(recurse=False)): + continue + + m._apply( + lambda t: torch.empty_like(t, device=materialization_device) if t.is_meta else t, + recurse=False, + ) + init_context = ( + mp_policy.model_init_context(m) if mp_policy is not None else nullcontext() + ) + with init_context: + if hasattr(m, "reset_parameters"): + m.reset_parameters() + elif hasattr(m, "_reset_parameters"): + m._reset_parameters() + else: + raise ValueError( + f"Module {name} contains meta parameters but cannot reset them" + ) + + if mesh is not None and mesh.size() > 1: + dp_group = mesh.get_group() + src_rank = torch.distributed.get_global_rank(dp_group, 0) + for param in self.parameters(): + if param.is_meta or isinstance(param, DTensor): + continue + torch.distributed.broadcast(param.data, src=src_rank, group=dp_group) + + def _init_fsdp_state( + self, enable_unshard_prefetch, enable_async_reduce_grad, bucket_allocator: BucketAllocator + ): + """Initialize FSDP state and mark nested FSDP modules as non-root. + + Important: This must be called BEFORE any forward/backward pass runs. + Re-initializing while child FSDPModules are actively unsharded + (mid-forward or mid-backward) will corrupt their state. The safety + check below enforces that constraint. + """ + named_forward_modules = [ + (name, child) for name, child in self.named_modules() if isinstance(child, FSDPModule) + ] + forward_order = [child for name, child in named_forward_modules] + + # Safety check: no child FSDPModule must be in an active state. + # - unshard_done_events[id(child)] non-None → unsharded, not yet resharded + # - reduce_grad_buckets[id(child)] non-empty → reduce-scatter in flight + # Re-initializing _fsdp_root_context while a child is in either state + # would overwrite its shared state mid-pass. + for child_module in forward_order: + if child_module is self: + continue + if not hasattr(child_module, "_fsdp_root_context"): + continue + ctx = child_module._fsdp_root_context + if ctx is None: + continue + if ctx.unshard_done_events.get(id(child_module)) is not None: + raise RuntimeError( + "_init_fsdp_state cannot be called while a child FSDPModule " + "is still unsharded. All children must be resharded before " + "re-initializing FSDP state." + ) + if ctx.reduce_grad_buckets.get(id(child_module)): + raise RuntimeError( + "_init_fsdp_state cannot be called while a child FSDPModule " + "has pending reduce-scatter operations. All children must have " + "completed gradient reduction before re-initializing FSDP state." + ) + + root_context = _FSDPRootContext( + ag_stream=( + torch.cuda.Stream() if enable_unshard_prefetch else torch.cuda.current_stream() + ), + rs_stream=( + torch.cuda.Stream() if enable_async_reduce_grad else torch.cuda.current_stream() + ), + forward_order=forward_order, + reduce_grad_buckets={id(module): [] for module in forward_order}, + unshard_done_events={id(module): None for module in forward_order}, + enable_unshard_prefetch=enable_unshard_prefetch, + enable_async_reduce_grad=enable_async_reduce_grad, + _reversed_order=list(reversed(forward_order)), + bucket_allocator=bucket_allocator, + ) + setattr(self, "_fsdp_state", _FSDPState()) + setattr(self, "_fsdp_root_context", root_context) + + module_idx = 0 + for name, module in named_forward_modules: + for param_group in module._fsdp_param_groups: + param_group.set_allocator(root_context.bucket_allocator) + + if module is not self: + module._fsdp_state._is_root = False + setattr(module, "_fsdp_root_context", root_context) + + setattr(module, "_fsdp_module_idx", module_idx) + setattr(module, "_fsdp_module_name", name) + module_idx += 1 + + def unshard(self, async_op: bool = False, bwd_pass: bool = False): + """ + Unshard parameters by all-gathering from the sharded buffer. + + This is called pre-forward to make parameters available for + computation. After unsharding, each param.data points to + the full (unsharded) tensor. + """ + torch.cuda.nvtx.range_push("MFSDP unshard") + ctx = self._fsdp_root_context + stream = ctx.ag_stream if async_op else torch.cuda.current_stream() + + if async_op: + # Synchronize ag_stream with current_stream to guarantee that main-stream + # writes to parameter data are visible before the all-gather kernel reads them + # on ag_stream. Without this barrier, stale or partially-written parameter + # shards may be gathered, causing convergence divergence. + stream.wait_stream(torch.cuda.current_stream()) + + # Unshard this module and optionally prefetch next modules in the forward/backward pass + if async_op: + prefetch_modules = ctx.get_prefetch_next_modules(self, bwd_pass=bwd_pass) + else: + prefetch_modules = [] + for module in [self] + prefetch_modules: + if all( + param_group.has_unsharded_weight_buffers(bwd_pass=bwd_pass) + for param_group in module._fsdp_param_groups + ): + continue + if bwd_pass and id(module) in ctx.backward_done_modules: + continue # Skip prefetch for modules whose backward is already done + + # Unshard parameters for this module + for param_names, param_group in module._named_param_groups: + # Optional NaN checking for debugging + if getattr(module, "_enable_nan_checks", False): + for name, dist_param in zip(param_names, param_group.dist_params): + assert not torch.isnan( + dist_param._local_tensor + ).any(), f"NaN detected in dist param for parameter {name}" + + with torch.cuda.stream(stream): + param_group.unshard(bwd_pass=bwd_pass) + + # Record event to track when unshard is done for this module + if async_op: + event = stream.record_event() + ctx.unshard_done_events[id(module)] = event + + # Ensure unshard is complete before forward. + # The event is NOT cleared here — it persists as a "currently unsharded" + # flag and is only cleared by reshard(). This prevents redundant + # all-gathers during activation recompute and prefetch re-entry. + if ctx.unshard_done_events[id(self)] is not None: + ctx.unshard_done_events[id(self)].wait() + + # Replace module parameters with unsharded versions + for param_names, param_group in self._named_param_groups: + for name, param in zip(param_names, param_group.params): + _replace_module_parameter(self, name, param) + + # Optional NaN checking for debugging + if getattr(self, "_enable_nan_checks", False): + for name, param in zip(param_names, param_group.params): + assert not torch.isnan(param).any(), f"NaN detected in parameter {name}" + + torch.cuda.nvtx.range_pop() + + def reshard(self): + """Reshard parameters by replacing with sharded DTensors.""" + torch.cuda.nvtx.range_push("MFSDP reshard") + ctx = self._fsdp_root_context + for param_names, param_group in self._named_param_groups: + param_group.reshard() + for name, dist_param in zip(param_names, param_group.dist_params): + _replace_module_parameter(self, name, dist_param) + ctx.unshard_done_events[id(self)] = None # Clear unshard event for this module + torch.cuda.nvtx.range_pop() + + def _wait_for_previous_async_reduce_grad(self): + """Release older async reduce buffers in backward order.""" + ctx = self._fsdp_root_context + if not ctx.enable_async_reduce_grad: + return + + backward_order = list(reversed(ctx.forward_order)) + for i, module in enumerate(backward_order): + if i - 2 >= 0: + buckets = ctx.reduce_grad_buckets[id(backward_order[i - 2])] + while len(buckets) > 0: + event, param_group = buckets.pop() + event.wait() + param_group.release_grad_buffer() + if module is self: + break + + def reduce_grad(self, async_op: bool = False): + """ + Reduce gradients across data-parallel ranks. + + This is called post-backward to: + 1. Copy gradients to main gradient buffer + 2. Perform all-reduce or reduce-scatter + 3. Install reduced gradients to distributed parameters + """ + torch.cuda.nvtx.range_push("MFSDP reduce_grad") + ctx = self._fsdp_root_context + stream = ctx.rs_stream if async_op else torch.cuda.current_stream() + + # Handle pending reduce events before this module to release buffers promptly. + self._wait_for_previous_async_reduce_grad() + + # Perform reduction for this module + for param_names, param_group in self._named_param_groups: + if not param_group.requires_grad: + continue + + # NaN check before reduction + if getattr(self, "_enable_nan_checks", False): + for name, param in zip(param_names, param_group.params): + if param.grad is not None: + assert not torch.isnan( + param.grad + ).any(), f"NaN in parameter grad for {name}" + + # Copy .grad -> main grad buffer on main stream (fast memcpy). + # When gradient_accumulation_fusion is active for FSDP params, the backward + # kernel writes directly into main_grad (weight.main_grad = get_main_grad() in + # layers.py) and sets grad_added_to_main_grad=True. In that case we must NOT + # zero or overwrite main_grad; discard the dummy .grad tensor if present. + for name, param in zip(param_names, param_group.params): + if getattr(param, "grad_added_to_main_grad", False): + if param.grad is not None: + del param.grad + elif param.grad is None: + main_grad = param.get_main_grad() + if hasattr(param, "main_grad") and param.main_grad is not None: + if param.main_grad.data_ptr() != main_grad.data_ptr(): + main_grad.copy_(param.main_grad.detach()) + else: + main_grad.zero_() + else: + main_grad = param.get_main_grad() + main_grad.copy_(param.grad.detach()) + del param.grad + + if async_op: + # ---- Overlapped path ---- + # Switch to rs_stream for the reduce-scatter kernel + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + param_group.reduce_grad() + else: + # ---- Non-overlapped path ---- + # Reduce gradients immediately and release grad buffer + param_group.reduce_grad() + param_group.release_grad_buffer() + + # Install reduced gradients to distributed parameters + for name, param, dist_param, dist_grad in zip( + param_names, param_group.params, param_group.dist_params, param_group.dist_grads + ): + if not param.requires_grad: + continue + if param_group.mp_policy.use_decoupled_grad: + setattr(dist_param, "decoupled_grad", dist_grad) + if dist_param.grad is not None: + del dist_param.grad + else: + if dist_grad is not None: + with torch.cuda.stream(stream): + dist_grad = dist_grad.to(dist_param.dtype) + setattr(dist_param, "grad", dist_grad) + if hasattr(dist_param, "decoupled_grad"): + dist_param.decoupled_grad = None + + if async_op: + event = stream.record_event() + ctx.reduce_grad_buckets[id(self)].append((event, param_group)) + + # NaN check after reduction + if getattr(self, "_enable_nan_checks", False): + for name, dist_grad in zip(param_names, param_group.dist_grads): + if dist_grad is not None: + assert not torch.isnan( + dist_grad._local_tensor + ).any(), f"NaN in dist grad for parameter {name}" + + torch.cuda.nvtx.range_pop() + + @torch.no_grad() + def _scale_gradients(self, scaling_factor: float): + """Scale optimizer-facing gradients by a factor.""" + ctx = self._fsdp_root_context + torch.cuda.current_stream().wait_stream(ctx.rs_stream) + for _, child in self.named_modules(): + if not isinstance(child, FSDPModule): + continue + for param_group in child._fsdp_param_groups: + for dist_param in param_group.dist_params: + grad = getattr(dist_param, "decoupled_grad", None) + if grad is None: + grad = dist_param.grad + if grad is None: + continue + if isinstance(grad, DTensor): + grad._local_tensor.mul_(scaling_factor) + else: + grad.mul_(scaling_factor) + + def _zero_grad_buffer(self): + """Zero the gradient buffer for all parameter groups.""" + for _, child in self.named_modules(): + if not isinstance(child, FSDPModule): + continue + for param_group in child._fsdp_param_groups: + if param_group.main_grad_buffer is not None: + param_group.main_grad_buffer.data.zero_() + param_group.release_grad_buffer() + for dist_param in param_group.dist_params: + if dist_param.grad is not None: + del dist_param.grad + if hasattr(dist_param, "decoupled_grad"): + dist_param.decoupled_grad = None + + def _copy_main_weights_to_model_weights(self): + """Copy main weight buffer to model weight buffer.""" + for child in self.modules(): + if not isinstance(child, FSDPModule): + continue + for param_group in child._fsdp_param_groups: + param_group.copy_main_weights_to_model_weights() + + def _compute_per_param_norms(self) -> Dict[str, Dict[str, float]]: + """ + Compute per-parameter L2 norms for params and grads. + + Returns {param_name: {"param_norm": float, "grad_norm": float}}. + Local squared norms are all-reduced across the DP group. + """ + results = {} + dp_group = None + + for module_name, child in self.named_modules(): + if not isinstance(child, FSDPModule): + continue + for param_names, param_group in child._named_param_groups: + if dp_group is None and param_group.dp_group is not None: + dp_group = param_group.dp_group + for param_name, dist_param, dist_grad in zip( + param_names, param_group.dist_params, param_group.dist_grads + ): + full_name = f"{module_name}.{param_name}" if module_name else param_name + results[full_name] = {"param_norm": 0.0, "grad_norm": 0.0} + if dist_param._local_tensor.numel() > 0: + results[full_name]["param_norm"] = ( + dist_param._local_tensor.float().norm(p=2).item() ** 2 + ) + if dist_grad is not None and dist_grad._local_tensor.numel() > 0: + results[full_name]["grad_norm"] = ( + dist_grad._local_tensor.float().norm(p=2).item() ** 2 + ) + + if dp_group is not None: + for param_name in results: + for key in ("param_norm", "grad_norm"): + value = torch.tensor([results[param_name][key]], device="cuda") + torch.distributed.all_reduce(value, group=dp_group) + results[param_name][key] = value.sqrt().item() + + return results + + def _log_per_param_norms(self, iteration: int, prefix: str = ""): + """Log per-parameter param and gradient L2 norms on rank 0.""" + norms = self._compute_per_param_norms() + if torch.distributed.get_rank() != 0: + return + for param_name in sorted(norms.keys()): + param_norm = norms[param_name]["param_norm"] + grad_norm = norms[param_name]["grad_norm"] + logger.info( + f"{prefix} iter={iteration} param={param_name} " + f"param_norm={param_norm:.6f} grad_norm={grad_norm:.6f}" + ) + + def _log_parameter_groups(self): + """Print a compact summary of rewrite-path FSDP parameter groups.""" + + def _fmt_dtype(dtype: torch.dtype) -> str: + short = { + torch.float32: "fp32", + torch.float16: "fp16", + torch.bfloat16: "bf16", + torch.int64: "i64", + torch.int32: "i32", + torch.uint8: "u8", + } + return short.get(dtype, str(dtype).removeprefix("torch.")) + + def _elem_size(dtype: torch.dtype) -> int: + return { + torch.float32: 4, + torch.float16: 2, + torch.bfloat16: 2, + torch.int64: 8, + torch.int32: 4, + torch.uint8: 1, + }.get(dtype, 1) + + def _mb(num_bytes: int | float) -> str: + return f"{num_bytes / 1_000_000:.2f} MB" + + rank = torch.distributed.get_rank() + lines = [f"FSDP parameter groups (rank {rank})"] + group_idx = 0 + total_model_elems = 0 + total_comm = 0 + total_pad = 0 + + for module_name, child in self.named_modules(): + if not isinstance(child, FSDPModule): + continue + for param_names, param_group in child._named_param_groups: + param_shapes = [p.shape for p in param_group.params] + numel = sum(s.numel() for s in param_shapes) + total_model_elems += numel + dp_size = torch.distributed.get_world_size(param_group.dp_group) + + buffer_entries = [] + group_pad = 0 + group_comm = 0 + for buffer_label, buffer in ( + ("W", param_group.model_weight_buffer), + ("MW", param_group.main_weight_buffer), + ("G", param_group.main_grad_buffer), + ): + if buffer is None: + continue + global_size = buffer.buffer_index.bucket_meta.size + elem_size = _elem_size(buffer.dtype) + group_pad += max(0, global_size - numel) * elem_size + group_comm += global_size * elem_size + dist_flag = "D" if buffer.is_distributed else "R" + buffer_entries.append( + f"{buffer_label}[{_fmt_dtype(buffer.dtype)}:{buffer.data_size}:{dist_flag}]" + ) + total_pad += group_pad + total_comm += group_comm + + lines.append( + f"- {module_name} #{group_idx} dp={dp_size} " + f"strategy={param_group.sharding_strategy} " + f"chunk_factor={param_group.chunk_size_factor}" + ) + lines.append( + f" {numel:,} elems x {_fmt_dtype(param_group.dtype)} " + f"comm={_mb(group_comm)} pad={_mb(group_pad)} " + f"{' '.join(buffer_entries)}" + ) + for param_name, param, param_shape in zip( + param_names, param_group.params, param_shapes + ): + dist_idx = param_group.param_idx.get(param) + offset_info = "" + if param_group.model_weight_buffer is not None and dist_idx is not None: + item_index = ( + param_group.model_weight_buffer.buffer_index.item_index_map.get( + dist_idx + ) + ) + if item_index is not None: + offset_info = f" @{item_index.global_data_index:,}+{item_index.size:,}" + lines.append(f" {param_name:50s} {str(tuple(param_shape)):24s}{offset_info}") + group_idx += 1 + + lines.append( + f"Summary: {group_idx} groups, {total_model_elems:,} model elems, " + f"comm={_mb(total_comm)}, pad={_mb(total_pad)}" + ) + logger.info("\n".join(lines)) + + def _set_nan_check(self, enable_nan_checks: bool): + """Enable or disable NaN checking.""" + for _, child in self.named_modules(): + if not isinstance(child, FSDPModule): + continue + setattr(child, "_enable_nan_checks", enable_nan_checks) + + if enable_nan_checks: + for name, param in self.named_parameters(): + if isinstance(param, DTensor): + param_data = param.data._local_tensor + else: + param_data = param.data + assert not torch.isnan(param_data).any(), f"NaN detected in parameter {name}" + for child in self.modules(): + if not isinstance(child, FSDPModule): + continue + for param_group in child._fsdp_param_groups: + for param in param_group.params: + wbuf = param_group.model_weight_buffer + param_data = wbuf.get_item(param_group.param_idx[param], as_shard=False) + assert not torch.isnan( + param_data + ).any(), "NaN detected in model weight buffer" + + +def _get_module_fsdp_param_groups( + module: nn.Module, + mp_policy: FullyShardMixedPrecisionPolicy, + allocator: BucketAllocator, + mesh: Optional[DeviceMesh] = None, + ignored_params: Optional[set[nn.Parameter]] = None, + gradient_scaling_factor: Optional[float] = None, +) -> List[ParameterGroup]: + """ + Group module parameters by (device, dtype, requires_grad) and create ParameterGroups. + + Parameters are grouped because they share the same buffer management + and sharding strategy. Each group gets its own DataParallelBuffer. + """ + param_groups = {} + + for param in module.parameters(): + if ignored_params is not None and param in ignored_params: + continue + + # The policy owns dtype-sensitive grouping, including FP8/MXFP8 tensors + # whose logical dtype may differ from their communication payload. + param_dtype = mp_policy.group_key_dtype(param) + param_attrs = (param.device, param_dtype, param.requires_grad) + if param_attrs not in param_groups: + param_groups[param_attrs] = [] + param_groups[param_attrs].append(param) + + # Create ParameterGroup for each group + fsdp_param_groups = [] + for i, params in enumerate(param_groups.values()): + fsdp_param_groups.append( + ParameterGroup( + params, + mesh=mesh, + param_group_id=ParamGroupIdx(id(module), i), + mp_policy=mp_policy, + gradient_scaling_factor=gradient_scaling_factor, + allocator=allocator, + ) + ) + + return fsdp_param_groups diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/fully_shard.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/fully_shard.py new file mode 100644 index 00000000000..b2446d3253f --- /dev/null +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/fully_shard.py @@ -0,0 +1,113 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Public fully_shard API for Megatron-FSDP2. + +The implementation is split across: +- fsdp_module.py: FSDPModule behavior and runtime state +- hooks.py: forward/backward hook registration +""" + +from typing import Callable, Optional + +import torch.nn as nn +from torch.distributed.device_mesh import DeviceMesh +from torch.distributed.tensor.placement_types import Shard + +from .allocator import StorageFreeingBucketAllocator, TracePoolAllocator +from .fsdp_module import FSDPModule +from .hooks import ( + _register_backward_hook, + _register_backward_pre_hook, + _register_fine_grained_forward_pre_hooks, + _register_forward_hook, + _register_forward_pre_hook, +) +from .mixed_precision import FullyShardMixedPrecisionPolicy +from .utils import _init_default_fully_shard_mesh + +__all__ = ["FSDPModule", "fully_shard"] + + +def fully_shard( + module: nn.Module, + *, + mesh: Optional[DeviceMesh] = None, + reshard_after_forward: Optional[bool | int] = None, # TODO: implement + shard_placement_fn: Optional[ + Callable[[nn.Parameter], Optional[Shard]] + ] = None, # TODO: implement + mp_policy: Optional[FullyShardMixedPrecisionPolicy] = None, + offload_policy: Optional["OffloadPolicy"] = None, # TODO: implement + ignored_params: Optional[set[nn.Parameter]] = None, + # --- Megatron-FSDP specific options --- + enable_unshard_prefetch: bool = True, + enable_async_reduce_grad: bool = True, + gradient_scaling_factor: Optional[float] = None, + enable_trace_pool: bool = False, +) -> nn.Module: + """ + Wrap a module with FSDP sharding semantics. + + This function: + 1. Converts the module class to FSDPModule dynamically (mixin pattern) + 2. Groups parameters by (device, dtype, requires_grad) + 3. Creates ParameterGroup for each group with dedicated buffers + 4. Registers forward/backward hooks for unshard/reshard/reduce + 5. Replaces module parameters with DTensor representations + """ + if isinstance(module, FSDPModule): + raise ValueError( + "The input module has already been fully sharded. " + "Please do not call fully_shard on the same module more than once." + ) + mesh = mesh or _init_default_fully_shard_mesh() + + if mp_policy is None: + mp_policy = FullyShardMixedPrecisionPolicy() + + cls = module.__class__ + new_cls = type(f"FSDP{cls.__name__}", (FSDPModule, cls), {}) + module.__class__ = new_cls + + bucket_allocator = ( + TracePoolAllocator() if enable_trace_pool else StorageFreeingBucketAllocator() + ) + + module._init_named_param_groups( + mesh, + ignored_params, + mp_policy=mp_policy, + bucket_allocator=bucket_allocator, + gradient_scaling_factor=gradient_scaling_factor, + ) + module._init_fsdp_state( + enable_unshard_prefetch=enable_unshard_prefetch, + enable_async_reduce_grad=enable_async_reduce_grad, + bucket_allocator=bucket_allocator, + ) + module._init_param_main_grad_func() + + if mp_policy.fine_grained_forward_hooks_required(module._fsdp_param_groups): + _register_fine_grained_forward_pre_hooks(module) + else: + _register_forward_pre_hook(module) + _register_forward_hook(module) + _register_backward_pre_hook(module) + _register_backward_hook(module) + + module.reshard() + + return module diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/mcore_fsdp_checkpoint_design.md b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/mcore_fsdp_checkpoint_design.md new file mode 100644 index 00000000000..ffbeb1fd245 --- /dev/null +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/mcore_fsdp_checkpoint_design.md @@ -0,0 +1,1009 @@ +# Megatron FSDP v2 Checkpoint Design + +## 1. Overview + +Megatron FSDP v2 (`use_megatron_fsdp_v2=True`) wraps model parameters as PyTorch +`DTensor` objects sharded across the data-parallel dimension. This document describes +the checkpoint save/load architecture and how it integrates with MCore's +`DistributedOptimizer` and DCP (`torch.distributed.checkpoint`). + +### Goals + +- Save and load Megatron FSDP v2 model + optimizer state via DCP. +- Support `fsdp_dtensor` as the canonical checkpoint format. +- Enable online checkpoint conversion from legacy Megatron formats (ND-parallel, + Megatron FSDP v1 baseline) to Megatron FSDP v2. +- Handle uneven DTensor sharding (parameters not evenly divisible by DP size). + +### Non-Goals (current scope) + +- Async checkpoint save/load for Megatron FSDP v2. +- Cross-node checkpoint resharding (different DP topology on load). + +### Feature Support Matrix + +| Source Format | Target Format | Model | Optimizer | Notes | +|---------------|--------------|-------|-----------|-------| +| MFSDP v2 | MFSDP v2 | ✓ | ✓ | Full round-trip and cross-setting (`optim_grads` ↔ `optim_grads_params`) | +| MFSDP v1 baseline | MFSDP v2 | ✓ | ✓ | Both `fsdp_dtensor` format, canonical key matching | +| ND-parallel (`torch_dist`) | MFSDP v2 | ✓ | ✓ (``fully_reshardable`` only) | Online conversion via `_load_torch_dist_into_megatron_fsdp_v2` in `checkpointing.py`. Expert weights are split from flattened multi-expert tensors. Optimizer states (`exp_avg`, `exp_avg_sq`) are loaded into V2's name-based format. Hi-precision FP32 optimizer param copies are used for model weights when available. ``dp_reshardable`` format is not supported (bucket-based layout incompatible with name-based V2 format). | +| ND-parallel (`torch`) | MFSDP v2 | ✗ | ✗ | Not supported (different serialization) | + +--- + +## 2. Background: Two FSDP Paths + +MCore wraps FSDP through `FullyShardedDataParallel` (in `mcore_fsdp_adapter.py`). +There are two code paths: + +| Path | Flag | Inner module | Model state dict | +|------|------|-------------|------------------| +| **Legacy Megatron FSDP** | `use_megatron_fsdp=True`, `use_megatron_fsdp_v2=False` | `MegatronFSDP` | `state_dict()` with DTensor hooks | +| **Megatron FSDP v2** | `use_megatron_fsdp=True`, `use_megatron_fsdp_v2=True` | `FSDPModule` (DTensor-native) | `state_dict()` (DTensors natively) | + +The `fsdp_dtensor` checkpoint format (`--ckpt-format fsdp_dtensor`) is the required +format for both paths. It uses DCP directly, storing each parameter as a `DTensor`. + +--- + +## 3. Architecture + +### 3.1 Component Map + +| Module | Location | Responsibility | +|--------|----------|----------------| +| `uneven_dtensor.py` | `megatron/core/distributed/fsdp/src/megatron_fsdp/` | `get_state_dict`, `preprocess_state_dict_for_uneven_dtensor`, chunk metadata for uneven DTensors | +| `fsdp_dtensor_checkpoint.py` | `megatron/core/transformer/` | SWiGLU split, GDN split, expert key remapping, FP8 cleanup | +| `distrib_optimizer.py` | `megatron/core/optimizer/` | `state_dict()`, `load_state_dict()`, `sharded_state_dict()`, `sharded_param_state_fsdp_dtensor()` | +| `checkpointing.py` | `megatron/training/` | High-level save/load orchestration, `_build_megatron_fsdp_v2_state_dict`, `preprocess_fsdp_dtensor_state_dict()` | +| `checkpoint.py` | `megatron/core/distributed/fsdp/` | `MegatronFSDPStateful`, `_apply_mcore_postprocess`, `_split_fused_params_v2` (SwiGLU+GDN), `_build_dtensor_optim_sd`, `_preprocess_and_verify_v2_state_dict`, `_build_torch_dist_to_v2_map`, `load_torch_dist_into_fsdp_v2` | +| `mcore_fsdp_adapter.py` | `megatron/core/distributed/fsdp/` | Routes to v1 `MegatronFSDP` or Megatron FSDP v2 `fully_shard` | + +### 3.1.1 `checkpoint.py` Functions (Megatron FSDP v2) + +| Function | Description | +|----------|-------------| +| `MegatronFSDPStateful` | ``Stateful`` wrapper implementing DCP protocol. ``state_dict()`` calls ``_get_state_dict`` from ``uneven_dtensor`` (attaches uneven DTensor chunk metadata), then ``_apply_mcore_postprocess`` (SwiGLU/GDN split, FP8 cleanup, expert remapping). ``load_state_dict()`` uses PyTorch's ``set_state_dict``. | +| `_preprocess_and_verify_v2_state_dict` | Build a shadow optimizer state dict with DTensors sharing storage with original plain tensors (dual-dict pattern). Verifies ``__create_chunk_list__`` and ``__create_write_items__`` metadata on all model and optimizer DTensors. Returns canonical ``(v2_by_canonical, v2_optim_state)`` maps. | +| `_build_torch_dist_to_v2_map` | Build mapping from torch_dist metadata keys to v2 DTensors. Returns ``(regular_model, hi_prec_model, optim_keys, optim_matched)``. | +| `_build_dtensor_optim_sd(raw_opt_sd, model)` | Wrap plain-tensor optimizer states as uneven DTensors. Returns a copy (does not mutate original). Uses ``_maybe_wrap_as_uneven_dtensor`` internally. | +| `_maybe_wrap_as_uneven_dtensor(tensor, dist_param)` | Wrap a single plain tensor as an uneven DTensor if it matches the parameter's local shard shape; otherwise return unchanged. Shared by ``_build_dtensor_optim_sd`` and ``_preprocess_and_verify_v2_state_dict``. | +| `load_torch_dist_into_fsdp_v2` | Entry point for online checkpoint conversion from legacy ``torch_dist`` format to ``fsdp_dtensor``. Five-phase pipeline: preprocess/verify, build name mapping, DCP load, expert params, verify. | +| `add_module_prefix(state_dict)` | Add ``module.`` prefix to all state dict keys. Megatron FSDP v2 lacks ``MegatronFSDP`` wrapper so ``model.state_dict()`` keys have no prefix; this aligns with Megatron's checkpoint format. | +| `strip_module_prefix(state_dict)` | Remove ``module.`` prefix from state dict keys. Inverse of ``add_module_prefix``, used when loading checkpoint back into FSDP v2 model. | +| `get_model_state_dict(model)` | Get model state dict with ``module.`` prefix. Auto-detects whether prefix is already present; adds it if missing. | +| `get_optimizer_state_dict(optimizer, is_loading)` | Get optimizer state dict via Path A (``sharded_param_state_fsdp_dtensor``). Delegates to ``optimizer.sharded_state_dict()`` with ``fsdp_dtensor`` sharding type. Returns ``{"state": ..., "param_to_group_meta": ...}``. | +| `handle_fp8_extra_state_case(model_sd)` | Remove ``._extra_state`` keys from model state dict (FP8 artifact cleanup). | +| `handle_experts_in_state_dict(model_sd, num_experts)` | Rename expert parameter keys for expert-parallel sharding. | +| `handle_swiglu_in_state_dict_v2(model, model_sd, opt_sd)` | Thin wrapper: precomputes ``layer_glu`` map, delegates to ``_split_fused_params_v2`` with SwiGLU detector and ``_w``/``_v`` suffix format. | +| `handle_gdn_in_state_dict_v2(model, model_sd, opt_sd)` | Thin wrapper: delegates to ``_split_fused_params_v2`` with ``_match_gdn_key`` detector and ``.query``/``.key``/``.value`` suffix format. | +| `_split_fused_params_v2(model, model_sd, opt_sd, detector, key_fmt, tag)` | Unified fused-parameter splitting skeleton. Iterates model and optimizer state dicts, calls the ``detector(key, dtensor, model) -> (sizes, names, dim)`` callback to identify fused tensors, splits via ``split_dtensor``, and renames keys via ``key_fmt(key, sub_name)``. Shared by SwiGLU and GDN. | +| `_is_swiglu_key(key)` | Check whether a key matches any SwiGLU fc1 pattern. | +| `_match_gdn_key(key, dtensor)` | GDN detector: returns ``(sizes, names, dim)`` for fused QKV projections, or ``None``. | +| `_detect_glu_layers(model)` | Return ``{layer_path: gated_linear_unit}`` for all TransformerLayers. | +| `_model_has_module_prefix(model)` | Detect whether model's ``named_parameters()`` keys already carry ``module.`` prefix. | +| `normalize_torch_dist_key(key)` | Normalize a torch_dist checkpoint key to v2 canonical form. Maps ``transformer_layer`` → ``mtp_model_layer``. | +| `reverse_normalize_torch_dist_key(key)` | Reverse the v2 canonical key back to torch_dist naming (``mtp_model_layer`` → ``transformer_layer``). Used when constructing DCP load paths that must match torch_dist storage paths. | + + +### 3.2 Current Save Flow + +> **Note:** This describes the current `generate_state_dict`-based flow. For FSDP v2, +> ``_build_megatron_fsdp_v2_state_dict`` (Section 12) replaces this in the v2 code path. + +``` +save_checkpoint() + | + +-- 1. generate_state_dict() + | | + | +-- model[i].state_dict_for_save_checkpoint() + | | Returns DTensor state dict. + | | + | +-- optimizer.sharded_state_dict(state_dict, metadata={'distrib_optim_sharding_type': 'fsdp_dtensor'}) + | | Returns {"state": {name: opt_state}, "param_to_group_meta": {...}} + | | via sharded_param_state_fsdp_dtensor(). + | | + | +-- rng_state, scheduler state, etc. + | + +-- 2. preprocess_fsdp_dtensor_state_dict() + | - handle_fp8_extra_state_case + | - handle_swiglu_in_state_dict (model + optimizer) + | - handle_experts_in_state_dict (EP key remapping) + | - preprocess_state_dict_for_uneven_dtensor + | + +-- 3. torch.distributed.checkpoint.save(state_dict, storage_writer) +``` + +### 3.3 Current Load Flow + +> **Note:** This describes the current `generate_state_dict`-based flow. For FSDP v2, +> ``_build_megatron_fsdp_v2_state_dict`` (Section 12) replaces this in the v2 code path. + +``` +load_checkpoint() + | + +-- 1. Build sharded_state_dict via generate_state_dict() + | Same structure as save. For optimizer: is_loading=True triggers + | _init_optimizer_states_with_dummy_values() to create placeholder states. + | + +-- 2. _load_base_checkpoint() + | - preprocess_fsdp_dtensor_state_dict() + | - torch.distributed.checkpoint.load_state_dict(state_dict, reader) + | + +-- 3. Post-load application + - ddp_model[i].load_state_dict(state_dict['model'], strict) + - optimizer.load_state_dict(state_dict['optimizer']) + - RNG states, scheduler, etc. +``` + +--- + +## 4. Model State Dict + +Megatron FSDP v2 uses `model.state_dict()` which returns a dict of `DTensor` values. +The keys follow the Megatron `module.` prefix convention +(e.g., `module.embedding.word_embeddings.weight`). + +### 4.1 `state_dict_for_save_checkpoint` — Adapter (Phase 1 Bridge) + +**Current status:** Set to `not_implemented_op` for v2 (line 384 of `mcore_fsdp_adapter.py`). + +**Phase 1 fix:** Wire to `model.state_dict()`. This is a **temporary bridge** until +Phase 2 replaces `generate_state_dict` with `get_state_dict` from `uneven_dtensor.py`. +Once Path B is the primary path, this wiring is no longer needed and will be removed +in Phase 3. + +```python +# In _init_with_fully_shard(), replace: +self.module.state_dict_for_save_checkpoint = not_implemented_op +self.state_dict_for_save_checkpoint = not_implemented_op +# With: +self.module.state_dict_for_save_checkpoint = lambda *args, **kwargs: module.state_dict() +self.state_dict_for_save_checkpoint = lambda *args, **kwargs: module.state_dict() +``` + +### 4.2 `load_state_dict` — Adapter + +**Current:** Calls `super().load_state_dict()` at line 398-399 which delegates through +the FSDP framework. **Status: OK.** + +```python +if self.ddp_config.use_megatron_fsdp_v2: + super().load_state_dict(state_dict, strict=strict) + return +``` + +--- + +## 5. Optimizer State Dict + +There are **two** paths for obtaining the optimizer state dict, depending on context. + +### 5.1 Path A: Megatron Training Loop (`checkpointing.py`) + +Uses `optimizer.sharded_state_dict(model_state_dict, sharding_type="fsdp_dtensor")`, +which calls `sharded_param_state_fsdp_dtensor()`. Returns: + +```python +{ + "state": { + "": {"exp_avg": torch.Tensor, "exp_avg_sq": torch.Tensor, ...}, + ... + }, + "param_to_group_meta": { + "": {"lr": ..., "weight_decay": ...}, + ... + } +} +``` + +**Important:** ``FusedAdam`` (and similar NVIDIA optimizers) store ``exp_avg`` / +``exp_avg_sq`` as **plain tensors** matching the parameter's local DTensor shard, +NOT as DTensors. This is because the optimizer kernel operates on the local data +directly. Path A returns these plain tensors as-is. + +For DCP compatibility, these plain tensors must be wrapped as DTensors via +``_wrap_optim_states_as_dtensors`` before DCP save (see Section 5.10), and +unwrapped back via ``_unwrap_optim_states_from_dtensors`` after DCP load. + +This is the primary path for Megatron-integrated training. + +### 5.2 Path B: Standalone `get_state_dict()` (`uneven_dtensor.py`) + +Uses PyTorch's native `torch.distributed.checkpoint.state_dict.get_state_dict()`, +which calls `optimizer.state_dict()` internally. For `DistributedOptimizer` with +Megatron FSDP, `state_dict()` returns the inner optimizer's full state dict directly. + +This path is used by: +- `checkpoint.py` (``MegatronFSDPStateful`` wrapper and MCore post-processing helpers) +- `test_mcore_checkpoint.py` (checkpoint save/load and online format conversion tests) + +### 5.3 `DistributedOptimizer.__init__` — FSDP Short-Circuit + +When `use_megatron_fsdp=True` or `use_megatron_fsdp_v2=True`, `__init__()` returns +early (line 543) without setting up buffer ranges, gbuf mappings, or shard slicing. +Megatron FSDP manages weight/gradient memory directly. + +### 5.4 `DistributedOptimizer.state_dict()` — FSDP Branch + +Returns the **full** inner optimizer state dict because: +- FSDP manages parameter state as DTensors (no separate `save_parameter_state`). +- PyTorch's `get_state_dict()` expects `optimizer.state_dict()` to include state. +- The `sharded_param_state_fsdp_dtensor()` path handles Megatron-specific key + remapping separately. + +```python +if self.ddp_config.use_megatron_fsdp or self.ddp_config.use_megatron_fsdp_v2: + return self.optimizer.state_dict() +``` + +**Why not strip `"state"` like the non-FSDP path?** The legacy non-FSDP path strips +the `"state"` key from `state_dict()` and stores optimizer parameter states in a +separate `param_state` checkpoint. This is necessary because the non-FSDP path manages +parameters in contiguous gradient buffers requiring manual sharding. For Megatron FSDP, +DCP handles DTensor sharding natively, so splitting is unnecessary. + +### 5.5 `DistributedOptimizer.load_state_dict()` — FSDP Branch + +Converts name-based `param_to_group_meta` back to tensor-based `param_groups`, then +delegates to the inner optimizer: + +```python +if self.ddp_config.use_megatron_fsdp or self.ddp_config.use_megatron_fsdp_v2: + if "param_to_group_meta" in state_dict: + state_dict["param_groups"] = self._param2group_meta_to_param_groups(...) + self.optimizer.load_state_dict(state_dict) + return +``` + +### 5.6 `DistributedOptimizer.sharded_state_dict()` — FSDP Branch + +Only `sharding_type="fsdp_dtensor"` is supported. Calls +`sharded_param_state_fsdp_dtensor()` which: +1. Optionally initializes optimizer states with dummy values (for loading). +2. Maps tensor keys to parameter name strings via `_param_name()`. +3. Returns `{"state": ..., "param_to_group_meta": ...}`. + +**Why this works for Megatron FSDP v2:** The v2 path uses the same `self.optimizer` +(a standard Torch optimizer like AdamW) and the same `_param_name` mapping. DTensor +parameters are correctly identified by name. The state dict keys are the same as the +checkpoint's model state dict keys, so DCP can match them. + +### 5.8 `sharding_type="fsdp_dtensor"` Rationale + +The `fsdp_dtensor` checkpoint format signals to MCore's `checkpointing.py` to: +1. Use DCP for all I/O. +2. Store parameters as DTensors (each carries its own sharding metadata). +3. Store optimizer state as a flat `{param_name: state}` dict. + +This is the only format compatible with Megatron FSDP because the non-DCP formats +(`torch`, `torch_dist`) use gather/scatter patterns that assume contiguous gradient +buffers, which don't exist in the FSDP path. + +### 5.10 DTensor Wrapping / Unwrapping Layer (Path A) + +Path A must add a wrapping/unwrapping layer to bridge the gap between +FusedAdam's plain-tensor optimizer states and DCP's DTensor requirement. + +**Why wrapping is needed on save:** + +DCP builds a global save plan by inspecting the state dict. Sharded data must +be DTensors so DCP knows about the sharding (mesh, placements, per-rank chunk +metadata). Plain tensors with the same FQN on every rank are treated as +non-sharded (replicated) data, which causes ``"item.index.fqn not in md"`` +errors when ranks hold different local shards. + +**Why unwrapping is needed on load:** + +After DCP loads DTensors into the skeleton, the optimizer state dict contains +DTensor values. But FusedAdam's ``load_state_dict`` → ``set_scaled_state`` +calls ``state[state_name].copy_(unscaled_state)`` where ``state[state_name]`` +is a plain tensor (FusedAdam's internal storage). Passing a DTensor causes +``RuntimeError: aten.copy_.default got mixed torch.Tensor and DTensor``. + +**The wrapping functions (in ``checkpoint.py``):** + +- ``_build_dtensor_optim_sd(raw_opt_sd, model)`` — wraps every plain-tensor optimizer + state that matches a parameter's local shard as an uneven DTensor. Returns a copy; + original is not mutated. +- ``_maybe_wrap_as_uneven_dtensor(tensor, dist_param)`` — wraps a single tensor. Shared + by ``_build_dtensor_optim_sd`` and ``_preprocess_and_verify_v2_state_dict``. Returns + the tensor unchanged if it is already a DTensor or the shape doesn't match. + +**Placement in the flow:** + +:: + + Save: get_optim_state_dict (plain) → _wrap_optim_states_as_dtensors → DCP save + Load: DCP load (DTensors) → _unwrap_optim_states_from_dtensors → optimizer.load_state_dict + +**Note on the baseline (v1) path:** + +The baseline path does NOT need these functions because the existing conversion +tests (``test_mcore_checkpoint.py``) never save optimizer states to DCP — they +save only model state dicts via ``dcp_save({"model": source_sd}, ...)``. If +optimizer state saving were added to the baseline path, the same wrapping/ +unwrapping layer would be required. + +### 5.11 Dual-Dict Pattern: Plain Tensors → DTensor Wrapping → Restore + +The main training loop (``checkpointing.py``) uses a **dual-dict pattern** to +avoid the need for ``_unwrap_optim_states_from_dtensors``: + +:: + + _build_megatron_fsdp_v2_state_dict() + | + +-- state_dict["optimizer"] = plain tensors (from get_optimizer_state_dict) + | + save_checkpoint() / _load_base_checkpoint() + | + +-- (load only) raw_optimizer_state_dict = state_dict["optimizer"].copy() + | captures plain-tensor dict before wrapping + | + +-- _wrap_optim_states_as_dtensors(state_dict, model) + | wraps in-place: plain → DTensor (shared storage via make_uneven_dtensor) + | + +-- DCP save / DCP load + | writes through DTensors → plain tensors auto-updated via shared storage + | + +-- (load only) state_dict["optimizer"] = raw_optimizer_state_dict + restores plain-tensor dict for optimizer.load_state_dict + +Key properties: + +* **Wrapping is deferred** — ``_build_megatron_fsdp_v2_state_dict`` keeps + optimizer states as plain tensors. Wrapping as DTensors happens later, in + ``save_checkpoint`` (before DCP save) or ``_load_base_checkpoint`` (after + raw copy, before DCP load). +* **No unwrapping needed** — the ``raw_optimizer_state_dict.copy()`` at + line 1526 of ``checkpointing.py`` (existing code) captures the plain-tensor + state dict. After DCP load writes through the DTensors (which share storage + with the plain tensors), the raw dict is restored. The raw dict already has + the correct loaded values — no ``.to_local()`` conversion needed. +* **split_optimizer=False** — ``_apply_mcore_postprocess`` only splits model + keys (SwiGLU ``_w``/``_v``). Optimizer keys stay as canonical parameter + names, which ``DistributedOptimizer.load_state_dict`` can match. + +--- + +## 6. Proposed: `uneven_dtensor.get_state_dict` as Primary Path + +### 6.1 Motivation + +The current `generate_state_dict()` in `checkpointing.py` mixes three concerns: +1. Building state dict structure per PP chunk +2. Format-specific branching (`torch_dist` vs `fsdp_dtensor` vs legacy) +3. Dispatching to optimizer-specific APIs (`sharded_state_dict` with sharding type) + +For Megatron FSDP v2, we can do better. `get_state_dict()` from `uneven_dtensor.py` +wraps PyTorch's native `torch.distributed.checkpoint.state_dict.get_state_dict()` +with uneven DTensor preprocessing. It produces correct DTensor state dicts for both +model and optimizer in a single call. MCore-specific post-processing (FP8 cleanup, +SwiGLU split, expert key remapping) can be applied as a separate layer. + +### 6.2 Proposed Architecture + +``` +┌──────────────────────────────────────────────────────────────┐ +│ Layer 1: Megatron FSDP State Dict │ +│ uneven_dtensor.get_state_dict(model, optimizer) │ +│ → model_state_dict: {name: DTensor} │ +│ → optimizer_state_dict: {state: ..., param_groups: ...} │ +│ Internally calls PyTorch's get_state_dict + preprocesses │ +│ uneven DTensor chunk metadata for DCP serialization. │ +│ Handles: DTensor serialization, model/optimizer state dict │ +│ separation, uneven DTensor chunk metadata │ +└──────────────────────────┬───────────────────────────────────┘ + │ +┌──────────────────────────▼───────────────────────────────────┐ +│ Layer 2: MCore Post-Processing │ +│ preprocess_fsdp_dtensor_state_dict() │ +│ → handle_fp8_extra_state_case() # strip FP8 artifacts │ +│ → handle_swiglu_in_state_dict() # split fc1 into _w/_v │ +│ → handle_experts_in_state_dict() # EP key remapping │ +│ (uneven DTensor is already handled by Layer 1) │ +└──────────────────────────┬───────────────────────────────────┘ + │ +┌──────────────────────────▼───────────────────────────────────┐ +│ Layer 3: DCP I/O │ +│ torch.distributed.checkpoint.save(state_dict, writer) │ +│ torch.distributed.checkpoint.load_state_dict(state_dict, │ +│ reader) │ +└──────────────────────────────────────────────────────────────┘ +``` + +### 6.3 Save Flow (Proposed) + +```python +def save_checkpoint_fsdp_v2(model_chunks, optimizer, ...): + # Layer 1: State dicts via uneven_dtensor.get_state_dict + from megatron.core.distributed.fsdp.src.megatron_fsdp.uneven_dtensor import ( + get_state_dict, + ) + + for pp_rank, model_chunk in enumerate(model_chunks): + model_sd, optim_sd = get_state_dict( + model=model_chunk, + optimizers=optimizer, + ) + # model_sd keys: "module.embedding.weight", "module.layers.0.fc.weight", ... + # optim_sd keys: "state", "param_groups" + # Uneven DTensor chunk metadata is already attached by get_state_dict. + + state_dict = { + f"model{pp_rank}" if len(model_chunks) > 1 else "model": model_sd, + "optimizer": optim_sd, + "rng_state": ..., + } + + # Layer 2: MCore post-processing (FP8, SwiGLU, expert keys) + preprocess_fsdp_dtensor_state_dict(args, state_dict, model_chunks[0]) + + # Layer 3: DCP save + torch.distributed.checkpoint.save(state_dict, checkpoint_id=str(ckpt_dir)) +``` + +### 6.4 Load Flow (Proposed) + +> **Important:** `get_state_dict` is designed for save — it gathers current +> state into a dict. For loading, we need a **skeleton** (dict with correct keys and +> tensor shapes but empty/placeholder values) that DCP fills in-place. The optimizer +> states must be pre-allocated BEFORE building the skeleton, so the skeleton has the +> right structure for DCP to write into. + +```python +def load_checkpoint_fsdp_v2(model_chunks, optimizer, ...): + from megatron.core.distributed.fsdp.src.megatron_fsdp.uneven_dtensor import ( + get_state_dict, + ) + + # Pre-allocate: optimizer states must exist before get_state_dict + # can produce a skeleton with the right structure. + # Equivalent to the current is_loading=True path in sharded_param_state_fsdp_dtensor. + optimizer._init_optimizer_states_with_dummy_values() + + for pp_rank, model_chunk in enumerate(model_chunks): + # Build skeleton — get_state_dict produces a dict with the right keys + # and tensor shapes, plus uneven DTensor chunk metadata for DCP. + # Values are stale (skeleton only needs structure). + model_sd, optim_sd = get_state_dict( + model=model_chunk, + optimizers=optimizer, + ) + + state_dict = { + f"model{pp_rank}" if len(model_chunks) > 1 else "model": model_sd, + "optimizer": optim_sd, + } + + # MCore post-processing (on skeleton, before DCP load) + state_dict["_model"] = model_chunks[0] + preprocess_fsdp_dtensor_state_dict(args, state_dict, model_chunks[0]) + + # DCP load fills DTensors in-place. + torch.distributed.checkpoint.load_state_dict( + state_dict=state_dict, + storage_reader=FileSystemReader(checkpoint_id), + planner=DefaultLoadPlanner(allow_partial_load=not strict), + ) + + # After DCP load, model DTensors are already updated in-place. + # For v2, load_state_dict is a no-op for DTensor params but required + # for non-DTensor state (buffers, extra_state) and strict key checking. + model_chunk.load_state_dict(state_dict["model"], strict=False) + optimizer.load_state_dict(state_dict["optimizer"]) +``` + +### 6.5 Key Design Decisions + +**Why `uneven_dtensor.get_state_dict` instead of `model.state_dict_for_save_checkpoint` + `optimizer.sharded_state_dict`?** + +| Aspect | Current (Path A) | Proposed (Path B) | +|--------|-----------------|-------------------| +| Model state dict | `model.state_dict_for_save_checkpoint()` — `not_implemented_op` for v2 | `get_state_dict(model, optimizer)` — works natively with DTensors | +| Optimizer state dict | `optimizer.sharded_state_dict(state_dict, sharding_type="fsdp_dtensor")` → `sharded_param_state_fsdp_dtensor()` | `get_state_dict(model, optimizer)` → calls `optimizer.state_dict()` internally | +| PP handling | `checkpointing.py` iterates model chunks manually | `get_state_dict` handles via `submodules`; caller iterates PP chunks | +| Loading skeleton | `is_loading=True` → `_init_optimizer_states_with_dummy_values()` | `get_state_dict` with pre-allocated optimizer states | +| Uneven DTensor | Done inside `preprocess_fsdp_dtensor_state_dict` | Done inside `get_state_dict` (via `preprocess_state_dict_for_uneven_dtensor`) — no separate step needed | + +**Why use `uneven_dtensor.get_state_dict` instead of raw `torch.distributed.checkpoint.state_dict.get_state_dict`?** + +`uneven_dtensor.get_state_dict` adds uneven DTensor chunk metadata in the same call. +With the raw PyTorch version, `preprocess_state_dict_for_uneven_dtensor` must be called +separately, and the caller must ensure it is not double-called (once by the wrapper, +once by `preprocess_fsdp_dtensor_state_dict`). Using `uneven_dtensor.get_state_dict` +consolidates this into a single well-tested entry point. + +**Why separate the layers?** + +1. **Testability** — Each layer can be tested independently. State dict correctness + can be verified with simple models. MCore post-processing can be tested with + synthetic state dicts. + +2. **Reduced coupling** — `checkpointing.py` no longer needs to know about + `sharding_type` or `sharded_param_state_fsdp_dtensor`. It just calls + `get_state_dict`, applies post-processing, and hands to DCP. + +3. **Online conversion is trivial** — The same `get_state_dict` + post-processing + pipeline works for both save and online conversion (loading a legacy checkpoint into + a v2 model). The only difference is the DCP operation (save vs load). + +### 6.6 Prerequisites + +For this approach to work, the following must be in place: + +1. **`DistributedOptimizer.state_dict()` returns inner state dict** — Already implemented + (Section 5.4). `get_state_dict` calls `optimizer.state_dict()` internally via + PyTorch's native `get_state_dict`. + +2. **`DistributedOptimizer.load_state_dict()` handles inner state dict** — Already + implemented (Section 5.5). After DCP loads into the skeleton, we call + `optimizer.load_state_dict()`. + +3. **Optimizer states pre-allocated before load** — DCP does in-place loading into + the skeleton's tensor buffers. The skeleton must have correctly-sized optimizer + state tensors (exp_avg, exp_avg_sq) before DCP writes into them. Currently this + is done by `_init_optimizer_states_with_dummy_values()` (called by + `sharded_param_state_fsdp_dtensor` with `is_loading=True`). With Path B, this + method must be called explicitly before `get_state_dict`. + +4. **Model `load_state_dict` works** — Already implemented (Section 4.2). + +5. **Post-processing functions are FSDP-agnostic** — `handle_fp8_extra_state_case`, + `handle_swiglu_in_state_dict`, `handle_experts_in_state_dict` operate on dicts, not + on model internals. They work for both DTensor and regular tensor state dicts. + +6. **Uneven DTensor preprocessing handled by `get_state_dict`** — `uneven_dtensor.get_state_dict` + calls `preprocess_state_dict_for_uneven_dtensor` on both model and optimizer state + dicts. `preprocess_fsdp_dtensor_state_dict` also calls it. To avoid double-calling, + `preprocess_fsdp_dtensor_state_dict` should skip the uneven DTensor step when the + state dict already went through `get_state_dict`. + +### 6.7 Migration Path + +| Phase | Description | +|-------|-------------| +| **Phase 1** (now) | Keep current `generate_state_dict` path. Add `use_megatron_fsdp_v2` guards to `distrib_optimizer.py`. Wire `state_dict_for_save_checkpoint` in adapter. | +| **Phase 2** | Add `get_state_dict`-based path (from `uneven_dtensor.py`) as an alternative code path in `checkpointing.py`, gated by `use_megatron_fsdp_v2`. Verify round-trip correctness. | +| **Phase 3** | Deprecate `sharded_param_state_fsdp_dtensor` and the `is_loading` skeleton pattern for v2. Remove `state_dict_for_save_checkpoint` wiring from adapter. | + +--- + +## 7. Online Checkpoint Conversion + +### 7.1 Problem + +Users may have checkpoints from legacy Megatron formats (ND-parallel with +`DistributedOptimizer`, or Megatron FSDP v1 baseline) and want to load them into a +Megatron FSDP v2 model. Key structures differ: + +| Format | Model Keys | Optimizer Keys | +|--------|-----------|----------------| +| ND-parallel | `module.layer.weight` | Tensor-keyed (by param tensor id) | +| Megatron FSDP v1 baseline | `module.layer.weight` | Tensor-keyed | +| Megatron FSDP v2 | `module.layer.weight` | String-keyed (by param name) | + +### 7.2 Solution: `_load_torch_dist_into_megatron_fsdp_v2` + +The function `_load_torch_dist_into_megatron_fsdp_v2` in `checkpointing.py` is the entry +point for online conversion from `torch_dist` to `fsdp_dtensor` format. It is called +from `_load_global_dist_base_checkpoint` when `use_megatron_fsdp_v2` is set and the +source checkpoint uses `torch_dist` format. + +The conversion proceeds in five phases (implemented in `load_torch_dist_into_fsdp_v2`): + +#### Phase 1 — Preprocess & Verify v2 State Dict + +``_preprocess_and_verify_v2_state_dict`` builds canonical maps of v2 model and +optimizer state entries. Plain-tensor optimizer states are wrapped as uneven +DTensors sharing storage with the originals (dual-dict pattern, see Section 5.11). +Both ``__create_chunk_list__`` and ``__create_write_items__`` metadata are verified +on all model and optimizer DTensors. + +#### Phase 2 — DCP Key Mapping + +``_build_torch_dist_to_v2_map`` iterates torch_dist metadata keys and matches them +against canonical v2 entries: regular model weights, hi-precision (``param``) +optimizer copies, and optimizer state tensors (``exp_avg``, ``exp_avg_sq``). + +Metadata keys from the torch_dist checkpoint are canonicalized for +**matching** against v2 entries, while the original torch_dist storage paths are kept +as DCP state-dict keys so they match the checkpoint metadata verbatim: + +- **Model weights** (``model.``): the ``model.`` prefix is stripped and + shard suffixes (``/shard_X_Y`` on ``_extra_state`` entries) are removed. The + remaining name is canonicalized via ``normalize_torch_dist_key`` + (``transformer_layer`` → ``mtp_model_layer``) **only for matching** + against the v2 model's canonical keys. The **original** torch_dist name + (without the ``model.`` prefix) is stored as the DCP load key. + +- **Hi-precision optimizer copies** (``optimizer.state.param.``): the + ``optimizer.state.param.`` prefix is stripped. The param name is canonicalized + and ``module.`` prefix is stripped for matching; the original torch_dist name is + used as the DCP key. When both a regular model weight and a hi-prec optimizer + copy map to the same v2 DTensor, the hi-prec copy takes priority. + +- **Optimizer state tensors** (``optimizer.state.exp_avg.``, + ``optimizer.state.exp_avg_sq.``): the state key and param name are + extracted; the param name is canonicalized and ``module.`` prefix stripped for + matching against the v2 optimizer's ``state`` dict + (``v2_optim_state[canonical_name][state_key]``). The full original + torch_dist key (``optimizer.state.exp_avg.original_name``) is used as the DCP + load key. + +#### Phase 3 — Single DCP Load + +A single ``dcp.load`` call loads all matched tensors: + +```python +mapped_sd = { + "model": { + "decoder.layers.0.weight": v2_dtensor, # regular weights + ... + }, + "optimizer": { + "state": { + "exp_avg": { + "decoder.layers.0.weight": v2_opt_dtensor, + ... + }, + "exp_avg_sq": { ... }, + "param": { # hi-prec model copies + "decoder.layers.0.weight": v2_dtensor, + ... + }, + } + } +} +dcp.load(state_dict=mapped_sd, checkpoint_id=..., planner=DefaultLoadPlanner(allow_partial_load=True)) +``` + +The DCP state dict mirrors the torch_dist checkpoint directory structure: +regular weights under ``model.``, and all optimizer-related tensors (including +hi-precision parameter copies) under ``optimizer.state.*``. + +After loading, hi-precision model copies are merged back into the `model` subtree +so the model state dict is complete. + +#### Phase 4 — Load Fused Layer / Expert Params by Slicing + +Torch_dist stores multi-layer and multi-expert tensors as a single fused tensor (e.g., +``decoder.layers.self_attention.linear_qkv.weight`` of shape ``(num_layers, H, 3*W)``). +FSDP v2 stores each layer as an individual DTensor +(``decoder.layers.0.self_attention.linear_qkv.weight`` of shape ``(H, 3*W)``). +DCP cannot split one tensor into many, so Phase 4 handles this in `load_torch_dist_into_fsdp_v2`. + +Three tensor formats are supported: + +1. **Regular fused** — shape ``(num_layers, ...)`` → sliced per layer index. +2. **GroupedMLP experts** — shape ``(num_layers, num_experts, ...)`` → sliced per + ``(layer_idx, expert_idx)``. +3. **SequentialMLP experts** — shape ``(num_global_experts, ...)`` → sliced per global + expert index (EP-aware, mapped to local expert via + ``ep_rank * num_local + local_expert_idx``). + +To avoid OOM when loading a large fused tensor (e.g., GroupedMLP with many layers and +experts), the fused buffer is created as a **Shard(0) DTensor** across the DP device +mesh: + +```python +device_mesh = example_val.device_mesh +flat = torch.distributed.tensor.empty( + fused_shape, dtype=example_val.dtype, + device_mesh=device_mesh, placements=[Shard(0)], +) +dcp.load(state_dict={td_key: flat}, ...) +flat = redistribute_uneven_dtensor_to_replicated(flat).to_local() +``` + +- ``torch.distributed.tensor.empty(..., placements=[Shard(0)])`` distributes memory + across DP ranks so no single rank allocates the full fused tensor. +- ``redistribute_uneven_dtensor_to_replicated`` gathers shards into a full local + tensor for per-layer/per-expert slicing. +- After slicing, ``__create_chunk_list__`` metadata is used to copy each rank's + DP-shard chunk with correct local offsets into the destination v2 DTensor. + +#### Phase 5 — Strictness Verification + +After all four loading phases complete, the v2 model's parameter names are compared +against the union of all loaded entries. ``_extra_state`` entries are excluded (they +are FP8/FP4 metadata that reinitializes on load). Optimizer state entries not matched +to torch_dist data are also reported. Any unmatched parameter triggers a +``RuntimeError``, ensuring no weights are silently skipped. + +### 7.3 Key Normalization Helpers + +Located in `megatron/core/distributed/fsdp/checkpoint.py`: + +| Function | Direction | Transforms | +|----------|-----------|------------| +| `normalize_torch_dist_key` | torch_dist → v2 | ``transformer_layer`` → ``mtp_model_layer`` | +| `reverse_normalize_torch_dist_key` | v2 → torch_dist | ``mtp_model_layer`` → ``transformer_layer`` | + +These are used both in `_load_torch_dist_into_megatron_fsdp_v2` (for matching DCP +keys) and in Phase 4 of `load_torch_dist_into_fsdp_v2` (for constructing storage paths +that match the torch_dist checkpoint). + +### 7.4 Strictness Checks + +After all phases complete (see Phase 5 above), two strictness checks run: + +1. **Model parameter coverage:** every v2 model canonical parameter name must appear + in the set of loaded entries (regular model + hi-prec + expert params). + ``._extra_state`` entries are excluded (FP8/FP4 metadata that reinitializes). + +2. **Optimizer state coverage:** every v2 optimizer state parameter must have been + matched to at least one torch_dist state tensor. Unmatched entries trigger a + ``RuntimeError``. + +--- + +## 8. Edge Cases + +### 8.1 Multi-Optimizer (MoE: Expert + Non-Expert) + +MoE models use a `ChainedOptimizer` wrapping two `DistributedOptimizer` instances +(one for expert params, one for non-expert params). The proposed Path B must handle +this by building separate state dicts per sub-optimizer: + +```python +# For ChainedOptimizer (expert + non-expert): +from megatron.core.distributed.fsdp.src.megatron_fsdp.uneven_dtensor import get_state_dict + +state_dict["optimizer_expert"] = get_state_dict(model, optimizers=optimizer.expert) +state_dict["optimizer_non_expert"] = get_state_dict(model, optimizers=optimizer.non_expert) +``` + +The current `generate_state_dict` path handles this by calling `optimizer.sharded_state_dict()` +on the `ChainedOptimizer`, which internally dispatches to each sub-optimizer. + +### 8.2 Frozen Sub-Models (`optimizer is None`) + +When sub-models are frozen, the `DistributedOptimizer` is a stub (`is_stub_optimizer=True`). +`state_dict()` and `load_state_dict()` should be no-ops. The checkpoint should only contain +model weights, not optimizer state. This needs to be tested. + +### 8.3 `param_to_group_meta` Under Path B + +**Current Path A**: `sharded_param_state_fsdp_dtensor()` converts `self.optimizer.param_groups` +(tensor-keyed) to `param_to_group_meta` (name-keyed) for the checkpoint, and back on load. + +**Path B**: `get_state_dict` calls `optimizer.state_dict()` which returns the inner +optimizer's param groups. These are **already in the standard PyTorch format** — no +conversion needed. The `param_to_group_meta` key format is a Megatron artifact that exists +only because Path A manually constructs the state dict. With Path B, `get_state_dict` +handles this natively, so `param_to_group_meta` is no longer needed. + +This means `DistributedOptimizer.load_state_dict()` for Path B should accept the raw +`self.optimizer.state_dict()` format directly, bypassing the `param_to_group_meta` +conversion path. The current FSDP branch in `load_state_dict()` already handles this: +if `"param_to_group_meta"` is absent, it just calls `self.optimizer.load_state_dict(state_dict)`. + +--- + +## 9. Uneven DTensor Handling + +### 9.1 Problem + +When a parameter's size is not evenly divisible by the DP world size, each rank +holds a different-sized local shard. Standard DCP assumes uniform shard sizes. + +### 9.2 Solution: Chunk Metadata Patching + +`preprocess_state_dict_for_uneven_dtensor()` walks the state dict, finds all +`DTensor` values, and calls `update_uneven_dtensor_chunk_metadata()` on each. +This function: + +1. Gathers chunk metadata from all ranks via `all_gather_object`. +2. Computes global offsets and sizes for each rank's shard. +3. Patches the DTensor with `__create_chunk_list__` and `__create_write_items__` + closures that DCP uses for serialization. + +This is called on **both** model and optimizer state dicts. + +--- + +## 10. DTensor Attribute Propagation + +When loading a checkpoint into a Megatron FSDP v2 model, parameters are DTensors. +Certain attributes (e.g., `is_embedding_parameter`, `allreduce`) set on the original +`nn.Parameter` objects by upstream layers (e.g., TE) must be propagated to the +DTensor wrappers. This is handled in `mcore_fsdp_adapter.py` at lines 307-331. +Missing attributes cause optimizer misclassification (`_get_param_groups`) and wrong +gradient scaling. + +See `design.md` (Pitfall: Attribute Propagation) for the full list of attributes +and their consumers. + +--- + +## 11. Key Differences from Legacy Path + +| Aspect | Legacy Megatron FSDP | Megatron FSDP v2 | +|--------|---------------------|-------------------| +| Parameter representation | `MegatronFSDP`-managed DTensors | Native `FSDPModule` DTensors | +| Model wrapper | `MegatronFSDP` (adds `module.` prefix in state dict) | No `MegatronFSDP` wrap — `fully_shard` applied directly | +| Model state dict keys | Keys have `module.` prefix from `MegatronFSDP` wrapper | Keys lack `module.` prefix; `add_module_prefix()` used for checkpoint alignment | +| Model `state_dict_for_save_checkpoint` | `model.state_dict()` with `state_dict_pre_hook` | `model.state_dict()` (already DTensors) | +| Optimizer buffer management | Megatron FSDP managed | Standard Torch optimizer managed | +| Gradient buffer | Megatron FSDP `param_and_grad_buffer` | None (FSDP v2 handles internally) | +| Load model state dict | `module.load_state_dict(custom, strict)` with `_load_state_dict_post_hook` | `super().load_state_dict(state_dict, strict)` — after `strip_module_prefix()` | +| Zero gradient | `model_chunk.zero_grad_buffer()` | `model_chunk._zero_grad_buffer()` | + +### 11.1 Model State Dict ``module.`` Prefix Alignment + +**Problem:** Legacy Megatron FSDP wraps the model in a ``MegatronFSDP`` class which +stores the model as ``self.module``. This means ``MegatronFSDP.state_dict()`` produces +keys with a ``module.`` prefix (e.g., ``module.embedding.word_embeddings.weight``). +Megatron FSDP v2 applies ``fully_shard`` directly without a ``MegatronFSDP`` wrapper, +so the raw model's ``state_dict()`` produces keys **without** the prefix (e.g., +``embedding.word_embeddings.weight``). + +**Solution:** ``checkpoint.py`` provides two post-processing functions: + +- ``add_module_prefix(state_dict)`` — adds ``module.`` prefix to all keys before + saving, aligning with Megatron's checkpoint format. +- ``strip_module_prefix(state_dict)`` — removes ``module.`` prefix after loading, + aligning with the FSDP v2 model's expected key format. + +``get_model_state_dict(model)`` auto-detects whether the prefix is present and adds +it if missing. ``load_checkpoint`` strips the prefix back before calling +``model.load_state_dict()``. + +**Note on optimizer keys:** ``sharded_param_state_fsdp_dtensor()`` uses +``model_chunk.named_parameters()`` to derive parameter names. When +``model_chunks`` are ``FullyShardedDataParallel`` instances (which store the model +as ``self.module``), the returned names already carry the ``module.`` prefix. This +ensures model and optimizer state dict keys are consistent in the checkpoint. + +--- + +## 12. Implementation Checklist + +### Phase 1: `distrib_optimizer.py` — FSDP v2 Guard Propagation + +- [x] `__init__`: Guard early return with `use_megatron_fsdp or use_megatron_fsdp_v2` +- [x] `state_dict()`: Return `self.optimizer.state_dict()` for FSDP paths +- [x] `load_state_dict()`: Guard direct load path with `use_megatron_fsdp_v2` +- [x] `sharded_state_dict()`: Guard `fsdp_dtensor` enforcement with `use_megatron_fsdp_v2` +- [x] `sharded_param_state_fsdp_dtensor()`: Update assertion to accept both flags + +### Phase 1: `mcore_fsdp_adapter.py` — Model State Dict + +- [x] Wire `state_dict_for_save_checkpoint` to `module.state_dict()` in + `_init_with_fully_shard()` (replace `not_implemented_op` at line 383-384) +- [x] Wire `self.state_dict_for_save_checkpoint` to `self.state_dict()` in + `_init_with_fully_shard()` (replace `not_implemented_op` at line 383-384) + +### Phase 2: Torch-Native `get_state_dict` Path (Path B — `MegatronFSDPStateful`) + +- [x] Add `_is_megatron_fsdp_v2()` helper — checks ``ddp_config.use_megatron_fsdp_v2`` + on the ``FullyShardedDataParallel`` adapter, with fallback to ``isinstance(model[0], FSDPModule)`` + for models that went through ``fully_shard()`` without the adapter wrapper +- [x] Add `_build_megatron_fsdp_v2_state_dict()` to `checkpointing.py` using + ``MegatronFSDPStateful`` (which internally uses ``get_state_dict`` from + ``uneven_dtensor`` for both model and optimizer, then applies MCore + post-processing) +- [x] Wire save path: when `_is_megatron_fsdp_v2(model)`, use `_build_megatron_fsdp_v2_state_dict` +- [x] Wire load path: same condition, pre-allocate optimizer states via + `_init_optimizer_states_with_dummy_values()`, then use `_build_megatron_fsdp_v2_state_dict` +- [x] Skip `preprocess_fsdp_dtensor_state_dict` for v2 in both save and load paths + (post-processing already handled by ``_apply_mcore_postprocess`` inside + ``MegatronFSDPStateful.state_dict()``) +- [ ] Handle PP: iterate model chunks, build per-chunk state dicts +- [ ] Handle multi-optimizer (ChainedOptimizer: expert + non-expert optimizers) + +### Phase 2: Path A — Standalone Save/Load (`checkpoint.py`) + +- [ ] Implement ``save_checkpoint(model, ckpt_dir, optimizer, args)`` — standalone DCP save + helper for external usage (not currently needed by the main training loop). + +### Path A Save/Load Flow (Final) + +:: + + _build_megatron_fsdp_v2_state_dict() + | + +-- get_model_state_dict(model[0]) # model DTensors + +-- preprocess_state_dict_for_uneven_dtensor(model_sd) # chunk metadata on params + +-- get_optimizer_state_dict(optimizer) # plain tensors (NO wrapping here) + +-- _apply_mcore_postprocess(split_optimizer=False) # model split, optimizer canonical + +-- preprocess_state_dict_for_uneven_dtensor(full_sd) # metadata on split model DTensors + +-- Scheduler, Rerun, RNG + | + save_checkpoint() # SAVE path + +-- _wrap_optim_states_as_dtensors(...) # wrap plain → DTensor for DCP save plan + +-- DCP save + | + _load_base_checkpoint() # LOAD path + +-- raw_optimizer_state_dict = state_dict["optimizer"].copy() # capture plain dict + +-- _wrap_optim_states_as_dtensors(...) # wrap plain → DTensor for DCP load + +-- DCP load (fills DTensors → plain tensors updated via shared storage) + +-- state_dict["optimizer"] = raw_optimizer_state_dict # restore plain dict + +-- optimizer.load_state_dict(state_dict["optimizer"]) # receives plain tensors + +### Testing + +- [x] Round-trip test: Save via MCore `save_checkpoint`, load via `setup_model_and_optimizer` → `load_checkpoint` +- [x] Optimizer state dict round-trip: save → load → verify values + (``test_megatron_fsdp_v2_round_trip`` now validates both model and optimizer) +- [ ] Cross-format optimizer state conversion: ND-parallel → Megatron FSDP v2 +- [ ] Unskip and fix hanging `get_state_dict` tests in `test_fully_shard.py` +- [ ] Uneven DTensor sharding with non-divisible parameter sizes +- [ ] Frozen parameter handling in `get_state_dict` +- [ ] PP + v2 checkpoint round-trip test +- [ ] Online conversion test with torch-native Path B + +--- + +## 13. Testing Strategy + +### Active Tests + +| Test | File | Status | +|------|------|--------| +| Online convert: ND-parallel → Megatron FSDP v2 | `test_mcore_checkpoint.py` | Active | +| Online convert: FSDP v1 baseline → Megatron FSDP v2 | `test_mcore_checkpoint.py` | Active | +| Megatron FSDP v2 → v2 round-trip | `test_mcore_checkpoint.py` | Active | +| `MegatronFSDPStateful` wrapper | `checkpoint.py` | Active | +| `get_state_dict` strict DTensor assert | `test_fully_shard.py` | Active | +| SWiGLU/expert key transforms | `test_fsdp_dtensor_checkpoint.py` | Active | +| Megatron FSDP v2 E2E training | `test_mcore_nd_parallel.py` | Active | + +### Known-Gap Tests + +| Test | File | Status | +|------|------|--------| +| `get_state_dict` basic | `test_fully_shard.py` | Skipped (hangs) | +| `get_state_dict` nested FSDP | `test_fully_shard.py` | Skipped (hangs) | +| `get_state_dict` LLM scenario | `test_fully_shard.py` | Skipped (hangs) | +| `get_state_dict` frozen params | `test_fully_shard.py` | Skipped (hangs) | + +### Test Gaps + +- Optimizer state dict round-trip for Megatron FSDP v2. +- Cross-format optimizer state conversion (ND-parallel optimizer → Megatron FSDP v2). +- Uneven DTensor sharding with non-divisible parameter sizes. +- Frozen parameter handling in `get_state_dict`. + +--- + +## 14. Debugging Checklist + +1. **Model state dict has DTensors?** — All parameters should be `DTensor` after + `fully_shard()`: + ```python + from torch.distributed.tensor import DTensor + for name, p in model.named_parameters(): + assert isinstance(p, DTensor), f"{name}: expected DTensor, got {type(p)}" + ``` + +2. **Checkpoint save succeeds?** — DCP writes to disk. Check for + `FileSystemWriter` or `FileSystemWriterAsync` in logs. + +3. **Checkpoint load succeeds?** — Verify that `dcp_load()` returns without errors. + Set `strict_fsdp_dtensor_load=True` for strict key matching. + +4. **Model parameters match after load?** — Compare source and loaded state dicts + using `_state_dict_to_full_tensor` helper. + +5. **Optimizer state restored?** — Check that `optimizer.state` is non-empty and + contains expected keys (`exp_avg`, `exp_avg_sq`, `step`). + +6. **NaN after load?** — Common causes: + - Missing `allreduce` attribute propagation (expert params misclassified) + - Missing `overwrite_main_grad=True` for wgrad fusion (gradient doubling) + - Wrong `gradient_accumulation_fusion` setting + +--- + +## 15. Future Work + +- **Async checkpoint**: Integrate `FileSystemWriterAsync` for non-blocking saves. +- **Cross-topology resharding**: Support loading checkpoints saved with a different + DP world size (requires DCP's resharding planner). +- **Unified `set_state_dict`**: Add a `set_state_dict()` wrapper in + `uneven_dtensor.py` that handles preprocessing inverse. +- **Complete skipped tests**: Debug and unskip the `test_fully_shard.py` checkpoint + tests that currently hang. diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/nvfp4_design.md b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/nvfp4_design.md new file mode 100644 index 00000000000..90097c55716 --- /dev/null +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/nvfp4_design.md @@ -0,0 +1,366 @@ +# NVFP4 Support in Megatron FSDP v2 — Design Document + +## 1. Overview + +This document proposes adding NVFP4 primary-weights support to Megatron FSDP v2 +(the `fully_shard` API path under `megatron_fsdp/v2/`). The existing +non-FSDP DDP path (`DistributedDataParallel` + `DistributedOptimizer`) already +supports `--fp4-param-gather` — this design mirrors that support into the v2 +FSDP path, following the same patterns used for FP8 support. + +## 2. Current State + +### 2.1 What already works (non-FSDP DDP path) + +- Model weights initialised as `NVFP4Tensor` (packed uint8: 2 FP4 values/byte, `numel_packed = numel_logical // 2`). +- `_ParamAndGradBuffer` maintains dual index maps: + - `param_index_map` — full-numel offsets for grads + - `nvfp4_packed_param_index_map` — packed offsets for the param buffer +- All-gather communicates packed uint8 bytes; gradients use full-numel FP32. +- After optimizer step: `quantize_nvfp4_param_shard()` calls TE's + `quantize_master_weights` to cast FP32 → NVFP4 in-place in the packed buffer. + +### 2.2 What already works (FSDP v2 FP8 path) + +- `FullyShardMixedPrecisionPolicy` (in `v2/mixed_precision.py`) handles + FP8 / MXFP8 detection, buffer dtype, raw-data extraction, post-unshard + processing, post-reshard cache invalidation, and main→model quantization. +- `ParameterGroup` creates 3–4 `DataParallelBuffer` instances: + `model_weight_buffer` (uint8 for FP8), optional `transpose_weight_buffer`, + `main_weight_buffer` (FP32), and `main_grad_buffer` (FP32). +- All buffers share the same `BufferIndex` layout (same params, same shapes). + +### 2.3 What is missing + +No NVFP4 path exists in any file under `megatron_fsdp/`. The v1 `MegatronFSDP` +wrapper also lacks it. The `DistributedDataParallelConfig` for FSDP +(`megatron_fsdp/distributed_data_parallel_config.py`) has no `fp4_param_gather` +field, and `mcore_fsdp_adapter.py` has no translation of `fp4_param_gather` into +the v2 mixed-precision policy. + +## 3. Key Difference from FP8: Packed Storage + +| Aspect | FP8 (Float8Tensor) | NVFP4 (NVFP4Tensor) | +|--------|--------------------|----------------------| +| Storage dtype | `uint8` | `uint8` (packed) | +| Bytes per logical element | 1 | 0.5 | +| Packed numel | `N` (same as logical) | `N // 2` | +| Internal raw attr | `_rowwise_data` / `_data` | `_rowwise_data` | +| Transpose/columnwise | MXFP8 has `_columnwise_data` | None | +| Transpose cache | Yes (needs invalidation) | No | +| Quantization API | `cast_master_weights_to_fp8` | `quantize_master_weights` (same TE function, different tensor type) | + +The **packed storage** is the central design challenge: the `model_weight_buffer` +needs a `BufferIndex` with **packed shapes** (`shape[-1] // 2`), while +`main_weight_buffer` and `main_grad_buffer` need **full shapes**. + +## 4. Design Decisions + +### 4.1 Separate BufferIndex per buffer (no shared layout) + +Each `DataParallelBuffer` already owns its own `BufferIndex`. For NVFP4 +param groups, the model-weight buffer will use packed shapes in its index, +while main-weight and main-grad buffers use the standard full shapes. + +**Why this is clean**: `param_idx` (mapping `param → position_in_list`) stays +consistent across all three buffers because it only depends on parameter +ordering. Each buffer's `BufferIndex` translates the same `param_idx` into +different byte offsets appropriate for that buffer's dtype/shape. + +### 4.2 Policy-driven shape transform + +The `FullyShardMixedPrecisionPolicy` already owns dtype and raw-data decisions. +We extend it with an `nvfp4` sub-policy (analogous to `fp8`) that: + +- Returns `torch.uint8` from `model_weight_buffer_dtype()`. +- Returns packed shapes from a new `get_param_storage_shapes()`. +- Returns `_rowwise_data` from `get_param_data()`. +- Calls `post_all_gather_processing()` in `post_unshard()`. +- Invokes `quantize_master_weights` (via `quantize_nvfp4_param_shard`) in + `copy_main_weights_to_model_weights()`. + +### 4.3 No transpose/cache management needed + +NVFP4 has no columnwise data or transpose cache. `needs_transpose_weight_buffer()` +already returns `False` for non-MXFP8 tensors. No changes needed there. + +## 5. Detailed Implementation Plan + +### 5.1 `megatron_fsdp/v2/mixed_precision.py` + +**Add NVFP4 detection** (following the FP8 pattern at lines 28–63): + +```python +try: + from transformer_engine.pytorch.tensor.nvfp4_tensor import NVFP4Tensor as NVFP4_TENSOR_CLASS + HAVE_TE_NVFP4 = True +except Exception: + NVFP4_TENSOR_CLASS = None + HAVE_TE_NVFP4 = False +``` + +**Add `FullyShardNVFP4Policy` dataclass** (analogous to `FullyShardFP8Policy`): + +```python +@dataclass(frozen=True) +class FullyShardNVFP4Policy: + enabled: bool = False + recipe: Optional[str] = None +``` + +**Extend `FullyShardMixedPrecisionPolicy`** with `nvfp4` field: + +```python +@dataclass(frozen=True) +class FullyShardMixedPrecisionPolicy: + ... + fp4_param_gather: bool = False + fp4_recipe: Optional[str] = None + nvfp4: Optional[FullyShardNVFP4Policy] = field(default=None, repr=False) + ... +``` + +**`__post_init__`**: initialize `nvfp4` from flat fields or vice versa (same pattern +as FP8, lines 160–172). + +**New / modified methods**: + +| Method | Change | +|--------|--------| +| `model_init_context()` | Enter `fp8_model_init(NVFP4BlockScaling)` when NVFP4 enabled | +| `group_key_dtype()` | Return `("quantized", "NVFP4Tensor", recipe)` for NVFP4 params | +| `is_nvfp4_param()` | New: `isinstance(tensor, NVFP4_TENSOR_CLASS)` | +| `model_weight_buffer_dtype()` | Return `torch.uint8` for NVFP4 (same as FP8) | +| `get_param_storage_shapes()` | **New**: Return packed shapes for NVFP4, original shapes otherwise | +| `get_param_data()` | Return `tensor._rowwise_data` (packed) for NVFP4 | +| `bind_unsharded_param()` | Set `_rowwise_data` to all-gathered buffer view (same pattern as FP8 `_data`) | +| `get_high_precision_value()` | Use `dequantize()` or preserved init val | +| `post_unshard()` | Call `post_all_gather_processing()` for NVFP4 params (same as FP8 rowwise) | +| `post_reshard()` | No-op for NVFP4 (no transpose cache to clear) | +| `copy_main_weights_to_model_weights()` | Call `quantize_nvfp4_param_shard()` (from `fp4_utils.py`) | +| `storage_tensors_to_free()` | Free `_rowwise_data` for NVFP4 (no transpose) | +| `needs_transpose_weight_buffer()` | Already returns `False` for non-MXFP8 — no change | +| `fine_grained_forward_hooks_required()` | Return `True` if any NVFP4 params present | +| `weight_buffers_for_unshard()` | Return `[model_weight_buffer]` (no transpose) | +| `validate_param_group()` | Works as-is (grouping key already differentiates NVFP4) | + +### 5.2 `megatron_fsdp/v2/param_group.py` + +**`_init_buffers()`** — unchanged. Buffer creation passes the shared logical +`chunk_size_factor` through `_create_buffer()`. Shape resolution is handled +in `DataParallelBuffer.__init__()` and `BufferIndex.compact()`. + +### 5.3 `megatron_fsdp/v2/dp_buffer.py` + +**`BufferIndex.compact(factor, compact_shapes)`** — **new method** that proportionally +scales all indices for packed storage: + +```python +def compact(self, factor: float, compact_shapes: List[torch.Size]) -> None: + for item_id, item in self.item_index_map.items(): + new_map[item_id] = ItemIndex( + global_data_index=int(item.global_data_index * factor), + size=int(item.size * factor), + item_id=item.item_id, + shape=compact_shapes[item_id], + ) + self.item_index_map = new_map + self.bucket_meta = BucketMeta( + global_data_index=0, + size=int(self.bucket_meta.size * factor), + items=list(new_map.values()), + ) + self.shard_meta = self._build_shard_meta(...) +``` + +For NVFP4, ``factor = 0.5`` and ``compact_shapes`` come from +``get_param_storage_shapes()``. This preserves the proportional item-offset +mapping between buffers while eliminating fragment-binning waste. + +**`DataParallelBuffer.__init__()`** — always builds with logical shapes, then compacts +NVFP4 weight buffers: + +```python +_logical_shapes = [p.shape for p in params] +self.buffer_index = BufferIndex( + param_shapes=_logical_shapes, + chunk_size_factor=chunk_size_factor, + ..., +) +if buffer_role in ("model_weight", "transpose_weight") and any(mp_policy.is_nvfp4_param(p) for p in params): + self.buffer_index.compact(0.5, mp_policy.get_param_storage_shapes(params)) +``` + +Main-weight and main-grad buffers are never compacted — they keep the +logical layout and the original ``chunk_size_factor``. + +### 5.4 `megatron_fsdp/v2/fully_shard.py` + +No changes needed — parameter replacement and hook registration are dtype-agnostic. + +### 5.5 `megatron_fsdp/v2/hooks.py` + +No changes needed. + +### 5.6 `megatron_fsdp/v2/fsdp_module.py` + +No changes to `FSDPModule`. The parameter grouping logic in +`_get_module_fsdp_param_groups()` already calls `mp_policy.group_key_dtype()` +which will differentiate NVFP4 params from others. + +### 5.7 `megatron_fsdp/distributed_data_parallel_config.py` + +Add `fp4_param_gather: bool = False` field (mirroring `fp8_param_gather` at line 37): + +```python +fp4_param_gather: bool = False +"""If true, keep the compute param in fp4 (do not use any other intermediate dtype) and + perform the param all-gather in fp4.""" +``` + +### 5.8 `megatron/core/distributed/distributed_data_parallel_config.py` + +Already has `fp4_param_gather: bool = False` at line 75. No change needed. + +### 5.9 `mcore_fsdp_adapter.py` + +**`_init_with_fully_shard()`** — pass FP4 configuration into the policy: + +```python +fully_shard_mp_policy = FullyShardMixedPrecisionPolicy( + ... + nvfp4=FullyShardNVFP4Policy( + enabled=ddp_config.fp4_param_gather, + recipe=config.fp4_recipe, + ), +) +``` + +Also pass `fp4_param_gather` to the v1 `MegatronFSDP` path if desired (similar +to FP8 wiring at line 257–261). + +### 5.10 `megatron/training/arguments.py` + +The `--fp4-param-gather` argument already exists (line 1816). Verify it is +propagated into the FSDP config when `--use-megatron-fsdp` is set. Currently +`DistributedDataParallelConfig` is auto-populated from `args` fields (line 1457 +in `training.py`), so adding `fp4_param_gather` to the FSDP sub-config dataclass +(step 5.7) will make it flow automatically. + +### 5.11 `megatron/core/fp4_utils.py` + +No changes needed. Existing functions (`is_nvfp4tensor`, `quantize_nvfp4_param_shard`, +`get_nvfp4_rowwise_packed_shape`) will be imported by the modified mixed_precision.py. + +## 6. Data Flow Summary + +``` +[Model Init] fp8_model_init(NVFP4BlockScaling) + → TE creates NVFP4Tensor params (packed uint8, shape[-1] // 2) + │ + ▼ +[FSDP v2 ParameterGroup._init_buffers()] + model_weight_buffer: DataParallelBuffer(uint8, packed shapes, distributed) + main_weight_buffer: DataParallelBuffer(fp32, full shapes, distributed) + main_grad_buffer: DataParallelBuffer(fp32, full shapes, distributed) + │ + ▼ +[Forward: unshard()] + 1. All-gather packed uint8 from model_weight_buffer shards + 2. bind_unsharded_param(): rebind NVFP4Tensor._rowwise_data → full buffer view + 3. post_unshard(): post_all_gather_processing() rebuilds TE caches + │ + ▼ +[Backward: reduce_grad()] + 1. Copy param.grad → main_grad_buffer (fp32, full numel) + 2. reduce_scatter_tensor() on main_grad_buffer + │ + ▼ +[Optimizer step] + Optimizer updates main_weight_buffer (fp32) + │ + ▼ +[copy_main_weights_to_model_weights()] + quantize_nvfp4_param_shard(): fp32 main → packed NVFP4 model buffer + (same TE quantize_master_weights API used by non-FSDP path) +``` + +## 7. Files to Modify + +| File | Change | +|------|--------| +| `megatron_fsdp/v2/mixed_precision.py` | Add `FullyShardNVFP4Policy`, NVFP4 detection, all policy methods | +| `megatron_fsdp/v2/param_group.py` | Pass logical `chunk_size_factor` (computed from ``p.shape``) to all buffers | +| `megatron_fsdp/v2/dp_buffer.py` | Add `BufferIndex.compact()`; all buffers build with logical shapes, only NVFP4 weight buffers are compacted | +| `megatron_fsdp/distributed_data_parallel_config.py` | Add `fp4_param_gather` field | +| `mcore_fsdp_adapter.py` | Wire `fp4_param_gather` + `fp4_recipe` into policy | +| `megatron_fsdp/v2/__init__.py` | Export new NVFP4 policy class | + +Files that do **not** need changes: `fp4_utils.py`, `hooks.py`, `fully_shard.py`, +`fsdp_module.py` (except possibly fine-grained hook gating). + +## 8. Risks and Edge Cases + +1. **Packed buffer alignment**: The `BufferIndex._build_layout()` pads to + `dp_world_size * chunk_size_factor`. `chunk_size_factor` is computed from + logical `p.shape[1:].numel()` so that it works correctly for all buffers + (model-weight, main-weight, main-grad). Only the model-weight buffer uses + packed shapes in its `BufferIndex`; the main buffers use logical shapes. + +2. **Empty shards**: NVFP4 params with odd inner dimensions are rejected by + TE (assertion in `get_nvfp4_rowwise_packed_shape`). No special handling needed. + +3. **Mixed NVFP4 + non-NVFP4 params in same group**: Prevented by + `group_key_dtype()` returning different keys. + +4. **Gradient accumulation fusion**: The `overwrite_main_grad` and + `grad_added_to_main_grad` flags set in `pre_backward_hook` work on + `ParameterGroup` params regardless of their storage dtype. + +5. **Checkpoint save/load**: The `state_dict` hooks and DTensor-based checkpoint + path in `mcore_fsdp_checkpoint_design.md` should work transparently since + the DTensor views point into the main_weight_buffer (fp32), not the packed + model-weight buffer. + +## 9. BufferIndex API (Standardized) + +The `BufferIndex` class provides three coordinate domains for querying item +positions. All method names use a consistent `_range` suffix and clear domain +prefixes: + +| Method | Domain | Returns | Description | +|--------|--------|---------|-------------| +| `_get_item_self_range(item_id)` | **Item-self** | `(start, end)` | Offset within the item itself (0 = start of item). What portion of this item falls in the current rank's shard. | +| `_get_item_local_range(item_id, *, as_shard=False)` | **Local buffer** | `(start, end)` | Byte offsets within `self.data` (the local GPU buffer). Where to read/write. The `as_shard` flag forces shard-intersection computation even for non-distributed buffers. | +| `_get_item_global_range(item_id)` | **Global buffer** | `(global_offset, size)` | Position and size in the full logical (unsharded) buffer. Same on all ranks. | + +`_get_item_local_range` absorbs the old `_get_item_local_index` and +`_get_item_local_shard_index`. The `as_shard` kwarg unifies what was +previously the `only_shard` flag on `get_item()`: + +```python +def get_item(self, item_id, *, as_shard=False): + start, end = self.buffer_index._get_item_local_range(item_id, as_shard=as_shard) + return self.data[start:end] +``` + +## 10. DataParallelBuffer.summon_full_params (planned) + +> **Not yet implemented.** Context manager that temporarily unshards a +> buffer, then automatically reshards on exit: + +```python +@contextmanager +def summon_full_params(self, async_op=False): + need_unshard = self.is_distributed and not self.is_unsharded() + if need_unshard: + self.unshard(async_op=async_op) + try: + yield + finally: + if need_unshard: + self.reshard() +``` + +Only unshards/reshards if the buffer is actually distributed and not +already unsharded. Non-distributed buffers are a transparent no-op. diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/utils.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/utils.py new file mode 100644 index 00000000000..bdebdc48564 --- /dev/null +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/utils.py @@ -0,0 +1,83 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from dataclasses import dataclass +from typing import Callable + +import torch +import torch.nn as nn +from torch.distributed import distributed_c10d +from torch.distributed.tensor import DeviceMesh, init_device_mesh + + +@dataclass(frozen=True, slots=True) +class ParamGroupIdx: + """Immutable identifier for a ParameterGroup: (module_id, index).""" + + module_id: int + index: int + + +class RegisterFSDPBackwardFunction(torch.autograd.Function): + """ + Autograd Function for registering post-backward hooks. + + This Function simply passes inputs through in forward, but its + backward calls the post_backward_hook to perform reshard and + gradient reduction after gradients are computed. + """ + + @staticmethod + def forward(ctx, post_backward: Callable, *inputs: torch.Tensor): + ctx.post_backward = post_backward + return inputs + + @staticmethod + def backward(ctx, *grads: torch.Tensor): + ctx.post_backward() + return (None,) + grads + + +def _replace_module_parameter(module: nn.Module, name: str, new_param: nn.Parameter): + """ + Replace a module's parameter while preserving module hierarchy. + + Example: + If name="layers.0.linear1.weight", this finds module.layers[0].linear1 + and replaces its weight parameter. + """ + parts = name.split(".") + parent = module + for part in parts[:-1]: + parent = getattr(parent, part) + setattr(parent, parts[-1], new_param) + + +def _init_default_fully_shard_mesh() -> DeviceMesh: + """Default to global CUDA mesh if possible else global CPU mesh.""" + if not distributed_c10d.is_initialized(): + distributed_c10d.init_process_group() + + default_pg = distributed_c10d._get_default_group() # still private + if torch.cuda.is_available(): + device = torch.device("cuda") + else: + device = torch.device("cpu") + mesh = DeviceMesh.from_group( + [default_pg], + device_type=device.type, + mesh=torch.distributed.get_process_group_ranks(default_pg), + mesh_dim_names=("dp",), + ) + return mesh diff --git a/megatron/core/optimizer/distrib_optimizer.py b/megatron/core/optimizer/distrib_optimizer.py index 0104c99f01d..276b11a2f5e 100644 --- a/megatron/core/optimizer/distrib_optimizer.py +++ b/megatron/core/optimizer/distrib_optimizer.py @@ -678,8 +678,8 @@ def __init__( return self.is_stub_optimizer = False - if self.ddp_config.use_megatron_fsdp: - # Megatron-FSDP will manage optimizer weights and gradients. + if self.ddp_config.use_megatron_fsdp or self.ddp_config.use_megatron_fsdp_v2: + # Megatron-FSDP / FSDP v2 will manage optimizer weights and gradients. return # Model grad buffer ranges. @@ -788,7 +788,15 @@ def state_dict(self): the standard model/RNG checkpoint file. The parameter and dependent optimizer state (e.g., exp_avg, exp_avg_sq) are stored in a separate checkpoint file by calling 'save_parameter_state()'. + + For Megatron-FSDP v2, the full inner optimizer state dict is + returned directly because FSDP manages parameter state as DTensors and + the standard PyTorch DCP / get_state_dict() APIs expect a complete + optimizer state dict. """ + if self.ddp_config.use_megatron_fsdp and self.ddp_config.use_megatron_fsdp_v2: + return self.optimizer.state_dict() + inner_state_dict = self.optimizer.state_dict() state_dict = {} @@ -872,9 +880,9 @@ def load_state_dict(self, state_dict): - state_order : The index of a parameter within the shared parameter list. """ - if self.ddp_config.use_megatron_fsdp: - # When using Megatron-FSDP, directly load the optimizer state - # into the wrapped optimizer. + if self.ddp_config.use_megatron_fsdp or self.ddp_config.use_megatron_fsdp_v2: + # When using Megatron-FSDP / FSDP v2, directly load the optimizer + # state into the wrapped optimizer. if "param_to_group_meta" in state_dict: state_dict["param_groups"] = self._param2group_meta_to_param_groups( state_dict["param_to_group_meta"], self.optimizer.param_groups @@ -1398,8 +1406,11 @@ def _param_name(self, param: torch.nn.Parameter) -> str: "Ensure that each model chunk has unique parameter names." ) name_to_param.update(_name_to_param) - num_experts = self.model_chunks[0].config.num_moe_experts if self.model_chunks else None - name_to_param = handle_experts_in_state_dict(name_to_param, num_experts) + if not self.ddp_config.use_megatron_fsdp_v2: + num_experts = ( + self.model_chunks[0].config.num_moe_experts if self.model_chunks else None + ) + name_to_param = handle_experts_in_state_dict(name_to_param, num_experts) self.param_to_name = {param: name for name, param in name_to_param.items()} assert ( param in self.param_to_name @@ -1460,12 +1471,14 @@ def sharded_state_dict( ) # Handle FSDP DistributedOptimizer States - if self.ddp_config.use_megatron_fsdp and sharding_type != "fsdp_dtensor": + if ( + self.ddp_config.use_megatron_fsdp or self.ddp_config.use_megatron_fsdp_v2 + ) and sharding_type != "fsdp_dtensor": raise NotImplementedError( f"sharding_type {sharding_type} is not supported with Megatron FSDP." ) if sharding_type == "fsdp_dtensor": - # Megatron-FSDP custom sharded state dict construction. + # Megatron-FSDP / FSDP v2 custom sharded state dict construction. state_dict = self.sharded_param_state_fsdp_dtensor(is_loading) return state_dict @@ -1584,7 +1597,7 @@ def sharded_param_state_fsdp_dtensor(self, is_loading: bool = False): Sharded state dict where each parameter is a separate PyTorch DTensor. """ assert ( - self.ddp_config.use_megatron_fsdp + self.ddp_config.use_megatron_fsdp or self.ddp_config.use_megatron_fsdp_v2 ), "fsdp_dtensor sharding type is only supported with Megatron FSDP." # Initialize optimizer states with dummy values if loading. diff --git a/megatron/core/optimizer/optimizer_config.py b/megatron/core/optimizer/optimizer_config.py index 5fbf6710f76..1dbf8d51d45 100644 --- a/megatron/core/optimizer/optimizer_config.py +++ b/megatron/core/optimizer/optimizer_config.py @@ -336,6 +336,9 @@ class OptimizerConfig: overlap_param_gather_with_optimizer_step: bool = False """If true, overlap param all-gather of first bucket with optimizer step.""" + use_megatron_fsdp: bool = False + """If true, use Megatron-FSDP and include its custom gradient clipping implementation.""" + ####################### # Optimizer Offload ####################### diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index 423a51f061d..8749daafd2a 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -1211,6 +1211,9 @@ def validate_args(args, defaults={}): ): raise ValueError("MXFP8 with inference optimized layers requires FlashInfer >= 0.6.4") + if getattr(args, 'use_megatron_fsdp_v2', False): + args.use_megatron_fsdp = True + if args.inference_dynamic_batching_sampling_backend == 'flashinfer': try: import flashinfer # noqa: F401 @@ -3888,6 +3891,13 @@ def _add_distributed_args(parser): help='Manually register the FSDP communication buffers to NCCL user buffer.' 'This option is only effective when use-megatron-fsdp and use-nccl-ub is set.', ) + group.add_argument( + '--use-megatron-fsdp-v2', + action='store_true', + dest='use_megatron_fsdp_v2', + help='Use PyTorch fully shard API for FSDP implementation. ' + 'This option is only effective when use-megatron-fsdp is set. ', + ) group.add_argument( '--create-all-gather-group', action='store_true', diff --git a/megatron/training/checkpointing.py b/megatron/training/checkpointing.py index 87102e112b2..1fd6e6710f1 100644 --- a/megatron/training/checkpointing.py +++ b/megatron/training/checkpointing.py @@ -52,6 +52,9 @@ from megatron.core.distributed.fsdp.src.megatron_fsdp.uneven_dtensor import ( preprocess_state_dict_for_uneven_dtensor, ) + from megatron.core.distributed.fsdp.checkpoint import ( + _apply_mcore_postprocess + ) from megatron.core.transformer.fsdp_dtensor_checkpoint import ( handle_experts_in_state_dict, handle_fp8_extra_state_case, @@ -1178,6 +1181,17 @@ def maybe_save_dataloader_state(train_iterator, iteration, dataloader_save_path) torch.save(dataloader_save_dict, data_state_save_path) +def _is_megatron_fsdp_v2(model): + """Check if model uses Megatron FSDP v2 (use_megatron_fsdp_v2 flag).""" + from megatron.core.distributed.fsdp.src.megatron_fsdp.v2 import FSDPModule + first_model = model[0] if isinstance(model, (list, tuple)) else model + for m in first_model.modules(): + if isinstance(m, FSDPModule): + return True + + return False + + def generate_state_dict( args, model, @@ -1216,6 +1230,9 @@ def generate_state_dict( ) else: # torch, torch_dcp, fsdp_dtensor model_sd = model[i].state_dict_for_save_checkpoint() + if args.use_megatron_fsdp_v2 and args.ckpt_format == "fsdp_dtensor": + from megatron.core.distributed.fsdp.checkpoint import _propagate_chunk_metadata_to_state_dict + _propagate_chunk_metadata_to_state_dict(model[i], model_sd) state_dict[key] = model_sd @@ -1264,6 +1281,9 @@ def generate_state_dict( def preprocess_fsdp_dtensor_state_dict(args, raw_state_dict, model): + if _is_megatron_fsdp_v2(model): + return _apply_mcore_postprocess(raw_state_dict, args, model) + state_dict = raw_state_dict.copy() handle_fp8_extra_state_case(state_dict["model"]) if args.swiglu: @@ -1441,6 +1461,13 @@ def _load_global_dist_base_checkpoint( state_dict = dist_checkpointing.load_common_state_dict(checkpoint_name) return state_dict, checkpoint_name, release, CheckpointType.GLOBAL + if args.use_megatron_fsdp_v2: + checkpoint_name = find_checkpoint_rank_0(load_dir, iteration, release) + state_dict = _load_torch_dist_into_megatron_fsdp_v2( + args, checkpoint_name, sharded_state_dict, strict=True + ) + return state_dict, checkpoint_name, release, CheckpointType.FSDP_DTENSOR + if sharded_state_dict is None: assert not args.auto_detect_ckpt_format and not args.use_dist_ckpt, ( args.auto_detect_ckpt_format, @@ -1903,7 +1930,7 @@ def load_checkpoint( model = unwrap_model(ddp_model) ckpt_format = args.ckpt_format - if args.auto_detect_ckpt_format or ckpt_format == "torch_dist": + if args.auto_detect_ckpt_format or ckpt_format in ("torch_dist", "fsdp_dtensor"): state_dict, checkpoint_name, release, ckpt_type = _load_base_checkpoint( load_dir, args, rank0=True, checkpointing_context=checkpointing_context ) @@ -2012,6 +2039,16 @@ def load_checkpoint( f" Please use `--ckpt-fully-parallel-save` flag during checkpoint saving." ) + if ( + args.use_megatron_fsdp_v2 + and sharded_sd_metadata['distrib_optim_sharding_type'] == 'dp_reshardable' + ): + raise RuntimeError( + "Megatron FSDP v2 does not support checkpoint conversion from " + "distrib_optim_sharding_type=dp_reshardable. " + "Please re-save the checkpoint with --dist-ckpt-optim-fully-reshardable." + ) + # Check if fully parallel load is compatible with sharding type if ( args.ckpt_fully_parallel_load diff --git a/megatron/training/config/training_config.py b/megatron/training/config/training_config.py index 93fc889b0bb..8b202ada035 100644 --- a/megatron/training/config/training_config.py +++ b/megatron/training/config/training_config.py @@ -268,6 +268,9 @@ class LoggerConfig: log_params_norm: bool = False """If set, calculate and log parameters norm.""" + log_per_param_norm: bool = False + """If set, log per-parameter norm with per-param granularity (requires --log-params-norm).""" + log_throughput: bool = False """If set, calculate and log throughput per GPU.""" diff --git a/megatron/training/training.py b/megatron/training/training.py index 34e65337703..52e14130510 100644 --- a/megatron/training/training.py +++ b/megatron/training/training.py @@ -3610,6 +3610,14 @@ def trace_handler(p): if args.log_params_norm: params_norm = calc_params_l2_norm(model) + if args.log_per_param_norm: + model_chunks = model if isinstance(model, list) else [model] + for model_chunk in model_chunks: + try: + if hasattr(model_chunk, 'log_per_param_norms'): + model_chunk.log_per_param_norms(iteration, prefix="[PER-PARAM]") + except Exception as e: + print(f"[PER-PARAM] logging failed for {model_chunk}: {e}") if optimizer is not None: learning_rate = get_canonical_lr_for_logging(optimizer.param_groups) else: diff --git a/tests/unit_tests/distributed/megatron_fsdp/v2/conftest.py b/tests/unit_tests/distributed/megatron_fsdp/v2/conftest.py new file mode 100644 index 00000000000..49ca9e349cc --- /dev/null +++ b/tests/unit_tests/distributed/megatron_fsdp/v2/conftest.py @@ -0,0 +1,15 @@ +import os +import sys + + +def _redirect_nonzero_ranks(): + rank = int(os.getenv("RANK", "0")) + redirect_nonzero_ranks = os.getenv("MFSDP_UT_REDIRECT_NONZERO_RANKS", "1").lower() == "1" + if rank != 0 and redirect_nonzero_ranks: + os.makedirs("pytest_ranks", exist_ok=True) + f = open(f"pytest_ranks/rank{rank}.out", "w", buffering=1) + sys.stdout = f + sys.stderr = f + + +_redirect_nonzero_ranks() diff --git a/tests/unit_tests/distributed/megatron_fsdp/v2/test_allocator.py b/tests/unit_tests/distributed/megatron_fsdp/v2/test_allocator.py new file mode 100644 index 00000000000..28ef81f21fa --- /dev/null +++ b/tests/unit_tests/distributed/megatron_fsdp/v2/test_allocator.py @@ -0,0 +1,99 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for TemporaryBucketAllocator. Pure CPU, no torch.distributed.""" + +import sys +from pathlib import Path + +import pytest +import torch + +sys.path.insert(0, str(Path(__file__).parents[2])) +from megatron.core.distributed.fsdp.src.megatron_fsdp.v2.allocator import ( + Bucket, + TemporaryBucketAllocator, +) + + +def _run_allocator_tests(allocator: TemporaryBucketAllocator) -> None: + """Three-phase test covering allocate, free, and re-allocate for any allocator. + + Phase 1 -- allocate: + - Returned bucket has correct size, dtype, and is tracked internally. + - Duplicate allocate on the same id returns the same object (no realloc). + - Different ids produce independent buckets. + + Phase 2 -- free: + - Freed id is removed from internal tracking. + - Underlying tensor storage is resized to 0. + - Freeing a non-existent id is silently ignored. + - Freeing one id does not affect others. + + Phase 3 -- re-allocate after free: + - Re-allocating a freed id returns a usable bucket with correct size. + """ + + # ---- Phase 1: allocate ---- + b0 = allocator.allocate( + param_group_id=0, size=1024, dtype=torch.float32, device=torch.device("cpu") + ) + assert b0.data.numel() == 1024 + assert b0.data.dtype == torch.float32 + assert 0 in allocator.buckets + + b0_again = allocator.allocate( + param_group_id=0, size=1024, dtype=torch.float32, device=torch.device("cpu") + ) + assert b0_again is b0 + assert b0_again.data.data_ptr() == b0.data.data_ptr() + + b1 = allocator.allocate( + param_group_id=1, size=512, dtype=torch.bfloat16, device=torch.device("cpu") + ) + assert b1 is not b0 + assert b1.data.numel() == 512 + assert b1.data.dtype == torch.bfloat16 + + # ---- Phase 2: free ---- + tensor_ref = b0.data + allocator.free(0) + assert 0 not in allocator.buckets + assert tensor_ref._typed_storage()._size() == 0 + + assert 1 in allocator.buckets + assert b1.data.numel() == 512 + + allocator.free(999) + + # ---- Phase 3: re-allocate ---- + b0_new = allocator.allocate( + param_group_id=0, size=1024, dtype=torch.float32, device=torch.device("cpu") + ) + assert b0_new.data.numel() == 1024 + assert 0 in allocator.buckets + + # cleanup + allocator.free(0) + allocator.free(1) + + +class TestTemporaryBucketAllocator: + + def test_full_lifecycle(self): + _run_allocator_tests(TemporaryBucketAllocator()) + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/tests/unit_tests/distributed/megatron_fsdp/v2/test_fully_shard.py b/tests/unit_tests/distributed/megatron_fsdp/v2/test_fully_shard.py new file mode 100644 index 00000000000..138bd9aa4fc --- /dev/null +++ b/tests/unit_tests/distributed/megatron_fsdp/v2/test_fully_shard.py @@ -0,0 +1,762 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for the Megatron-FSDP v2 ``fully_shard`` API, ``FSDPModule``, +and checkpoint (``get_state_dict``). + +Covers: +- Basic fully_shard: class mutation, hooks, reshard on init +- Multi-layer LLM-style nesting (embedding → transformer layers → lm_head) +- Multimodal-style: separate encoders with partial freezing +- Partially frozen training (requires_grad=False on some params) +- Nested FSDP (expert-in-layer pattern) +- Mixed precision policies (fp32 main params, fp32 grad reduce) +- ignored_params +- enable_unshard_prefetch / enable_async_reduce_grad feature flags +- Forward/backward lifecycle correctness +- get_state_dict / preprocess_state_dict_for_uneven_dtensor +- Double-shard safety (reject re-wrap) + +Run with: + torchrun --nproc_per_node=2 -m pytest \\ + tests/unit_tests/distributed/megatron_fsdp/v2/test_fully_shard.py -v + +Single-GPU tests: + pytest tests/unit_tests/distributed/megatron_fsdp/v2/test_fully_shard.py -v \\ + -k "test_double_shard_rejected or test_no_params_module" +""" + +import sys +from pathlib import Path + +import pytest +import torch +import torch.nn as nn +from torch.distributed.checkpoint.state_dict import StateDictOptions +from torch.distributed.checkpoint.state_dict import get_state_dict as torch_get_state_dict + +sys.path.insert(0, str(Path(__file__).parents[2])) +from megatron.core.distributed.fsdp.src.megatron_fsdp.uneven_dtensor import ( + get_state_dict, + preprocess_state_dict_for_uneven_dtensor, +) +from megatron.core.distributed.fsdp.src.megatron_fsdp.v2.fsdp_module import FSDPModule +from megatron.core.distributed.fsdp.src.megatron_fsdp.v2.fully_shard import fully_shard +from megatron.core.distributed.fsdp.src.megatron_fsdp.v2.mixed_precision import ( + FullyShardMixedPrecisionPolicy, +) + +# ------------------------------------------------------------------ # +# Distributed environment (NCCL session-scoped) +# ------------------------------------------------------------------ # + + +@pytest.fixture(scope="session", autouse=True) +def dist_env(): + if not torch.distributed.is_initialized(): + torch.distributed.init_process_group(backend="nccl") + rank = torch.distributed.get_rank() + device = torch.device(f"cuda:{rank % torch.cuda.device_count()}") + torch.cuda.set_device(device) + yield + if torch.distributed.is_initialized(): + torch.distributed.destroy_process_group() + + +def _rank(): + return torch.distributed.get_rank() + + +def _world_size(): + return torch.distributed.get_world_size() + + +def _device(): + return torch.device(f"cuda:{_rank() % torch.cuda.device_count()}") + + +# ------------------------------------------------------------------ # +# Mock models for different application scenarios +# ------------------------------------------------------------------ # + + +class SimpleMLP(nn.Module): + """Single linear layer with optional bias.""" + + def __init__(self, hidden=64, bias=True): + super().__init__() + self.fc = nn.Linear(hidden, hidden, bias=bias) + + def forward(self, x): + return self.fc(x) + + +class TinyLLM(nn.Module): + """Simulates an LLM: embedding → block of layers → lm_head. + + Structure:: + embedding (nn.Embedding) → layers (nn.ModuleList of SimpleMLP) → lm_head (nn.Linear) + """ + + def __init__(self, vocab=128, hidden=64, num_layers=3): + super().__init__() + self.embedding = nn.Embedding(vocab, hidden) + self.layers = nn.ModuleList([SimpleMLP(hidden) for _ in range(num_layers)]) + self.lm_head = nn.Linear(hidden, vocab) + self.norm = nn.LayerNorm(hidden) + + def forward(self, x): + h = self.embedding(x) + for layer in self.layers: + h = layer(h) + return self.lm_head(self.norm(h)) + + +class MultimodalModel(nn.Module): + """Simulates a multimodal model with separate vision/text encoders. + + Structure:: + vision_encoder (nn.Linear) — may be frozen + text_encoder (nn.Linear) — trainable + fusion (nn.Linear) — trainable + """ + + def __init__(self, hidden=64): + super().__init__() + self.vision_encoder = nn.Linear(hidden, hidden) + self.text_encoder = nn.Linear(hidden, hidden) + self.fusion = nn.Linear(hidden * 2, hidden) + + def forward(self, img, txt): + v = self.vision_encoder(img) + t = self.text_encoder(txt) + return self.fusion(torch.cat([v, t], dim=-1)) + + +class ExpertBlock(nn.Module): + """Simulates an MoE expert: two linear layers.""" + + def __init__(self, hidden=64, ffn_hidden=128): + super().__init__() + self.fc1 = nn.Linear(hidden, ffn_hidden) + self.fc2 = nn.Linear(ffn_hidden, hidden) + + def forward(self, x): + return self.fc2(torch.relu(self.fc1(x))) + + +class MOETransformerLayer(nn.Module): + """Simulates a transformer layer with MoE: attention → MoE experts. + + Structure:: + attn (nn.Linear) → experts (ExpertBlock) → norm (nn.LayerNorm) + """ + + def __init__(self, hidden=64, ffn_hidden=128): + super().__init__() + self.attn = nn.Linear(hidden, hidden) + self.experts = ExpertBlock(hidden, ffn_hidden) + self.norm = nn.LayerNorm(hidden) + + def forward(self, x): + h = self.attn(x) + h = self.experts(h) + return self.norm(h + x) + + +# ------------------------------------------------------------------ # +# Helpers +# ------------------------------------------------------------------ # + + +def _forward_backward(model, x): + """Run forward + backward and return loss.""" + out = model(x) + loss = out.sum() + loss.backward() + return loss.item() + + +def _assert_dtensor_params(module): + """Assert all parameters in the module (and any nested FSDPModules) are DTensors.""" + from torch.distributed.tensor import DTensor + + for name, p in module.named_parameters(): + assert isinstance(p, DTensor), ( + f"Parameter '{name}' should be a DTensor after fully_shard, " f"got {type(p).__name__}" + ) + + +def _assert_original_params_unchanged(module, originals): + """After fully_shard, the original (pre-fully_shard) param OBJECTS should + still be the same Python objects (identity check), but their .data may have + been freed (empty tensor).""" + for name, p in module.named_parameters(): + assert ( + p is originals[name] + ), f"Original param object for '{name}' was replaced; expected identity match." + + +def _count_fsdp_modules(module): + """Return number of FSDPModule instances in the module tree.""" + return sum(1 for m in module.modules() if isinstance(m, FSDPModule)) + + +# ------------------------------------------------------------------ # +# 1. Basic fully_shard — class mutation, hooks, reshard on init +# ------------------------------------------------------------------ # + + +class TestFullyShardBasic: + def test_module_class_becomes_fsdp(self): + """fully_shard should dynamically convert the module class to a FSDPModule mixin.""" + torch.manual_seed(42) + model = SimpleMLP(64).to(_device()) + original_cls = model.__class__ + wrapped = fully_shard(model) + assert wrapped is model # returns same object + assert isinstance(wrapped, FSDPModule) + assert FSDPModule in type(wrapped).__mro__ + assert original_cls in type(wrapped).__mro__ + + def test_params_are_dtensor_after_reshard(self): + """After fully_shard, module.reshard() is called, so params must be DTensors.""" + torch.manual_seed(42) + model = SimpleMLP(64).to(_device()) + fully_shard(model) + _assert_dtensor_params(model) + + def test_forward_without_errors(self): + """A simple forward pass after fully_shard should succeed.""" + torch.manual_seed(42) + model = SimpleMLP(64).to(_device()) + fully_shard(model) + x = torch.randn(2, 64, device=_device()) + out = model(x) + assert out.shape == (2, 64) + + def test_forward_backward_no_nan(self): + """Forward + backward should produce finite loss and gradients.""" + torch.manual_seed(42) + model = SimpleMLP(64).to(_device()) + fully_shard(model) + x = torch.randn(2, 64, device=_device()) + loss = _forward_backward(model, x) + assert not torch.isnan(torch.tensor(loss)), "Loss is NaN" + assert not torch.isinf(torch.tensor(loss)), "Loss is Inf" + + @pytest.mark.parametrize( + "enable_unshard_prefetch,enable_async_reduce_grad", + [(False, False), (False, True), (True, False), (True, True)], + ) + def test_feature_flags(self, enable_unshard_prefetch, enable_async_reduce_grad): + """All combinations of overlap flags should work.""" + torch.manual_seed(42) + model = SimpleMLP(64).to(_device()) + fully_shard( + model, + enable_unshard_prefetch=enable_unshard_prefetch, + enable_async_reduce_grad=enable_async_reduce_grad, + ) + x = torch.randn(2, 64, device=_device()) + out = model(x) + loss = out.sum() + loss.backward() + assert not torch.isnan(torch.tensor(loss.item())) + + +# ------------------------------------------------------------------ # +# 2. Scenarios — LLM-style nesting (embedding + layers + lm_head) +# ------------------------------------------------------------------ # + + +class TestLLMScenario: + def test_llm_full_shard_root(self): + """Shard the root TinyLLM — all params go into one FSDP module.""" + torch.manual_seed(42) + model = TinyLLM(vocab=128, hidden=64, num_layers=2).to(_device()) + fully_shard(model) + assert _count_fsdp_modules(model) == 1 + _assert_dtensor_params(model) + + x = torch.randint(0, 128, (4, 8), device=_device()) + loss = _forward_backward(model, x) + assert not torch.isnan(torch.tensor(loss)) + + def test_llm_per_layer_shard(self): + """Shard each transformer layer individually (typical FSDP setup).""" + torch.manual_seed(42) + model = TinyLLM(vocab=128, hidden=64, num_layers=2).to(_device()) + # Shard each child layer separately + for layer in model.layers: + fully_shard(layer) + # Shard the root (embedding + lm_head + norm covered) + fully_shard(model) + + assert _count_fsdp_modules(model) == 3 # 2 layers + root + _assert_dtensor_params(model) + + x = torch.randint(0, 128, (4, 8), device=_device()) + loss = _forward_backward(model, x) + assert not torch.isnan(torch.tensor(loss)) + + +# ------------------------------------------------------------------ # +# 3. Multimodal scenario — separate encoders + partial freezing +# ------------------------------------------------------------------ # + + +class TestMultimodalScenario: + def test_frozen_vision_encoder(self): + """Freeze vision encoder params; they should NOT be sharded (no grad).""" + torch.manual_seed(42) + model = MultimodalModel(hidden=64).to(_device()) + # Freeze vision encoder + for p in model.vision_encoder.parameters(): + p.requires_grad = False + fully_shard(model) + + x_img = torch.randn(2, 64, device=_device(), requires_grad=True) + x_txt = torch.randn(2, 64, device=_device(), requires_grad=True) + out = model(x_img, x_txt) + loss = out.sum() + loss.backward() + + # Frozen params should have no grad after reduce_grad + for name, p in model.named_parameters(): + if "vision_encoder" in name: + assert not p.requires_grad, f"Frozen param {name} should have requires_grad=False" + + assert not torch.isnan(torch.tensor(loss.item())) + + def test_frozen_params_in_own_group(self): + """Frozen params are included but grouped separately (different requires_grad).""" + torch.manual_seed(42) + model = MultimodalModel(hidden=64).to(_device()) + for p in model.vision_encoder.parameters(): + p.requires_grad = False + fully_shard(model) + + # Frozen params should be in a group with requires_grad=False + has_frozen_group = False + for param_group in model._fsdp_param_groups: + if not param_group.requires_grad: + has_frozen_group = True + for p in param_group.params: + assert ( + not p.requires_grad + ), "Param in frozen group should have requires_grad=False" + assert has_frozen_group, "Frozen params should be in their own param group" + + def test_mixed_frozen_and_trainable(self): + """Some parts frozen, some trainable — all sharded together.""" + torch.manual_seed(42) + model = MultimodalModel(hidden=64).to(_device()) + # Freeze only vision, text and fusion stay trainable + for p in model.vision_encoder.parameters(): + p.requires_grad = False + fully_shard(model) + + x_img = torch.randn(2, 64, device=_device(), requires_grad=True) + x_txt = torch.randn(2, 64, device=_device(), requires_grad=True) + + # Should run without error + out = model(x_img, x_txt) + loss = out.sum() + loss.backward() + assert not torch.isnan(torch.tensor(loss.item())) + + +# ------------------------------------------------------------------ # +# 4. Nested FSDP — expert-in-layer (EDP pattern) +# ------------------------------------------------------------------ # + + +class TestNestedFSDP: + def test_nested_expert_in_layer(self): + """Shard experts inside layer, then shard layer, then root — EDP pattern.""" + torch.manual_seed(42) + device = _device() + + class Model(nn.Module): + def __init__(self): + super().__init__() + self.layer = MOETransformerLayer(64, 128) + self.head = nn.Linear(64, 10) + + def forward(self, x): + return self.head(self.layer(x)) + + model = Model().to(device) + + # Step 1: shard the expert (nested FSDP) + model.layer.experts = fully_shard(model.layer.experts) + # Step 2: shard the layer (will detect nested FSDP) + model.layer = fully_shard(model.layer) + # Step 3: shard the root + model = fully_shard(model) + + assert _count_fsdp_modules(model) == 3 # experts, layer, root + + x = torch.randn(2, 64, device=device, requires_grad=True) + loss = _forward_backward(model, x) + assert not torch.isnan(torch.tensor(loss)) + + def test_nested_ignored_params_are_skipped(self): + """Nested FSDP module params must be in the parent's ignored_params.""" + torch.manual_seed(42) + device = _device() + model = MOETransformerLayer(64, 128).to(device) + + expert = fully_shard(model.experts) + model.experts = expert # rebind (fully_shard returns same object) + model = fully_shard(model) + + # The outer FSDP (model) should have a parameter group that does NOT + # include the inner FSDP's (expert) params + expert_param_ids = set(id(p) for p in expert.parameters()) + for _, param_group in model._named_param_groups: + for p in param_group.params: + assert ( + id(p) not in expert_param_ids + ), "Nested FSDP param leaked into parent param group" + + def test_nested_forward_backward(self): + """Nested FSDP forward+backward produces correct loss pattern.""" + torch.manual_seed(42) + device = _device() + model = MOETransformerLayer(64, 128).to(device) + + model.experts = fully_shard(model.experts) + model = fully_shard(model) + + # Run twice — second pass should work correctly (buffer reuse) + for _ in range(2): + x = torch.randn(2, 64, device=device, requires_grad=True) + out = model(x) + loss = out.sum() + loss.backward() + assert not torch.isnan(torch.tensor(loss.item())) + + +# ------------------------------------------------------------------ # +# 5. Mixed precision policies +# ------------------------------------------------------------------ # + + +class TestMixedPrecision: + def test_main_params_fp32(self): + """With fp32 main params, main_weight_buffer should be created.""" + torch.manual_seed(42) + mp_policy = FullyShardMixedPrecisionPolicy(main_params_dtype=torch.float32) + model = SimpleMLP(64).to(_device()).bfloat16() + fully_shard(model, mp_policy=mp_policy) + + # Verify main_weight_buffer exists and is fp32 + for param_group in model._fsdp_param_groups: + if param_group.main_weight_buffer is not None: + assert ( + param_group.main_weight_buffer.dtype == torch.float32 + ), f"Expected fp32 main weight buffer, got {param_group.main_weight_buffer.dtype}" + + def test_main_params_none(self): + """With no main_params_dtype, no main_weight_buffer should be created.""" + torch.manual_seed(42) + mp_policy = FullyShardMixedPrecisionPolicy(main_params_dtype=None) + model = SimpleMLP(64).to(_device()) + fully_shard(model, mp_policy=mp_policy) + + for param_group in model._fsdp_param_groups: + assert ( + param_group.main_weight_buffer is None + ), "main_weight_buffer should be None when main_params_dtype is None" + + def test_fp32_grad_reduce(self): + """grad_reduce_in_fp32=True should use fp32 gradient communication.""" + torch.manual_seed(42) + mp_policy = FullyShardMixedPrecisionPolicy( + main_grads_dtype=torch.float32, grad_comm_dtype=torch.float32 + ) + model = SimpleMLP(64).to(_device()).bfloat16() + fully_shard(model, mp_policy=mp_policy) + + x = torch.randn(2, 64, device=_device(), dtype=torch.bfloat16, requires_grad=True) + out = model(x) + loss = out.sum() + loss.backward() + assert not torch.isnan(torch.tensor(loss.item())) + + +# ------------------------------------------------------------------ # +# 6. ignored_params +# ------------------------------------------------------------------ # + + +class TestIgnoredParams: + def test_ignored_params_excluded_from_groups(self): + """Params passed as ignored_params should not appear in FSDP groups.""" + torch.manual_seed(42) + model = SimpleMLP(64).to(_device()) + # Pre-identify the param to ignore before wrapping + ignored = {model.fc.weight} + fully_shard(model, ignored_params=ignored) + + for param_group in model._fsdp_param_groups: + for p in param_group.params: + assert p is not model.fc.weight, "Ignored param leaked into group" + + def test_ignored_param_stays_on_module(self): + """Ignored param should remain as a regular nn.Parameter on the module.""" + torch.manual_seed(42) + model = SimpleMLP(64).to(_device()) + original_weight = model.fc.weight + ignored = {model.fc.weight} + fully_shard(model, ignored_params=ignored) + + # After fully_shard and reshard, the ignored weight should still be + # the original nn.Parameter (not a DTensor) on the module + from torch.distributed.tensor import DTensor + + assert not isinstance( + model.fc.weight, DTensor + ), "Ignored param should not be converted to DTensor" + assert model.fc.weight is original_weight, "Ignored param identity changed" + + +# ------------------------------------------------------------------ # +# 7. Forward/backward lifecycle correctness +# ------------------------------------------------------------------ # + + +class TestLifecycle: + def test_params_unsharded_during_forward(self): + """During forward, model parameters should be in unsharded state (full tensors).""" + torch.manual_seed(42) + model = SimpleMLP(64).to(_device()) + + captured_shapes = [] + + def hook(module, inp): + # At this point, the pre-forward hook has already unsharded + for name, p in module.named_parameters(): + captured_shapes.append((name, p.data.shape)) + + model.register_forward_pre_hook(hook) + fully_shard(model) + + # Before forward, params should be DTensors (reshard called at init) + from torch.distributed.tensor import DTensor + + for _, p in model.named_parameters(): + assert isinstance(p, DTensor), "Params should be DTensors after init reshard" + + x = torch.randn(2, 64, device=_device()) + model(x) + + # Inside the forward pre-hook, params should be full tensors + for name, shape in captured_shapes: + assert shape == torch.Size([64, 64]) or shape == torch.Size( + [64] + ), f"Param {name} has wrong shape during forward: {shape}" + + def test_params_resharded_after_forward(self): + """After forward, model parameters should be resharded back to DTensors.""" + torch.manual_seed(42) + model = SimpleMLP(64).to(_device()) + fully_shard(model) + + from torch.distributed.tensor import DTensor + + x = torch.randn(2, 64, device=_device()) + model(x) + + for _, p in model.named_parameters(): + assert isinstance(p, DTensor), "Params should be DTensors after forward reshard" + + def test_params_unsharded_during_backward(self): + """During backward, model parameters should be unsharded.""" + torch.manual_seed(42) + model = SimpleMLP(64).to(_device()) + fully_shard(model) + + from torch.distributed.tensor import DTensor + + captured_dtensor = [] + + def grad_hook(grad): + for _, p in model.named_parameters(): + captured_dtensor.append(isinstance(p, DTensor)) + + x = torch.randn(2, 64, device=_device(), requires_grad=True) + out = model(x) + out.register_hook(grad_hook) + loss = out.sum() + loss.backward() + + # During backward (grad_hook fires during backward pass), params should NOT be DTensor + for is_dt in captured_dtensor: + assert not is_dt, "Params should be unsharded during backward" + + def test_params_resharded_after_backward(self): + """After full backward pass, params should be DTensors again.""" + torch.manual_seed(42) + model = SimpleMLP(64).to(_device()) + fully_shard(model) + + from torch.distributed.tensor import DTensor + + x = torch.randn(2, 64, device=_device(), requires_grad=True) + out = model(x) + loss = out.sum() + loss.backward() + + for _, p in model.named_parameters(): + assert isinstance(p, DTensor), "Params should be DTensors after backward reshard" + + +# ------------------------------------------------------------------ # +# 8. Safety — double-shard rejection +# ------------------------------------------------------------------ # + + +class TestSafety: + def test_double_shard_rejected(self): + """Calling fully_shard on an already-wrapped module should raise ValueError.""" + torch.manual_seed(42) + model = SimpleMLP(64).to(_device()) + fully_shard(model) + with pytest.raises(ValueError, match="already been fully sharded"): + fully_shard(model) + + def test_no_params_module_ok(self): + """fully_shard on a module with no parameters should succeed (no-op).""" + model = nn.Sequential().to(_device()) + wrapped = fully_shard(model) + assert isinstance(wrapped, FSDPModule) + assert _count_fsdp_modules(wrapped) == 1 + + +# ------------------------------------------------------------------ # +# 9. Checkpoint — get_state_dict and preprocess_state_dict_for_uneven_dtensor +# ------------------------------------------------------------------ # + + +class TestCheckpoint: + @pytest.mark.skip(reason="Hangs. Debug in progress.") + def test_get_state_dict_returns_dicts(self): + """get_state_dict should return model and optimizer state dicts.""" + torch.manual_seed(42) + model = SimpleMLP(64).to(_device()) + fully_shard(model) + optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3) + + # Run one step to populate optimizer state + x = torch.randn(2, 64, device=_device()) + out = model(x) + loss = out.sum() + loss.backward() + optimizer.step() + + model_sd, opt_sd = get_state_dict(model, optimizer) + assert isinstance(model_sd, dict) + assert isinstance(opt_sd, dict) + assert len(model_sd) > 0, "Model state dict should not be empty" + + @pytest.mark.skip(reason="Hangs. Debug in progress.") + def test_get_state_dict_nested_fsdp(self): + """get_state_dict should work with nested FSDP modules.""" + torch.manual_seed(42) + device = _device() + model = MOETransformerLayer(64, 128).to(device) + model.experts = fully_shard(model.experts) + model = fully_shard(model) + optimizer = torch.optim.SGD(model.parameters(), lr=0.01) + + x = torch.randn(2, 64, device=device) + out = model(x) + loss = out.sum() + loss.backward() + optimizer.step() + + model_sd, opt_sd = get_state_dict(model, optimizer) + assert len(model_sd) > 0 + + @pytest.mark.skip(reason="Hangs. Debug in progress.") + def test_preprocess_state_dict_adds_metadata(self): + """preprocess_state_dict_for_uneven_dtensor should add chunk metadata.""" + torch.manual_seed(42) + model = SimpleMLP(64).to(_device()) + fully_shard(model) + opt = torch.optim.SGD(model.parameters(), lr=0.0) + + # Build a raw state dict via torch's state_dict + sd = torch_get_state_dict( + model, opt, options=StateDictOptions(full_state_dict=True, cpu_offload=True) + ) + preprocess_state_dict_for_uneven_dtensor(sd) + + # Check that the state dict still contains parameter data + assert len(sd) > 0 + + def test_get_state_dict_strict_all_dtensor(self): + """get_state_dict should assert all params are DTensors.""" + torch.manual_seed(42) + model = SimpleMLP(64).to(_device()) + # DON'T call fully_shard — params are NOT DTensors + optimizer = torch.optim.AdamW(model.parameters()) + + with pytest.raises(AssertionError, match="Expected all parameters to be DTensors"): + get_state_dict(model, optimizer) + + @pytest.mark.skip(reason="Hangs. Debug in progress.") + def test_get_state_dict_llm_scenario(self): + """Full LLM forward-backward-checkpoint cycle should work.""" + torch.manual_seed(42) + device = _device() + model = TinyLLM(vocab=128, hidden=64, num_layers=2).to(device) + for layer in model.layers: + fully_shard(layer) + fully_shard(model) + optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3) + + x = torch.randint(0, 128, (4, 8), device=device) + out = model(x) + loss = out.sum() + loss.backward() + optimizer.step() + + model_sd, opt_sd = get_state_dict(model, optimizer) + assert len(model_sd) > 0 + assert len(opt_sd) > 0 + + @pytest.mark.skip(reason="Hangs. Debug in progress.") + def test_get_state_dict_with_frozen_params(self): + """get_state_dict should work with mixed frozen/trainable params.""" + torch.manual_seed(42) + device = _device() + model = MultimodalModel(hidden=64).to(device) + for p in model.vision_encoder.parameters(): + p.requires_grad = False + fully_shard(model) + optimizer = torch.optim.AdamW([p for p in model.parameters() if p.requires_grad], lr=1e-3) + + x_img = torch.randn(2, 64, device=device, requires_grad=True) + x_txt = torch.randn(2, 64, device=device, requires_grad=True) + out = model(x_img, x_txt) + loss = out.sum() + loss.backward() + optimizer.step() + + model_sd, opt_sd = get_state_dict(model, optimizer) + assert len(model_sd) > 0 diff --git a/tests/unit_tests/distributed/megatron_fsdp/v2/test_mcore_checkpoint.py b/tests/unit_tests/distributed/megatron_fsdp/v2/test_mcore_checkpoint.py new file mode 100644 index 00000000000..812886a1441 --- /dev/null +++ b/tests/unit_tests/distributed/megatron_fsdp/v2/test_mcore_checkpoint.py @@ -0,0 +1,567 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +import logging +import shutil +from pathlib import Path + +import pytest +import torch +from torch.testing import assert_close + +import megatron.core.parallel_state as mpu +from megatron.core.utils import is_torch_min_version +from tests.unit_tests.distributed.megatron_fsdp.utils import ( + make_gpt_mock_data_iterator, + make_moe_args_model_and_optimizer, + pretrain_forward_backward, + set_manual_seed, +) +from tests.unit_tests.test_utilities import Utils + +logger = logging.getLogger(__name__) + +SHARED_TMP_DIR = "/tmp/pytest-shared-tmp" + +try: + from transformer_engine.pytorch.fp8 import check_nvfp4_support + + _NVFP4_AVAILABLE, _NVFP4_SKIP_REASON = check_nvfp4_support() +except Exception: + _NVFP4_AVAILABLE = False + _NVFP4_SKIP_REASON = "NVFP4 support unavailable" + + +# ------------------------------------------------------------------ +# Helpers +# ------------------------------------------------------------------ + + +def _state_dict_to_full_tensor(sd): + from torch.distributed.tensor import DTensor + + from megatron.core.distributed.fsdp.src.megatron_fsdp.uneven_dtensor import ( + uneven_dtensor_to_full_tensor, + ) + + out = {} + for k, v in sd.items(): + if isinstance(v, DTensor): + out[k] = uneven_dtensor_to_full_tensor(v) + else: + out[k] = v + return out + + +def _normalize_key(key: str) -> str: + while key.startswith("module."): + key = key[len("module.") :] + return key + + +def _get_model_from_chunks(model_chunks): + if isinstance(model_chunks, list): + return model_chunks[0] + return model_chunks + + +def _assert_model_match(source_full, loaded_full): + nonempty = False + for s_key, s_val in source_full.items(): + if s_key.endswith("._extra_state"): + continue # Skip _extra_state buffers which is runtime-dependent and not guaranteed to match + canonical = _normalize_key(s_key) + matched_key = None + for l_key in loaded_full: + if _normalize_key(l_key) == canonical: + matched_key = l_key + break + assert ( + matched_key is not None + ), f"Key {s_key} (canonical: {canonical}) not found in loaded state dict" + l_val = loaded_full[matched_key] + if s_val is None and l_val is None: + continue + assert ( + s_val is not None and l_val is not None + ), f"One of source or loaded value for {s_key} is None while the other is not" + assert ( + s_val.shape == l_val.shape + ), f"Shape mismatch for {s_key}: {s_val.shape} vs {l_val.shape}" + if s_val.numel() > 0: + nonempty = True + else: + continue + assert_close( + s_val, + l_val, + rtol=0, + atol=0, + msg=f"Value mismatch for {s_key}, s_val: {s_val}, l_val: {l_val}", + ) + + world_size = torch.distributed.get_world_size() + all_nonempty = [False] * world_size + torch.distributed.all_gather_object(all_nonempty, nonempty) + assert any(all_nonempty), "All ranks had empty model state after load." + + +def _optim_state_to_full(optim_sd_or_optim, model): + """Wrap optimizer states as DTensors or unflatten flat format and gather to full tensors. + + Accepts either a sharded_state_dict (fsdp_dtensor format) or a + DistributedOptimizer instance (nd fully_reshardable format). + """ + from torch.distributed.tensor import DTensor + + from megatron.core.distributed.fsdp.checkpoint import _build_dtensor_optim_sd + from megatron.core.distributed.fsdp.src.megatron_fsdp.uneven_dtensor import ( + uneven_dtensor_to_full_tensor, + ) + + # ND-parallel source: use get_parameter_state_dp_zero directly + if not isinstance(optim_sd_or_optim, dict): + return _nd_optim_state_to_full(optim_sd_or_optim, model) + + optim_sd = optim_sd_or_optim + if "param_state_sharding_type" in optim_sd: + return _flat_optim_state_to_full(optim_sd, model) + + wrapped = {"optimizer": _build_dtensor_optim_sd(optim_sd, model)} + out = {} + for param_name, param_states in wrapped["optimizer"]["state"].items(): + out[param_name] = {} + for state_key, state_val in param_states.items(): + if isinstance(state_val, DTensor): + out[param_name][state_key] = uneven_dtensor_to_full_tensor(state_val) + else: + out[param_name][state_key] = state_val + return out + + +def _nd_optim_state_to_full(optim, model): + """Get full optimizer state from an ND-parallel DistributedOptimizer. + + Uses ``get_parameter_state_dp_zero`` (the same method as + ``sharded_param_state_fully_reshardable``) to gather all-gathered + world tensors, then unflattens them per-parameter using + ``param_index_map``. Handles ``ChainedOptimizer`` (MoE) by + iterating all inner ``DistributedOptimizer`` instances. + """ + param_to_name = {} + for name, p in model.named_parameters(): + if p.requires_grad: + param_to_name[p] = name + + out = {} + inner_optims = getattr(optim, "chained_optimizers", [optim]) + for inner_optim in inner_optims: + dp_zero = inner_optim.get_parameter_state_dp_zero( + use_gloo_comm=False, return_on_all_ranks=True + ) + for gbuf_idx, gbuf_range_maps in enumerate(inner_optim.gbuf_ranges): + buffer = inner_optim.buffers[gbuf_idx] + for dtype, gbuf_range_map_for_all_buckets in gbuf_range_maps.items(): + world_tensors = dp_zero[gbuf_idx][dtype] + for model_param, (start, end, _) in buffer.param_index_map.items(): + if model_param not in param_to_name: + continue + param_name = param_to_name[model_param] + out[param_name] = { + state_key: world_tensors[state_key][start:end] + .reshape(model_param.shape) + .clone() + .cuda() + for state_key in world_tensors + if state_key in ("param", "exp_avg", "exp_avg_sq") + } + return out + + +def _flat_optim_state_to_full(optim_sd, model): + """Convert fully_reshardable optimizer sd to {param_name: {state_key: tensor}}. + + The ``fully_reshardable`` format stores optimizer state as a nested + dict: ``optim_sd["param_state"]`` = ``{0: {"param": ShardedTensor, + "exp_avg": ShardedTensor, ...}, 1: {...}, ...}``. This function + extracts the full tensors from the ShardedTensor wrappers and maps + integer indices to parameter names by matching the ``"param"`` + tensor shape against ``model.named_parameters()``. + """ + from torch.distributed.checkpoint.stateful import ShardedTensor + + param_state = optim_sd.get("param_state", {}) + if not param_state: + return {} + + # Build shape → param_name lookup + param_by_shape = {} + for name, p in model.named_parameters(): + if p.requires_grad: + param_by_shape.setdefault(p.shape, []).append(name) + + out = {} + used_names = set() + for idx, states in sorted(param_state.items()): + # Find matching param name via the "param" tensor shape + param_tensor = None + if "param" in states: + st = states["param"] + if isinstance(st, ShardedTensor): + shards = st.local_shards() + if shards: + param_tensor = shards[0].tensor + param_name = None + if param_tensor is not None: + shape = param_tensor.shape + candidates = [n for n in param_by_shape.get(shape, []) if n not in used_names] + if candidates: + param_name = candidates[0] + used_names.add(param_name) + if param_name is None: + param_name = f"param_{idx}" + + out[param_name] = {} + for state_key, st in states.items(): + if isinstance(st, ShardedTensor): + shards = st.local_shards() + out[param_name][state_key] = shards[0].tensor if shards else st + else: + out[param_name][state_key] = st + + return out + + +def _assert_optim_match(source_optim_full, loaded_optim_full, ignore_param=True): + for param_name, source_states in source_optim_full.items(): + canonical = _normalize_key(param_name) + matched_param = None + for l_param in loaded_optim_full: + if _normalize_key(l_param) == canonical: + matched_param = l_param + break + assert matched_param is not None, ( + f"Optimizer param {param_name} (canonical: {canonical}) " + f"not found in loaded optimizer state" + ) + loaded_states = loaded_optim_full[matched_param] + for state_key, s_val in source_states.items(): + if ignore_param and state_key == "param": + continue + assert ( + state_key in loaded_states + ), f"Optimizer state '{state_key}' for param {param_name} not found after load" + l_val = loaded_states[state_key] + if s_val is None and l_val is None: + continue + if isinstance(s_val, torch.Tensor) and s_val.numel() > 0: + assert ( + s_val.shape == l_val.shape + ), f"Optimizer state shape mismatch for {param_name}.{state_key}" + assert_close( + s_val, + l_val, + atol=0, + rtol=0, + msg=f"Optimizer state value mismatch for {param_name}.{state_key}, s_val: {s_val}, l_val: {l_val}", + ) + + +# ================================================================== +# Test class +# ================================================================== + + +class TestMegatronFsdpV2Checkpoint: + + @staticmethod + def _init_model_and_optimizer(seed=42, **kwargs): + VOCAB_SIZE = kwargs.get("vocab_size", 100) + MAX_SEQ_LEN = kwargs.get("seq_length", 128) + MICRO_BATCH_SIZE = kwargs.get("micro_batch_size", 2) + GLOBAL_BATCH_SIZE = kwargs.get("global_batch_size", 32) + NUM_TRAINING_STEPS = kwargs.get("train_iters", 2) + TP = kwargs.get("TP", 1) + PP = kwargs.get("PP", 1) + VPP = kwargs.get("VPP", None) + EP = kwargs.get("EP", 1) + ETP = kwargs.get("ETP", 1) + OUTER_DP = kwargs.get("OUTER_DP", 1) + + Utils.initialize_model_parallel( + tensor_model_parallel_size=TP, + pipeline_model_parallel_size=PP, + expert_model_parallel_size=EP, + expert_tensor_parallel_size=ETP, + num_distributed_optimizer_instances=OUTER_DP, + ) + set_manual_seed(seed) + + model_chunks, optim = make_moe_args_model_and_optimizer( + ut_filename="test_mcore_checkpoint.py", + micro_batch_size=MICRO_BATCH_SIZE, + global_batch_size=GLOBAL_BATCH_SIZE, + vocab_size=VOCAB_SIZE, + padded_vocab_size=VOCAB_SIZE, + seq_length=MAX_SEQ_LEN, + max_position_embeddings=MAX_SEQ_LEN, + sequence_parallel=TP > 1, + tensor_model_parallel_size=TP, + pipeline_model_parallel_size=PP, + num_layers_per_virtual_pipeline_stage=VPP, + train_iters=NUM_TRAINING_STEPS, + **kwargs, + ) + from megatron.training.training import get_optimizer_param_scheduler + + opt_param_scheduler = get_optimizer_param_scheduler(optim) + return model_chunks, optim, opt_param_scheduler + + @staticmethod + def _training_loop(seed=42, **kwargs): + model_chunks, optim, scheduler = TestMegatronFsdpV2Checkpoint._init_model_and_optimizer( + seed=seed, **kwargs + ) + MICRO_BATCH_SIZE = kwargs.get("micro_batch_size", 2) + GLOBAL_BATCH_SIZE = kwargs.get("global_batch_size", 32) + MAX_SEQ_LEN = kwargs.get("seq_length", 128) + DP_GROUP = mpu.get_data_parallel_group() + + data_iterator = make_gpt_mock_data_iterator( + dp_group=DP_GROUP, + vocab_size=kwargs.get("vocab_size", 100), + sequence_length=MAX_SEQ_LEN, + batch_size=MICRO_BATCH_SIZE, + num_samples=GLOBAL_BATCH_SIZE * kwargs.get("train_iters", 2), + ) + for _ in range(kwargs.get("train_iters", 2)): + optim.zero_grad() + pretrain_forward_backward( + model=model_chunks, + data_iterator=data_iterator, + sequence_length=MAX_SEQ_LEN, + micro_batch_size=MICRO_BATCH_SIZE, + num_micro_batches=GLOBAL_BATCH_SIZE // MICRO_BATCH_SIZE // DP_GROUP.size(), + ) + optim.step() + + if "save" in kwargs: + from megatron.training.checkpointing import save_checkpoint as mcore_save_checkpoint + + mcore_save_checkpoint( + iteration=0, + model=model_chunks, + optimizer=optim, + opt_param_scheduler=scheduler, + num_floating_point_operations_so_far=0, + ) + + model = _get_model_from_chunks(model_chunks) + return model, model.state_dict(), optim, scheduler + + @classmethod + def setup_class(cls): + Utils.initialize_model_parallel() + + @classmethod + def teardown_class(cls): + Utils.destroy_model_parallel() + + # ================================================================== + # Online checkpoint conversion + # ================================================================== + + # ---- Shared target config (always MFSDP v2) ---- + _TARGET_BASE = dict( + use_megatron_fsdp=True, + use_megatron_fsdp_v2=True, + init_model_with_meta_device=True, + ckpt_format="fsdp_dtensor", + gradient_accumulation_fusion=False, + overlap_param_gather=True, + overlap_grad_reduce=True, + recompute_granularity="full", + recompute_method="uniform", + recompute_num_layers=1, + fp8_param_gather=False, + ) + + @pytest.mark.skipif( + not is_torch_min_version("2.4.0"), + reason="Requires DTensor and DeviceMesh support (PyTorch >= 2.4.0).", + ) + @pytest.mark.parametrize("nd_topology", [pytest.param({"EP": 2}, id="EP2")]) + @pytest.mark.parametrize( + "source_type, source_configs, target_configs", + [ + # ---- ND-parallel → MFSDP v2 ---- + pytest.param( + "nd", + dict( + distrib_optim_sharding_type="fully_reshardable", + dist_ckpt_optim_fully_reshardable=True, + ), + dict(data_parallel_sharding_strategy="optim_grads_params"), + id="nd_fully_reshardable_to_v2", + ), + # ---- MFSDP v2 → MFSDP v2 (round-trip) ---- + pytest.param( + "v2", + dict(data_parallel_sharding_strategy="optim_grads_params"), + dict(data_parallel_sharding_strategy="optim_grads_params"), + id="v2_rt_optim_grads_params", + ), + pytest.param( + "v2", + dict(data_parallel_sharding_strategy="optim_grads"), + dict(data_parallel_sharding_strategy="optim_grads"), + id="v2_rt_optim_grads", + ), + # ---- MFSDP v2 → MFSDP v2 (cross-setting) ---- + pytest.param( + "v2", + dict(data_parallel_sharding_strategy="optim_grads_params"), + dict(data_parallel_sharding_strategy="optim_grads"), + id="v2_x_optim_grads_params_to_optim_grads", + ), + pytest.param( + "v2", + dict(data_parallel_sharding_strategy="optim_grads"), + dict(data_parallel_sharding_strategy="optim_grads_params"), + id="v2_x_optim_grads_to_optim_grads_params", + ), + # ---- MFSDP v2 → MFSDP v2 (round-trip, NVFP4) ---- + pytest.param( + "v2", + dict( + data_parallel_sharding_strategy="optim_grads_params", + fp4="e2m1", + fp4_recipe="nvfp4", + fp4_param_gather=True, + bf16=True, + ), + dict( + data_parallel_sharding_strategy="optim_grads_params", + fp4="e2m1", + fp4_recipe="nvfp4", + fp4_param_gather=True, + bf16=True, + ), + marks=pytest.mark.skipif(not _NVFP4_AVAILABLE, reason=_NVFP4_SKIP_REASON), + id="v2_rt_nvfp4_optim_grads_params", + ), + # ---- MFSDP v1 baseline → MFSDP v2 ---- + pytest.param( + "v1", + dict(data_parallel_sharding_strategy="optim_grads_params"), + dict(data_parallel_sharding_strategy="optim_grads_params"), + id="v1_to_v2_optim_grads_params", + ), + pytest.param( + "v1", + dict(data_parallel_sharding_strategy="optim_grads"), + dict(data_parallel_sharding_strategy="optim_grads"), + id="v1_to_v2_optim_grads", + ), + pytest.param( + "v1", + dict(data_parallel_sharding_strategy="optim"), + dict(data_parallel_sharding_strategy="optim_grads_params"), + id="v1_optim_to_v2_optim_grads_params", + ), + ], + ) + def test_checkpoint_online_convert( + self, request, nd_topology, source_type, source_configs, target_configs + ): + """ + Train a source model (ND-parallel, MFSDP v1, or MFSDP v2) with + *source_configs*, save via ``save_checkpoint``, load into an + MFSDP v2 model with *target_configs*, and verify both model and + optimizer state. + """ + if source_type == "v1": + pytest.skip("v1 checkpoint format not available") + + if source_type == "v2": + src_sharding = source_configs.get( + "data_parallel_sharding_strategy", "optim_grads_params" + ) + tgt_sharding = target_configs.get( + "data_parallel_sharding_strategy", "optim_grads_params" + ) + if src_sharding != "optim_grads_params" or tgt_sharding != "optim_grads_params": + pytest.skip( + "Only 'optim_grads_params' sharding strategy is supported for v2 source" + ) + + # ---- Build source config ---- + if source_type == "v2": + src_base = dict(self._TARGET_BASE, auto_detect_ckpt_format=True) + elif source_type == "v1": + src_base = dict( + use_megatron_fsdp=True, + init_model_with_meta_device=True, + ckpt_format="fsdp_dtensor", + gradient_accumulation_fusion=False, + ) + elif source_type == "nd": + src_base = dict( + use_distributed_optimizer=True, fp8_param_gather=False, ckpt_format="torch_dist" + ) + else: + raise ValueError(f"Unknown source_type: {source_type}") + + # ND-parallel → V2: optimizer format is incompatible (bucket-based vs name-based). + src_sharding_type = source_configs.get("distrib_optim_sharding_type", "fsdp_dtensor") + supports_optim = src_sharding_type != "dp_reshardable" + + save_config = {**nd_topology, **src_base, **source_configs} + tgt_configs = {**nd_topology, **self._TARGET_BASE, **target_configs} + if not supports_optim: + save_config["no_save_optim"] = True + tgt_configs["no_load_optim"] = True + + test_id = request.node.name.replace("[", "_").replace("]", "").replace("/", "_") + ckpt_base = Path(SHARED_TMP_DIR) / TestMegatronFsdpV2Checkpoint.__name__ / test_id + + # ---- Train + save with source config ---- + source_model, source_sd, source_optim, _ = TestMegatronFsdpV2Checkpoint._training_loop( + save=str(ckpt_base), save_interval=1, **save_config + ) + source_full = _state_dict_to_full_tensor(source_sd) + + if supports_optim: + if source_type == "nd": + source_optim_full = _optim_state_to_full(source_optim, source_model) + else: + source_optim_sd = source_optim.sharded_state_dict( + model_sharded_state_dict=( + source_model.sharded_state_dict() if source_type not in ["v1", "v2"] else {} + ), + metadata={"distrib_optim_sharding_type": src_sharding_type}, + ) + source_optim_full = _optim_state_to_full(source_optim_sd, source_model) + Utils.destroy_model_parallel() + + # ---- Load with target config (always MFSDP v2) ---- + v2_model_chunks, loaded_optim, _ = TestMegatronFsdpV2Checkpoint._init_model_and_optimizer( + load=str(ckpt_base), **tgt_configs + ) + v2_model = _get_model_from_chunks(v2_model_chunks) + + # ---- Verify model ---- + loaded_full = _state_dict_to_full_tensor(v2_model.state_dict()) + _assert_model_match(source_full, loaded_full) + + # ---- Verify optimizer (FSDP source types only) ---- + if supports_optim: + loaded_optim_sd = loaded_optim.sharded_state_dict( + metadata={"distrib_optim_sharding_type": "fsdp_dtensor"} + ) + loaded_optim_full = _optim_state_to_full(loaded_optim_sd, v2_model) + _assert_optim_match(source_optim_full, loaded_optim_full) + + Utils.destroy_model_parallel() + if torch.distributed.get_rank() == 0: + shutil.rmtree(ckpt_base, ignore_errors=True) + torch.distributed.barrier() diff --git a/tests/unit_tests/distributed/megatron_fsdp/v2/test_mcore_nd_parallel.py b/tests/unit_tests/distributed/megatron_fsdp/v2/test_mcore_nd_parallel.py new file mode 100644 index 00000000000..3984733eba3 --- /dev/null +++ b/tests/unit_tests/distributed/megatron_fsdp/v2/test_mcore_nd_parallel.py @@ -0,0 +1,312 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +import copy + +import pytest +import torch +from torch.testing import assert_close + +import megatron.core.parallel_state as mpu +from megatron.core.distributed.fsdp.src.megatron_fsdp.mixed_precision import HAVE_TE_MXFP8TENSOR +from megatron.core.distributed.fsdp.src.megatron_fsdp.v2.mixed_precision import ( + HAVE_TE_NVFP4, + HAVE_TE_NVFP4_RECIPE, +) +from megatron.core.utils import is_torch_min_version +from tests.unit_tests.distributed.megatron_fsdp.utils import ( + make_gpt_mock_data_iterator, + make_moe_args_model_and_optimizer, + pretrain_forward_backward, + set_manual_seed, +) +from tests.unit_tests.test_utilities import Utils + + +@pytest.fixture(scope="class") +def ref_cache(): + """ + Shared read/write cache for an class. + Keys: arbitrary strings, values: anything (tensors, dicts, etc.). + """ + return {} + + +class TestMegatronFSDPE2E: + + @staticmethod + def _training_loop(seed=42, **kwargs): + """ + Run a small deterministic (optional) training loop using a mocked MoE/GPT model and optimizer. + This helper initializes model-parallel state, creates a model and optimizer via + make_moe_args_model_and_optimizer, constructs a mock GPT data iterator, and runs + NUM_TRAINING_STEPS iterations of forward/backward/optimization. Losses from each + training step are collected and returned. + Args: + seed (int, optional): RNG seed for reproducibility. Default: 42. + **kwargs: Configuration overrides (all optional). Recognized keys: + - vocab_size (int): Vocabulary size for the mock model. Default: 100. + - seq_length (int): Sequence length used for the mock data. Default: 128. + - micro_batch_size (int): Per-microbatch size. Default: 2. + - global_batch_size (int): Global batch size across data-parallel ranks. Default: 32. + - train_iters (int): Number of training iterations to run. Default: 20. + - tensor_model_parallel_size (int): Tensor model parallel world size. Default: 1. + - pipeline_model_parallel_size (int): Pipeline model parallel world size. Default: 1. + - num_layers_per_virtual_pipeline_stage (int or None): Virtual pipeline configuration. + - expert_model_parallel_size (int): Expert model parallel size for MoE. Default: 1. + - expert_tensor_parallel_size (int): Expert tensor parallel size for MoE. Default: 1. + - num_distributed_optimizer_instances (int): Number of distributed optimizer instances. Default: 1. + Returns: + list: A list of length train_iters containing the per-step language-model loss values + (the value appended from output[-1] each iteration). Loss objects are returned as produced + by the training utilities (typically tensors or scalars). + Side effects: + - Calls Utils.initialize_model_parallel(...) and Utils.destroy_model_parallel(). + - Sets global RNG state via set_manual_seed(seed). + - Constructs models/optimizers via make_moe_args_model_and_optimizer and a data iterator + via make_gpt_mock_data_iterator. + - Runs optimizer.zero_grad(), pretrain_forward_backward(...), and optim.step() repeatedly. + - Calculates the number of micro-batches per step as: + global_batch_size // micro_batch_size // data_parallel_world_size. + This requires that global_batch_size be divisible by micro_batch_size * data_parallel_world_size. + Raises: + ValueError: If batch-size arithmetic or other setup assumptions (e.g., divisibility) are violated. + """ + # Configuration parameters with defaults + VOCAB_SIZE = kwargs.pop("vocab_size", 100) + MAX_SEQ_LEN = kwargs.pop("seq_length", 128) + MICRO_BATCH_SIZE = kwargs.pop("micro_batch_size", 2) + GLOBAL_BATCH_SIZE = kwargs.pop("global_batch_size", 32) + NUM_TRAINING_STEPS = kwargs.pop("train_iters", 20) + TP = kwargs.pop("TP", 1) + PP = kwargs.pop("PP", 1) + VPP = kwargs.pop("VPP", None) + EP = kwargs.pop("EP", 1) + ETP = kwargs.pop("ETP", 1) + OUTER_DP = kwargs.pop("OUTER_DP", 1) + + # Initialize model parallel groups + Utils.initialize_model_parallel( + tensor_model_parallel_size=TP, + pipeline_model_parallel_size=PP, + expert_model_parallel_size=EP, + expert_tensor_parallel_size=ETP, + num_distributed_optimizer_instances=OUTER_DP, + ) + DP_GROUP = mpu.get_data_parallel_group() + + # Set manual seed for reproducibility + set_manual_seed(seed) + + # Create model and optimizer + model_chunks, optim = make_moe_args_model_and_optimizer( + ut_filename="test_mcore_fully_sharded_data_parallel.py", + micro_batch_size=MICRO_BATCH_SIZE, + global_batch_size=GLOBAL_BATCH_SIZE, + vocab_size=VOCAB_SIZE, + padded_vocab_size=VOCAB_SIZE, + seq_length=MAX_SEQ_LEN, + sequence_parallel=TP > 1, + tensor_model_parallel_size=TP, + pipeline_model_parallel_size=PP, + num_layers_per_virtual_pipeline_stage=VPP, + train_iters=NUM_TRAINING_STEPS, + **kwargs, + ) + if kwargs.get("use_megatron_fsdp", False) and kwargs.get( + "use_precision_aware_optimizer", False + ): + assert ( + not optim.optimizer.master_weights + ), "Megatron-FSDP should not use FusedAdam master weights." + + # Prepare data iterator + data_iterator = make_gpt_mock_data_iterator( + dp_group=DP_GROUP, + vocab_size=VOCAB_SIZE, + sequence_length=MAX_SEQ_LEN, + batch_size=MICRO_BATCH_SIZE, + num_samples=GLOBAL_BATCH_SIZE * NUM_TRAINING_STEPS, + ) + + outputs = [] + + # Training loop + for _ in range(NUM_TRAINING_STEPS): + optim.zero_grad() + output = pretrain_forward_backward( + model=model_chunks, + data_iterator=data_iterator, + sequence_length=MAX_SEQ_LEN, + micro_batch_size=MICRO_BATCH_SIZE, + num_micro_batches=GLOBAL_BATCH_SIZE // MICRO_BATCH_SIZE // DP_GROUP.size(), + ) + optim.step() + + # Collect loss + outputs.append(output[-1]) + + Utils.destroy_model_parallel() + + return outputs + + @pytest.mark.skipif( + not is_torch_min_version("2.4.0"), reason="Test needs to be updated for torch >= 2.4.0" + ) + @pytest.mark.parametrize("nd_topology", [pytest.param({"EP": 2}, id="EP2")]) + @pytest.mark.parametrize( + "spec_configs", + [ + pytest.param( + dict( + data_parallel_sharding_strategy="optim_grads_params", + recompute_granularity="full", + recompute_method="uniform", + recompute_num_layers=1, + overlap_param_gather=True, + overlap_grad_reduce=True, + use_megatron_fsdp_v2=True, + gradient_accumulation_fusion=True, + ), + id="optim_grads_params_double_buffer", + ), + pytest.param( + dict( + bf16=True, + data_parallel_sharding_strategy="optim_grads_params", + fp8="e4m3", + fp8_param_gather=True, + fp8_recipe="mxfp8", + moe_grouped_gemm=True, + overlap_param_gather=True, + overlap_grad_reduce=True, + use_megatron_fsdp_v2=True, + ), + id="optim_grads_params_mxfp8_param_gather", + ), + pytest.param( + dict( + bf16=True, + data_parallel_sharding_strategy="optim_grads_params", + fp4="e2m1", + fp4_recipe="nvfp4", + fp4_param_gather=True, + main_grads_dtype="fp32", + main_params_dtype="fp32", + overlap_param_gather=True, + overlap_grad_reduce=True, + use_megatron_fsdp_v2=True, + ), + id="optim_grads_params_nvfp4_param_gather", + ), + ], + ) + def test_compatible_with_nd_parallel(self, ref_cache, nd_topology, spec_configs): + if spec_configs.get("fp8_recipe") == "mxfp8" and ( + not torch.cuda.is_available() + or torch.cuda.get_device_capability()[0] < 10 + or not HAVE_TE_MXFP8TENSOR + ): + pytest.skip("Requires PyTorch & CUDA device with TE MXFP8Tensor support") + + if spec_configs.get("fp4_param_gather"): + if not torch.cuda.is_available(): + pytest.skip("CUDA is required for NVFP4") + if not (HAVE_TE_NVFP4 and HAVE_TE_NVFP4_RECIPE): + pytest.skip("NVFP4 requires Transformer Engine >= 2.7.0.dev0") + try: + from transformer_engine.pytorch.fp8 import check_nvfp4_support + + is_nvfp4_available, reason = check_nvfp4_support() + if not is_nvfp4_available: + pytest.skip("NVFP4 not available: " + reason) + except ImportError: + pytest.skip("NVFP4 support check requires Transformer Engine >= 2.7.0.dev0") + + reference_kind = "distopt" + ref_cache_key = ( + reference_kind, + tuple(sorted(nd_topology.items())), + tuple(sorted((key, repr(value)) for key, value in spec_configs.items())), + ) + if ref_cache_key not in ref_cache: + reference_spec_configs = copy.deepcopy(spec_configs) + reference_spec_configs["use_megatron_fsdp_v2"] = False + reference_spec_configs["gradient_accumulation_fusion"] = False + reference_spec_configs["fp8_param_gather"] = False + ref_cache[ref_cache_key] = TestMegatronFSDPE2E._training_loop( + use_distributed_optimizer=True, **nd_topology, **reference_spec_configs + ) + + fsdp_spec_configs = copy.deepcopy(spec_configs) + fsdp_spec_configs.setdefault("gradient_accumulation_fusion", False) + outputs = TestMegatronFSDPE2E._training_loop( + use_megatron_fsdp=True, + init_model_with_meta_device=True, + ckpt_format="fsdp_dtensor", + **nd_topology, + **fsdp_spec_configs, + ) + reference_outputs = ref_cache[ref_cache_key] + + if torch.distributed.get_rank() == 0: + for step, (output, ref_output) in enumerate(zip(outputs, reference_outputs)): + loss = output["lm loss"] + ref_loss = ref_output["lm loss"] + assert_close( + loss, + ref_loss, + atol=0, + rtol=0.05, + msg=( + f"Loss mismatch at step {step}, FSDP Loss = {loss.detach().item()}, " + f"Reference Loss = {ref_loss.detach().item()}" + f", Compare = {compare_losses(loss.detach().item(), ref_loss.detach().item())}" + f", outputs = {outputs}, reference_outputs = {reference_outputs}" + ), + ) + + +def compare_losses(loss_a: float, loss_b: float, reference: str = "b"): + """ + Compare two loss values with absolute and relative differences. + + Parameters + ---------- + loss_a : float + First loss value (e.g., baseline model). + loss_b : float + Second loss value (e.g., new model). + reference : {"a", "b"}, default "b" + Which loss to treat as the reference when computing the + relative difference. If "b", relative diff is vs loss_b; + if "a", vs loss_a. + + Returns + ------- + dict with keys: + "abs_diff" : float + |loss_a - loss_b| + "rel_diff" : float + |loss_a - loss_b| / reference_loss + "better" : str + "a" if loss_a < loss_b, "b" if loss_b < loss_a, "equal" otherwise. + """ + abs_diff = abs(loss_a - loss_b) + + if reference == "a": + ref = loss_a + else: + ref = loss_b + + if ref == 0: + rel_diff = float("inf") # or None, depending on your preference + else: + rel_diff = abs_diff / ref + + if loss_a < loss_b: + better = "a" + elif loss_b < loss_a: + better = "b" + else: + better = "equal" + + return {"abs_diff": abs_diff, "rel_diff": rel_diff, "better": better} diff --git a/tests/unit_tests/distributed/megatron_fsdp/v2/test_param_group.py b/tests/unit_tests/distributed/megatron_fsdp/v2/test_param_group.py new file mode 100644 index 00000000000..ca2e80e3acc --- /dev/null +++ b/tests/unit_tests/distributed/megatron_fsdp/v2/test_param_group.py @@ -0,0 +1,313 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Multi-process integration tests for ParameterGroup. + +Builds a layer with mixed-dtype params (bf16 + uint8), splits them into +two ParameterGroups, and validates shard / unshard / reshard / reduce_grad +across all four sharding strategies. + +Run with: + torchrun --nproc_per_node=4 -m pytest megatron.core.distributed.fsdp.src.megatron_fsdp.tests.test_param_group -v +""" + +import sys +from pathlib import Path + +import pytest +import torch +import torch.nn as nn + +sys.path.insert(0, str(Path(__file__).parents[2])) +from megatron.core.distributed.fsdp.src.megatron_fsdp.v2.mixed_precision import ( + FullyShardMixedPrecisionPolicy, +) +from megatron.core.distributed.fsdp.src.megatron_fsdp.v2.param_group import ParameterGroup +from megatron.core.distributed.fsdp.src.megatron_fsdp.v2.utils import ParamGroupIdx + +# ------------------------------------------------------------------ # +# Process group — init once per pytest session, shared by all tests +# ------------------------------------------------------------------ # + + +@pytest.fixture(scope="session", autouse=True) +def dist_env(): + """Initialize NCCL process group once and tear down at session end.""" + if not torch.distributed.is_initialized(): + torch.distributed.init_process_group(backend="nccl") + rank = torch.distributed.get_rank() + device = torch.device(f"cuda:{rank % torch.cuda.device_count()}") + torch.cuda.set_device(device) + yield + torch.distributed.destroy_process_group() + + +# ------------------------------------------------------------------ # +# Test model — contains bf16 and uint8 (simulated fp8) params +# ------------------------------------------------------------------ # + + +class MixedDtypeLayer(nn.Module): + """A toy layer with large matrices, small biases, and quantized projections.""" + + def __init__(self): + super().__init__() + self.linear1 = nn.Linear(16, 32, bias=False) # bf16, shape [512, 256] + self.linear2 = nn.Linear(32, 16, bias=True) # bf16, shape [256, 512] + bias [256] + self.norm = nn.LayerNorm(16) # bf16, weight [256] + bias [256] + self.quant_proj = nn.Linear(16, 8, bias=False) # uint8, shape [128, 256] + self.quant_gate = nn.Linear(16, 4, bias=False) # uint8, shape [64, 256] + + +# ------------------------------------------------------------------ # +# Helpers +# ------------------------------------------------------------------ # + + +def _build_groups(strategy): + """Create two ParameterGroups (bf16 + uint8) and call init_buffers. + + Returns (groups, originals, dp_group, rank, world_size, device) where + `originals[i]` is a list of cloned param tensors before any sharding. + """ + rank = torch.distributed.get_rank() + device = torch.device(f"cuda:{rank % torch.cuda.device_count()}") + dp_group = torch.distributed.group.WORLD + + # Fixed seed so all ranks start with identical weights + torch.manual_seed(42) + layer = MixedDtypeLayer() + + # Split params by dtype + bf16_params, uint8_params = [], [] + for name, p in layer.named_parameters(): + if "quant" in name: + uint8_params.append( + nn.Parameter(p.data.to(torch.uint8).to(device), requires_grad=False) + ) + else: + bf16_params.append(nn.Parameter(p.data.to(torch.bfloat16).to(device))) + + # Build one ParameterGroup per dtype, each with its own param_group_id + groups, originals = [], [] + for gid, params in enumerate([bf16_params, uint8_params]): + if not params: + continue + originals.append([p.detach().clone() for p in params]) + pg = ParameterGroup( + params=params, + param_group_id=ParamGroupIdx(0, gid), + mp_policy=FullyShardMixedPrecisionPolicy(), + mesh=None, + sharding_strategy=strategy, + ) + groups.append(pg) + return groups, originals, dp_group, rank, torch.distributed.get_world_size(), device + + +def _flags(s): + """Return (has_model_weight_buf, has_grad_buf, weight_distributed, grad_distributed). + + - has_model_weight_buf: whether model_weight_buffer is created (False for no_shard) + - has_grad_buf: always True when requires_grad, across all strategies + - weight_distributed: only True for optim_grads_params (full FSDP) + - grad_distributed: True for optim_grads and optim_grads_params + """ + return { + "no_shard": (False, True, False, False), + "optim": (True, True, False, False), + "optim_grads": (True, True, False, True), + "optim_grads_params": (True, True, True, True), + }[s] + + +# ------------------------------------------------------------------ # +# PyTorch reference — thin wrappers around torch.distributed +# ------------------------------------------------------------------ # + + +class Ref: + @staticmethod + def all_gather(shard, group): + ws = torch.distributed.get_world_size(group) + out = torch.empty(shard.numel() * ws, dtype=shard.dtype, device=shard.device) + torch.distributed.all_gather_into_tensor(out, shard, group=group) + return out + + @staticmethod + def reduce_scatter(full, group): + ws = torch.distributed.get_world_size(group) + ss = full.numel() // ws + out = torch.empty(ss, dtype=full.dtype, device=full.device) + torch.distributed.reduce_scatter_tensor(out, full, group=group) + return out + + @staticmethod + def all_reduce(t, group): + torch.distributed.all_reduce(t, group=group) + return t + + +# ------------------------------------------------------------------ # +# Part 1: init_buffers — verify buffer creation and shard correctness +# ------------------------------------------------------------------ # + + +@pytest.mark.parametrize("strategy", ["no_shard", "optim", "optim_grads", "optim_grads_params"]) +def test_init_buffers(strategy): + if strategy != "optim_grads_params": + pytest.skip( + f"Only optim_grads_params strategy creates distributed buffers, skipping {strategy} test." + ) + + groups, originals, dp_group, rank, ws, device = _build_groups(strategy) + has_wbuf, _, w_dist, g_dist = _flags(strategy) + + for pg, orig in zip(groups, originals): + # -- model_weight_buffer -- + if has_wbuf: + assert pg.model_weight_buffer is not None + wbuf = pg.model_weight_buffer + assert wbuf.is_distributed == w_dist + + # Per-param check: get_item should return this rank's portion of + # the original param. A param may span shard boundaries, so the + # returned slice can be shorter than the full param or even empty. + for i, p in enumerate(orig): + item = wbuf.get_item(i) + if w_dist: + s, e = wbuf.buffer_index._get_item_self_range(i) + expected = p.flatten()[s:e] + else: + expected = p.flatten() + assert torch.equal(item, expected) + else: + # no_shard: model_weight_buffer should not exist + assert pg.model_weight_buffer is None + + # -- main_grad_buffer -- + if pg.requires_grad: + assert pg.main_grad_buffer is not None + assert pg.main_grad_buffer.is_distributed == g_dist + # Grad buffer must be zero-initialized + assert torch.all(pg.main_grad_buffer.data == 0) + + torch.distributed.barrier() + + +# ------------------------------------------------------------------ # +# Part 2: unshard + reshard — verify all-gather and cleanup +# ------------------------------------------------------------------ # + + +@pytest.mark.parametrize("strategy", ["no_shard", "optim", "optim_grads", "optim_grads_params"]) +def test_unshard_reshard(strategy): + if strategy != "optim_grads_params": + pytest.skip( + f"Only optim_grads_params strategy creates distributed buffers, skipping {strategy} test." + ) + + groups, originals, dp_group, rank, ws, device = _build_groups(strategy) + _, _, w_dist, _ = _flags(strategy) + + for pg, orig in zip(groups, originals): + wbuf = pg.model_weight_buffer + if wbuf is None: + # no_shard: nothing to unshard + continue + + shard_before = wbuf.data.clone() + unsharded, work = wbuf.unshard(async_op=False) + + if not w_dist: + # Non-distributed: unshard returns self.data directly, no comm + assert unsharded is wbuf.data + assert work is None + else: + # Distributed: after all-gather, every param should be fully + # recoverable from the unsharded buffer at its global offset + for i, p in enumerate(orig): + off, sz = wbuf.buffer_index._get_item_global_range(i) + recovered = unsharded[off : off + sz] + assert torch.equal(recovered, p.flatten()) + + # Reshard: release temporary buffer, persistent shard must be intact + wbuf.reshard() + if w_dist: + assert wbuf._unsharded_buffer is None + torch.set_printoptions( + threshold=float("inf"), # never switch to summarised view + edgeitems=None, # show all edge items + linewidth=200 # optional, for wide tensors + ) + assert torch.equal(wbuf.data, shard_before), ( + f"wbuf.data={wbuf.data}\n" + f"shard_before={shard_before}\n" + ) + + torch.distributed.barrier() + + +# ------------------------------------------------------------------ # +# Part 3: reduce_grad — verify all-reduce / reduce-scatter +# ------------------------------------------------------------------ # + + +@pytest.mark.parametrize("strategy", ["no_shard", "optim", "optim_grads", "optim_grads_params"]) +def test_reduce_grad(strategy): + if strategy != "optim_grads_params": + pytest.skip( + f"Only optim_grads_params strategy creates distributed buffers, skipping {strategy} test." + ) + + groups, _, dp_group, rank, ws, device = _build_groups(strategy) + _, _, _, g_dist = _flags(strategy) + + for pg in groups: + gbuf = pg.main_grad_buffer + if gbuf is None: + # uint8 group has requires_grad=False, so no grad buffer + continue + + if not g_dist: + # Non-distributed: each rank fills with (rank+1), then all-reduce + # should produce sum across all ranks + gbuf.data.fill_(float(rank + 1)) + ref = torch.full_like(gbuf.data, float(rank + 1)) + Ref.all_reduce(ref, dp_group) + gbuf.reduce_grad() + assert torch.equal(gbuf.data, ref) + else: + # Distributed: fill a full-size temp buffer with (rank+1), + # reduce-scatter it, compare against PyTorch ref + full_size = gbuf.buffer_index.bucket_meta.size + full = torch.full((full_size,), float(rank + 1), dtype=gbuf.dtype, device=device) + + # Pre-populate the allocator so reduce_grad sees the data + bucket = gbuf.allocator.allocate( + key=gbuf.alloc_key, size=full_size, dtype=gbuf.dtype, device=device + ) + bucket.data.copy_(full) + + ref_shard = Ref.reduce_scatter(full.clone(), dp_group) + + gbuf.data.zero_() + gbuf.reduce_grad() + + # Only compare the shard region of self.data + sm = gbuf.buffer_index.shard_meta + actual = gbuf.data[sm.local_data_index : sm.local_data_index + sm.size] + assert torch.equal(actual, ref_shard) + + torch.distributed.barrier() From 646ea0ad8cfe69c1dfd3013ba08774a579bb4e70 Mon Sep 17 00:00:00 2001 From: jianbinc Date: Fri, 29 May 2026 09:20:34 +0800 Subject: [PATCH 03/28] docs: polish v2 README, add TP design doc, activation checkpointing support - 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) --- examples/megatron_fsdp/fsdp_toy.py | 24 +- .../fsdp/src/megatron_fsdp/v2/README.md | 166 +++++++++--- .../fsdp/src/megatron_fsdp/v2/fsdp_module.py | 3 + .../src/megatron_fsdp/v2/tp_support_design.md | 244 ++++++++++++++++++ .../megatron_fsdp/v2/test_fully_shard.py | 208 ++++++++++++++- 5 files changed, 612 insertions(+), 33 deletions(-) create mode 100644 megatron/core/distributed/fsdp/src/megatron_fsdp/v2/tp_support_design.md diff --git a/examples/megatron_fsdp/fsdp_toy.py b/examples/megatron_fsdp/fsdp_toy.py index 6e7299f4411..72c80b39337 100644 --- a/examples/megatron_fsdp/fsdp_toy.py +++ b/examples/megatron_fsdp/fsdp_toy.py @@ -36,8 +36,14 @@ def __init__(self, dim: int): super().__init__() self.linear1 = nn.Linear(dim, dim) self.linear2 = nn.Linear(dim, dim) + self._use_activation_checkpointing = False def forward(self, x): + if self._use_activation_checkpointing: + return torch.utils.checkpoint.checkpoint(self._forward_impl, x, use_reentrant=False) + return self._forward_impl(x) + + def _forward_impl(self, x): return self.linear2(torch.relu(self.linear1(x))) @@ -52,6 +58,10 @@ def forward(self, x): x = layer(x) return self.out(x) + def enable_activation_checkpointing(self): + for layer in self.layers: + layer._use_activation_checkpointing = True + # ----------------------- # Distributed init / mesh @@ -72,10 +82,15 @@ def build_fsdp_model( dim: int, n_layers: int, use_megatron_fsdp: bool, + use_activation_checkpointing: bool = False, ) -> Tuple["FSDPModule", torch.distributed.device_mesh.DeviceMesh]: if use_megatron_fsdp: - from megatron.core.distributed.fsdp.src.megatron_fsdp.uneven_dtensor import get_state_dict - from megatron.core.distributed.fsdp.src.megatron_fsdp.v2 import FSDPModule, fully_shard + try: + from megatron_fsdp.uneven_dtensor import get_state_dict + from megatron_fsdp.v2 import FSDPModule, fully_shard + except ImportError: + from megatron.core.distributed.fsdp.src.megatron_fsdp.uneven_dtensor import get_state_dict + from megatron.core.distributed.fsdp.src.megatron_fsdp.v2 import FSDPModule, fully_shard sys.modules["get_state_dict"] = get_state_dict else: from torch.distributed.fsdp import FSDPModule, fully_shard @@ -83,6 +98,9 @@ def build_fsdp_model( mesh = init_distributed() model = ToyModel(dim=dim, n_layers=n_layers).to("cuda") + if use_activation_checkpointing: + model.enable_activation_checkpointing() + # Example: per-layer sharding for layer in model.layers: fully_shard(layer, mesh=mesh) @@ -233,6 +251,7 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--ckpt-interval", type=int, default=20) parser.add_argument("--log-interval", type=int, default=5) parser.add_argument("--use-megatron-fsdp", action="store_true", help="Use Megatron-FSDP instead of PyTorch FSDP2") + parser.add_argument("--activation-checkpoint", action="store_true", help="Enable activation checkpointing on transformer layers") return parser.parse_args() @@ -243,6 +262,7 @@ def main() -> None: dim=args.model_dim, n_layers=args.n_layers, use_megatron_fsdp=args.use_megatron_fsdp, + use_activation_checkpointing=args.activation_checkpoint, ) optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr) diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/README.md b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/README.md index f3a9e90e05d..49b7663e832 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/README.md +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/README.md @@ -1,21 +1,25 @@ -# Megatron FSDP2 (fully_shard_v2) +# Megatron FSDP v2 -This directory contains the **fully_shard_v2** implementation — a PyTorch FSDP2-compatible API layer for Megatron Core. +This directory contains the **Megatron FSDP v2** implementation — a PyTorch FSDP2-compatible API layer for Megatron Core. ## Architecture ``` v2/ ├── README.md # This file -├── __init__.py # Public exports -├── fully_shard.py # Public fully_shard() API -├── fsdp_module.py # FSDPModule runtime state and methods +├── __init__.py # Public exports (FSDPModule, fully_shard, mixed precision policies) +├── fully_shard.py # Public fully_shard() API entry point +├── fsdp_module.py # FSDPModule runtime state (unshard/reshard/reduce_grad) ├── hooks.py # Forward/backward hook registration ├── param_group.py # ParameterGroup — groups params with shared buffers ├── dp_buffer.py # DataParallelBuffer — flat buffer management -├── allocator.py # BucketAllocator implementations -├── utils.py # Internal utility functions -└── design.md # Detailed design documentation (overlap, memory, sync) +├── allocator.py # BucketAllocator (Temporary, StorageFreeing, TracePool) +├── mixed_precision.py # MixedPrecisionPolicy, FP8Policy, NVFP4Policy +├── utils.py # Internal utilities (mesh init, backward Function) +├── design.md # Overlap, memory, and synchronization design +├── nvfp4_design.md # NVFP4 primary-weights design +├── mcore_fsdp_checkpoint_design.md # Checkpoint save/load and format conversion design +└── tp_support_design.md # Tensor Parallelism support plan (future) ``` For the broader `megatron_fsdp` package layout, see the parent directory. @@ -32,6 +36,23 @@ Wraps a module with FSDP sharding semantics: 4. Registers forward/backward hooks for unshard/reshard/reduce 5. Replaces module parameters with `DTensor` representations +### FullyShardMixedPrecisionPolicy + +Controls parameter/gradient dtypes and communication precision: + +| Policy | `main_params_dtype` | `main_grads_dtype` | Notes | +|--------|---------------------|---------------------|-------| +| `FullyShardMixedPrecisionPolicy` | `None` or `fp32` | `None` | Base policy (bf16 param, fp32 reduce) | +| `FullyShardFP8Policy` | varies | varies | MXFP8 mixed precision | +| `FullyShardNVFP4Policy` | `fp4` | `fp32` | NVFP4 primary weights | + +```python +from megatron_fsdp.v2 import fully_shard, FullyShardFP8Policy + +mp_policy = FullyShardFP8Policy() +fully_shard(model, mp_policy=mp_policy) +``` + ### FSDPModule Mixin class added to wrapped modules. Methods: @@ -79,45 +100,114 @@ See the parent directory `..` for `uneven_dtensor.py` which provides: | `optim` | No | No | **Not yet supported** | Like ZeRO-1: shard optimizer states only | | `optim_grads` | No | Yes | **Not yet supported** | Like ZeRO-2: shard optimizer states + gradients | -> **FIXME:** `no_shard`, `optim`, and `optim_grads` sharding strategies are not yet supported in v2. -> Currently only `optim_grads_params` is fully implemented and tested. -> These strategies will be added in a follow-up change. +## Known Limitations + +### Parallelism + +- **Tensor Parallelism (TP):** Not supported. v2 currently operates on a 1D + DP-only DeviceMesh. Parameters that are already partitioned by TP layers + (e.g., `ColumnParallelLinear`, `RowParallelLinear`) are not correctly handled. + See [tp_support_design.md](tp_support_design.md) for the planned design. +- **Hybrid Sharding (HSDP):** Not supported. v2 does not yet support an outer + DP dimension for hybrid (inter-node + intra-node) sharding. +- **Context Parallelism (CP):** Not supported. +- **Expert Parallelism (EP):** Not supported in v2's param group logic. MoE + expert parameters are not currently handled via EP submeshes. + +### Sharding Strategies + +Only `optim_grads_params` (ZeRO-3 equivalent) is implemented. `no_shard` +(DDP-like), `optim` (ZeRO-1), and `optim_grads` (ZeRO-2) are planned but +not yet available. + +### `fully_shard()` API Parameters + +The following parameters are accepted in the function signature but are **not +yet implemented** (marked `TODO`): + +- `reshard_after_forward` — no-op +- `shard_placement_fn` — no-op; all params use `Shard(0)` on the DP dimension +- `offload_policy` — no-op; CPU offloading is not supported + +### Hardware & Platform + +- **GPU only.** CUDA devices only. CPU, XPU, and ROCm are not tested or supported. +- **NVFP4** (`mixed_precision.py`): The non-distributed quantization path is + not implemented (`FIXME` at `mixed_precision.py:686`). Distributed NVFP4 + (reduce-scatter + quantize) is functional. + +### Checkpointing + +- **Async checkpoint:** Not supported for the v2 path. +- **`dp_reshardable` checkpoints:** Loading from `dp_reshardable` format is + not supported. Re-save checkpoints with `--ckpt-fully-parallel-save` first. ## Integration with Megatron -`FullyShardedDataParallel` in `mcore_fsdp_adapter.py` supports two code paths: +There are two ways to use Megatron FSDP v2: + +### Option A: Through Megatron Core (MCore) + +Set `--use-megatron-fsdp-v2` in your training arguments. The adapter +(`mcore_fsdp_adapter.py`) will automatically route to the v2 `fully_shard` +path for model sharding. -1. **Legacy path** (`use_fully_shard_api=False`): Uses `MegatronFSDP` from the original megatron_fsdp module -2. **fully_shard path** (`use_fully_shard_api=True`): Uses PyTorch or Megatron's `fully_shard_v2` +```bash +python pretrain_gpt.py \ + --use-megatron-fsdp-v2 \ + --use-megatron-fsdp \ + ... +``` -### Using fully_shard_v2 +This is the recommended path for Megatron-LM training workflows. + +### Option B: Standalone (without Megatron-LM) + +Import `fully_shard` directly from the `megatron_fsdp.v2` package: ```python -from megatron.core.distributed.fsdp.mcore_fsdp_adapter import FullyShardedDataParallel - -# In your config: -ddp_config.data_parallel_sharding_strategy = "optim_grads_params" -ddp_config.use_fully_shard_api = True - -# During model setup: -model = FullyShardedDataParallel( - config=transformer_config, - ddp_config=ddp_config, - module=model, -) +from megatron_fsdp.v2 import FSDPModule, fully_shard + +model = MyModel().cuda() +fully_shard(model) ``` +### Installation + +**Pre-release (current):** v2 has not been released as a standalone package yet. +Clone and install via `PYTHONPATH`: + +```bash +git clone -b mfsdp_refactor https://github.com/shjwudp/Megatron-LM.git +export PYTHONPATH=$PWD/Megatron-LM/megatron/core/distributed/fsdp/src:$PYTHONPATH +``` + +**After release** (once `megatron-fsdp` is published to PyPI): + +```bash +pip install megatron-fsdp +``` + +The standalone import (`from megatron_fsdp.v2 import fully_shard`) will work +without any Megatron-LM dependency once the package is released. + ## Toy Example See `examples/megatron_fsdp/fsdp_toy.py` for a standalone example showing: - Basic model wrapping with `fully_shard()` - Training loop with gradient accumulation +- Activation checkpointing (`--activation-checkpoint`) - Distributed checkpointing with `torch.distributed.checkpoint` ```bash torchrun --nproc_per_node=2 examples/megatron_fsdp/fsdp_toy.py \ --model-dim 512 --n-layers 2 --batch-size 4 + +# With activation checkpointing and Megatron-FSDP +torchrun --nproc_per_node=2 examples/megatron_fsdp/fsdp_toy.py \ + --model-dim 512 --n-layers 2 --batch-size 4 \ + --use-megatron-fsdp --activation-checkpoint ``` ## Gotchas / Pitfalls @@ -128,6 +218,24 @@ torchrun --nproc_per_node=2 examples/megatron_fsdp/fsdp_toy.py \ ## Unit Tests ```bash -# Run FSDP2 API tests -pytest -xvs tests/unit_tests/distributed/megatron_fsdp/test_mcore_fsdp_fully_shard_v2_api.py +# Run all v2 unit tests (requires 2 GPUs) +torchrun --nproc_per_node=2 -m pytest \ + tests/unit_tests/distributed/megatron_fsdp/v2/ -v -x + +# Run specific test files +torchrun --nproc_per_node=2 -m pytest \ + tests/unit_tests/distributed/megatron_fsdp/v2/test_fully_shard.py -v + +torchrun --nproc_per_node=2 -m pytest \ + tests/unit_tests/distributed/megatron_fsdp/v2/test_mcore_checkpoint.py -v + +# Single-GPU tests +pytest tests/unit_tests/distributed/megatron_fsdp/v2/ -v \ + -k "test_double_shard_rejected or test_no_params_module or test_get_state_dict_strict" ``` + +## Acknowledgements + +This work was designed and implemented by **Jianbin Chang**\*, **Tong Liu**\*, **Cory Ye** and **Jingyue Wu**. + +**\*** Primary code developers. diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/fsdp_module.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/fsdp_module.py index ad831a56dff..3901afe44aa 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/fsdp_module.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/fsdp_module.py @@ -297,6 +297,9 @@ def _materialize_meta_module( f"Module {name} contains meta parameters but cannot reset them" ) + # Move materialized parameters to the same target device (e.g., GPU) + m.to(materialization_device) + if mesh is not None and mesh.size() > 1: dp_group = mesh.get_group() src_rank = torch.distributed.get_global_rank(dp_group, 0) diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/tp_support_design.md b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/tp_support_design.md new file mode 100644 index 00000000000..5f5e828e610 --- /dev/null +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/tp_support_design.md @@ -0,0 +1,244 @@ +# Tensor Parallelism Support in Megatron FSDP v2 — Design Document + +## 1. Overview + +This document proposes adding Tensor Parallelism (TP) support to Megatron FSDP v2. +Currently v2 operates on a 1D DP-only DeviceMesh and does not handle parameters +that are already sharded by TP layers (e.g., `ColumnParallelLinear`, +`RowParallelLinear`). This document describes how to compose TP and FSDP sharding +in v2, referencing the proven patterns from v1. + +## 2. Background + +### 2.1 How v1 Handles TP + +Megatron FSDP v1 composes TP and FSDP using a **two-step DTensor construction** +in `make_fsdp_dtensor()` (`param_and_grad_buffer.py:4553–4712`): + +**Step 1 — TP-sharded DTensor.** If a parameter has a `_tensor_parallel_mode` +attribute (`"column"` or `"row"`), a TP-only DTensor is constructed first: + +```python +tp_mesh = dist_index.get_submesh("tp") +placements = [Shard(partition_dim)] # dim=0 for column, dim=1 for row +global_shape[partition_dim] *= tp_mesh.mesh.numel() +param = DTensor.from_local(local_tensor, tp_mesh, placements, shape=global_shape, ...) +``` + +**Step 2 — FSDP-sharded DTensor.** The TP DTensor is then wrapped in an +additional FSDP placement, composing a multi-dimensional DTensor: + +```python +device_mesh = dist_index.get_submesh(("dp_cp", "tp")) # 2D mesh +placements = [Shard(0), Shard(tp_dim)] # DP shard + TP shard +param = DTensor.from_local(param.to_local(), device_mesh, placements, ...) +``` + +The key insight: **FSDP v1 communicates gradients ONLY on the DP group**. +TP communication (all-reduce, all-gather, reduce-scatter) is handled +independently by the TP layers themselves (`tensor_parallel/layers.py`). +The 2D DTensor representation is purely for sharding semantics — the +communication operations are decoupled. + +### 2.2 Current v2 State + +v2 currently has no TP awareness: + +| Component | v2 Status | +|-----------|-----------| +| DeviceMesh | 1D DP only (`v2/utils.py:67–82`) | +| Placements | `[Shard(0)]` or `[Replicate()]` only (`v2/param_group.py:306`) | +| `fully_shard()` API | No TP mesh parameter; `shard_placement_fn` is TODO | +| Adapter mesh | `_init_dp_mesh()` creates DP-only mesh (`mcore_fsdp_adapter.py:688`) | +| TP annotation | Not consumed; `_annotate_tensor_parallelism` exists in adapter | +| Gradient ops | Reduce-scatter on DP group only (correct, no change needed) | + +### 2.3 Why TP + FSDP Composition Matters + +When TP is active, each TP rank holds a shard of the weight. For example, with +`tensor_model_parallel_size=4`: + +- `ColumnParallelLinear.weight`: shape `[out_per_partition, in]` — sharded along dim 0 +- `RowParallelLinear.weight`: shape `[out, in_per_partition]` — sharded along dim 1 + +v2 currently wraps these **partial** tensors with `[Shard(0)]` on a DP mesh, +which misrepresents the true global shape: the DTensor reports the TP-partial +shape as "global," and checkpoint save/load produces incorrect shard metadata. + +Correct behavior: the DTensor should carry a 2D `(dp, tp)` mesh with placements +`[Shard(0), Shard(tp_dim)]` and the full global shape. + +## 3. Proposed Design + +### 3.1 High-Level Approach + +Follow v1's two-step DTensor construction, adapted for v2's `ParameterGroup` and +`DataParallelBuffer` architecture: + +1. **Annotate:** Reuse `_annotate_tensor_parallelism` from v1 to mark params +2. **Construct 2D mesh:** Build `(dp, tp)` DeviceMesh in the adapter or directly + in `fully_shard()` +3. **Two-step DTensor:** In `ParameterGroup._init_buffers()`, wrap TP-sharded + params first, then apply FSDP shard +4. **No gradient changes:** FSDP reduce-scatter remains DP-only (correct) +5. **Checkpoint metadata:** Leverage existing uneven DTensor infrastructure + +### 3.2 Phase 1: Standalone `fully_shard()` with TP Mesh (Minimal Viable) + +Add a `tp_mesh` parameter to `fully_shard()`: + +```python +def fully_shard( + module, + *, + mesh: Optional[DeviceMesh] = None, + tp_mesh: Optional[DeviceMesh] = None, # NEW + ... +) -> nn.Module: +``` + +When `tp_mesh` is provided: +- **In `_init_named_param_groups()`:** For each parameter, if it belongs to a + TP-sharded module (detected via `_tensor_parallel_mode` or module type), + construct placements as `[Shard(0), Shard(tp_dim)]` on a combined 2D mesh +- **In `dp_buffer.py`:** `BufferIndex.shard_meta` must account for the TP-local + size (the "logical" size visible to this rank after TP sharding) vs. the + "storage" size (the FSDP shard size). The `unshard()` all-gather returns the + TP-local full tensor, not the true global tensor — this is correct because + TP layers operate on their local TP shard during forward/backward. + +**TP-local vs. global semantics:** + +``` +True global shape: [out, in] +TP-local shape: [out/tp, in] (column-parallel) +FSDP shard shape: [out/tp/dp, in] (sharded across DP ranks) + +DTensor shape: [out, in] ← global shape annotated +DTensor mesh: (dp, tp) +DTensor placements: [Shard(0), Shard(0)] ← DP shard then TP shard +``` + +The `DTensor.to_local()` returns the TP-local+FSDP-local shard — exactly the +data that PyTorch's autograd operates on during forward/backward. + +### 3.3 Phase 2: Integration with MCore Adapter + +In `mcore_fsdp_adapter.py`, extend `_init_with_fully_shard()` to: + +1. Call `_annotate_tensor_parallelism()` (already exists at line 434–502) +2. Build a 2D `(dp, tp)` DeviceMesh via `_get_dp_tp_mesh()` (already exists at + line 783) instead of the current `_init_dp_mesh()` +3. Pass the 2D mesh + tp_mesh to `fully_shard()` +4. The adapter's `_init_distributed_params()` path for v2 should use the + same two-step DTensor construction + +### 3.4 Phase 3: Bucketing with TP + +v1's `BucketingPolicy` does NOT consider TP when computing buffer sizes — it +only uses DP shard sizes. v2 should do the same: buffers are sized for the +TP-local parameter shape (after TP sharding), not the global shape. This is +correct because: + +- `unshard()` all-gathers to the TP-local full tensor (not global) +- `reduce_grad()` reduce-scatters from the TP-local full gradient (not global) +- TP layers independently handle their own communication + +### 3.5 Gradient Communication (No Change) + +v2's gradient reduce-scatter already operates on the DP group only — this is +correct and should not change for TP support. TP parameters participate in +TP-level all-reduce/reduce-scatter within the TP layer's forward/backward, and +in DP-level reduce-scatter within v2's `reduce_grad()`. The two are independent +and naturally compose. + +### 3.6 Checkpointing with TP + +The uneven DTensor infrastructure (`uneven_dtensor.py`) already supports +multi-dimensional shard metadata via `_shard_order`. v2's checkpoint path +(through `get_state_dict()` → `preprocess_state_dict_for_uneven_dtensor()`) +should work correctly once DTensors carry the 2D mesh, because: + +- Each DTensor knows its own shard placement and global shape +- DCP (`torch.distributed.checkpoint`) natively handles multi-dimensional + DTensor serialization +- Chunk metadata computation (`gather_and_compute_chunk_metadata()`) supports + arbitrary `_shard_order` values + +## 4. Implementation Checklist + +### Phase 1: Standalone TP (3–5 days) + +- [ ] Add `tp_mesh` parameter to `fully_shard()` signature +- [ ] In `ParameterGroup._init_buffers()`: detect TP params and construct 2D mesh + placements +- [ ] In `ParameterGroup`: add `_tp_partition_dim()` helper (delegate to + `get_mcore_tensor_parallel_partition_dim`) +- [ ] In `BufferIndex`: ensure `shard_meta` correctly reflects TP-local size +- [ ] Unit test: `test_tp_column_parallel()` — ColumnParallelLinear + fully_shard + forward/backward +- [ ] Unit test: `test_tp_row_parallel()` — RowParallelLinear + fully_shard + forward/backward +- [ ] Unit test: `test_tp_fsdp_loss_identity()` — compare FSDP+TP loss vs DDP baseline + +### Phase 2: MCore Adapter Integration (2–3 days) + +- [ ] Wire `_annotate_tensor_parallelism()` in adapter v2 path +- [ ] Build 2D `(dp, tp)` mesh in `_init_with_fully_shard()` +- [ ] Pass TP mesh and annotations through `fully_shard()` calls +- [ ] Integration test: `gpt3_mcore_tp2_pp1` config with `--use-megatron-fsdp-v2` + +### Phase 3: Checkpoint + Polish (2–3 days) + +- [ ] Verify `get_state_dict()` returns correct global shapes for TP+FSDP params +- [ ] Test checkpoint save/load round-trip with TP active +- [ ] Test online format conversion: torch_dist → fsdp_dtensor with TP +- [ ] Handle `force_sync_tp_duplicated_param` for replicated TP params (e.g., + RowParallel bias, LayerNorm) — broadcast from tp_rank=0 +- [ ] Update README with TP usage example + +## 5. Replicated TP Parameters + +Some parameters in TP modules are not sharded (e.g., `RowParallelLinear.bias`, +`LayerNorm.weight`). These have `_tensor_parallel_mode = "replicated"` in v1's +annotation system. For these params: + +- **DTensor type:** `Replicate()` on the TP dimension +- **Synchronization:** v1 broadcasts from `tp_rank=0` during init to ensure + consistency. v2 should do the same via `force_sync_tp_duplicated_param` +- **FSDP handling:** Normal FSDP shard — these params have the full unpartitioned + shape, so v2's default logic applies correctly + +Combined placements for replicated TP + FSDP: `[Shard(0), Replicate()]` + +## 6. Risks and Open Questions + +1. **`DTensor.full_tensor()` behavior with 2D mesh.** Calling `.full_tensor()` + on a `(dp, tp)` DTensor all-gathers across BOTH dimensions. v2's `unshard()` + should only all-gather across the DP dimension — never across TP. Solution: + use `DTensor.redistribute(placements=[Replicate(), Shard(tp_dim)])` instead + of `.full_tensor()`, or construct a 1D DP submesh for the all-gather. + +2. **Stream ordering with TP.** TP layers perform their own communication on the + default stream. v2's `ag_stream` / `rs_stream` operations must be correctly + ordered relative to TP comms. v1 handles this with `stream.wait_stream()` + barriers; v2's `_FSDPRootContext` already has this mechanism. + +3. **HSDP + TP (3D mesh).** v1 supports `(outer_dp, inner_dp, tp)` 3D meshes. + v2 should defer HSDP+TP support to a follow-up; the design is extensible. + +4. **Expert parallelism with TP.** MoE experts may have their own EP mesh + dimension. v1 handles this via `is_expert_parallel` parameter in + `get_submesh()`. v2 should similarly extend its param group logic. + +## 7. v1 Code References + +| v1 Component | File | Lines | +|-------------|------|-------| +| `make_fsdp_dtensor()` | `param_and_grad_buffer.py` | 4553–4712 | +| `_get_fsdp_tensor_spec()` | `param_and_grad_buffer.py` | 4477–4550 | +| `FSDPDistributedIndex` | `utils.py` | 475–673 | +| TP helper functions | `utils.py` | 819–850 | +| `_annotate_tensor_parallelism()` | `mcore_fsdp_adapter.py` | 434–502 | +| `_get_dp_tp_mesh()` | `mcore_fsdp_adapter.py` | 783 | +| 2D/3D mesh construction | `mcore_fsdp_adapter.py` | 504–633 | +| Gradient communication (DP-only) | `param_and_grad_buffer.py` | 3168–3268 | +| `_init_optimizer_named_parameters` (TP attr propagation) | `param_and_grad_buffer.py` | 2872–2926 | +| Chunk metadata with `_shard_order` | `uneven_dtensor.py` | 32–97 | diff --git a/tests/unit_tests/distributed/megatron_fsdp/v2/test_fully_shard.py b/tests/unit_tests/distributed/megatron_fsdp/v2/test_fully_shard.py index 138bd9aa4fc..94a0897115c 100644 --- a/tests/unit_tests/distributed/megatron_fsdp/v2/test_fully_shard.py +++ b/tests/unit_tests/distributed/megatron_fsdp/v2/test_fully_shard.py @@ -626,7 +626,211 @@ def test_params_resharded_after_backward(self): # ------------------------------------------------------------------ # -# 8. Safety — double-shard rejection +# 8. Activation checkpointing +# ------------------------------------------------------------------ # + + +class MLPWithCheckpointing(nn.Module): + """A multi-layer MLP that supports activation checkpointing on its blocks.""" + + def __init__(self, hidden=64, num_layers=3): + super().__init__() + self.layers = nn.ModuleList( + [nn.Sequential(nn.Linear(hidden, hidden), nn.ReLU(), nn.Linear(hidden, hidden)) + for _ in range(num_layers)] + ) + self._use_activation_checkpointing = False + + def forward(self, x): + for layer in self.layers: + if self._use_activation_checkpointing: + x = torch.utils.checkpoint.checkpoint(layer, x, use_reentrant=False) + else: + x = layer(x) + return x + + def enable_activation_checkpointing(self): + self._use_activation_checkpointing = True + + +class LargePerLayerModel(nn.Module): + """Multi-layer model with individually wrapped FSDP layers and optional + activation checkpointing support.""" + + def __init__(self, hidden=256, num_layers=4): + super().__init__() + self.layers = nn.ModuleList([ + nn.Sequential(nn.Linear(hidden, hidden), nn.ReLU(), nn.Linear(hidden, hidden)) + for _ in range(num_layers) + ]) + + def forward(self, x): + for layer in self.layers: + x = layer(x) + return x + + +class TestActivationCheckpointing: + def test_activation_checkpointing_forward_backward(self): + """Forward + backward with activation checkpointing should produce finite loss.""" + torch.manual_seed(42) + device = _device() + model = MLPWithCheckpointing(hidden=64, num_layers=4).to(device) + model.enable_activation_checkpointing() + + for layer in model.layers: + fully_shard(layer) + fully_shard(model) + _assert_dtensor_params(model) + + x = torch.randn(2, 64, device=device, requires_grad=True) + out = model(x) + loss = out.sum() + loss.backward() + + assert not torch.isnan(torch.tensor(loss.item())), "Loss is NaN" + assert not torch.isinf(torch.tensor(loss.item())), "Loss is Inf" + + def test_activation_checkpointing_multi_step(self): + """Multiple forward+backward steps with activation checkpointing should be stable.""" + torch.manual_seed(42) + device = _device() + model = MLPWithCheckpointing(hidden=64, num_layers=4).to(device) + model.enable_activation_checkpointing() + + for layer in model.layers: + fully_shard(layer) + fully_shard(model) + + losses = [] + for step in range(4): + torch.manual_seed(step) + x = torch.randn(2, 64, device=device, requires_grad=True) + out = model(x) + loss = out.sum() + loss.backward() + losses.append(loss.item()) + + for i, loss_val in enumerate(losses): + assert not torch.isnan(torch.tensor(loss_val)), f"Loss at step {i} is NaN" + assert not torch.isinf(torch.tensor(loss_val)), f"Loss at step {i} is Inf" + + def test_activation_checkpointing_with_overlap(self): + """Activation checkpointing should work with unshard_prefetch and async_reduce_grad.""" + torch.manual_seed(42) + device = _device() + model = MLPWithCheckpointing(hidden=128, num_layers=4).to(device) + model.enable_activation_checkpointing() + + for layer in model.layers: + fully_shard( + layer, + enable_unshard_prefetch=True, + enable_async_reduce_grad=True, + ) + fully_shard(model, enable_unshard_prefetch=True, enable_async_reduce_grad=True) + + x = torch.randn(2, 128, device=device, requires_grad=True) + out = model(x) + loss = out.sum() + loss.backward() + + assert not torch.isnan(torch.tensor(loss.item())) + + def test_activation_checkpointing_nested_fsdp(self): + """Activation checkpointing with nested FSDP (expert-in-layer) should work.""" + torch.manual_seed(42) + device = _device() + + class NestedCheckpointModel(nn.Module): + def __init__(self, hidden=64): + super().__init__() + self.attn = nn.Linear(hidden, hidden) + self.experts = nn.Sequential( + nn.Linear(hidden, hidden), nn.ReLU(), nn.Linear(hidden, hidden) + ) + self.norm = nn.LayerNorm(hidden) + + def forward(self, x): + h = self.attn(x) + if self._use_activation_checkpointing: + h = torch.utils.checkpoint.checkpoint(self.experts, h, use_reentrant=False) + else: + h = self.experts(h) + return self.norm(h + x) + + model = NestedCheckpointModel(hidden=64).to(device) + model._use_activation_checkpointing = True + model.experts = fully_shard(model.experts) + model = fully_shard(model) + + x = torch.randn(2, 64, device=device, requires_grad=True) + out = model(x) + loss = out.sum() + loss.backward() + + assert not torch.isnan(torch.tensor(loss.item())) + + def test_activation_checkpointing_disabled_vs_enabled_same_loss(self): + """With same inputs and no parameter updates, checkpointed and non-checkpointed + forward should produce the same output (checkpoint recompute is numerically transparent).""" + torch.manual_seed(42) + device = _device() + + model = MLPWithCheckpointing(hidden=64, num_layers=3).to(device) + + for layer in model.layers: + fully_shard(layer) + fully_shard(model) + + x = torch.randn(2, 64, device=device, requires_grad=True) + + # Forward without activation checkpointing + model._use_activation_checkpointing = False + out_no_ckpt = model(x) + + # Forward with activation checkpointing + torch.manual_seed(42) + x2 = torch.randn(2, 64, device=device, requires_grad=True) + model._use_activation_checkpointing = True + out_ckpt = model(x2) + + assert torch.allclose( + out_no_ckpt, out_ckpt, atol=1e-5 + ), "Checkpointing changed forward output (same inputs)" + + def test_activation_checkpointing_per_layer_shard_with_ckpt(self): + """Per-layer FSDP with activation checkpointing on each layer — full training step.""" + torch.manual_seed(42) + device = _device() + model = LargePerLayerModel(hidden=256, num_layers=6).to(device) + optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3) + + for layer in model.layers: + fully_shard(layer) + + fully_shard(model) + + # Use checkpoint on every other layer to test mixed use + def ckpt_forward(x): + for i, layer in enumerate(model.layers): + if i % 2 == 0: + x = torch.utils.checkpoint.checkpoint(layer, x, use_reentrant=False) + else: + x = layer(x) + return x + + x = torch.randn(4, 256, device=device, requires_grad=True) + out = ckpt_forward(x) + loss = out.sum() + loss.backward() + optimizer.step() + + assert not torch.isnan(torch.tensor(loss.item())) + + +# ------------------------------------------------------------------ # +# 9. Safety — double-shard rejection # ------------------------------------------------------------------ # @@ -648,7 +852,7 @@ def test_no_params_module_ok(self): # ------------------------------------------------------------------ # -# 9. Checkpoint — get_state_dict and preprocess_state_dict_for_uneven_dtensor +# 10. Checkpoint — get_state_dict and preprocess_state_dict_for_uneven_dtensor # ------------------------------------------------------------------ # From f50b7bbb74b5e924df4283946ece8bbaa6eb03c3 Mon Sep 17 00:00:00 2001 From: jianbinc Date: Fri, 29 May 2026 17:33:22 +0800 Subject: [PATCH 04/28] fix: checkpoint save in fsdp_toy, import naming, test cleanup - 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 --- examples/megatron_fsdp/fsdp_toy.py | 9 ++++++++- megatron/core/distributed/fsdp/checkpoint.py | 2 +- megatron/training/checkpointing.py | 3 ++- .../megatron_fsdp/v2/test_param_group.py | 20 ++++++------------- 4 files changed, 17 insertions(+), 17 deletions(-) diff --git a/examples/megatron_fsdp/fsdp_toy.py b/examples/megatron_fsdp/fsdp_toy.py index 72c80b39337..c519f224be8 100644 --- a/examples/megatron_fsdp/fsdp_toy.py +++ b/examples/megatron_fsdp/fsdp_toy.py @@ -17,6 +17,7 @@ import sys from pathlib import Path from typing import Tuple +import warnings import torch import torch.distributed as dist @@ -91,7 +92,6 @@ def build_fsdp_model( except ImportError: from megatron.core.distributed.fsdp.src.megatron_fsdp.uneven_dtensor import get_state_dict from megatron.core.distributed.fsdp.src.megatron_fsdp.v2 import FSDPModule, fully_shard - sys.modules["get_state_dict"] = get_state_dict else: from torch.distributed.fsdp import FSDPModule, fully_shard @@ -134,6 +134,13 @@ def __init__(self, model, optimizer=None): self.optimizer = optimizer def state_dict(self): + try: + from megatron_fsdp.uneven_dtensor import get_state_dict + except ImportError: + from megatron.core.distributed.fsdp.src.megatron_fsdp.uneven_dtensor import get_state_dict + except ImportError: + warnings.warn("Could not import get_state_dict from megatron_fsdp.uneven_dtensor.") + # this line automatically manages FSDP FQN's, as well as sets the default state dict type to FSDP.SHARDED_STATE_DICT model_state_dict, optimizer_state_dict = get_state_dict(self.model, self.optimizer) return { diff --git a/megatron/core/distributed/fsdp/checkpoint.py b/megatron/core/distributed/fsdp/checkpoint.py index 0c753a23a38..61d60220724 100644 --- a/megatron/core/distributed/fsdp/checkpoint.py +++ b/megatron/core/distributed/fsdp/checkpoint.py @@ -1177,7 +1177,7 @@ def _canonicalize_td_key(td_key, *, strip_model_prefix=True): return key -def load_torch_dist_into_fsdp_v2(args, checkpoint_name, v2_state_dict, strict=True): +def _load_torch_dist_into_megatron_fsdp_v2(args, checkpoint_name, v2_state_dict, strict=True): """Load a torch_dist checkpoint into a Megatron FSDP v2 skeleton via DCP. This is the entry point for online checkpoint conversion from diff --git a/megatron/training/checkpointing.py b/megatron/training/checkpointing.py index 1fd6e6710f1..feb592c7e79 100644 --- a/megatron/training/checkpointing.py +++ b/megatron/training/checkpointing.py @@ -53,7 +53,8 @@ preprocess_state_dict_for_uneven_dtensor, ) from megatron.core.distributed.fsdp.checkpoint import ( - _apply_mcore_postprocess + _apply_mcore_postprocess, + _load_torch_dist_into_megatron_fsdp_v2, ) from megatron.core.transformer.fsdp_dtensor_checkpoint import ( handle_experts_in_state_dict, diff --git a/tests/unit_tests/distributed/megatron_fsdp/v2/test_param_group.py b/tests/unit_tests/distributed/megatron_fsdp/v2/test_param_group.py index ca2e80e3acc..ead8dca61f1 100644 --- a/tests/unit_tests/distributed/megatron_fsdp/v2/test_param_group.py +++ b/tests/unit_tests/distributed/megatron_fsdp/v2/test_param_group.py @@ -63,11 +63,11 @@ class MixedDtypeLayer(nn.Module): def __init__(self): super().__init__() - self.linear1 = nn.Linear(16, 32, bias=False) # bf16, shape [512, 256] - self.linear2 = nn.Linear(32, 16, bias=True) # bf16, shape [256, 512] + bias [256] - self.norm = nn.LayerNorm(16) # bf16, weight [256] + bias [256] - self.quant_proj = nn.Linear(16, 8, bias=False) # uint8, shape [128, 256] - self.quant_gate = nn.Linear(16, 4, bias=False) # uint8, shape [64, 256] + self.linear1 = nn.Linear(16, 32, bias=False) # bf16, shape [16, 32] + self.linear2 = nn.Linear(32, 16, bias=True) # bf16, shape [32, 16] + bias [16] + self.norm = nn.LayerNorm(16) # bf16, weight [16] + bias [16] + self.quant_proj = nn.Linear(16, 8, bias=False) # uint8, shape [16, 8] + self.quant_gate = nn.Linear(16, 4, bias=False) # uint8, shape [16, 4] # ------------------------------------------------------------------ # @@ -246,15 +246,7 @@ def test_unshard_reshard(strategy): wbuf.reshard() if w_dist: assert wbuf._unsharded_buffer is None - torch.set_printoptions( - threshold=float("inf"), # never switch to summarised view - edgeitems=None, # show all edge items - linewidth=200 # optional, for wide tensors - ) - assert torch.equal(wbuf.data, shard_before), ( - f"wbuf.data={wbuf.data}\n" - f"shard_before={shard_before}\n" - ) + assert torch.equal(wbuf.data, shard_before) torch.distributed.barrier() From 8b4f6e89fbecd55e07e0721e963686017edd3e06 Mon Sep 17 00:00:00 2001 From: jianbinc Date: Fri, 29 May 2026 17:38:20 +0800 Subject: [PATCH 05/28] update README --- .../fsdp/src/megatron_fsdp/v2/README.md | 41 +++++++++---------- 1 file changed, 19 insertions(+), 22 deletions(-) diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/README.md b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/README.md index 49b7663e832..f60d7f5ee75 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/README.md +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/README.md @@ -1,6 +1,8 @@ # Megatron FSDP v2 -This directory contains the **Megatron FSDP v2** implementation — a PyTorch FSDP2-compatible API layer for Megatron Core. +> ⚠️ Prototype Implementation — Not Yet Production Ready + +Megatron FSDP v2 (M-FSDP2) inherits the performance of Megatron FSDP v1 and supports API drop-in replacement for FSDP2. ## Architecture @@ -36,20 +38,20 @@ Wraps a module with FSDP sharding semantics: 4. Registers forward/backward hooks for unshard/reshard/reduce 5. Replaces module parameters with `DTensor` representations -### FullyShardMixedPrecisionPolicy +### MixedPrecisionPolicy -Controls parameter/gradient dtypes and communication precision: +Controls parameter/gradient dtypes and communication precision. Available policies: -| Policy | `main_params_dtype` | `main_grads_dtype` | Notes | -|--------|---------------------|---------------------|-------| -| `FullyShardMixedPrecisionPolicy` | `None` or `fp32` | `None` | Base policy (bf16 param, fp32 reduce) | -| `FullyShardFP8Policy` | varies | varies | MXFP8 mixed precision | -| `FullyShardNVFP4Policy` | `fp4` | `fp32` | NVFP4 primary weights | +| Policy | Notes | +|--------|-------| +| `MixedPrecisionPolicy` | Base policy (bf16 param, fp32 main weight & main grad) | +| `FullyShardFP8Policy` | MXFP8 rowwise/colwise quantized weights | +| `FullyShardNVFP4Policy` | NVFP4 primary weights | ```python -from megatron_fsdp.v2 import fully_shard, FullyShardFP8Policy +from megatron_fsdp.v2 import fully_shard, MixedPrecisionPolicy -mp_policy = FullyShardFP8Policy() +mp_policy = MixedPrecisionPolicy(main_params_dtype=torch.float32) fully_shard(model, mp_policy=mp_policy) ``` @@ -61,7 +63,7 @@ Mixin class added to wrapped modules. Methods: |--------|-------------|---------| | `unshard()` | Pre-forward | All-gather params from sharded buffer | | `reshard()` | Post-forward, post-backward | Release unsharded buffer | -| `reduce_grad()` | Post-backward | All-reduce or reduce-scatter gradients | +| `reduce_grad()` | Post-backward / grad sync | Reduce-scatter gradients into optimizer-facing shards | ### DataParallelBuffer @@ -69,15 +71,15 @@ Flat buffer managing (a shard of) parameter/gradient data: - `unshard()` — all-gather to full tensor - `reshard()` — free temporary buffer -- `reduce_grad()` — all-reduce or reduce-scatter gradients +- `reduce_grad()` — reduce-scatter gradients into optimizer-facing shards - Uses `BufferIndex` to track parameter layout within the buffer ### ParameterGroup Groups parameters sharing the same (device, dtype, requires_grad): -- `model_weight_buffer` — stores sharded model weights -- `main_weight_buffer` — optional high-precision copy +- `model_weight_buffer` — stores compute weights; replicated for ZeRO-1/2 and sharded for ZeRO-3 +- `main_weight_buffer` — optional high-precision optimizer copy; sharded when optimizer state is sharded - `main_grad_buffer` — accumulates gradients before reduce - `dist_params` — DTensor views into the buffer @@ -110,7 +112,6 @@ See the parent directory `..` for `uneven_dtensor.py` which provides: See [tp_support_design.md](tp_support_design.md) for the planned design. - **Hybrid Sharding (HSDP):** Not supported. v2 does not yet support an outer DP dimension for hybrid (inter-node + intra-node) sharding. -- **Context Parallelism (CP):** Not supported. - **Expert Parallelism (EP):** Not supported in v2's param group logic. MoE expert parameters are not currently handled via EP submeshes. @@ -219,6 +220,7 @@ torchrun --nproc_per_node=2 examples/megatron_fsdp/fsdp_toy.py \ ```bash # Run all v2 unit tests (requires 2 GPUs) +TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD=1 \ torchrun --nproc_per_node=2 -m pytest \ tests/unit_tests/distributed/megatron_fsdp/v2/ -v -x @@ -226,16 +228,11 @@ torchrun --nproc_per_node=2 -m pytest \ torchrun --nproc_per_node=2 -m pytest \ tests/unit_tests/distributed/megatron_fsdp/v2/test_fully_shard.py -v +TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD=1 \ torchrun --nproc_per_node=2 -m pytest \ tests/unit_tests/distributed/megatron_fsdp/v2/test_mcore_checkpoint.py -v # Single-GPU tests pytest tests/unit_tests/distributed/megatron_fsdp/v2/ -v \ -k "test_double_shard_rejected or test_no_params_module or test_get_state_dict_strict" -``` - -## Acknowledgements - -This work was designed and implemented by **Jianbin Chang**\*, **Tong Liu**\*, **Cory Ye** and **Jingyue Wu**. - -**\*** Primary code developers. +``` \ No newline at end of file From 041dcc27d83f4ded00871c6046b6ac6242380ebb Mon Sep 17 00:00:00 2001 From: tongliu Date: Sun, 31 May 2026 22:17:04 +0800 Subject: [PATCH 06/28] feat(mfsdp-v2): ZeRO-1/2/3 sharding strategies, MixedPrecisionPolicy 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 --- .../distributed/fsdp/mcore_fsdp_adapter.py | 19 +- .../fsdp/src/megatron_fsdp/v2/README.md | 26 +- .../fsdp/src/megatron_fsdp/v2/__init__.py | 15 +- .../fsdp/src/megatron_fsdp/v2/allocator.py | 11 +- .../fsdp/src/megatron_fsdp/v2/design.md | 200 +++++++---- .../fsdp/src/megatron_fsdp/v2/dp_buffer.py | 139 +++++--- .../fsdp/src/megatron_fsdp/v2/fsdp_module.py | 90 +++-- .../fsdp/src/megatron_fsdp/v2/fully_shard.py | 15 +- .../fsdp/src/megatron_fsdp/v2/hooks.py | 12 +- .../src/megatron_fsdp/v2/mixed_precision.py | 88 ++++- .../fsdp/src/megatron_fsdp/v2/nvfp4_design.md | 10 +- .../fsdp/src/megatron_fsdp/v2/param_group.py | 46 +-- .../megatron_fsdp/v2/test_fully_shard.py | 10 +- .../v2/test_mcore_nd_parallel.py | 331 ++++++++++++++++++ .../megatron_fsdp/v2/test_param_group.py | 33 +- 15 files changed, 803 insertions(+), 242 deletions(-) diff --git a/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py b/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py index d78efb61c56..aac55076c72 100644 --- a/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py +++ b/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py @@ -52,11 +52,6 @@ MixedPrecisionPolicy, ) from megatron.core.distributed.fsdp.src.megatron_fsdp.v2 import FSDPModule - from megatron.core.distributed.fsdp.src.megatron_fsdp.v2.mixed_precision import ( - FullyShardFP8Policy, - FullyShardMixedPrecisionPolicy, - FullyShardNVFP4Policy, - ) HAVE_MEGATRON_FSDP = True except ImportError as import_megatron_fsdp_error: @@ -227,6 +222,11 @@ def _init_with_fully_shard( ): if ddp_config.use_megatron_fsdp: from megatron.core.distributed.fsdp.src.megatron_fsdp.v2 import fully_shard + from megatron.core.distributed.fsdp.src.megatron_fsdp.v2.mixed_precision import ( + FullyShardFP8Policy, + MixedPrecisionPolicy, + FullyShardNVFP4Policy, + ) else: from torch.distributed.fsdp import fully_shard @@ -242,7 +242,7 @@ def _init_with_fully_shard( edp_mesh = _init_dp_mesh(pg_collection, edp=True) dp_mesh = _init_dp_mesh(pg_collection, edp=False) - fully_shard_mp_policy = FullyShardMixedPrecisionPolicy( + fully_shard_mp_policy = MixedPrecisionPolicy( main_params_dtype=ddp_config.megatron_fsdp_main_params_dtype, main_grads_dtype=( torch.float32 @@ -269,6 +269,7 @@ def _init_with_fully_shard( "enable_unshard_prefetch": ddp_config.overlap_param_gather, "enable_async_reduce_grad": ddp_config.overlap_grad_reduce, "enable_trace_pool": ddp_config.fsdp_double_buffer, + "sharding_strategy": ddp_config.data_parallel_sharding_strategy, } if config.calculate_per_token_loss: gradient_scaling_factor = None @@ -362,15 +363,11 @@ def not_implemented_op(): self.start_param_sync = noop self.start_grad_sync = noop - def finish_grad_sync(force_all_reduce: Optional[bool] = False): - ctx = self.module._fsdp_root_context - torch.cuda.current_stream().wait_stream(ctx.rs_stream) - def synchronize_param_gather(): ctx = self.module._fsdp_root_context torch.cuda.current_stream().wait_stream(ctx.ag_stream) - self.finish_grad_sync = finish_grad_sync + self.finish_grad_sync = self.module.finish_grad_sync self.scale_gradients = self.module._scale_gradients self.zero_grad_buffer = self.module._zero_grad_buffer self.log_per_param_norms = self.module._log_per_param_norms diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/README.md b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/README.md index f60d7f5ee75..493a666af9c 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/README.md +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/README.md @@ -40,19 +40,41 @@ Wraps a module with FSDP sharding semantics: ### MixedPrecisionPolicy -Controls parameter/gradient dtypes and communication precision. Available policies: +Controls parameter/gradient dtypes and communication precision. | Policy | Notes | |--------|-------| -| `MixedPrecisionPolicy` | Base policy (bf16 param, fp32 main weight & main grad) | +| `MixedPrecisionPolicy` | Base policy; all fields default to ``None`` | | `FullyShardFP8Policy` | MXFP8 rowwise/colwise quantized weights | | `FullyShardNVFP4Policy` | NVFP4 primary weights | +**Key fields:** + +| Field | Default | Purpose | +|-------|---------|---------| +| `main_params_dtype` | ``None`` | Dtype for optimizer main-weight buffer. ``None`` = no separate buffer, optimizer mutates model weights directly. Set to ``torch.float32`` for quantized models (FP8/NVFP4) so the optimizer works on high-precision copies. | +| `main_grads_dtype` | ``None`` | Dtype for optimizer main-grad buffer. When ``None`` and ``use_decoupled_grad=False``, aligns with ``main_params_dtype``. Otherwise falls back to ``param.dtype``. | +| `grad_comm_dtype` | ``None`` | Dtype for gradient reduce-scatter communication. ``None`` = use ``main_grads_dtype``. | +| `use_decoupled_grad` | ``False`` | When ``False``, ``main_grads_dtype`` is inferred from ``main_params_dtype`` so the optimizer operates in a consistent precision context. | + ```python from megatron_fsdp.v2 import fully_shard, MixedPrecisionPolicy +# No separate main buffer — optimizer mutates model params directly +mp_policy = MixedPrecisionPolicy() +fully_shard(model, mp_policy=mp_policy) + +# fp32 optimizer precision for bf16 model mp_policy = MixedPrecisionPolicy(main_params_dtype=torch.float32) fully_shard(model, mp_policy=mp_policy) + +# FP8 mixed precision — fp32 main weights auto-created by adapter +from megatron_fsdp.v2 import FullyShardFP8Policy +mp_policy = MixedPrecisionPolicy( + main_params_dtype=torch.float32, + fp8=FullyShardFP8Policy(enabled=True) +) +fully_shard(model, mp_policy=mp_policy) ``` ### FSDPModule diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/__init__.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/__init__.py index 21879c4717a..0e78fc0e52c 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/__init__.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/__init__.py @@ -12,14 +12,9 @@ # See the License for the specific language governing permissions and # limitations under the License. -# FIXME: The following sharding strategies are not yet supported in FSDP v2: -# - no_shard -# - optim -# - optim_grads -# Currently only optim_grads_params is fully implemented and tested. -# We will add support for these strategies in a follow-up change. -# When using ParameterGroup directly or via fully_shard(), passing an unsupported -# sharding_strategy will raise NotImplementedError. See README.md for details. +# FSDP v2 currently supports ZeRO-1/2/3 strategies (`optim`, +# `optim_grads`, and `optim_grads_params`). `no_shard` is not implemented in +# this path yet. from ..uneven_dtensor import ( gather_and_compute_chunk_metadata, @@ -35,7 +30,7 @@ from .fully_shard import FSDPModule, fully_shard from .mixed_precision import ( FullyShardFP8Policy, - FullyShardMixedPrecisionPolicy, + MixedPrecisionPolicy, FullyShardNVFP4Policy, ) from .param_group import ParameterGroup @@ -44,7 +39,7 @@ "FSDPModule", "fully_shard", "FullyShardFP8Policy", - "FullyShardMixedPrecisionPolicy", + "MixedPrecisionPolicy", "FullyShardNVFP4Policy", "ParameterGroup", "BufferIndex", diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/allocator.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/allocator.py index 80ba21e750e..ff7e7eaa89f 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/allocator.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/allocator.py @@ -163,7 +163,8 @@ class TracePoolAllocator(BucketAllocator): in ``_slot_map``. After coloring, slots are laid out contiguously and a single - ``torch.empty`` per ``(dtype, device)`` group is allocated. + ``torch.empty`` per ``(dtype, device)`` group is allocated. If the + trace is empty, planning is a no-op and later cursor resets also no-op. **Phase 3 — Optimized** (after ``plan()``) @@ -297,7 +298,9 @@ def plan(self) -> int: Multiply by ``element_size(dtype)`` for bytes. """ assert self._phase == "trace", "plan() can only be called in trace phase" - assert len(self._trace) > 0, "empty trace — nothing to plan" + if len(self._trace) == 0: + self._phase = "optimized" + return 0 # ---- step 1: build intervals from alloc/free pairs ---- alloc_stack: Dict[AllocatorKey, List[int]] = {} # key -> [alloc_seq, ...] @@ -318,7 +321,9 @@ def plan(self) -> int: ) ) - assert len(intervals) > 0, "no paired alloc/free intervals found in trace" + if len(intervals) == 0: + self._phase = "optimized" + return 0 # ---- step 2 & 3: color and allocate ---- return self._assign_pool(intervals) diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/design.md b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/design.md index 0a6c53dd3bc..b4c5f9b0345 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/design.md +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/design.md @@ -9,8 +9,8 @@ | `fully_shard.py` | Public `fully_shard()` API and allocator selection | | `fsdp_module.py` | `FSDPModule`, `_FSDPRootContext`, `_FSDPState`, `unshard()`, `reshard()`, `reduce_grad()` | | `hooks.py` | Forward/backward hook registration and final callback | -| `param_group.py` | `ParameterGroup.unshard(async_op)`, `reduce_grad()`, `release_grad_buffer()`, `_init_buffers()` (memory optimization) | -| `dp_buffer.py` | `DataParallelBuffer.unshard(async_op)` (all-gather + `p.data` rebind), `reduce_grad()` (reduce-scatter + shard accumulation) | +| `param_group.py` | `ParameterGroup.unshard()`, `reduce_grad()`, `release_grad_buffer()`, `_init_buffers()` (memory optimization) | +| `dp_buffer.py` | `DataParallelBuffer.unshard()` (all-gather + `p.data` rebind), `reduce_grad()` (reduce-scatter + shard accumulation) | | `allocator.py` | `BucketAllocator` hierarchy: `TemporaryBucketAllocator`, `StorageFreeingBucketAllocator`, `TracePoolAllocator` — pooled memory for unsharded parameter and gradient buffers | | `mcore_fsdp_adapter.py` | `FullyShardedDataParallel.stop_communication()` — synchronizes ag_stream and rs_stream into main stream | @@ -46,7 +46,7 @@ class _FSDPRootContext: # module_id -> Event: signals when that module's all-gather is complete. # None means "not yet launched" or "already consumed by a wait". - # --- Reduce-grad overlap tracking --- + # --- Reduce-scatter grad overlap tracking --- reduce_grad_buckets: Dict[int, List[Tuple[torch.cuda.Event, ParameterGroup]]] # module_id -> [(event, param_group), ...] # Each entry: event signals RS complete; param_group holds the grad buffer. @@ -159,8 +159,9 @@ for module in [self] + prefetch: for _, param_group in module._named_param_groups: param_group.unshard() # → DataParallelBuffer.unshard(): - # allocate unsharded bucket, launch all_gather_into_tensor, - # rebind p.data → unsharded buffer slice (even before AG done!) + # for sharded buffers, allocate an unsharded bucket, launch + # all_gather_into_tensor, and rebind p.data → unsharded buffer slice + # (even before AG done!); replicated buffers bind directly to self.data. # NOTE: async_op is NOT passed; the stream context handles dispatch. if async_op: @@ -181,8 +182,8 @@ for param_names, param_group in self._named_param_groups: **Important: `p.data` rebind race.** `DataParallelBuffer.unshard()` rebinds `p.data` to the unsharded buffer slice -**inside the `with torch.cuda.stream(stream)` block**, before the all-gather completes (when -`async_op=True`). The memory is already allocated and the slice indices are correct; only the +**inside the `with torch.cuda.stream(stream)` block**, before the all-gather completes when +side-stream prefetch is enabled. The memory is already allocated and the slice indices are correct; only the NCCL fill is in-flight. The outer `unshard()` guards correctness by calling `event.wait()` before calling `_replace_module_parameter`, so the module's parameters are safe to read by the time the forward kernel uses them. @@ -226,7 +227,7 @@ ctx.unshard_done_events[id(self)] = None # reset so next iteration can prefet --- -## Feature 2: Reduce-Grad Overlap +## Feature 2: Reduce-Scatter Grad Overlap ### Hook entry point @@ -241,57 +242,58 @@ module.post_backward_issued = True ### `FSDPModule.reduce_grad(async_op)` ```python -stream = ctx.rs_stream if async_op else torch.cuda.current_stream() +def reduce_grad(self, async_op: bool = False): + stream = ctx.rs_stream if async_op else torch.cuda.current_stream() -# --- Step 1: Sliding drain — free grad buffers 2 positions back in backward order --- -if async_op: - backward_order = list(reversed(ctx.forward_order)) - for i, module in enumerate(backward_order): - if i - 2 >= 0: - for event, param_group in drain(ctx.reduce_grad_buckets[id(backward_order[i-2])]): - event.wait() - param_group.release_grad_buffer() - # → deletes param.main_grad views (prevents TE grad-accum-fusion leak) - # → DataParallelBuffer.reshard() (frees unsharded grad bucket) - if module is self: break - -# --- Step 2: Copy .grad → main_grad_buffer (on main stream, fast memcpy) --- -for param_names, param_group in self._named_param_groups: - if not param_group.requires_grad: continue - - for name, param in zip(param_names, param_group.params): - main_grad = param.get_main_grad() - if param.grad is None: - if not getattr(param, 'grad_added_to_main_grad', False): - main_grad.zero_() # no TE fusion: zero the slot - else: - main_grad.copy_(param.grad.detach()) # normal backward: copy .grad - del param.grad - - # --- Step 3: Reduce-scatter on rs_stream --- + # --- Step 1: Sliding drain — free grad buffers 2 positions back in backward order --- if async_op: - stream.wait_stream(torch.cuda.current_stream()) # ensure .grad copy is visible to rs_stream - with torch.cuda.stream(stream): - param_group.reduce_grad() - # → DataParallelBuffer.reduce_grad() (synchronous within this stream): - # fetch_unsharded_buffer() allocates full grad buffer - # reduce_scatter_tensor(output=grad_shard, input=full_grad) - # self.data[local_idx:...] += grad_shard - event = stream.record_event() - ctx.reduce_grad_buckets[id(self)].append((event, param_group)) - # param_group.release_grad_buffer() is NOT called here; deferred until drain/final CB - else: - param_group.reduce_grad() - param_group.release_grad_buffer() - - # --- Step 4: Install dist_grad on dist_param (runs in stream context) --- - for name, param, dist_param, dist_grad in zip( - param_names, param_group.params, param_group.dist_params, param_group.dist_grads - ): - if param.requires_grad and dist_grad is not None: + backward_order = list(reversed(ctx.forward_order)) + for i, module in enumerate(backward_order): + if i - 2 >= 0: + for event, param_group in drain(ctx.reduce_grad_buckets[id(backward_order[i-2])]): + event.wait() + param_group.release_grad_buffer() + # → deletes param.main_grad views (prevents TE grad-accum-fusion leak) + # → DataParallelBuffer.reshard() (frees unsharded grad bucket) + if module is self: break + + # --- Step 2: Copy .grad → main_grad_buffer (on main stream, fast memcpy) --- + for param_names, param_group in self._named_param_groups: + if not param_group.requires_grad: continue + + for name, param in zip(param_names, param_group.params): + main_grad = param.get_main_grad() + if param.grad is None: + if not getattr(param, 'grad_added_to_main_grad', False): + main_grad.zero_() # no TE fusion: zero the slot + else: + main_grad.copy_(param.grad.detach()) # normal backward: copy .grad + del param.grad + + # --- Step 3: Reduce-scatter on rs_stream --- + if async_op: + stream.wait_stream(torch.cuda.current_stream()) # ensure .grad copy is visible to rs_stream with torch.cuda.stream(stream): - dist_grad = dist_grad.to(dist_param.dtype) # dtype cast on rs_stream - setattr(dist_param, "grad", dist_grad) # Python ref, no GPU dependency + param_group.reduce_grad() + # → DataParallelBuffer.reduce_grad() (synchronous within this stream): + # fetch_unsharded_buffer() allocates full grad buffer + # reduce_scatter_tensor(output=grad_shard, input=full_grad) + # self.data[local_idx:...] += grad_shard + event = stream.record_event() + ctx.reduce_grad_buckets[id(self)].append((event, param_group)) + # param_group.release_grad_buffer() is NOT called here; deferred until drain/final CB + else: + param_group.reduce_grad() + param_group.release_grad_buffer() + + # --- Step 4: Install dist_grad on dist_param (runs in stream context) --- + for name, param, dist_param, dist_grad in zip( + param_names, param_group.params, param_group.dist_params, param_group.dist_grads + ): + if param.requires_grad and dist_grad is not None: + with torch.cuda.stream(stream): + dist_grad = dist_grad.to(dist_param.dtype) # dtype cast on rs_stream + setattr(dist_param, "grad", dist_grad) # Python ref, no GPU dependency ``` **Key design point — `DataParallelBuffer.reduce_grad()` has no `async_op` parameter.** @@ -364,6 +366,64 @@ def _post_backward_final_callback(root_state, root_module): --- +## ZeRO-1 and ZeRO-2 Workflow + +`optim` and `optim_grads` keep compute weights replicated but still expose +optimizer-facing DTensor shards through `dist_params`. + +### ZeRO-1 (`optim`) + +1. Forward and backward read replicated `model_weight_buffer`; no parameter + all-gather is needed in the steady state. +2. Backward writes local gradients into the replicated `main_grad_buffer`. + Post-backward reduce-scatter skips `optim` groups, so gradients remain full + replicas across local gradient accumulation. +3. `finish_grad_sync()` performs one delayed reduce-scatter for each `optim` + grad buffer. The reduce-scatter output is written into this rank's virtual + shard, which is what the optimizer consumes through `dist_grads`. +4. The optimizer updates this rank's sharded `main_weight_buffer` view. After + `copy_main_weights_to_model_weights()`, the next forward refreshes the + replicated compute weights from those updated shards. + +### ZeRO-2 (`optim_grads`) + +1. Forward and backward also read replicated `model_weight_buffer`; no parameter + all-gather is needed in the steady state. +2. Backward writes gradients into a temporary full grad buffer returned by + `main_grad_buffer.fetch_unsharded_buffer()`. +3. The post-backward hook reduce-scatters that temporary full buffer and + accumulates the result into the persistent sharded `main_grad_buffer.data`. + With overlap enabled, this reduce-scatter is launched on `ctx.rs_stream` and + the normal sliding drain/final callback releases the temporary buffer after + its event completes. +4. `finish_grad_sync()` only waits for outstanding `rs_stream` work for + `optim_grads`; it does not launch another reduce-scatter. +5. The optimizer updates this rank's sharded `main_weight_buffer` view. The next + forward refreshes replicated compute weights the same way as ZeRO-1. + +### Replicated Weight Refresh + +For ZeRO-1/2, `copy_main_weights_to_model_weights()` marks the replicated +`DataParallelBuffer` dirty when `main_weight_buffer` is sharded and +`model_weight_buffer` is replicated. The next normal unshard for that buffer +calls `DataParallelBuffer.unshard()`, which refreshes any dirty replicated +buffer before compute: + +1. Non-FP8 weights copy this rank's updated main-weight shard into the matching + slice of the replicated model-weight buffer. +2. FP8 weights quantize the local FP32 main-weight shard into the local FP8 + model-weight shard first; MXFP8 marks the transpose buffer dirty as well. +3. `DataParallelBuffer.unshard(bind_params=...)` sees the dirty flag and gathers + the updated shards directly into the full replicated compute buffer on every + rank, then clears the flag. The same call can bind params to `self.data` for + the current compute phase. + +The rowwise/model buffer is refreshed on forward unshard. For MXFP8, the +transpose buffer is refreshed on backward unshard, where +`weight_buffers_for_unshard(..., bwd_pass=True)` selects it. + +--- + ## Feature 3: Activation Recomputation (Gradient Checkpointing) ### Problem @@ -536,18 +596,30 @@ final_callback: No `async_op` parameter is needed. The method is purely synchronous within the calling stream: ```python -def reduce_grad(self): - if not self.is_distributed: +def reduce_grad(self, grad_comm_dtype=None): + sm = self.buffer_index.shard_meta + local_grad_shard = self.data[sm.local_data_index : sm.local_data_index + sm.size] + + if not self.is_distributed and self.sharding_strategy == "no_shard": torch.distributed.all_reduce(self.data, group=self.dp_group) return - full_grad = self.fetch_unsharded_buffer() # allocates bucket if not present - sm = self.buffer_index.shard_meta - grad_shard = full_grad[sm.bucket_data_index : sm.bucket_data_index + sm.size] + + if self.is_distributed: + full_grad = self.fetch_unsharded_buffer() # temporary full grad buffer + input_buffer = full_grad + output_offset = sm.bucket_data_index + accumulate_output = True + else: + input_buffer = self.data # ZeRO-1 replicated accumulation buffer + output_offset = sm.local_data_index + accumulate_output = False + grad_shard = input_buffer[output_offset : output_offset + sm.size] torch.distributed.reduce_scatter_tensor( - output=grad_shard, input=full_grad, group=self.dp_group + output=grad_shard, input=input_buffer, group=self.dp_group ) - # Accumulate into persistent shard — supports multi-micro-batch grad accumulation - self.data[sm.local_data_index : sm.local_data_index + sm.size] += grad_shard + if accumulate_output: + # ZeRO-2/3 accumulate into persistent shard for micro-batch grad accumulation. + local_grad_shard += grad_shard ``` The caller (`FSDPModule.reduce_grad`) provides the stream context; `DataParallelBuffer` diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/dp_buffer.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/dp_buffer.py index 313fbcc3e02..2010f873d05 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/dp_buffer.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/dp_buffer.py @@ -20,7 +20,7 @@ import torch from .allocator import BucketAllocator, TemporaryBucketAllocator -from .mixed_precision import FullyShardMixedPrecisionPolicy +from .mixed_precision import MixedPrecisionPolicy from .utils import ParamGroupIdx logger = logging.getLogger(__name__) @@ -355,7 +355,7 @@ def __init__( device: torch.device, dp_group: torch.distributed.ProcessGroup, param_group_id: ParamGroupIdx, - mp_policy: FullyShardMixedPrecisionPolicy, + mp_policy: MixedPrecisionPolicy, *, allocator: Optional[BucketAllocator] = None, buffer_role: str = "model_weight", @@ -375,6 +375,7 @@ def __init__( self.alloc_key = (param_group_id, buffer_role) self.mp_policy = mp_policy self.is_distributed = is_distributed + self.sharding_strategy = sharding_strategy self.gradient_scaling_factor = gradient_scaling_factor dp_rank = torch.distributed.get_rank(dp_group) @@ -408,6 +409,8 @@ def __init__( self.data: Optional[torch.Tensor] = None self._unsharded_buffer: Optional[torch.Tensor] = None + # Set when a replicated buffer only has this rank's updated shard. + self._dirty = False # ------------------------------------------------------------------ # # Public API @@ -534,48 +537,61 @@ def get_item(self, item_id: int, *, as_shard: bool = False) -> torch.Tensor: def is_unsharded(self) -> bool: """Return whether this buffer currently has a full unsharded view.""" - return not self.is_distributed or self._unsharded_buffer is not None + return not self._dirty and ( + not self.is_distributed or self._unsharded_buffer is not None + ) @torch.no_grad() def unshard( - self, async_op: bool = True - ) -> Tuple[torch.Tensor, Optional[torch.distributed.Work]]: + self, + bind_params: bool = True, + ) -> torch.Tensor: """All-gather the full buffer from all shards and bind parameter storage. For non-distributed buffers self.data is already full, so - (self.data, None) is returned directly. + self.data is returned directly. If a replicated buffer only has this + rank's updated shard, the shard is all-gathered into self.data first. """ - work = None - if self.is_distributed: - bucket = self.allocator.allocate( - key=self.alloc_key, - size=self.buffer_index.bucket_meta.size, - dtype=self.dtype, - device=self.device, - ) - self._unsharded_buffer = bucket.data + assert self.data is not None + shard_buffer = None + if self._dirty: + assert not self.is_distributed, "dirty unshard requires a replicated buffer" + full_buffer = self.data sm = self.buffer_index.shard_meta - shard = self.data[sm.local_data_index : sm.local_data_index + sm.size] - - work = torch.distributed.all_gather_into_tensor( - output_tensor=self._unsharded_buffer, - input_tensor=shard, - group=self.dp_group, - async_op=async_op, - ) - if self._unsharded_buffer.is_cuda: - # The temporary bucket may be released from another stream before this - # collective finishes; record the producer stream for allocator safety. - self._unsharded_buffer.record_stream(torch.cuda.current_stream()) + shard_buffer = self.data[sm.local_data_index : sm.local_data_index + sm.size] + elif self.is_distributed: + if self._unsharded_buffer is None: + bucket = self.allocator.allocate( + key=self.alloc_key, + size=self.buffer_index.bucket_meta.size, + dtype=self.dtype, + device=self.device, + ) + self._unsharded_buffer = bucket.data + sm = self.buffer_index.shard_meta + shard_buffer = self.data[sm.local_data_index : sm.local_data_index + sm.size] full_buffer = self._unsharded_buffer else: full_buffer = self.data - # Bind the full buffer to the params according to the layout. - self._bind_buffer_to_params(full_buffer) + if shard_buffer is not None: + torch.distributed.all_gather_into_tensor( + output_tensor=full_buffer, + input_tensor=shard_buffer, + group=self.dp_group, + ) + if self.is_distributed and full_buffer.is_cuda: + # Temporary all-gather buckets may be released from another stream before + # the collective finishes; record the producer stream for allocator safety. + full_buffer.record_stream(torch.cuda.current_stream()) + if self._dirty: + self._dirty = False + + if bind_params: + self._bind_buffer_to_params(full_buffer) - return (full_buffer, work) + return full_buffer def _bind_buffer_to_params(self, buffer: torch.Tensor) -> None: """Bind the given buffer to the params according to the layout.""" @@ -627,11 +643,13 @@ def fetch_buffer(self, *, as_shard: bool = False) -> torch.Tensor: @torch.no_grad() def reduce_grad(self, grad_comm_dtype: Optional[torch.dtype] = None): - """Reduce gradients across the data-parallel group. + """Reduce gradients into the optimizer-facing local shard. - For distributed buffers: reduce-scatter the full gradient into each - rank's shard, then accumulate into self.data. - For non-distributed buffers: all-reduce in-place. + For distributed buffers, this reduce-scatters a temporary full gradient + and accumulates the result into the persistent local shard. For + replicated buffers, this reduce-scatters the full accumulation buffer + once into this rank's virtual shard for ZeRO-1 optimizer consumption. + For no-shard buffers, this all-reduces the full gradient buffer. If grad_comm_dtype differs from self.dtype, communicate with a temporary casted tensor and cast the reduced result back before accumulation. """ @@ -647,34 +665,55 @@ def reduce_grad(self, grad_comm_dtype: Optional[torch.dtype] = None): op = torch.distributed.ReduceOp.SUM prescale = True - if not self.is_distributed: - comm_data = ( + sm = self.buffer_index.shard_meta + local_grad_shard = self.data[sm.local_data_index : sm.local_data_index + sm.size] + + if not self.is_distributed and self.sharding_strategy == "no_shard": + comm_input = ( self.data if grad_comm_dtype == self.dtype else self.data.to(grad_comm_dtype) ) if prescale: - comm_data.mul_(self.gradient_scaling_factor) - torch.distributed.all_reduce(comm_data, group=self.dp_group, op=op) - if comm_data is not self.data: - self.data.copy_(comm_data.to(self.dtype)) + comm_input.mul_(self.gradient_scaling_factor) + torch.distributed.all_reduce(comm_input, group=self.dp_group, op=op) + if grad_comm_dtype != self.dtype: + self.data.copy_(comm_input.to(self.dtype)) return - full_grad = self.fetch_buffer() - comm_input = full_grad if grad_comm_dtype == self.dtype else full_grad.to(grad_comm_dtype) - if full_grad.is_cuda: - # Keep temporary reduce-scatter buffers tied to the stream that uses them. - full_grad.record_stream(torch.cuda.current_stream()) + if self.is_distributed: + # ZeRO-2/3 (optim_grads/optim_grads_params): ``self.data`` is the + # persistent local grad shard. The full grad buffer is temporary, + # assembled only for this reduce-scatter, and the RS result is + # accumulated into ``local_grad_shard`` for gradient accumulation. + input_buffer = self.fetch_buffer() + output_offset = sm.bucket_data_index + accumulate_output = True + if input_buffer.is_cuda: + # Keep temporary reduce-scatter buffers tied to the stream that uses them. + input_buffer.record_stream(torch.cuda.current_stream()) + else: + # ZeRO-1 (optim): ``self.data`` is the replicated full grad + # accumulation buffer. The optimizer consumes only this rank's + # virtual shard, so the one delayed RS writes directly into that + # slice instead of accumulating into a separate shard buffer. + input_buffer = self.data + output_offset = sm.local_data_index + accumulate_output = False + + comm_input = ( + input_buffer if grad_comm_dtype == self.dtype else input_buffer.to(grad_comm_dtype) + ) if prescale: comm_input.mul_(self.gradient_scaling_factor) - - sm = self.buffer_index.shard_meta - local_grad_shard = self.data[sm.local_data_index : sm.local_data_index + sm.size] - reduced_grad_shard = comm_input[sm.bucket_data_index : sm.bucket_data_index + sm.size] + reduced_grad_shard = comm_input[output_offset : output_offset + sm.size] torch.distributed.reduce_scatter_tensor( output=reduced_grad_shard, input=comm_input, group=self.dp_group, op=op ) - local_grad_shard += reduced_grad_shard + if accumulate_output: + local_grad_shard += reduced_grad_shard + elif grad_comm_dtype != self.dtype: + local_grad_shard.copy_(reduced_grad_shard.to(self.dtype)) def check_all_fsdp_buffers(module) -> bool: diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/fsdp_module.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/fsdp_module.py index 3901afe44aa..c38a041cf32 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/fsdp_module.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/fsdp_module.py @@ -25,7 +25,7 @@ from torch.distributed.tensor import DTensor from .allocator import BucketAllocator -from .mixed_precision import FullyShardMixedPrecisionPolicy +from .mixed_precision import MixedPrecisionPolicy from .param_group import ParameterGroup from .utils import ParamGroupIdx, _replace_module_parameter @@ -102,9 +102,9 @@ class _FSDPRootContext: # ------------------------------------------------------------------ # Reduce-scatter (gradient sync) tracking # ------------------------------------------------------------------ - reduce_grad_buckets: Dict[int, List[Tuple[torch.cuda.Event, "ParameterGroup"]]] = field( - default_factory=dict - ) + reduce_grad_buckets: Dict[ + int, List[Tuple[torch.cuda.Event, "ParameterGroup"]] + ] = field(default_factory=dict) """ Maps module_id -> list of (event, parameter_group) tuples. @@ -171,16 +171,17 @@ class FSDPModule(nn.Module): methods for managing parameter sharding state: - unshard(): All-gather parameters before forward - reshard(): Release unsharded buffer after forward - - reduce_grad(): Reduce-scatter gradients after backward + - reduce_grad(): Reduce gradients after backward """ def _init_named_param_groups( self, mesh: Optional[DeviceMesh], ignored_params: Optional[set], - mp_policy: FullyShardMixedPrecisionPolicy, + mp_policy: MixedPrecisionPolicy, bucket_allocator: BucketAllocator, gradient_scaling_factor: Optional[float] = None, + sharding_strategy: str = "optim_grads_params", ): """ Initialize parameter groups and build param name mapping. @@ -212,6 +213,7 @@ def _init_named_param_groups( ignored_params=ignored_params, allocator=bucket_allocator, gradient_scaling_factor=gradient_scaling_factor, + sharding_strategy=sharding_strategy, ) setattr(self, "_fsdp_param_groups", fsdp_param_groups) @@ -262,7 +264,7 @@ def _materialize_meta_module( self, ignored_modules: set, mesh: Optional[DeviceMesh] = None, - mp_policy: Optional[FullyShardMixedPrecisionPolicy] = None, + mp_policy: Optional[MixedPrecisionPolicy] = None, ): """ Materialize meta parameters to actual device and initialize. @@ -483,7 +485,7 @@ def reduce_grad(self, async_op: bool = False): This is called post-backward to: 1. Copy gradients to main gradient buffer - 2. Perform all-reduce or reduce-scatter + 2. Perform gradient reduction 3. Install reduced gradients to distributed parameters """ torch.cuda.nvtx.range_push("MFSDP reduce_grad") @@ -517,10 +519,11 @@ def reduce_grad(self, async_op: bool = False): del param.grad elif param.grad is None: main_grad = param.get_main_grad() - if hasattr(param, "main_grad") and param.main_grad is not None: - if param.main_grad.data_ptr() != main_grad.data_ptr(): - main_grad.copy_(param.main_grad.detach()) - else: + param_main_grad = getattr(param, "main_grad", None) + if ( + param_main_grad is None + or param_main_grad.data_ptr() != main_grad.data_ptr() + ): main_grad.zero_() else: main_grad = param.get_main_grad() @@ -540,22 +543,26 @@ def reduce_grad(self, async_op: bool = False): param_group.release_grad_buffer() # Install reduced gradients to distributed parameters - for name, param, dist_param, dist_grad in zip( - param_names, param_group.params, param_group.dist_params, param_group.dist_grads - ): - if not param.requires_grad: - continue - if param_group.mp_policy.use_decoupled_grad: - setattr(dist_param, "decoupled_grad", dist_grad) - if dist_param.grad is not None: - del dist_param.grad - else: - if dist_grad is not None: - with torch.cuda.stream(stream): - dist_grad = dist_grad.to(dist_param.dtype) - setattr(dist_param, "grad", dist_grad) - if hasattr(dist_param, "decoupled_grad"): - dist_param.decoupled_grad = None + with torch.cuda.stream(stream): + for name, param, dist_param, dist_grad in zip( + param_names, + param_group.params, + param_group.dist_params, + param_group.dist_grads, + ): + if not param.requires_grad: + continue + if param_group.mp_policy.use_decoupled_grad: + setattr(dist_param, "decoupled_grad", dist_grad) + if dist_param.grad is not None: + del dist_param.grad + else: + assert ( + dist_grad is None or dist_param.dtype == dist_grad.dtype + ), f"{name} Dist param dtype {dist_param.dtype} does not match dist grad dtype {dist_grad.dtype}" + setattr(dist_param, "grad", dist_grad) + if hasattr(dist_param, "decoupled_grad"): + dist_param.decoupled_grad = None if async_op: event = stream.record_event() @@ -571,6 +578,29 @@ def reduce_grad(self, async_op: bool = False): torch.cuda.nvtx.range_pop() + @torch.no_grad() + def finish_grad_sync(self, force_all_reduce: Optional[bool] = False): + """Finish optimizer-facing gradient synchronization for this iteration.""" + ctx = self._fsdp_root_context + for _, child in self.named_modules(): + if not isinstance(child, FSDPModule): + continue + if any( + param_group.sharding_strategy == "optim" + for param_group in child._fsdp_param_groups + ): + # ZeRO-1 keeps gradients replicated during backward and performs + # exactly one reduce-scatter at the iteration grad-sync boundary. + child.reduce_grad(async_op=False) + for param_group in child._fsdp_param_groups: + for param, dist_grad in zip(param_group.params, param_group.dist_grads): + if param.requires_grad: + # v1 replaces module params with optimizer-facing distributed + # params after grad sync. v2 keeps compute params in the module, + # so mirror the reduced grad for shared finalizers. + param.main_grad = dist_grad + torch.cuda.current_stream().wait_stream(ctx.rs_stream) + @torch.no_grad() def _scale_gradients(self, scaling_factor: float): """Scale optimizer-facing gradients by a factor.""" @@ -790,11 +820,12 @@ def _set_nan_check(self, enable_nan_checks: bool): def _get_module_fsdp_param_groups( module: nn.Module, - mp_policy: FullyShardMixedPrecisionPolicy, + mp_policy: MixedPrecisionPolicy, allocator: BucketAllocator, mesh: Optional[DeviceMesh] = None, ignored_params: Optional[set[nn.Parameter]] = None, gradient_scaling_factor: Optional[float] = None, + sharding_strategy: str = "optim_grads_params", ) -> List[ParameterGroup]: """ Group module parameters by (device, dtype, requires_grad) and create ParameterGroups. @@ -827,6 +858,7 @@ def _get_module_fsdp_param_groups( mp_policy=mp_policy, gradient_scaling_factor=gradient_scaling_factor, allocator=allocator, + sharding_strategy=sharding_strategy, ) ) diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/fully_shard.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/fully_shard.py index b2446d3253f..ed7b59a7905 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/fully_shard.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/fully_shard.py @@ -35,7 +35,7 @@ _register_forward_hook, _register_forward_pre_hook, ) -from .mixed_precision import FullyShardMixedPrecisionPolicy +from .mixed_precision import MixedPrecisionPolicy from .utils import _init_default_fully_shard_mesh __all__ = ["FSDPModule", "fully_shard"] @@ -49,7 +49,7 @@ def fully_shard( shard_placement_fn: Optional[ Callable[[nn.Parameter], Optional[Shard]] ] = None, # TODO: implement - mp_policy: Optional[FullyShardMixedPrecisionPolicy] = None, + mp_policy: Optional[MixedPrecisionPolicy] = None, offload_policy: Optional["OffloadPolicy"] = None, # TODO: implement ignored_params: Optional[set[nn.Parameter]] = None, # --- Megatron-FSDP specific options --- @@ -57,6 +57,7 @@ def fully_shard( enable_async_reduce_grad: bool = True, gradient_scaling_factor: Optional[float] = None, enable_trace_pool: bool = False, + sharding_strategy: str = "optim_grads_params", ) -> nn.Module: """ Wrap a module with FSDP sharding semantics. @@ -76,15 +77,18 @@ def fully_shard( mesh = mesh or _init_default_fully_shard_mesh() if mp_policy is None: - mp_policy = FullyShardMixedPrecisionPolicy() + mp_policy = MixedPrecisionPolicy() cls = module.__class__ new_cls = type(f"FSDP{cls.__name__}", (FSDPModule, cls), {}) module.__class__ = new_cls - bucket_allocator = ( - TracePoolAllocator() if enable_trace_pool else StorageFreeingBucketAllocator() + use_trace_pool = enable_trace_pool and sharding_strategy in ( + "optim", + "optim_grads", + "optim_grads_params", ) + bucket_allocator = TracePoolAllocator() if use_trace_pool else StorageFreeingBucketAllocator() module._init_named_param_groups( mesh, @@ -92,6 +96,7 @@ def fully_shard( mp_policy=mp_policy, bucket_allocator=bucket_allocator, gradient_scaling_factor=gradient_scaling_factor, + sharding_strategy=sharding_strategy, ) module._init_fsdp_state( enable_unshard_prefetch=enable_unshard_prefetch, diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/hooks.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/hooks.py index 30a269e9437..2ecbe30ab06 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/hooks.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/hooks.py @@ -148,7 +148,11 @@ def post_backward(module: FSDPModule): ctx.backward_done_modules.add(id(module)) ctx._advance_backward_module() module.reshard() - module.reduce_grad(async_op=ctx.enable_async_reduce_grad) + if any( + param_group.sharding_strategy in ("optim_grads", "optim_grads_params") + for param_group in module._fsdp_param_groups + ): + module.reduce_grad(async_op=ctx.enable_async_reduce_grad) module.post_backward_issued = True @torch.compiler.disable @@ -225,7 +229,11 @@ def _post_backward_final_callback(root_state: _FSDPState, root_module: nn.Module if getattr(module, "post_backward_issued", False): continue module.reshard() - module.reduce_grad(async_op=ctx.enable_async_reduce_grad) + if any( + param_group.sharding_strategy in ("optim_grads", "optim_grads_params") + for param_group in module._fsdp_param_groups + ): + module.reduce_grad(async_op=ctx.enable_async_reduce_grad) for buckets in ctx.reduce_grad_buckets.values(): while len(buckets) > 0: event, param_group = buckets.pop() diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/mixed_precision.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/mixed_precision.py index ed0d44cab83..254ceb3d2dd 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/mixed_precision.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/mixed_precision.py @@ -176,7 +176,7 @@ class FullyShardNVFP4Policy: @dataclass(frozen=True) -class FullyShardMixedPrecisionPolicy: +class MixedPrecisionPolicy: """Mixed precision policy owned by the v2 ``fully_shard`` path.""" main_params_dtype: Optional[torch.dtype] = None @@ -403,22 +403,39 @@ def needs_transpose_weight_buffer(self, tensor: torch.Tensor) -> bool: return HAVE_TE_MXFP8 and isinstance(tensor, MXFP8Tensor) def main_params_dtype_for_param(self, tensor: torch.Tensor) -> Optional[torch.dtype]: - """Return the main-parameter dtype for a parameter group.""" - if self.is_fp8_param(tensor) and self.main_params_dtype is None: - return torch.float32 - if self.is_nvfp4_param(tensor) and self.main_params_dtype is None: - return torch.float32 + """Return the dtype for the optimizer main-weight buffer. + + Returns ``self.main_params_dtype`` unchanged. When this is ``None`` + no ``main_weight_buffer`` is allocated — the optimizer mutates the + model-weight buffer directly. Set to ``torch.float32`` when the + model uses quantized weights (FP8 / NVFP4) so the optimizer works on + high-precision copies. + """ return self.main_params_dtype def main_grads_dtype_for_param(self, tensor: torch.Tensor) -> torch.dtype: """Return the main-gradient dtype for a parameter group. - Defaults to float32 for stable gradient accumulation (avoids - precision loss in reduce-scatter and optimizer steps). + Resolution order (first match wins): + 1. Explicit ``main_grads_dtype`` in the policy — user override wins. + 2. When ``use_decoupled_grad`` is *disabled*, the optimizer's main-grad + buffer should match the main-param buffer dtype (the optimizer + writes gradients into the same context as the main params). + 3. Default — fall back to the parameter's own dtype. + + Only the first condition triggers an independent main-grad buffer + creation with a different dtype than the main params. Conditions 2 + and 3 re-use the dtype already chosen for the main-weight buffer. """ if self.main_grads_dtype is not None: return self.main_grads_dtype - return torch.float32 + + if not self.use_decoupled_grad: + main_param_dtype = self.main_params_dtype_for_param(tensor) + if main_param_dtype is not None: + return main_param_dtype + + return tensor.dtype def get_high_precision_value(self, tensor: torch.Tensor) -> torch.Tensor: """Return a high-precision value for initializing optimizer main weights.""" @@ -435,7 +452,17 @@ def get_high_precision_value(self, tensor: torch.Tensor) -> torch.Tensor: return tensor.dequantize() def post_unshard(self, params: List[torch.Tensor], bwd_pass: bool = False) -> None: - """Run post-unshard mixed precision processing for a parameter group.""" + """Run post-unshard processing for quantized parameters. + + After FSDP all-gathers the raw quantized data into the model-weight + buffer, Transformer Engine needs recipe-specific state rebuilt: + + - **FP8:** rebinds rowwise/columnwise data pointers so TE's forward + and backward kernels find the correct quantized payloads. + - **NVFP4:** on backward pass only, calls TE's post-all-gather + handler to restore the rowwise/colwise scale factors needed by + the bwd kernel. + """ fp8_params = [param for param in params if self.is_fp8_param(param)] nvfp4_params = [param for param in params if self.is_nvfp4_param(param)] @@ -486,7 +513,16 @@ def post_unshard(self, params: List[torch.Tensor], bwd_pass: bool = False) -> No post_all_gather_processing(valid_nvfp4) def post_reshard(self, params: List[torch.Tensor]) -> None: - """Run post-reshard mixed precision processing for a parameter group.""" + """Run post-reshard cleanup for quantized parameters. + + After FSDP releases the all-gathered buffer, TE's temporary views + into the unsharded data become invalid: + + - **FP8:** drops transpose/cache views unless ``keep_transpose_cache`` + is set (for fine-grained recompute scenarios where the backward + kernel may request transpose data inside the forward call stack). + - **NVFP4:** skips — NVFP4 teardown is handled separately. + """ if self.fp8.keep_transpose_cache: return for param in params: @@ -517,6 +553,10 @@ def copy_main_weights_to_model_weights( assert model_weight_buffer is not None, "main weights require a model-weight buffer" + dirty_replicated_buffer = ( + main_weight_buffer.is_distributed and not model_weight_buffer.is_distributed + ) + if self.is_nvfp4_param(params[0]): quantize_main_weights_to_nvfp4( params, param_idx, data_parallel_group, model_weight_buffer, main_weight_buffer @@ -524,7 +564,27 @@ def copy_main_weights_to_model_weights( return if not self.is_fp8_param(params[0]): - model_weight_buffer.data.copy_(main_weight_buffer.data) + if model_weight_buffer.is_distributed and not main_weight_buffer.is_distributed: + raise RuntimeError( + "Unsupported FSDP main/model weight buffer layout: " + "model weights are sharded but main weights are replicated." + ) + if model_weight_buffer.is_distributed == main_weight_buffer.is_distributed: + model_weight_buffer.data.copy_(main_weight_buffer.data) + else: + model_shard_meta = model_weight_buffer.buffer_index.shard_meta + main_shard_meta = main_weight_buffer.buffer_index.shard_meta + model_weight_buffer.data[ + model_shard_meta.local_data_index : model_shard_meta.local_data_index + + model_shard_meta.size + ].copy_( + main_weight_buffer.data[ + main_shard_meta.local_data_index : main_shard_meta.local_data_index + + main_shard_meta.size + ] + ) + if dirty_replicated_buffer: + model_weight_buffer._dirty = True return fp8_params = [] @@ -554,6 +614,10 @@ def copy_main_weights_to_model_weights( quantize_main_weights_to_fp8( fp8_params, main_params, start_offsets, data_parallel_group, model_param_shards ) + if dirty_replicated_buffer: + model_weight_buffer._dirty = True + if transpose_weight_buffer is not None: + transpose_weight_buffer._dirty = True def is_fp8_param(tensor: torch.Tensor) -> bool: diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/nvfp4_design.md b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/nvfp4_design.md index 90097c55716..0610731f271 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/nvfp4_design.md +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/nvfp4_design.md @@ -22,7 +22,7 @@ FSDP path, following the same patterns used for FP8 support. ### 2.2 What already works (FSDP v2 FP8 path) -- `FullyShardMixedPrecisionPolicy` (in `v2/mixed_precision.py`) handles +- `MixedPrecisionPolicy` (in `v2/mixed_precision.py`) handles FP8 / MXFP8 detection, buffer dtype, raw-data extraction, post-unshard processing, post-reshard cache invalidation, and main→model quantization. - `ParameterGroup` creates 3–4 `DataParallelBuffer` instances: @@ -69,7 +69,7 @@ different byte offsets appropriate for that buffer's dtype/shape. ### 4.2 Policy-driven shape transform -The `FullyShardMixedPrecisionPolicy` already owns dtype and raw-data decisions. +The `MixedPrecisionPolicy` already owns dtype and raw-data decisions. We extend it with an `nvfp4` sub-policy (analogous to `fp8`) that: - Returns `torch.uint8` from `model_weight_buffer_dtype()`. @@ -108,11 +108,11 @@ class FullyShardNVFP4Policy: recipe: Optional[str] = None ``` -**Extend `FullyShardMixedPrecisionPolicy`** with `nvfp4` field: +**Extend `MixedPrecisionPolicy`** with `nvfp4` field: ```python @dataclass(frozen=True) -class FullyShardMixedPrecisionPolicy: +class MixedPrecisionPolicy: ... fp4_param_gather: bool = False fp4_recipe: Optional[str] = None @@ -227,7 +227,7 @@ Already has `fp4_param_gather: bool = False` at line 75. No change needed. **`_init_with_fully_shard()`** — pass FP4 configuration into the policy: ```python -fully_shard_mp_policy = FullyShardMixedPrecisionPolicy( +fully_shard_mp_policy = MixedPrecisionPolicy( ... nvfp4=FullyShardNVFP4Policy( enabled=ddp_config.fp4_param_gather, diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/param_group.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/param_group.py index 2cb0648853e..6ed8455b660 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/param_group.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/param_group.py @@ -30,7 +30,7 @@ from ..uneven_dtensor import make_uneven_dtensor, update_uneven_dtensor_chunk_metadata from .allocator import BucketAllocator, TemporaryBucketAllocator, _free_storage from .dp_buffer import DataParallelBuffer -from .mixed_precision import FullyShardMixedPrecisionPolicy +from .mixed_precision import MixedPrecisionPolicy from .utils import ParamGroupIdx @@ -56,7 +56,7 @@ def __init__( params: List[torch.nn.Parameter], param_group_id: ParamGroupIdx, *, - mp_policy: FullyShardMixedPrecisionPolicy, + mp_policy: MixedPrecisionPolicy, mesh: Optional[DeviceMesh] = None, sharding_strategy: str = "optim_grads_params", gradient_scaling_factor: Optional[float] = None, @@ -83,16 +83,12 @@ def __init__( self.mesh = mesh self.dp_group = mesh.get_group() - # FIXME: no_shard, optim, and optim_grads sharding strategies are not yet supported in v2. - # Currently only optim_grads_params is fully implemented and tested. - # See README.md § "Sharding Strategies" for details. - # We will add support for these strategies in a follow-up change. - if sharding_strategy not in ("optim_grads_params",): + if sharding_strategy == "no_shard": raise NotImplementedError( - f"Sharding strategy '{sharding_strategy}' is not yet supported in FSDP v2. " - f"Currently only 'optim_grads_params' is implemented. " - f"We will add support for 'no_shard', 'optim', and 'optim_grads' in a follow-up change." + "Sharding strategy 'no_shard' is not yet supported in FSDP v2." ) + if sharding_strategy not in ("optim", "optim_grads", "optim_grads_params"): + raise ValueError(f"Unsupported sharding strategy: {sharding_strategy}") self.sharding_strategy = sharding_strategy self.param_group_id = param_group_id @@ -114,7 +110,6 @@ def __init__( self.hsdp_wbuf: Optional[DataParallelBuffer] = None self.hsdp_gbuf: Optional[DataParallelBuffer] = None self.hsdp_comm_gbuf: Optional[DataParallelBuffer] = None - # Initialize buffers and distributed parameters self._init_buffers() @@ -149,9 +144,6 @@ def _create_buffer( mp_policy=self.mp_policy, ) - # FIXME: The branching below currently only handles optim_grads_params since - # no_shard, optim, and optim_grads are gated by a NotImplementedError at init. - # When support for those strategies is added, the logic below must be validated. def _init_buffers(self) -> None: """ Initialize all buffers based on sharding strategy. @@ -215,35 +207,32 @@ def _init_buffers(self) -> None: # Create distributed parameter views self._init_dist_params() - def unshard(self, async_op: bool = False, bwd_pass: bool = False): + def unshard(self, bwd_pass: bool = False): """ Unshard model weights by all-gathering from sharded buffer. After unshard, parameters point to full unsharded storage. FP8 parameters rebind their TE raw payload instead of ``param.data``. """ - work = None for weight_buffer in self.mp_policy.weight_buffers_for_unshard( self.model_weight_buffer, self.transpose_weight_buffer, bwd_pass=bwd_pass ): - if weight_buffer is None: - continue - if weight_buffer.is_distributed and weight_buffer.is_unsharded(): - continue - _, weight_work = weight_buffer.unshard(async_op=async_op) - if work is None: - work = weight_work + if weight_buffer is not None: + weight_buffer.unshard(bind_params=True) self.mp_policy.post_unshard(self.params, bwd_pass=bwd_pass) - return work def has_unsharded_weight_buffers(self, bwd_pass: bool = False) -> bool: - """Return whether all weight buffers needed for this forward/backward phase are unsharded.""" + """Return whether this phase can skip launching another distributed unshard.""" for weight_buffer in self.mp_policy.weight_buffers_for_unshard( self.model_weight_buffer, self.transpose_weight_buffer, bwd_pass=bwd_pass ): if weight_buffer is None: continue + if not weight_buffer.is_distributed: + # Replicated buffers still need DataParallelBuffer.unshard() to + # rebind original parameter storage before the module uses it. + return False if not weight_buffer.is_unsharded(): return False return True @@ -271,8 +260,9 @@ def reduce_grad(self): """ Reduce gradients across DP ranks. - For distributed buffers: reduce-scatter the full gradient - For non-distributed buffers: all-reduce in-place + ZeRO-2/3 reduce-scatter sharded grad buffers during backward. + ZeRO-1 keeps grads replicated during backward and reduce-scatters + the replicated buffer once when the optimizer syncs. """ self.main_grad_buffer.reduce_grad(grad_comm_dtype=self.mp_policy.grad_comm_dtype) @@ -294,7 +284,7 @@ def _init_dist_params(self): Creates DTensor views of model weights and gradients based on sharding strategy: - "optim_grads_params": weights and grads sharded, full ZeRO-3 - "optim_grads": grads sharded, weights replicated (ZeRO-2) - - "optim": optimizer state sharding only, weights and grads replicated (ZeRO-1) + - "optim": grads accumulate replicated, optimizer consumes reduced shards - "no_shard": replicated, no sharding (DDP-equivalent) """ self.dist_params = [] diff --git a/tests/unit_tests/distributed/megatron_fsdp/v2/test_fully_shard.py b/tests/unit_tests/distributed/megatron_fsdp/v2/test_fully_shard.py index 94a0897115c..336d810a4dd 100644 --- a/tests/unit_tests/distributed/megatron_fsdp/v2/test_fully_shard.py +++ b/tests/unit_tests/distributed/megatron_fsdp/v2/test_fully_shard.py @@ -54,7 +54,7 @@ from megatron.core.distributed.fsdp.src.megatron_fsdp.v2.fsdp_module import FSDPModule from megatron.core.distributed.fsdp.src.megatron_fsdp.v2.fully_shard import fully_shard from megatron.core.distributed.fsdp.src.megatron_fsdp.v2.mixed_precision import ( - FullyShardMixedPrecisionPolicy, + MixedPrecisionPolicy, ) # ------------------------------------------------------------------ # @@ -459,7 +459,7 @@ class TestMixedPrecision: def test_main_params_fp32(self): """With fp32 main params, main_weight_buffer should be created.""" torch.manual_seed(42) - mp_policy = FullyShardMixedPrecisionPolicy(main_params_dtype=torch.float32) + mp_policy = MixedPrecisionPolicy(main_params_dtype=torch.float32) model = SimpleMLP(64).to(_device()).bfloat16() fully_shard(model, mp_policy=mp_policy) @@ -473,7 +473,7 @@ def test_main_params_fp32(self): def test_main_params_none(self): """With no main_params_dtype, no main_weight_buffer should be created.""" torch.manual_seed(42) - mp_policy = FullyShardMixedPrecisionPolicy(main_params_dtype=None) + mp_policy = MixedPrecisionPolicy(main_params_dtype=None) model = SimpleMLP(64).to(_device()) fully_shard(model, mp_policy=mp_policy) @@ -485,8 +485,8 @@ def test_main_params_none(self): def test_fp32_grad_reduce(self): """grad_reduce_in_fp32=True should use fp32 gradient communication.""" torch.manual_seed(42) - mp_policy = FullyShardMixedPrecisionPolicy( - main_grads_dtype=torch.float32, grad_comm_dtype=torch.float32 + mp_policy = MixedPrecisionPolicy( + grad_comm_dtype=torch.float32 ) model = SimpleMLP(64).to(_device()).bfloat16() fully_shard(model, mp_policy=mp_policy) diff --git a/tests/unit_tests/distributed/megatron_fsdp/v2/test_mcore_nd_parallel.py b/tests/unit_tests/distributed/megatron_fsdp/v2/test_mcore_nd_parallel.py index 3984733eba3..7ee9f648b5d 100644 --- a/tests/unit_tests/distributed/megatron_fsdp/v2/test_mcore_nd_parallel.py +++ b/tests/unit_tests/distributed/megatron_fsdp/v2/test_mcore_nd_parallel.py @@ -20,6 +20,10 @@ ) from tests.unit_tests.test_utilities import Utils +STRICT_LOSS_ATOL = 5e-3 +STRICT_PARAM_ATOL = 5e-3 +STRICT_PARAM_RTOL = 1e-3 + @pytest.fixture(scope="class") def ref_cache(): @@ -32,6 +36,77 @@ def ref_cache(): class TestMegatronFSDPE2E: + @staticmethod + def _normalize_param_name(name): + while name.startswith("module."): + name = name[len("module.") :] + return name + + @staticmethod + def _materialize_param_tensor(param): + from torch.distributed.tensor import DTensor + + from megatron.core.distributed.fsdp.src.megatron_fsdp.uneven_dtensor import ( + uneven_dtensor_to_full_tensor, + ) + from megatron.core.fp8_utils import dequantize_fp8_tensor, is_float8tensor + + tensor = param.detach() + if isinstance(tensor, DTensor): + tensor = uneven_dtensor_to_full_tensor(tensor) + elif is_float8tensor(tensor): + tensor = dequantize_fp8_tensor(tensor) + return tensor.detach().float().cpu() + + @staticmethod + def _capture_named_params(model_chunks): + # All ranks must enter DTensor gather collectives, but only rank 0 + # keeps CPU copies for comparison. + snapshots = {} + for chunk_idx, model_chunk in enumerate(model_chunks): + for name, param in model_chunk.named_parameters(): + tensor = TestMegatronFSDPE2E._materialize_param_tensor(param) + if torch.distributed.get_rank() == 0: + key = f"{chunk_idx}.{TestMegatronFSDPE2E._normalize_param_name(name)}" + snapshots[key] = tensor + return snapshots + + @staticmethod + def _assert_replicated_weight_buffers_match(model_chunks): + from megatron.core.distributed.fsdp.src.megatron_fsdp.v2.fsdp_module import FSDPModule + + for model_chunk in model_chunks: + for _, module in model_chunk.named_modules(): + if not isinstance(module, FSDPModule): + continue + for param_group in module._fsdp_param_groups: + if ( + param_group.model_weight_buffer is None + or param_group.model_weight_buffer.is_distributed + ): + continue + param_group.unshard(bwd_pass=False) + if param_group.transpose_weight_buffer is not None: + param_group.unshard(bwd_pass=True) + + for buffer_name, buffer in ( + ("model_weight_buffer", param_group.model_weight_buffer), + ("transpose_weight_buffer", param_group.transpose_weight_buffer), + ): + if buffer is None or buffer.is_distributed: + continue + gathered = [ + torch.empty_like(buffer.data) + for _ in range(torch.distributed.get_world_size(param_group.dp_group)) + ] + torch.distributed.all_gather(gathered, buffer.data, group=param_group.dp_group) + for group_rank, replica in enumerate(gathered): + assert torch.equal(buffer.data, replica), ( + f"Replicated {buffer_name} mismatch for " + f"param_group={param_group.param_group_id}, " + f"group_rank={group_rank}" + ) + @staticmethod def _training_loop(seed=42, **kwargs): """ @@ -82,6 +157,11 @@ def _training_loop(seed=42, **kwargs): EP = kwargs.pop("EP", 1) ETP = kwargs.pop("ETP", 1) OUTER_DP = kwargs.pop("OUTER_DP", 1) + capture_param_snapshots = kwargs.pop("capture_param_snapshots", False) + verify_replicated_weight_buffers = kwargs.pop( + "verify_replicated_weight_buffers", False + ) + return_dict = kwargs.pop("return_dict", capture_param_snapshots) # Initialize model parallel groups Utils.initialize_model_parallel( @@ -128,6 +208,7 @@ def _training_loop(seed=42, **kwargs): ) outputs = [] + param_snapshots = [] # Training loop for _ in range(NUM_TRAINING_STEPS): @@ -140,12 +221,23 @@ def _training_loop(seed=42, **kwargs): num_micro_batches=GLOBAL_BATCH_SIZE // MICRO_BATCH_SIZE // DP_GROUP.size(), ) optim.step() + if verify_replicated_weight_buffers: + TestMegatronFSDPE2E._assert_replicated_weight_buffers_match(model_chunks) # Collect loss outputs.append(output[-1]) + if capture_param_snapshots: + param_snapshots.append( + TestMegatronFSDPE2E._capture_named_params(model_chunks) + ) Utils.destroy_model_parallel() + if return_dict: + result = {"outputs": outputs} + if capture_param_snapshots: + result["param_snapshots"] = param_snapshots + return result return outputs @pytest.mark.skipif( @@ -264,6 +356,245 @@ def test_compatible_with_nd_parallel(self, ref_cache, nd_topology, spec_configs) ), ) + @pytest.mark.skipif( + not is_torch_min_version("2.4.0"), reason="Test needs to be updated for torch >= 2.4.0" + ) + @pytest.mark.parametrize( + "case", + [ + pytest.param( + dict( + strategy="optim", + precision_configs=dict(bf16=True), + reference_kind="distopt", + capture_param_snapshots=True, + ), + id="bf16-optim", + ), + pytest.param( + dict( + strategy="optim_grads", + precision_configs=dict(bf16=True), + reference_kind="distopt", + capture_param_snapshots=True, + ), + id="bf16-optim_grads", + ), + pytest.param( + dict( + strategy="optim_grads_params", + precision_configs=dict(bf16=True), + reference_kind="distopt", + capture_param_snapshots=True, + ), + id="bf16-optim_grads_params", + ), + pytest.param( + dict( + strategy="optim_grads_params", + precision_configs=dict( + bf16=True, + fp8="e4m3", + fp8_param_gather=True, + fp8_recipe="mxfp8", + main_grads_dtype="fp32", + main_params_dtype="fp32", + exp_avg_dtype="bf16", + exp_avg_sq_dtype="bf16", + moe_grouped_gemm=True, + use_precision_aware_optimizer=True, + ), + reference_kind="fsdp_v1", + capture_param_snapshots=False, + ), + id="mxfp8_param_gather-optim_grads_params", + ), + ], + ) + def test_strict_iter_equivalence_zero_strategies(self, ref_cache, case): + strategy = case["strategy"] + precision_configs = case["precision_configs"] + if precision_configs.get("fp8_recipe") == "mxfp8" and ( + not torch.cuda.is_available() + or torch.cuda.get_device_capability()[0] < 10 + or not HAVE_TE_MXFP8TENSOR + ): + pytest.skip("Requires PyTorch & CUDA device with TE MXFP8Tensor support") + if Utils.world_size < 2: + pytest.skip("Requires at least 2 distributed ranks for ZeRO sharding") + + common_configs = dict( + data_parallel_sharding_strategy=strategy, + train_iters=3, + seq_length=64, + micro_batch_size=1, + global_batch_size=8, + # Keep strict iter-equivalence on the ordinary model-init path. The + # FSDP v1/v2 meta-init paths materialize nested FSDP units in a + # different order, so they can legitimately start from different + # random initial weights even with the same seed. + init_model_with_meta_device=False, + gradient_accumulation_fusion=False, + overlap_param_gather=False, + overlap_grad_reduce=False, + verify_replicated_weight_buffers=strategy in ("optim", "optim_grads"), + **precision_configs, + ) + reference_kind = case["reference_kind"] + capture_param_snapshots = case["capture_param_snapshots"] + ref_cache_key = ( + "strict_iter_equivalence", + reference_kind, + strategy, + capture_param_snapshots, + tuple(sorted((key, repr(value)) for key, value in common_configs.items())), + ) + + if ref_cache_key not in ref_cache: + reference_configs = copy.deepcopy(common_configs) + if reference_kind == "fsdp_v1": + reference_configs["use_megatron_fsdp_v2"] = False + ref_cache[ref_cache_key] = TestMegatronFSDPE2E._training_loop( + use_megatron_fsdp=True, + ckpt_format="fsdp_dtensor", + capture_param_snapshots=capture_param_snapshots, + return_dict=True, + **reference_configs, + ) + else: + ref_cache[ref_cache_key] = TestMegatronFSDPE2E._training_loop( + use_distributed_optimizer=True, + capture_param_snapshots=capture_param_snapshots, + return_dict=True, + **reference_configs, + ) + + fsdp_configs = copy.deepcopy(common_configs) + fsdp_configs["use_megatron_fsdp_v2"] = True + actual = TestMegatronFSDPE2E._training_loop( + use_megatron_fsdp=True, + ckpt_format="fsdp_dtensor", + capture_param_snapshots=capture_param_snapshots, + return_dict=True, + **fsdp_configs, + ) + reference = ref_cache[ref_cache_key] + + if torch.distributed.get_rank() != 0: + return + + assert len(actual["outputs"]) == len(reference["outputs"]) + for step, (output, ref_output) in enumerate( + zip(actual["outputs"], reference["outputs"]) + ): + loss = output["lm loss"] + ref_loss = ref_output["lm loss"] + assert_close( + loss, + ref_loss, + atol=STRICT_LOSS_ATOL, + rtol=0, + msg=( + f"Loss mismatch at step {step}, strategy={strategy}, " + f"actual={loss.detach().item()}, reference={ref_loss.detach().item()}, " + f"compare={compare_losses(loss.detach().item(), ref_loss.detach().item())}" + ), + ) + + if not capture_param_snapshots: + return + + assert len(actual["param_snapshots"]) == len(reference["param_snapshots"]) + for step, (params, ref_params) in enumerate( + zip(actual["param_snapshots"], reference["param_snapshots"]) + ): + missing = sorted(set(ref_params) ^ set(params)) + assert not missing, ( + f"Parameter key mismatch at step {step}, strategy={strategy}: {missing[:20]}" + ) + for name in sorted(ref_params): + assert_close( + params[name], + ref_params[name], + atol=STRICT_PARAM_ATOL, + rtol=STRICT_PARAM_RTOL, + msg=( + f"Parameter mismatch at step {step}, strategy={strategy}, " + f"name={name}, actual_shape={tuple(params[name].shape)}, " + f"reference_shape={tuple(ref_params[name].shape)}" + ), + ) + + @pytest.mark.skipif( + not is_torch_min_version("2.4.0"), reason="Test needs to be updated for torch >= 2.4.0" + ) + @pytest.mark.parametrize( + "strategy,precision_configs", + [ + pytest.param( + strategy, + dict( + bf16=True, + fp8="e4m3", + fp8_param_gather=True, + fp8_recipe="mxfp8", + main_grads_dtype="fp32", + main_params_dtype="fp32", + exp_avg_dtype="bf16", + exp_avg_sq_dtype="bf16", + moe_grouped_gemm=True, + use_precision_aware_optimizer=True, + ), + id=f"mxfp8_param_gather-{strategy}", + ) + for strategy in ("optim", "optim_grads") + ], + ) + def test_zero_strategy_non_equivalent_precision_paths_run( + self, strategy, precision_configs + ): + """Exercise valid ZeRO paths that intentionally lack a strict reference. + + MXFP8 ZeRO-1/2 refreshes replicated quantized compute buffers after + sharded optimizer updates; v1 and v2 do not provide a strict multi-step + golden comparison for that replicated-weight quantization path. + """ + if precision_configs.get("fp8_recipe") == "mxfp8" and ( + not torch.cuda.is_available() + or torch.cuda.get_device_capability()[0] < 10 + or not HAVE_TE_MXFP8TENSOR + ): + pytest.skip("Requires PyTorch & CUDA device with TE MXFP8Tensor support") + if Utils.world_size < 2: + pytest.skip("Requires at least 2 distributed ranks for ZeRO sharding") + + outputs = TestMegatronFSDPE2E._training_loop( + use_megatron_fsdp=True, + use_megatron_fsdp_v2=True, + ckpt_format="fsdp_dtensor", + data_parallel_sharding_strategy=strategy, + train_iters=3, + seq_length=64, + micro_batch_size=1, + global_batch_size=8, + init_model_with_meta_device=False, + gradient_accumulation_fusion=False, + overlap_param_gather=False, + overlap_grad_reduce=False, + **precision_configs, + ) + + if torch.distributed.get_rank() != 0: + return + + assert len(outputs) == 3 + for step, output in enumerate(outputs): + loss = output["lm loss"] + assert torch.isfinite(loss), ( + f"Non-finite loss at step {step}, strategy={strategy}, " + f"precision={precision_configs}" + ) + def compare_losses(loss_a: float, loss_b: float, reference: str = "b"): """ diff --git a/tests/unit_tests/distributed/megatron_fsdp/v2/test_param_group.py b/tests/unit_tests/distributed/megatron_fsdp/v2/test_param_group.py index ead8dca61f1..dfd0372413b 100644 --- a/tests/unit_tests/distributed/megatron_fsdp/v2/test_param_group.py +++ b/tests/unit_tests/distributed/megatron_fsdp/v2/test_param_group.py @@ -31,7 +31,7 @@ sys.path.insert(0, str(Path(__file__).parents[2])) from megatron.core.distributed.fsdp.src.megatron_fsdp.v2.mixed_precision import ( - FullyShardMixedPrecisionPolicy, + MixedPrecisionPolicy, ) from megatron.core.distributed.fsdp.src.megatron_fsdp.v2.param_group import ParameterGroup from megatron.core.distributed.fsdp.src.megatron_fsdp.v2.utils import ParamGroupIdx @@ -108,7 +108,7 @@ def _build_groups(strategy): pg = ParameterGroup( params=params, param_group_id=ParamGroupIdx(0, gid), - mp_policy=FullyShardMixedPrecisionPolicy(), + mp_policy=MixedPrecisionPolicy(), mesh=None, sharding_strategy=strategy, ) @@ -228,12 +228,11 @@ def test_unshard_reshard(strategy): continue shard_before = wbuf.data.clone() - unsharded, work = wbuf.unshard(async_op=False) + unsharded = wbuf.unshard() if not w_dist: # Non-distributed: unshard returns self.data directly, no comm assert unsharded is wbuf.data - assert work is None else: # Distributed: after all-gather, every param should be fully # recoverable from the unsharded buffer at its global offset @@ -272,29 +271,31 @@ def test_reduce_grad(strategy): # uint8 group has requires_grad=False, so no grad buffer continue - if not g_dist: - # Non-distributed: each rank fills with (rank+1), then all-reduce - # should produce sum across all ranks + if strategy == "no_shard": + # No-shard: each rank fills with (rank+1), then all-reduce should + # produce the sum across all ranks. gbuf.data.fill_(float(rank + 1)) ref = torch.full_like(gbuf.data, float(rank + 1)) Ref.all_reduce(ref, dp_group) gbuf.reduce_grad() assert torch.equal(gbuf.data, ref) else: - # Distributed: fill a full-size temp buffer with (rank+1), - # reduce-scatter it, compare against PyTorch ref + # ZeRO-1/2/3: reduce-scatter a full gradient buffer and compare + # this rank's optimizer-facing shard against the PyTorch reference. full_size = gbuf.buffer_index.bucket_meta.size full = torch.full((full_size,), float(rank + 1), dtype=gbuf.dtype, device=device) - # Pre-populate the allocator so reduce_grad sees the data - bucket = gbuf.allocator.allocate( - key=gbuf.alloc_key, size=full_size, dtype=gbuf.dtype, device=device - ) - bucket.data.copy_(full) - ref_shard = Ref.reduce_scatter(full.clone(), dp_group) - gbuf.data.zero_() + if g_dist: + # Pre-populate the allocator so reduce_grad sees the full temp buffer. + bucket = gbuf.allocator.allocate( + key=gbuf.alloc_key, size=full_size, dtype=gbuf.dtype, device=device + ) + bucket.data.copy_(full) + gbuf.data.zero_() + else: + gbuf.data.copy_(full) gbuf.reduce_grad() # Only compare the shard region of self.data From 2874c91bd21ca6a68681d8c98c489879aaa5668c Mon Sep 17 00:00:00 2001 From: jianbinc Date: Sun, 31 May 2026 21:45:25 +0800 Subject: [PATCH 07/28] feat(mfsdp-v2): sync module states after checkpoint load, harden uneven 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 --- megatron/core/distributed/fsdp/checkpoint.py | 4 +- .../fsdp/src/megatron_fsdp/uneven_dtensor.py | 11 +- .../fsdp/src/megatron_fsdp/v2/README.md | 2 - .../fsdp/src/megatron_fsdp/v2/fsdp_module.py | 11 + megatron/training/checkpointing.py | 8 + tools/checkpoint/checkpoint_inspector.py | 262 +++++++++++++++--- 6 files changed, 245 insertions(+), 53 deletions(-) diff --git a/megatron/core/distributed/fsdp/checkpoint.py b/megatron/core/distributed/fsdp/checkpoint.py index 61d60220724..98059f1bf12 100644 --- a/megatron/core/distributed/fsdp/checkpoint.py +++ b/megatron/core/distributed/fsdp/checkpoint.py @@ -285,7 +285,7 @@ def _split_fused_params_v2( optimizer_state_dict: Optional optimizer state dict (a copy is returned). detector(key, dtensor, model) -> ``(sizes, names, dim)`` or ``None``. key_fmt(key, sub_name) -> new_key. - tag: Log tag (e.g. ``"SwiGLU v2"``). + tag: Log tag (e.g. ``"SwiGLU"``). state_keys: Optimizer state keys to split (default: ``exp_avg``, ``exp_avg_sq``). Returns: @@ -395,7 +395,7 @@ def detector(key, dtensor, _model): optimizer_state_dict, detector, lambda k, s: f"{k}_{s}", - "SwiGLU v2", + "SwiGLU", ) diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/uneven_dtensor.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/uneven_dtensor.py index d42c10494e7..1efdeec2376 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/uneven_dtensor.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/uneven_dtensor.py @@ -395,8 +395,7 @@ def split_dtensor( Split a DTensor into smaller DTensors along a specified dimension. This function handles uneven sharding correctly by computing chunk metadata - for every output DTensor without extra collective operations. A single call - to ``gather_and_compute_chunk_metadata`` is made upfront; all subsequent + for every output DTensor without extra collective operations. All subsequent per-split metadata is derived locally from that result using pure integer arithmetic. @@ -491,10 +490,10 @@ def split_dtensor( split_points.append(split_points[-1] + size) # One collective call — result reused for all splits below. - if not hasattr(dtensor._local_tensor, "__create_chunk_list__"): - chunk_meta = gather_and_compute_chunk_metadata(dtensor) - else: - chunk_meta = dtensor.__create_chunk_list__()[0] + assert hasattr(dtensor._local_tensor, "__create_chunk_list__"), ( + "DTensor local tensor is missing chunk metadata." + ) + chunk_meta = dtensor._local_tensor.__create_chunk_list__()[0] chunk_slice = slice(chunk_meta.offsets[dim], chunk_meta.offsets[dim] + chunk_meta.sizes[dim]) local_offset = chunk_meta.offsets[dim] local_tensor = dtensor.to_local() diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/README.md b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/README.md index 493a666af9c..3fb04b2f2f4 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/README.md +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/README.md @@ -134,8 +134,6 @@ See the parent directory `..` for `uneven_dtensor.py` which provides: See [tp_support_design.md](tp_support_design.md) for the planned design. - **Hybrid Sharding (HSDP):** Not supported. v2 does not yet support an outer DP dimension for hybrid (inter-node + intra-node) sharding. -- **Expert Parallelism (EP):** Not supported in v2's param group logic. MoE - expert parameters are not currently handled via EP submeshes. ### Sharding Strategies diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/fsdp_module.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/fsdp_module.py index c38a041cf32..884c51d5275 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/fsdp_module.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/fsdp_module.py @@ -162,6 +162,10 @@ def get_prefetch_next_modules( raise AssertionError("Current module not found in forward module order") + def get_root_module(self): + """Return the root FSDP module associated with this context.""" + return self.forward_order[0] if self.forward_order else None + class FSDPModule(nn.Module): """ @@ -817,6 +821,13 @@ def _set_nan_check(self, enable_nan_checks: bool): param_data ).any(), "NaN detected in model weight buffer" + def get_root_module(self): + """Return the root FSDP module associated with this module.""" + return self._fsdp_root_context.get_root_module() + + def _sync_module_states_after_load(self): + self._copy_main_weights_to_model_weights() + def _get_module_fsdp_param_groups( module: nn.Module, diff --git a/megatron/training/checkpointing.py b/megatron/training/checkpointing.py index feb592c7e79..1324782e17e 100644 --- a/megatron/training/checkpointing.py +++ b/megatron/training/checkpointing.py @@ -2223,6 +2223,14 @@ def load_model_state_dict(module, state_dict, strict: bool): load_return = module.load_state_dict(state_dict, strict=False) print(f"load_return: {load_return}") + # Megatron FSDP v2 checkpoints require an extra step to sync module states after loading. + if _is_megatron_fsdp_v2(module): + from megatron.core.distributed.fsdp.src.megatron_fsdp.v2 import FSDPModule + for m in module.modules(): + if isinstance(m, FSDPModule): + root_module = m.get_root_module() + root_module._sync_module_states_after_load() + # Model. if not skip_load_to_model_and_opt: if len(ddp_model) == 1: diff --git a/tools/checkpoint/checkpoint_inspector.py b/tools/checkpoint/checkpoint_inspector.py index 7c946cb5ac3..fd9d5cbe0c5 100644 --- a/tools/checkpoint/checkpoint_inspector.py +++ b/tools/checkpoint/checkpoint_inspector.py @@ -218,7 +218,32 @@ def print_tensor(checkpoint_dir, key): torch.distributed.checkpoint.load( state_dict, storage_reader=reader, planner=DefaultLoadPlanner() ) - print(state_dict, state_dict[key].shape, state_dict[key]._local_tensor.shape) + + dtensor = state_dict[key] + local_data = dtensor._local_tensor.float() + local_norm = torch.norm(local_data) + local_norm_sq = local_norm.pow(2) + if dist.is_initialized() and dist.get_world_size() > 1: + global_norm_sq = local_norm_sq.clone() + dist.all_reduce(global_norm_sq, op=dist.ReduceOp.SUM) + global_norm = torch.sqrt(global_norm_sq) + else: + global_norm = local_norm + + click.echo(click.style(f"L2 Norm: {global_norm.item():.6f}", fg="green")) + click.echo(click.style(f" Global shape: {dtensor.shape}", fg="white")) + click.echo(click.style(f" Local shape: {dtensor._local_tensor.shape}", fg="white")) + + click.echo(click.style(f"\nLocal shard data (rank {dist.get_rank()}):", fg="cyan")) + if local_data.numel() <= 100: + click.echo(local_data) + else: + flat = local_data.flatten() + half = flat.numel() // 2 + click.echo(f" First {half}: {flat[:half]}") + click.echo(f" Last {flat.numel() - half}: {flat[half:]}") + + dist.destroy_process_group() def check_gpu_memory(threshold=0.9): @@ -1002,82 +1027,233 @@ def modify_state_dict(input_dir, output_dir, op, enable_msc): ) -def _compare_two_checkpoint(checkpoint_1, checkpoint_2): +def _compare_two_checkpoint( + checkpoint_1, + checkpoint_2, + memory_threshold: float = 0.7, + batch_max_bytes: int = 4 * 1024**3, + atol: float = 1e-8, + rtol: float = 1e-5, +): + """Compare two distributed checkpoints with batched loading and memory awareness. + + Parameters + ---------- + checkpoint_1, checkpoint_2 : Path + Paths to the two checkpoint directories. + memory_threshold : float + GPU memory utilisation ratio (0–1) that triggers CPU offload. + Default 0.7 (70 %). + batch_max_bytes : int + Maximum number of bytes to load in a single dcp.load call. + Default 4 GiB. + atol, rtol : float + Absolute / relative tolerance for torch.allclose. + """ + device_mesh = DeviceMesh.from_group(dist.group.WORLD, device_type="cuda") + reader_1 = FileSystemReader(checkpoint_1) metadata_1 = reader_1.read_metadata() - reader_2 = FileSystemReader(checkpoint_2) metadata_2 = reader_2.read_metadata() keys_1 = set(metadata_1.state_dict_metadata.keys()) keys_2 = set(metadata_2.state_dict_metadata.keys()) - click.echo(click.style("Comparing checkpoints...", fg="blue")) - - # Compare keys - missing_in_1 = keys_2 - keys_1 - missing_in_2 = keys_1 - keys_2 - common_keys = keys_1 & keys_2 + rank0_echo(f"Checkpoint 1: {len(keys_1)} keys") + rank0_echo(f"Checkpoint 2: {len(keys_2)} keys") - click.echo(click.style("Keys missing in checkpoint 1:", fg="red")) - for key in missing_in_1: - click.echo(click.style(f" - {key}", fg="red")) + # Phase 1 — key set comparison (no I/O) + _report_keys("Keys only in checkpoint 1", keys_1 - keys_2, "yellow") + _report_keys("Keys only in checkpoint 2", keys_2 - keys_1, "yellow") - click.echo(click.style("Keys missing in checkpoint 2:", fg="red")) - for key in missing_in_2: - click.echo(click.style(f" - {key}", fg="red")) + # Collect tensor keys with matching metadata for batch load + common_keys = sorted(keys_1 & keys_2) + metadata_mismatches = [] + batch_candidates = [] - # Compare common keys - click.echo(click.style("Common keys in both checkpoints:", fg="green")) + total_bytes = 0 for key in common_keys: - meta_1 = metadata_1.state_dict_metadata[key] - meta_2 = metadata_2.state_dict_metadata[key] + m1 = metadata_1.state_dict_metadata[key] + m2 = metadata_2.state_dict_metadata[key] - if not isinstance(meta_1, TensorStorageMetadata): + if not isinstance(m1, TensorStorageMetadata): + continue + if not isinstance(m2, TensorStorageMetadata): continue - if meta_1.size != meta_2.size or meta_1.properties.dtype != meta_2.properties.dtype: - click.echo( - click.style( - f" - {key} (metadata differ) meta_1: {meta_1}, meta_2: {meta_2}", fg="red" - ) + if m1.size != m2.size or m1.properties.dtype != m2.properties.dtype: + metadata_mismatches.append((key, m1, m2)) + continue + + size_bytes = m1.size.numel() * m1.properties.dtype.itemsize + total_bytes += size_bytes + batch_candidates.append((key, m1.size, m1.properties.dtype, size_bytes)) + + if metadata_mismatches: + rank0_echo(f"[Metadata mismatch] {len(metadata_mismatches)} keys:") + for key, m1, m2 in metadata_mismatches[:20]: + rank0_echo( + f" {key} | ckpt1: {m1.size} {m1.properties.dtype}" + f" | ckpt2: {m2.size} {m2.properties.dtype}" ) - else: - value_1 = torch.empty(meta_1.size, dtype=meta_1.properties.dtype) - value_2 = value_1.clone() + if len(metadata_mismatches) > 20: + rank0_echo(f" ... and {len(metadata_mismatches) - 20} more") - dcp.load({key: value_1}, storage_reader=reader_1, planner=DefaultLoadPlanner()) - dcp.load({key: value_2}, storage_reader=reader_2, planner=DefaultLoadPlanner()) + rank0_echo( + f"[Phase 2] Load & compare {len(batch_candidates)} matching tensors " + f"({total_bytes / 1e9:.2f} GB total across ranks)" + ) - if not torch.allclose(value_1, value_2, atol=1e-8, rtol=1e-5): - click.echo( - click.style( - f" - {key} (values differ) value_1: {value_1}, value_2: {value_2}", fg="red" + # Phase 2 — batched load & compare + mismatched_keys = [] + matched_count = 0 + skipped_non_recv = 0 + batch_start = 0 + batch_idx = 0 + + while batch_start < len(batch_candidates): + batch_bytes = 0 + batch_end = batch_start + while batch_end < len(batch_candidates): + if batch_bytes + batch_candidates[batch_end][3] > batch_max_bytes and batch_end > batch_start: + break + batch_bytes += batch_candidates[batch_end][3] + batch_end += 1 + + batch_keys = batch_candidates[batch_start:batch_end] + rank0_echo( + f" Batch {batch_idx} | {batch_start}-{batch_end - 1} " + f"({len(batch_keys)} keys, {batch_bytes / 1e9:.2f} GB)" + ) + + # Build state dicts for this batch + sd_1 = {} + sd_2 = {} + for key, shape, dtype, _ in batch_keys: + empty_tensor_1 = torch.distributed.tensor.empty( + shape, dtype=dtype, device_mesh=device_mesh, placements=[Shard(0)] + ) + empty_tensor_2 = torch.distributed.tensor.empty( + shape, dtype=dtype, device_mesh=device_mesh, placements=[Shard(0)] + ) + sd_1[key] = empty_tensor_1 + sd_2[key] = empty_tensor_2 + + # Load both checkpoints in one call each + dcp.load(sd_1, storage_reader=reader_1, planner=DefaultLoadPlanner()) + dcp.load(sd_2, storage_reader=reader_2, planner=DefaultLoadPlanner()) + + # Compare and count mismatches per key + for key, shape, dtype, _ in batch_keys: + v1 = sd_1[key] + v2 = sd_2[key] + + if v1._local_tensor.numel() == 0 or v2._local_tensor.numel() == 0: + skipped_non_recv += 1 + continue + + if not torch.allclose(v1._local_tensor, v2._local_tensor, atol=atol, rtol=rtol): + diff = (v1._local_tensor.float() - v2._local_tensor.float()).abs() + mismatched_keys.append( + ( + key, + diff.max().item(), + diff.mean().item(), + (diff > atol + rtol * v2._local_tensor.float().abs()).sum().item(), ) ) + else: + matched_count += 1 + + # Free batch tensors and manage GPU memory + del sd_1, sd_2 + if check_gpu_memory(memory_threshold): + gc.collect() + torch.cuda.empty_cache() + + batch_start = batch_end + batch_idx += 1 + + # Phase 3 — report + if skipped_non_recv > 0: + rank0_echo(f"[Skipped] {skipped_non_recv} keys with empty local shard on this rank") + rank0_echo( + f"[Result] {matched_count} matched, {len(mismatched_keys)} mismatched", + ) + if metadata_mismatches: + rank0_echo(f"[Result] {len(metadata_mismatches)} metadata mismatches") + + if mismatched_keys: + rank0_echo("Mismatched keys (max_abs_err | mean_abs_err | elems_mismatched):") + for key, max_err, mean_err, n_mis in mismatched_keys: + rank0_echo(f" {key} | {max_err:.6e} | {mean_err:.6e} | {n_mis}") + else: + rank0_echo("All matching tensors are bitwise-identical within tolerance.") + + +def _report_keys(title, keys, color): + if not keys: + return + rank0_echo(f"[{title}] {len(keys)}:") + for k in sorted(keys)[:20]: + rank0_echo(f" {k}") + if len(keys) > 20: + rank0_echo(f" ... and {len(keys) - 20} more") @cli.command() @click.argument("checkpoint_1", type=click.Path(exists=True)) @click.argument("checkpoint_2", type=click.Path(exists=True)) @click.option("--enable-msc", is_flag=True, help="Enable MultiStorageClient feature.") -def compare_two_checkpoint(checkpoint_1, checkpoint_2, enable_msc): +@click.option( + "--memory-threshold", + type=float, + default=0.7, + help="GPU memory utilisation ratio (0-1) that triggers CPU offload. Default 0.7.", +) +@click.option( + "--batch-max-bytes", + type=int, + default=4 * 1024**3, + help="Max bytes to load per batch. Default 4 GiB.", +) +@click.option("--atol", type=float, default=1e-8, help="Absolute tolerance for allclose.") +@click.option("--rtol", type=float, default=1e-5, help="Relative tolerance for allclose.") +def compare_two_checkpoint( + checkpoint_1, + checkpoint_2, + enable_msc, + memory_threshold, + batch_max_bytes, + atol, + rtol, +): """ - Compare two checkpoints. + Compare two distributed checkpoints with batched, memory-aware loading. + + Supports torchrun for DTensor (sharded) checkpoints: + + \\b + torchrun --nproc_per_node=8 checkpoint_inspector.py \\ + compare-two-checkpoint /ckpt/a /ckpt/b """ init_process_group(f"compare_two_checkpoint from {checkpoint_1} to {checkpoint_2}") if not enable_msc: MultiStorageClientFeature.disable() - _compare_two_checkpoint(Path(checkpoint_1), Path(checkpoint_2)) + _compare_two_checkpoint( + Path(checkpoint_1), + Path(checkpoint_2), + memory_threshold=memory_threshold, + batch_max_bytes=batch_max_bytes, + atol=atol, + rtol=rtol, + ) - click.echo( - click.style( - f"Comparison between {checkpoint_1} and {checkpoint_2} completed.", - fg="green", - bold=True, - ) + rank0_echo( + f"Comparison between {checkpoint_1} and {checkpoint_2} completed.", ) From 44565574d26a80998b6034cb5f8fe21c6f112cce Mon Sep 17 00:00:00 2001 From: tongliu Date: Wed, 3 Jun 2026 03:16:05 -0700 Subject: [PATCH 08/28] Add Megatron FSDP v2 Muon optimizer --- .../fsdp/src/megatron_fsdp/v2/dp_buffer.py | 74 ++++ .../fsdp/src/megatron_fsdp/v2/param_group.py | 138 +++++++ megatron/core/optimizer/__init__.py | 146 +++++++- .../core/optimizer/fully_shard_v2_muon.py | 354 ++++++++++++++++++ megatron/training/arguments.py | 21 +- .../megatron_fsdp/v2/test_dist_muon.py | 224 +++++++++++ 6 files changed, 950 insertions(+), 7 deletions(-) create mode 100644 megatron/core/optimizer/fully_shard_v2_muon.py create mode 100644 tests/unit_tests/distributed/megatron_fsdp/v2/test_dist_muon.py diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/dp_buffer.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/dp_buffer.py index 2010f873d05..c118051666f 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/dp_buffer.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/dp_buffer.py @@ -61,6 +61,79 @@ def __init__( self.shard_meta = self._build_shard_meta( self.bucket_meta, is_distributed, dp_world_size, dp_rank ) + self.item_root_ranks = self._build_item_root_ranks() + + def _build_item_root_ranks(self) -> Dict[int, int]: + """Assign each item a root rank, greedily balancing Newton-Schulz load. + + Only 2D matrices run NS (on their root rank), and that cost dominates the + single-root Muon step. NS flops scale ~ ``numel * min(dim)`` per matrix + (constant #steps dropped). We use Longest-Processing-Time greedy: assign + items heaviest-first to the currently least-loaded *holder* rank. + + Properties: + * Items with a single holder (not split across a shard boundary) are + forced onto their owner -> zero gather/scatter, comm unchanged. + * Only split items have a real choice; among holders we pick the + least-loaded, tie-broken toward the largest holder so an indifferent + balance still minimises the gather/scatter volume (old behaviour). + * Non-2D items get cost 0 (their root is never used by Muon). + + Must be deterministic and identical on every rank (gather/scatter must + agree on the root): inputs (param_shapes, shard_size, dp_world_size) are + rank-independent, and the sort/min tie-breaks are fully ordered by item_id. + """ + shard_size = self.bucket_meta.size // self.dp_world_size + load = [0] * self.dp_world_size + root_ranks: Dict[int, int] = {} + split_items = [] + + # Pass 1: a param living entirely on one rank is rooted there -- no choice, + # zero gather/scatter. Account its NS cost first so the split params in + # pass 2 balance on top of the complete fixed (non-split) load per rank. + for item_id, idx in self.item_index_map.items(): + item_start = idx.global_data_index + item_end = item_start + idx.size + overlaps = [ + max(0, min(item_end, (r + 1) * shard_size) - max(item_start, r * shard_size)) + for r in range(self.dp_world_size) + ] + holders = [r for r, o in enumerate(overlaps) if o > 0] + ns_cost = idx.size * min(idx.shape) if len(idx.shape) == 2 else 0 + if len(holders) == 1: + root_ranks[item_id] = holders[0] + load[holders[0]] += ns_cost + else: + split_items.append((ns_cost, item_id, holders, overlaps)) + + # Pass 2: only split params have a choice. LPT (heaviest first; ties by + # item_id for cross-rank determinism) -> assign each to its least-loaded + # holder (ties -> largest holder = min gather/scatter volume -> lowest rank). + for ns_cost, item_id, holders, overlaps in sorted(split_items, key=lambda e: (-e[0], e[1])): + best = min(holders, key=lambda r: (load[r], -overlaps[r], r)) + root_ranks[item_id] = best + load[best] += ns_cost + + # TODO(muon-debug): temporary one-time dump of the NS root assignment so we + # can see whether split-param roots actually land on different ranks. Drop + # once the load-balance behaviour is confirmed. + if self.dp_rank == 0: + twod = [iid for iid, idx in self.item_index_map.items() if len(idx.shape) == 2] + if twod: + roots_per_rank = [0] * self.dp_world_size + for iid in twod: + roots_per_rank[root_ranks[iid]] += 1 + split_dump = [ + (iid, tuple(self.item_index_map[iid].shape), holders, root_ranks[iid]) + for _, iid, holders, _ in sorted(split_items, key=lambda e: (-e[0], e[1])) + ] + logger.warning( + "[muon-root] pg=%s ws=%d: %d 2D items (%d split). " + "NS-load/rank=%s 2D-roots/rank=%s | split(item,shape,holders,root)=%s", + self.param_group_id, self.dp_world_size, len(twod), len(split_items), + load, roots_per_rank, split_dump, + ) + return root_ranks # ------------------------------------------------------------------ # # Layout construction (global, rank-independent) @@ -204,6 +277,7 @@ def pack_fragments( size=_pad_if_needed(data_index), items=list(item_index_map.values()), ) + return item_index_map, bucket_meta # ------------------------------------------------------------------ # diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/param_group.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/param_group.py index 6ed8455b660..eb39c469131 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/param_group.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/param_group.py @@ -82,6 +82,8 @@ def __init__( assert mesh.ndim == 1, "Only 1D mesh is supported" self.mesh = mesh self.dp_group = mesh.get_group() + self._dp_rank = torch.distributed.get_rank(self.dp_group) + self._dp_world_size = torch.distributed.get_world_size(self.dp_group) if sharding_strategy == "no_shard": raise NotImplementedError( @@ -277,6 +279,90 @@ def release_grad_buffer(self): del param.main_grad self.main_grad_buffer.reshard() + @torch.no_grad() + def unshard_grad_to_root_async(self, param): + """Gather this param's grad to its assigned root rank (async, non-blocking). + + All ranks must call this collectively (after reduce_grad), in the same param + order. The cross-rank gather is issued as non-blocking P2P + (``batch_isend_irecv``), which runs on the NCCL backend's own stream and so + overlaps the caller's compute; the root's own-piece copy is a local memcpy on + the current stream. Returns ``(full_grad_or_None, reqs)`` — the full grad on + root (valid only after ``reqs`` complete, so the caller waits them before NS), + None on others; the returned tensors must stay alive until ``reqs`` complete. + """ + idx = self.param_idx[param] + root, src, dst = self._grad_gather_plans[idx] + gbuf = self.main_grad_buffer + ops = [] + if self._dp_rank == root: + full_grad = torch.empty(param.shape, dtype=gbuf.dtype, device=gbuf.device).flatten() + for r in range(self._dp_world_size): + src_lo, src_hi = src[r] + dst_lo, dst_hi = dst[r] + if dst_hi <= dst_lo: + continue # rank r holds no part of this param + if r == root: + full_grad[dst_lo:dst_hi].copy_(gbuf.data[src_lo:src_hi]) # local + else: + ops.append( + torch.distributed.P2POp( + torch.distributed.irecv, full_grad[dst_lo:dst_hi], r, self.dp_group + ) + ) + reqs = torch.distributed.batch_isend_irecv(ops) if ops else [] + return full_grad.view(param.shape), reqs + else: + src_lo, src_hi = src[self._dp_rank] + if src_hi <= src_lo: + return None, [] + ops.append( + torch.distributed.P2POp( + torch.distributed.isend, gbuf.data[src_lo:src_hi], root, self.dp_group + ) + ) + return None, torch.distributed.batch_isend_irecv(ops) + + @torch.no_grad() + def scatter_grad_from_root_async(self, param, data=None): + """Scatter this param's transformed grad from root back (async, non-blocking). + + All ranks must call this collectively, in the same param order. P2P runs on + the NCCL backend's stream (overlaps compute). Root copies its own piece + locally and isends the rest; each non-root holder irecvs its piece straight + into its grad-buffer shard (safe: the matching gather isend that read that + shard was enqueued earlier on the same NCCL stream, so it completes before + this recv writes). Returns the in-flight P2P Works (empty if uninvolved). + """ + idx = self.param_idx[param] + root, src, dst = self._grad_gather_plans[idx] + gbuf = self.main_grad_buffer + ops = [] + if self._dp_rank == root: + full = data.to(dtype=gbuf.dtype).flatten() + for r in range(self._dp_world_size): + src_lo, src_hi = src[r] + dst_lo, dst_hi = dst[r] + if dst_hi <= dst_lo: + continue # rank r holds no part of this param + if r == root: + gbuf.data[src_lo:src_hi].copy_(full[dst_lo:dst_hi]) # local + else: + ops.append( + torch.distributed.P2POp( + torch.distributed.isend, full[dst_lo:dst_hi], r, self.dp_group + ) + ) + else: + src_lo, src_hi = src[self._dp_rank] + if src_hi > src_lo: + ops.append( + torch.distributed.P2POp( + torch.distributed.irecv, gbuf.data[src_lo:src_hi], root, self.dp_group + ) + ) + return torch.distributed.batch_isend_irecv(ops) if ops else [] + def _init_dist_params(self): """ Initialize distributed parameter views (DTensors) into the buffers. @@ -318,6 +404,11 @@ def _init_dist_params(self): # Mark as FSDP parameter for special handling setattr(param, "__fsdp_param__", True) setattr(dist_param, "__fsdp_param__", True) + # Back-references so an optimizer that only sees the optimizer-facing + # dist_param (e.g. FullyShardV2Muon) can reach this ParameterGroup and + # the original param it indexes (for per-param gather/scatter). + dist_param._fsdp_param_group = self + dist_param._fsdp_orig_param = param self.dist_params.append(dist_param) # Update dist_param chunk metadata for checkpointing and debugging. @@ -345,3 +436,50 @@ def _init_dist_params(self): ) else: self.dist_grads.append(None) + + # Per-param send/recv plans from BufferIndex (identical on all ranks, so + # no metadata comm). Entry is (root, src, dst): + # root — DP rank that assembles the full grad. + # src[r] — (lo, hi) of rank r's piece in its LOCAL grad buffer. + # dst[r] — (lo, hi) of that piece in the ASSEMBLED full grad on root. + # Same elements live at gbuf.data[src[r]] (rank r) and full_grad[dst[r]] + # (root); unshard moves src->dst, scatter moves dst->src. Empty range + # (lo==hi) => rank r owns nothing (runtime skips via lo (lo-shard_start, hi-shard_start) (gbuf.data is the shard); + # replicated buffer (ZeRO-1) -> (lo, hi) (gbuf.data is the full buffer). + # Only src[self._dp_rank] ever indexes gbuf.data. + # + # Example (distributed): 4 ranks, shard_size=100, param [150,380). + # rank 0: shard [0,100) ∩ param = ∅ -> src=(0,0) dst=(0,0) + # rank 1: shard [100,200) ∩ [150,200) -> src=(50,100) dst=(0,50) + # rank 2: shard [200,300) ∩ [200,300) -> src=(0,100) dst=(50,150) + # rank 3: shard [300,400) ∩ [300,380) -> src=(0,80) dst=(150,230) + # (rank 1's src starts at 50 because the param begins 50 into + # its shard; its dst starts at 0 because it is the param's head.) + bi = self.main_grad_buffer.buffer_index + is_distributed = self.main_grad_buffer.is_distributed + shard_size = bi.bucket_meta.size // self._dp_world_size + self._grad_gather_plans: List[tuple] = [] + for idx in range(len(self.params)): + item = bi.item_index_map[idx] + root = bi.item_root_ranks[idx] + param_start = item.global_data_index + param_end = param_start + item.size + + src, dst = [], [] + for r in range(self._dp_world_size): + shard_start = r * shard_size + shard_end = shard_start + shard_size + lo = max(param_start, shard_start) + hi = min(param_end, shard_end) + if hi <= lo: + src.append((0, 0)) + dst.append((0, 0)) + else: + if is_distributed: + src.append((lo - shard_start, hi - shard_start)) + else: + src.append((lo, hi)) + dst.append((lo - param_start, hi - param_start)) + self._grad_gather_plans.append((root, src, dst)) diff --git a/megatron/core/optimizer/__init__.py b/megatron/core/optimizer/__init__.py index 5787049cd8c..7a9d10fb2d3 100644 --- a/megatron/core/optimizer/__init__.py +++ b/megatron/core/optimizer/__init__.py @@ -856,6 +856,129 @@ def _get_megatron_emerging_optimizer( return ChainedOptimizer(results) +def _build_megatron_fsdp_v2_muon_optimizer( + config: OptimizerConfig, + model_chunks: List[MegatronModule], + all_dense_model_chunks: List, + config_overrides: Dict, + mp_group, + dp_cp_group, + dp_cp_group_gloo, + data_parallel_group_idx, + intra_dist_opt_group, + distributed_optimizer_instance_id, + pg_collection, +) -> MegatronOptimizer: + """Build Muon (2D matrices) + Adam (rest) for Megatron-FSDP v2. + + EXPERIMENTAL / UNVALIDATED (Phase 3). Routes 2D matrix params to the + single-root :class:`FullyShardV2Muon` and every other param to the standard + FSDP Adam, combined via :class:`ChainedOptimizer`. Called only from + ``get_megatron_optimizer`` when ``optimizer='muon'`` + ``use_megatron_fsdp``, + so no other optimizer path is affected. The Muon/Adam split reuses the same + ``_is_nonlinear_or_embedding`` config override as the non-FSDP emerging path. + + Needs validation on a GPU node (grad-clip semantics for Muon, momentum + checkpointing, and passing full FSDP buffers to an Adam built over only the + non-matrix param subset). + """ + from .fully_shard_v2_muon import FullyShardV2Muon, FullyShardV2MuonOptimizer + + # Tag params and route non-linear/embedding params to Adam (same mechanism + # as the non-FSDP emerging path in _get_megatron_emerging_optimizer). + # _default_param_overrides_factory is defined regardless of whether the + # emerging_optimizers package is installed (FullyShardV2Muon has a built-in + # Newton-Schulz fallback), so we don't index _EMERGING_OPTIMIZERS here. + from .emerging_optimizers import _default_param_overrides_factory + + for model_chunk in model_chunks: + for name, param in model_chunk.named_parameters(): + if not param.requires_grad: + continue + if 'experts' in name and 'shared' not in name: + param.expert_tp = True + if 'linear_qkv.weight' in name and len(param.shape) == 2: + param.is_qkv = True + config_overrides.update(_default_param_overrides_factory()) + + muon_kwargs = dict( + lr=config.lr, + momentum=getattr(config, 'muon_momentum', 0.95), + nesterov=getattr(config, 'muon_nesterov', False), + weight_decay=config.weight_decay, + num_ns_steps=getattr(config, 'muon_num_ns_steps', 5), + coefficient_type=getattr(config, 'muon_coefficient_type', 'quintic'), + scale_mode=getattr(config, 'muon_scale_mode', 'spectral'), + extra_scale_factor=getattr(config, 'muon_extra_scale_factor', 1.0), + fp32_matmul_prec=getattr(config, 'muon_fp32_matmul_prec', 'medium'), + # TODO(phase3): wire pg_collection/tp_mode for tensor-parallel NS (Phase 2). + pg_collection=None, + ) + + adam_config = copy.copy(config) + adam_config.optimizer = 'adam' + + optimizers: List[MegatronOptimizer] = [] + model_chunk_offset = 0 + for chunk_list in all_dense_model_chunks: + # --- Adam for the non-matrix params (standard FSDP path) --- + adam_groups, buffers = _get_param_groups_and_buffers( + chunk_list, + model_chunk_offset=model_chunk_offset, + config=adam_config, + config_overrides=config_overrides, + filter_fn=lambda g: g.get('optimizer') == 'adam', + buffer_name='buffers', + ) + if adam_groups: + adam_part = _get_megatron_optimizer_based_on_param_groups( + config=adam_config, + model_chunks=chunk_list, + param_groups=adam_groups, + per_model_buffers=buffers, + model_parallel_group=mp_group, + data_parallel_group=dp_cp_group, + data_parallel_group_gloo=dp_cp_group_gloo, + data_parallel_group_idx=data_parallel_group_idx, + intra_dist_opt_group=intra_dist_opt_group, + distributed_optimizer_instance_id=distributed_optimizer_instance_id, + pg_collection=pg_collection, + ) + if ( + not USING_PYTORCH_OPTIMIZER + and config.use_precision_aware_optimizer + and getattr(adam_part.optimizer, "master_weights", None) is not None + ): + setattr(adam_part.optimizer, "master_weights", False) + # ChainedOptimizer asserts sub-optimizers share one config object. + adam_part.config = config + optimizers.append(adam_part) + + # --- Muon for the 2D matrix params (single-root FSDP v2) --- + # Symmetric with Adam: pass the muon bucket of optimizer-facing params + # (dist_params). FullyShardV2Muon resolves each back to its FSDP + # ParameterGroup via the dist_param back-references. + muon_groups, _ = _get_param_groups_and_buffers( + chunk_list, + model_chunk_offset=model_chunk_offset, + config=config, + config_overrides=config_overrides, + filter_fn=lambda g: g.get('optimizer') != 'adam', + buffer_name='buffers', + ) + if muon_groups: + muon = FullyShardV2Muon(muon_groups, **muon_kwargs) + optimizers.append( + FullyShardV2MuonOptimizer(muon, config, model_chunks=chunk_list) + ) + + model_chunk_offset += 1 + + if len(optimizers) == 1: + return optimizers[0] + return ChainedOptimizer(optimizers) + + def get_megatron_optimizer( config: OptimizerConfig, model_chunks: List[MegatronModule], @@ -895,7 +1018,12 @@ def get_megatron_optimizer( # TODO: the standard and emerging optimizer paths handle pg_collection differently; # unify them so both use a single pg_collection-based flow. - if config.optimizer not in ('adam', 'sgd'): + # Megatron-FSDP routes emerging optimizers (e.g. Muon) through its own branch + # below (it needs the FSDP process groups); only the non-FSDP case is handled here. + use_megatron_fsdp = getattr( + getattr(model_chunks[0], 'ddp_config', None), 'use_megatron_fsdp', False + ) + if config.optimizer not in ('adam', 'sgd') and not use_megatron_fsdp: return _get_megatron_emerging_optimizer( config=config, model_chunks=model_chunks, @@ -939,6 +1067,22 @@ def get_megatron_optimizer( model_chunk_offset = 0 ddp_config = model_chunks[0].ddp_config # Use the first model chunk's DDP config if ddp_config.use_megatron_fsdp: + # Emerging optimizers (e.g. Muon): route matrices -> FullyShardV2Muon and + # the rest -> Adam, here where the FSDP process groups are in scope. + if config.optimizer not in ('adam', 'sgd'): + return _build_megatron_fsdp_v2_muon_optimizer( + config=config, + model_chunks=model_chunks, + all_dense_model_chunks=all_dense_model_chunks, + config_overrides=config_overrides, + mp_group=mp_group, + dp_cp_group=dp_cp_group, + dp_cp_group_gloo=intra_dp_cp_group_gloo, + data_parallel_group_idx=model_parallel_rank, + intra_dist_opt_group=intra_dist_opt_group, + distributed_optimizer_instance_id=distributed_optimizer_instance_id, + pg_collection=pg_collection, + ) # For no_shard, gradients are replicated across DP ranks after all-reduce, so grad stats # should only be reduced over TP/PP (model_parallel_group) to avoid inflating the norm. effective_intra_dist_opt_group = ( diff --git a/megatron/core/optimizer/fully_shard_v2_muon.py b/megatron/core/optimizer/fully_shard_v2_muon.py new file mode 100644 index 00000000000..d7120e6a80f --- /dev/null +++ b/megatron/core/optimizer/fully_shard_v2_muon.py @@ -0,0 +1,354 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""DP-distributed Muon optimizer for Megatron-FSDP v2 (single-root, async-pipelined). + +Muon orthogonalizes the gradient of 2D matrix parameters via Newton-Schulz (NS). +NS needs the *full* gradient matrix, but under FSDP the gradient is sharded over +the data-parallel (DP) group. This optimizer reconstructs each parameter's full +gradient on a single, load-balanced "root" DP rank (chosen by the parameter +group's ``BufferIndex.item_root_ranks``), runs NS there *once*, and scatters the +orthogonalized gradient back — rather than all-gathering the full gradient to +every rank and running NS redundantly. Momentum and the weight update stay fully +local (elementwise on each rank's shard), so optimizer state remains sharded. + +Two classes, following the "MegatronOptimizer wraps a torch optimizer" pattern: + +- :class:`FullyShardV2Muon` is a ``torch.optim.Optimizer``. Its ``orthogonalize`` + uses ``emerging_optimizers``' production Newton-Schulz (``newton_schulz_tp`` + + ``get_muon_scale_factor``) WHEN that optional package is installed, and falls + back to a self-contained vanilla quintic NS otherwise — so it runs (and can be + unit-tested) even without the package. Its ``step()`` does the single-root + gather/scatter + sharded momentum/update. +- :class:`FullyShardV2MuonOptimizer` is the thin ``MegatronOptimizer`` adapter + that lets the above slot into ``ChainedOptimizer`` alongside Adam. + +Per managed 2D matrix parameter, each ``step()`` runs three phases: + +1. Local momentum on this rank's gradient shard (sharded optimizer state). +2. Gather the (post-momentum) full gradient to the parameter's root rank + (``ParameterGroup.unshard_grad_to_root_async``). +3a. Root only: ``self.orthogonalize(param, full_grad)`` (NS + scale), then scatter + back (``ParameterGroup.scatter_grad_from_root_async``). +3b. Local elementwise weight update on this rank's master shard. + +The gather/NS/scatter run as an async pipeline on the FSDP reduce-scatter stream: +each parameter's gather is prefetched and the previous parameter's scatter is left +in flight so they overlap the current parameter's Newton-Schulz, while the +collective ops are still issued in the same parameter order on every rank. + +Scope: tensor-parallel size 1 (orthogonalize runs the non-TP NS path). Real +tensor-parallel NS + QKV split and distributed checkpointing of the sharded +momentum are later phases. Needs validation on a GPU node. +""" + +from typing import List + +import torch + +from .emerging_optimizers import HAVE_EMERGING_OPTIMIZERS +from .optimizer import MegatronOptimizer + +if HAVE_EMERGING_OPTIMIZERS: + from .emerging_optimizers import get_muon_scale_factor, newton_schulz_tp + + +@torch.no_grad() +def _vanilla_newton_schulz(grad: torch.Tensor, steps: int = 5, eps: float = 1e-7) -> torch.Tensor: + """Self-contained quintic Newton-Schulz orthogonalization (fp32). + + Fallback used when ``emerging_optimizers`` is not installed; matches the + canonical Muon ``zeropower`` iteration. Operates on the shorter dimension by + transposing when rows > cols. + """ + assert grad.ndim == 2, f"newton_schulz expects a 2D matrix, got {tuple(grad.shape)}" + a, b, c = 3.4445, -4.7750, 2.0315 + x = grad.to(torch.float32) + x = x / (x.norm() + eps) + transposed = x.size(0) > x.size(1) + if transposed: + x = x.t() + for _ in range(steps): + gram = x @ x.t() + x = a * x + (b * gram + c * (gram @ gram)) @ x + if transposed: + x = x.t() + return x + + +def _vanilla_muon_scale(rows: int, cols: int) -> float: + """Muon spectral scale factor ``max(1, rows / cols) ** 0.5`` (fallback).""" + return max(1.0, rows / cols) ** 0.5 + + +class FullyShardV2Muon(torch.optim.Optimizer): + """Single-root distributed Muon for Megatron-FSDP v2 (a torch optimizer). + + Args: + params: optimizer-facing params (FSDP-v2 ``dist_param`` DTensors) or param + groups — the same form every torch/Megatron optimizer takes. Each + dist_param carries ``_fsdp_param_group`` / ``_fsdp_orig_param`` + back-references (set in ``ParameterGroup._init_dist_params``) used to + drive the per-param gather/scatter. + lr / momentum / nesterov / weight_decay: Muon update hyperparameters. + num_ns_steps / coefficient_type / scale_mode / extra_scale_factor: NS + + scaling config used by ``orthogonalize``. + fp32_matmul_prec / pg_collection / tp_mode: reserved for the Phase-2 + tensor-parallel NS path (currently TP size 1; accepted and ignored). + """ + + def __init__( + self, + params, + *, + lr: float, + momentum: float = 0.95, + nesterov: bool = True, + weight_decay: float = 0.01, + num_ns_steps: int = 5, + coefficient_type: str = "quintic", + scale_mode: str = "spectral", + extra_scale_factor: float = 1.0, + fp32_matmul_prec: str = "medium", + pg_collection=None, + tp_mode: str = "duplicated", + ) -> None: + super().__init__(params, dict(lr=lr, weight_decay=weight_decay)) + + # Muon update hyperparameters owned by this class. lr is NOT stored — it + # lives in param_groups[0]["lr"] (where an LR scheduler writes), read in step. + self._momentum_coef = momentum + self._nesterov = nesterov + self._weight_decay = weight_decay + # NS + scaling config for orthogonalize(). + self._num_ns_steps = num_ns_steps + self._coefficient_type = coefficient_type + self._scale_mode = scale_mode + self._extra_scale_factor = extra_scale_factor + + # Resolve the managed 2D matrix params to their FSDP ParameterGroup via + # the dist_param back-references, and PRE-ALLOCATE the (sharded) momentum + # buffer for each one here at init (keeps step() free of allocation / + # lazy-init on the hot path). Iterating self.param_groups in a fixed order + # keeps the collective send/recv in phases 2/3a aligned across ranks. + # Each entry: (ParameterGroup, idx, original param, momentum shard). + self._managed = [] + for group in self.param_groups: + for dist_param in group["params"]: + pg = getattr(dist_param, "_fsdp_param_group", None) + orig = getattr(dist_param, "_fsdp_orig_param", None) + if pg is None or orig is None: + continue # not an FSDP-v2 param + if pg.sharding_strategy == "no_shard": + raise NotImplementedError( + "FullyShardV2Muon does not support the 'no_shard' strategy." + ) + if pg.main_grad_buffer is None: + continue + if ( + orig.requires_grad + and orig.dim() == 2 + and not getattr(orig, "is_embedding_or_output_parameter", False) + ): + idx = pg.param_idx[orig] + grad_shard = pg.main_grad_buffer.get_item(idx, as_shard=True) + momentum_shard = torch.zeros_like(grad_shard) + self._managed.append((pg, idx, orig, momentum_shard)) + + if not self._managed: + raise ValueError("FullyShardV2Muon got no 2D matrix parameters to manage.") + + + @torch.no_grad() + def orthogonalize(self, param, grad: torch.Tensor) -> torch.Tensor: + """Newton-Schulz orthogonalize + Muon spectral scale of a 2D gradient. + + Uses emerging_optimizers' production NS when available, else the built-in + vanilla NS fallback. (Phase 1 runs the non-TP path; ``param`` is accepted + for parity with the eventual TP / QKV-split path.) + """ + g = grad.to(torch.float32) + if HAVE_EMERGING_OPTIMIZERS: + orth = newton_schulz_tp( + g, + steps=self._num_ns_steps, + coefficient_type=self._coefficient_type, + tp_group=None, + partition_dim=None, + tp_mode="duplicated", + ) + scale = get_muon_scale_factor(g.size(-2), g.size(-1), mode=self._scale_mode) + else: + orth = _vanilla_newton_schulz(g, steps=self._num_ns_steps) + scale = _vanilla_muon_scale(g.size(-2), g.size(-1)) + return orth * (scale * self._extra_scale_factor) + + @torch.no_grad() + def step(self, closure=None): + """Run one Muon update over all managed 2D matrix params. + + Per param: (1) local momentum on this rank's grad shard; (2) gather the full + grad to the param's root rank; (3a) root orthogonalizes (NS + scale) and + scatters the result back into every holder's grad shard; (3b) local + decoupled-weight-decay update of this rank's master shard. Momentum/update + stay sharded so optimizer state is sharded. + + The gather/scatter are async (non-blocking ``batch_isend_irecv``); the NCCL + backend runs them on its own stream, so they overlap the Newton-Schulz + compute on the main stream — no explicit stream/event management needed. It + is a 2-stage pipeline: we prefetch the next param's gather so it overlaps the + current param's NS, and leave each scatter in flight so it overlaps the next + param's NS (NS dominates, so one gather in flight is enough to hide the comm); + the scatters are drained at the end. ``work.wait()`` + orders compute-after-comm; the backend orders comm-after-compute (a send + waits for the grad/NS that produced its buffer). P2P is issued in the same + param order on every rank, so it is deadlock-free. + + NOTE: the async pipeline is UNVALIDATED on GPU — needs a multi-rank run to + confirm overlap and numerical parity with a serial reference. + """ + lr = self.param_groups[0]["lr"] # set by the LR scheduler + items = self._managed + n = len(items) + + # Phase 1: in-place momentum on every rank's grad shard (view into gbuf), so + # the gathers below see a fully post-momentum grad. Cheap elementwise. + for pg, idx, _, momentum in items: + grad_shard = pg.main_grad_buffer.get_item(idx, as_shard=True) + if grad_shard.numel() > 0: + momentum.mul_(self._momentum_coef).add_(grad_shard) + if self._nesterov: + grad_shard.add_(momentum, alpha=self._momentum_coef) + else: + grad_shard.copy_(momentum) + + # Phase 2/3a: 2-stage pipeline — prefetch the next param's gather so it + # overlaps the current param's NS; fire each scatter and don't wait, drain + # at the end. + inflight = [] # all in-flight P2P Works + tensors kept alive until drained + + pg, _, param, _ = items[0] + pending = pg.unshard_grad_to_root_async(param) # kick off the first gather + + for i in range(n): + pg, _, param, _ = items[i] + full, gather_reqs = pending + if i + 1 < n: # prefetch next gather before this NS + npg, _, nparam, _ = items[i + 1] + pending = npg.unshard_grad_to_root_async(nparam) + + if full is not None: # this rank is param i's root + for req in gather_reqs: + req.wait() # only the root waits its gather (NS needs it) + orth = self.orthogonalize(param, full).to(full.dtype) + inflight.append((pg.scatter_grad_from_root_async(param, orth), orth, full)) + else: # non-root holder / uninvolved + # Fire-and-forget: its grad isend (gather) + the scatter recv straight + # into its grad shard. Don't block here; drained below. + inflight.append((gather_reqs, None, None)) + inflight.append((pg.scatter_grad_from_root_async(param, None), None, None)) + + # Drain all in-flight P2P. The scatter recvs have written each rank's grad + # shard in place, so Phase 3b can read it directly. + for reqs, _orth, _full in inflight: + for req in reqs: + req.wait() + + # Phase 3b: decoupled-WD update of each rank's master shard from the (now + # orthogonalized) grad shard, then refresh model weights from master. + for pg, idx, _, _ in items: + wbuf = pg.main_weight_buffer if pg.main_weight_buffer is not None else pg.model_weight_buffer + weight_shard = wbuf.get_item(idx, as_shard=True) + if weight_shard.numel() == 0: + continue + orth_shard = pg.main_grad_buffer.get_item(idx, as_shard=True) + if self._weight_decay != 0.0: + weight_shard.mul_(1.0 - lr * self._weight_decay) + weight_shard.add_(orth_shard.to(weight_shard.dtype), alpha=-lr) + + refreshed = set() + for pg, _, _, _ in items: + if id(pg) in refreshed or pg.main_weight_buffer is None: + continue + pg.copy_main_weights_to_model_weights() + refreshed.add(id(pg)) + + +class FullyShardV2MuonOptimizer(MegatronOptimizer): + """MegatronOptimizer adapter that lets FullyShardV2Muon slot into ChainedOptimizer. + + Presents the MegatronOptimizer interface around the inner FullyShardV2Muon + (a torch-style optimizer that updates the FSDP buffers directly). The actual + Muon update runs in ``step_with_ready_grads`` (-> ``self.optimizer.step()``). + + NOTE (UNVALIDATED): the inner Muon optimizer updates its 2D matrix params + end-to-end on the FSDP buffers (gather -> NS -> scatter -> update), so those + params are intentionally excluded from the ChainedOptimizer's global gradient + clipping / zero-counting / grad-norm (``get_parameters`` and + ``get_main_grads_for_grad_norm`` return ``[]``). Orthogonalized Muon updates + are already self-normalized, so this matches common practice; revisit if a + joint Muon+Adam grad-norm is desired. Checkpointing of the sharded momentum + is a TODO. This whole adapter needs validation on a GPU node. + """ + + def __init__(self, muon: "FullyShardV2Muon", config, model_chunks=None): + super().__init__(muon, config) + self.model_chunks = list(model_chunks) if model_chunks else [] + self.is_stub_optimizer = False + + # --- Excluded from the chained grad-norm / clip / zero-count machinery. --- + def get_parameters(self) -> List[torch.nn.Parameter]: + return [] + + def get_main_grads_for_grad_norm(self) -> List[torch.Tensor]: + return [] + + # --- Step contract used by ChainedOptimizer. --- + def prepare_grads(self) -> bool: + # FSDP gradients are already reduced into the grad buffers; no loss-scale + # inf/nan unscaling is applied here. Returns False = "no inf found". + return False + + @torch.no_grad() + def step_with_ready_grads(self) -> bool: + self.optimizer.step() + return True + + def step(self): + if self.prepare_grads(): + return False, None, None + return self.step_with_ready_grads(), None, None + + # --- Minimal remainder of the MegatronOptimizer interface. --- + def zero_grad(self, set_to_none: bool = True): + # FSDP zeroes its grad buffers via the module's own grad-sync path. + pass + + def get_loss_scale(self) -> torch.Tensor: + return torch.tensor([1.0], dtype=torch.float32, device=torch.cuda.current_device()) + + def reload_model_params(self, state_dict=None): + pass + + def state_dict(self): + # TODO: checkpoint the sharded Muon momentum buffers. + return {} + + def load_state_dict(self, state_dict): + # TODO: restore the sharded Muon momentum buffers. + pass + + def sharded_state_dict(self, model_sharded_state_dict, is_loading: bool = False, **kwargs): + # TODO: sharded checkpoint support for the Muon momentum. + return {} diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index 8749daafd2a..0b39188c930 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -1239,11 +1239,15 @@ def validate_args(args, defaults={}): args.use_distributed_optimizer = True # Optimizer step MXFP8 buffer operation that is not relevant or supported for Megatron-FSDP. args.reuse_grad_buf_for_mxfp8_param_ag = False - # Optimizer compatibility check. - assert args.optimizer in ( - 'sgd', - 'adam', - ), f"Megatron-FSDP does not support the {args.optimizer} optimizer yet." + # Optimizer compatibility check. Megatron-FSDP supports sgd/adam, plus the + # single-root distributed Muon (FullyShardV2Muon) on the v2 path only + # (the v2 ParameterGroup sets the dist_param back-references it needs). + assert args.optimizer in ('sgd', 'adam') or ( + args.optimizer == 'muon' and getattr(args, 'use_megatron_fsdp_v2', False) + ), ( + f"Megatron-FSDP does not support the {args.optimizer} optimizer " + "(muon requires --use-megatron-fsdp-v2)." + ) if ( args.data_parallel_sharding_strategy in ["optim_grads_params", "optim_grads"] @@ -1750,7 +1754,12 @@ def validate_args(args, defaults={}): # emerging optimizer check args.use_layer_wise_distributed_optimizer = False - if args.optimizer not in ('sgd', 'adam'): + # Megatron-FSDP handles emerging optimizers (e.g. muon) through its own + # optimizer factory (FullyShardV2Muon), NOT the LayerWiseDistributedOptimizer + # path below — so skip this whole block when Megatron-FSDP is enabled (it + # would otherwise disable use_distributed_optimizer and reject FSDP / the + # fsdp_dtensor ckpt format). muon+FSDP is gated to v2 by the assert above. + if args.optimizer not in ('sgd', 'adam') and not args.use_megatron_fsdp: if args.optimizer == 'dist_muon': warn_rank_0( "optimizer='dist_muon' is deprecated. " diff --git a/tests/unit_tests/distributed/megatron_fsdp/v2/test_dist_muon.py b/tests/unit_tests/distributed/megatron_fsdp/v2/test_dist_muon.py new file mode 100644 index 00000000000..336c7e04d7b --- /dev/null +++ b/tests/unit_tests/distributed/megatron_fsdp/v2/test_dist_muon.py @@ -0,0 +1,224 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Multi-process correctness tests for the single-root distributed Muon optimizer. + +Validates that ``FullyShardV2Muon`` (gather full grad to a load-balanced +root rank -> Newton-Schulz -> scatter back -> sharded weight update) produces the +same updated weights as a single-process reference Muon that runs the *same* +Newton-Schulz on the full gradient. The shared NS isolates the distribution logic +(gather/scatter/sharded momentum/sharded update) from NS numerics. + +Run with: + torchrun --nproc_per_node=4 -m pytest tests/unit_tests/distributed/megatron_fsdp/v2/test_dist_muon.py -v +""" + +import sys +from pathlib import Path + +import pytest +import torch +import torch.nn as nn + +sys.path.insert(0, str(Path(__file__).parents[2])) +from megatron.core.distributed.fsdp.src.megatron_fsdp.v2.mixed_precision import ( + FullyShardMixedPrecisionPolicy, +) +from megatron.core.distributed.fsdp.src.megatron_fsdp.v2.param_group import ParameterGroup +from megatron.core.distributed.fsdp.src.megatron_fsdp.v2.utils import ParamGroupIdx +from megatron.core.optimizer.fully_shard_v2_muon import FullyShardV2Muon + +STRATEGIES = ["optim", "optim_grads", "optim_grads_params"] + + +@pytest.fixture(scope="session", autouse=True) +def dist_env(): + """Initialize NCCL process group once and tear down at session end.""" + if not torch.distributed.is_initialized(): + torch.distributed.init_process_group(backend="nccl") + rank = torch.distributed.get_rank() + device = torch.device(f"cuda:{rank % torch.cuda.device_count()}") + torch.cuda.set_device(device) + yield + torch.distributed.destroy_process_group() + + +# ------------------------------------------------------------------ # +# Helpers +# ------------------------------------------------------------------ # + + +def _build_group(strategy, shapes): + """Build one fp32 ParameterGroup with 2D matrix params of the given shapes. + + fp32 master + fp32 grads so the update is computed in fp32 (tight tolerance). + Returns (pg, params, device). + """ + rank = torch.distributed.get_rank() + device = torch.device(f"cuda:{rank % torch.cuda.device_count()}") + + # Identical init weights on every rank. + torch.manual_seed(1234) + params = [nn.Parameter(torch.randn(*s, dtype=torch.float32, device=device)) for s in shapes] + + pg = ParameterGroup( + params=params, + param_group_id=ParamGroupIdx(0, 0), + mp_policy=FullyShardMixedPrecisionPolicy( + main_params_dtype=torch.float32, main_grads_dtype=torch.float32 + ), + mesh=None, + sharding_strategy=strategy, + ) + return pg, params, device + + +def _full_grads(shapes, device, step): + """Deterministic full gradients (identical across ranks), varying per step.""" + grads = [] + for i, s in enumerate(shapes): + g = torch.Generator(device="cpu").manual_seed(100 * step + i) + grads.append(torch.randn(*s, generator=g, dtype=torch.float32).to(device)) + return grads + + +def _plant_grads(pg, full_grads): + """Write each param's full grad into main_grad_buffer (per-rank shard).""" + for idx, g in enumerate(full_grads): + pg.main_grad_buffer.set_item(idx, g) + + +def _read_full_weight(pg, idx, shape): + """All-gather this param's full master weight and return it (reshaped, cloned).""" + mbuf = pg.main_weight_buffer + full = mbuf.unshard(bind_params=False) + off, sz = mbuf.buffer_index._get_item_global_range(idx) + out = full[off : off + sz].clone().view(shape) + mbuf.reshard() + return out + + +def _ref_muon_step(opt, param, w, g, state, *, lr, momentum, nesterov, weight_decay): + """Single-process reference Muon update on full tensors. + + Reuses the optimizer's own ``orthogonalize`` so the Newton-Schulz + scaling + path is identical by construction; the test then only validates the + distributed momentum / gather / scatter / sharded-update logic. + """ + buf = state.get("buf") + if buf is None: + buf = torch.zeros_like(g) + state["buf"] = buf + buf.mul_(momentum).add_(g) + pre_ns = g + momentum * buf if nesterov else buf.clone() + orth = opt.orthogonalize(param, pre_ns.to(torch.float32)).to(w.dtype) + w_new = w.clone() + if weight_decay != 0.0: + w_new.mul_(1.0 - lr * weight_decay) + w_new.add_(orth, alpha=-lr) + return w_new + + +# ------------------------------------------------------------------ # +# Core correctness: distributed Muon == single-process reference +# ------------------------------------------------------------------ # + + +@pytest.mark.parametrize("strategy", STRATEGIES) +def test_dist_muon_matches_reference(strategy): + # Shapes chosen so params land at different buffer offsets -> different + # load-balanced roots and shards that span rank boundaries unevenly. + shapes = [(3, 8), (7, 8), (5, 8), (9, 8), (12, 8)] + hp = dict(lr=0.1, momentum=0.9, nesterov=True, weight_decay=0.01, num_ns_steps=5) + + pg, params, device = _build_group(strategy, shapes) + + # Reference starts from the actual initialized master weights. + w_full = [_read_full_weight(pg, i, s) for i, s in enumerate(shapes)] + ref_state = [dict() for _ in shapes] + + opt = FullyShardV2Muon( + pg.dist_params, + lr=hp["lr"], + momentum=hp["momentum"], + nesterov=hp["nesterov"], + weight_decay=hp["weight_decay"], + num_ns_steps=hp["num_ns_steps"], + ) + + for step in range(3): + full_grads = _full_grads(shapes, device, step) + _plant_grads(pg, full_grads) + opt.step() + + # Advance the reference in lockstep (reusing opt.orthogonalize). + for i in range(len(shapes)): + w_full[i] = _ref_muon_step( + opt, params[i], w_full[i], full_grads[i], ref_state[i], + lr=hp["lr"], momentum=hp["momentum"], nesterov=hp["nesterov"], + weight_decay=hp["weight_decay"], + ) + + for i, s in enumerate(shapes): + got = _read_full_weight(pg, i, s) + # Tolerance reflects fp32 Newton-Schulz matmul non-determinism on GPU + # (the distributed root and the single-process reference run the SAME + # NS on the SAME assembled grad, so any real distribution bug would + # diverge by O(update magnitude), not ~1e-5). + assert torch.allclose(got, w_full[i], rtol=1e-3, atol=1e-4), ( + f"[{strategy}] param {i} shape {s} step {step}: max abs diff " + f"{(got - w_full[i]).abs().max().item():.3e}" + ) + + torch.distributed.barrier() + + +# ------------------------------------------------------------------ # +# Double-update guard: Muon only touches 2D matrix params +# ------------------------------------------------------------------ # + + +@pytest.mark.parametrize("strategy", STRATEGIES) +def test_dist_muon_skips_non_matrix_params(strategy): + rank = torch.distributed.get_rank() + device = torch.device(f"cuda:{rank % torch.cuda.device_count()}") + + torch.manual_seed(7) + matrix = nn.Parameter(torch.randn(8, 8, dtype=torch.float32, device=device)) + bias = nn.Parameter(torch.randn(8, dtype=torch.float32, device=device)) # 1D -> not Muon + params = [matrix, bias] + + pg = ParameterGroup( + params=params, + param_group_id=ParamGroupIdx(0, 0), + mp_policy=FullyShardMixedPrecisionPolicy( + main_params_dtype=torch.float32, main_grads_dtype=torch.float32 + ), + mesh=None, + sharding_strategy=strategy, + ) + + matrix_before = _read_full_weight(pg, 0, (8, 8)) + bias_before = _read_full_weight(pg, 1, (8,)) + + _plant_grads(pg, [torch.ones(8, 8, device=device), torch.ones(8, device=device)]) + FullyShardV2Muon(pg.dist_params, lr=0.5).step() + + matrix_after = _read_full_weight(pg, 0, (8, 8)) + bias_after = _read_full_weight(pg, 1, (8,)) + + assert torch.equal(bias_before, bias_after), "Muon must not modify the 1D bias param" + assert not torch.allclose(matrix_after, matrix_before), "Muon should update the 2D matrix" + + torch.distributed.barrier() From 9aedf56ab6e5092a3aa3a605d1933be0eb20ef47 Mon Sep 17 00:00:00 2001 From: tongliu Date: Thu, 4 Jun 2026 00:48:51 -0700 Subject: [PATCH 09/28] Clean up FullyShardV2Muon step() and docstrings 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 --- .../core/optimizer/fully_shard_v2_muon.py | 184 ++++++------------ 1 file changed, 60 insertions(+), 124 deletions(-) diff --git a/megatron/core/optimizer/fully_shard_v2_muon.py b/megatron/core/optimizer/fully_shard_v2_muon.py index d7120e6a80f..19db3fbcad2 100644 --- a/megatron/core/optimizer/fully_shard_v2_muon.py +++ b/megatron/core/optimizer/fully_shard_v2_muon.py @@ -12,46 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""DP-distributed Muon optimizer for Megatron-FSDP v2 (single-root, async-pipelined). - -Muon orthogonalizes the gradient of 2D matrix parameters via Newton-Schulz (NS). -NS needs the *full* gradient matrix, but under FSDP the gradient is sharded over -the data-parallel (DP) group. This optimizer reconstructs each parameter's full -gradient on a single, load-balanced "root" DP rank (chosen by the parameter -group's ``BufferIndex.item_root_ranks``), runs NS there *once*, and scatters the -orthogonalized gradient back — rather than all-gathering the full gradient to -every rank and running NS redundantly. Momentum and the weight update stay fully -local (elementwise on each rank's shard), so optimizer state remains sharded. - -Two classes, following the "MegatronOptimizer wraps a torch optimizer" pattern: - -- :class:`FullyShardV2Muon` is a ``torch.optim.Optimizer``. Its ``orthogonalize`` - uses ``emerging_optimizers``' production Newton-Schulz (``newton_schulz_tp`` + - ``get_muon_scale_factor``) WHEN that optional package is installed, and falls - back to a self-contained vanilla quintic NS otherwise — so it runs (and can be - unit-tested) even without the package. Its ``step()`` does the single-root - gather/scatter + sharded momentum/update. -- :class:`FullyShardV2MuonOptimizer` is the thin ``MegatronOptimizer`` adapter - that lets the above slot into ``ChainedOptimizer`` alongside Adam. - -Per managed 2D matrix parameter, each ``step()`` runs three phases: - -1. Local momentum on this rank's gradient shard (sharded optimizer state). -2. Gather the (post-momentum) full gradient to the parameter's root rank - (``ParameterGroup.unshard_grad_to_root_async``). -3a. Root only: ``self.orthogonalize(param, full_grad)`` (NS + scale), then scatter - back (``ParameterGroup.scatter_grad_from_root_async``). -3b. Local elementwise weight update on this rank's master shard. - -The gather/NS/scatter run as an async pipeline on the FSDP reduce-scatter stream: -each parameter's gather is prefetched and the previous parameter's scatter is left -in flight so they overlap the current parameter's Newton-Schulz, while the -collective ops are still issued in the same parameter order on every rank. - -Scope: tensor-parallel size 1 (orthogonalize runs the non-TP NS path). Real -tensor-parallel NS + QKV split and distributed checkpointing of the sharded -momentum are later phases. Needs validation on a GPU node. -""" from typing import List @@ -146,30 +106,29 @@ def __init__( self._managed = [] for group in self.param_groups: for dist_param in group["params"]: - pg = getattr(dist_param, "_fsdp_param_group", None) + param_group = getattr(dist_param, "_fsdp_param_group", None) orig = getattr(dist_param, "_fsdp_orig_param", None) - if pg is None or orig is None: + if param_group is None or orig is None: continue # not an FSDP-v2 param - if pg.sharding_strategy == "no_shard": + if param_group.sharding_strategy == "no_shard": raise NotImplementedError( "FullyShardV2Muon does not support the 'no_shard' strategy." ) - if pg.main_grad_buffer is None: + if param_group.main_grad_buffer is None: continue if ( orig.requires_grad and orig.dim() == 2 and not getattr(orig, "is_embedding_or_output_parameter", False) ): - idx = pg.param_idx[orig] - grad_shard = pg.main_grad_buffer.get_item(idx, as_shard=True) + idx = param_group.param_idx[orig] + grad_shard = param_group.main_grad_buffer.get_item(idx, as_shard=True) momentum_shard = torch.zeros_like(grad_shard) - self._managed.append((pg, idx, orig, momentum_shard)) + self._managed.append((param_group, idx, orig, momentum_shard)) if not self._managed: raise ValueError("FullyShardV2Muon got no 2D matrix parameters to manage.") - @torch.no_grad() def orthogonalize(self, param, grad: torch.Tensor) -> torch.Tensor: """Newton-Schulz orthogonalize + Muon spectral scale of a 2D gradient. @@ -198,34 +157,23 @@ def orthogonalize(self, param, grad: torch.Tensor) -> torch.Tensor: def step(self, closure=None): """Run one Muon update over all managed 2D matrix params. - Per param: (1) local momentum on this rank's grad shard; (2) gather the full - grad to the param's root rank; (3a) root orthogonalizes (NS + scale) and - scatters the result back into every holder's grad shard; (3b) local - decoupled-weight-decay update of this rank's master shard. Momentum/update - stay sharded so optimizer state is sharded. - - The gather/scatter are async (non-blocking ``batch_isend_irecv``); the NCCL - backend runs them on its own stream, so they overlap the Newton-Schulz - compute on the main stream — no explicit stream/event management needed. It - is a 2-stage pipeline: we prefetch the next param's gather so it overlaps the - current param's NS, and leave each scatter in flight so it overlaps the next - param's NS (NS dominates, so one gather in flight is enough to hide the comm); - the scatters are drained at the end. ``work.wait()`` - orders compute-after-comm; the backend orders comm-after-compute (a send - waits for the grad/NS that produced its buffer). P2P is issued in the same - param order on every rank, so it is deadlock-free. - - NOTE: the async pipeline is UNVALIDATED on GPU — needs a multi-rank run to - confirm overlap and numerical parity with a serial reference. + Per param, in order: + 1. momentum - update this rank's grad shard locally (stays sharded); + 2. gather - send the full grad to the param's root rank; + 3. NS - root orthogonalizes the full grad (Newton-Schulz); + 4. scatter - root sends each orthogonalized shard back to its holder; + 5. update - each rank applies a local decoupled-WD step to its master shard. + + Gather/scatter are async: all gathers (step 2) fire up front so the later ones + overlap the NS (step 3) of earlier params. """ + assert closure is None, "FullyShardV2Muon does not support a closure." lr = self.param_groups[0]["lr"] # set by the LR scheduler items = self._managed - n = len(items) - # Phase 1: in-place momentum on every rank's grad shard (view into gbuf), so - # the gathers below see a fully post-momentum grad. Cheap elementwise. - for pg, idx, _, momentum in items: - grad_shard = pg.main_grad_buffer.get_item(idx, as_shard=True) + # Phase 1: in-place momentum on each rank's grad shard, so gathers see post-momentum grad. + for param_group, idx, _, momentum in items: + grad_shard = param_group.main_grad_buffer.get_item(idx, as_shard=True) if grad_shard.numel() > 0: momentum.mul_(self._momentum_coef).add_(grad_shard) if self._nesterov: @@ -233,73 +181,61 @@ def step(self, closure=None): else: grad_shard.copy_(momentum) - # Phase 2/3a: 2-stage pipeline — prefetch the next param's gather so it - # overlaps the current param's NS; fire each scatter and don't wait, drain - # at the end. - inflight = [] # all in-flight P2P Works + tensors kept alive until drained - - pg, _, param, _ = items[0] - pending = pg.unshard_grad_to_root_async(param) # kick off the first gather - - for i in range(n): - pg, _, param, _ = items[i] - full, gather_reqs = pending - if i + 1 < n: # prefetch next gather before this NS - npg, _, nparam, _ = items[i + 1] - pending = npg.unshard_grad_to_root_async(nparam) - - if full is not None: # this rank is param i's root - for req in gather_reqs: - req.wait() # only the root waits its gather (NS needs it) - orth = self.orthogonalize(param, full).to(full.dtype) - inflight.append((pg.scatter_grad_from_root_async(param, orth), orth, full)) - else: # non-root holder / uninvolved - # Fire-and-forget: its grad isend (gather) + the scatter recv straight - # into its grad shard. Don't block here; drained below. - inflight.append((gather_reqs, None, None)) - inflight.append((pg.scatter_grad_from_root_async(param, None), None, None)) - - # Drain all in-flight P2P. The scatter recvs have written each rank's grad - # shard in place, so Phase 3b can read it directly. - for reqs, _orth, _full in inflight: + # Phase 2: fire ALL gathers up front so later ones overlap Phase-3 NS. + gather_work = [] # (param_group, param, full_grad_or_None, reqs) + for param_group, idx, param, _ in items: + full_grad, reqs = param_group.unshard_grad_to_root_async(param) + gather_work.append((param_group, param, full_grad, reqs)) + + # Phase 3: root waits its own gather, orthogonalizes, fires scatter; non-root fires recv. + scatter_work = [] # (reqs, orth_keepalive) + for param_group, param, full_grad, reqs in gather_work: + if full_grad is not None: # this rank is the param's root + for req in reqs: + req.wait() + orth = self.orthogonalize(param, full_grad).to(full_grad.dtype) + scatter_work.append((param_group.scatter_grad_from_root_async(param, orth), orth)) + else: # non-root holder / uninvolved + scatter_work.append((param_group.scatter_grad_from_root_async(param, None), None)) + + # Drain: non-root gather isends, then all scatters (recvs wrote each grad shard in place). + for _param_group, _param, _full_grad, reqs in gather_work: + for req in reqs: + req.wait() + for reqs, _orth in scatter_work: for req in reqs: req.wait() - # Phase 3b: decoupled-WD update of each rank's master shard from the (now - # orthogonalized) grad shard, then refresh model weights from master. - for pg, idx, _, _ in items: - wbuf = pg.main_weight_buffer if pg.main_weight_buffer is not None else pg.model_weight_buffer + # Phase 3b: decoupled-WD master-shard update from the orthogonalized grad, then refresh weights. + for param_group, idx, _, _ in items: + wbuf = ( + param_group.main_weight_buffer + if param_group.main_weight_buffer is not None + else param_group.model_weight_buffer + ) weight_shard = wbuf.get_item(idx, as_shard=True) if weight_shard.numel() == 0: continue - orth_shard = pg.main_grad_buffer.get_item(idx, as_shard=True) + orth_shard = param_group.main_grad_buffer.get_item(idx, as_shard=True) if self._weight_decay != 0.0: weight_shard.mul_(1.0 - lr * self._weight_decay) weight_shard.add_(orth_shard.to(weight_shard.dtype), alpha=-lr) refreshed = set() - for pg, _, _, _ in items: - if id(pg) in refreshed or pg.main_weight_buffer is None: + for param_group, _, _, _ in items: + if id(param_group) in refreshed or param_group.main_weight_buffer is None: continue - pg.copy_main_weights_to_model_weights() - refreshed.add(id(pg)) + param_group.copy_main_weights_to_model_weights() + refreshed.add(id(param_group)) class FullyShardV2MuonOptimizer(MegatronOptimizer): - """MegatronOptimizer adapter that lets FullyShardV2Muon slot into ChainedOptimizer. - - Presents the MegatronOptimizer interface around the inner FullyShardV2Muon - (a torch-style optimizer that updates the FSDP buffers directly). The actual - Muon update runs in ``step_with_ready_grads`` (-> ``self.optimizer.step()``). - - NOTE (UNVALIDATED): the inner Muon optimizer updates its 2D matrix params - end-to-end on the FSDP buffers (gather -> NS -> scatter -> update), so those - params are intentionally excluded from the ChainedOptimizer's global gradient - clipping / zero-counting / grad-norm (``get_parameters`` and - ``get_main_grads_for_grad_norm`` return ``[]``). Orthogonalized Muon updates - are already self-normalized, so this matches common practice; revisit if a - joint Muon+Adam grad-norm is desired. Checkpointing of the sharded momentum - is a TODO. This whole adapter needs validation on a GPU node. + """MegatronOptimizer adapter wrapping FullyShardV2Muon for ChainedOptimizer. + + The inner Muon optimizer updates its FSDP buffers end-to-end (gather -> NS -> + scatter -> update), so its params are excluded from the chained grad clip / + grad-norm / zero-count (``get_parameters`` / ``get_main_grads_for_grad_norm`` + return ``[]``) — orthogonalized updates are already self-normalized. """ def __init__(self, muon: "FullyShardV2Muon", config, model_chunks=None): From a2bf8a032379f59f6eb5d4130a2c2b9362beb6e3 Mon Sep 17 00:00:00 2001 From: tongliu Date: Thu, 4 Jun 2026 02:01:27 -0700 Subject: [PATCH 10/28] Rework distributed Muon UT: scales, nesterov, bit-exact reference 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 --- .../megatron_fsdp/v2/test_dist_muon.py | 237 +++++++----------- 1 file changed, 91 insertions(+), 146 deletions(-) diff --git a/tests/unit_tests/distributed/megatron_fsdp/v2/test_dist_muon.py b/tests/unit_tests/distributed/megatron_fsdp/v2/test_dist_muon.py index 336c7e04d7b..bd56b130a6a 100644 --- a/tests/unit_tests/distributed/megatron_fsdp/v2/test_dist_muon.py +++ b/tests/unit_tests/distributed/megatron_fsdp/v2/test_dist_muon.py @@ -12,13 +12,28 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Multi-process correctness tests for the single-root distributed Muon optimizer. +"""Multi-process correctness test for the single-root distributed Muon optimizer. -Validates that ``FullyShardV2Muon`` (gather full grad to a load-balanced -root rank -> Newton-Schulz -> scatter back -> sharded weight update) produces the -same updated weights as a single-process reference Muon that runs the *same* -Newton-Schulz on the full gradient. The shared NS isolates the distribution logic -(gather/scatter/sharded momentum/sharded update) from NS numerics. +``test_dist_muon_matches_reference`` validates ``FullyShardV2Muon`` against a +full-tensor reference in five steps, per parameter and per step: + + 1. construct a full gradient, identical on every rank; + 2. plant it into the sharded grad buffer (each rank keeps only its own slice); + 3. reference: run a full-tensor Muon update on every rank, reusing the + optimizer's own ``orthogonalize`` so the Newton-Schulz path is shared; + 4. distributed: ``opt.step()`` gathers each grad to its root rank, runs NS + there, scatters the result back, and updates each rank's shard; + 5. compare the two results bit for bit (``torch.equal``). + +Why exact equality: momentum and the weight update are elementwise, so they +commute with sharding; the gather/scatter just moves bytes; and both sides run +the SAME ``orthogonalize`` on the SAME assembled full gradient. The only +non-elementwise op is NS, and it sees identical input on same-arch GPUs. So any +divergence is a real bug in the gather / scatter / sharded-update logic. + +The param set also carries a trailing 1D bias (shape ``(8,)``): Muon manages only +2D matrices, so its reference stays frozen at init and the exact-match check thus +also asserts Muon never touched it. Run with: torchrun --nproc_per_node=4 -m pytest tests/unit_tests/distributed/megatron_fsdp/v2/test_dist_muon.py -v @@ -33,7 +48,7 @@ sys.path.insert(0, str(Path(__file__).parents[2])) from megatron.core.distributed.fsdp.src.megatron_fsdp.v2.mixed_precision import ( - FullyShardMixedPrecisionPolicy, + MixedPrecisionPolicy, ) from megatron.core.distributed.fsdp.src.megatron_fsdp.v2.param_group import ParameterGroup from megatron.core.distributed.fsdp.src.megatron_fsdp.v2.utils import ParamGroupIdx @@ -41,6 +56,17 @@ STRATEGIES = ["optim", "optim_grads", "optim_grads_params"] +# Each set is a list of param shapes at a different scale: a wide matrix +# (rows < cols), a tall one (rows > cols -> NS transposes), and a trailing 1D bias +# Muon must skip. Sizes are picked so that under grad sharding (ZeRO-2/3) each +# matrix straddles a shard boundary and is split across ranks, so the gather/scatter +# is actually exercised. +SHAPE_SETS = [ + pytest.param([(5, 8), (9, 8), (8,)], id="small"), # dims < 10 + pytest.param([(200, 256), (400, 256), (256,)], id="medium"), # dims in the hundreds + pytest.param([(1500, 2048), (3000, 2048), (2048,)], id="large"), # dims in the thousands +] + @pytest.fixture(scope="session", autouse=True) def dist_env(): @@ -54,51 +80,6 @@ def dist_env(): torch.distributed.destroy_process_group() -# ------------------------------------------------------------------ # -# Helpers -# ------------------------------------------------------------------ # - - -def _build_group(strategy, shapes): - """Build one fp32 ParameterGroup with 2D matrix params of the given shapes. - - fp32 master + fp32 grads so the update is computed in fp32 (tight tolerance). - Returns (pg, params, device). - """ - rank = torch.distributed.get_rank() - device = torch.device(f"cuda:{rank % torch.cuda.device_count()}") - - # Identical init weights on every rank. - torch.manual_seed(1234) - params = [nn.Parameter(torch.randn(*s, dtype=torch.float32, device=device)) for s in shapes] - - pg = ParameterGroup( - params=params, - param_group_id=ParamGroupIdx(0, 0), - mp_policy=FullyShardMixedPrecisionPolicy( - main_params_dtype=torch.float32, main_grads_dtype=torch.float32 - ), - mesh=None, - sharding_strategy=strategy, - ) - return pg, params, device - - -def _full_grads(shapes, device, step): - """Deterministic full gradients (identical across ranks), varying per step.""" - grads = [] - for i, s in enumerate(shapes): - g = torch.Generator(device="cpu").manual_seed(100 * step + i) - grads.append(torch.randn(*s, generator=g, dtype=torch.float32).to(device)) - return grads - - -def _plant_grads(pg, full_grads): - """Write each param's full grad into main_grad_buffer (per-rank shard).""" - for idx, g in enumerate(full_grads): - pg.main_grad_buffer.set_item(idx, g) - - def _read_full_weight(pg, idx, shape): """All-gather this param's full master weight and return it (reshaped, cloned).""" mbuf = pg.main_weight_buffer @@ -109,116 +90,80 @@ def _read_full_weight(pg, idx, shape): return out -def _ref_muon_step(opt, param, w, g, state, *, lr, momentum, nesterov, weight_decay): - """Single-process reference Muon update on full tensors. - - Reuses the optimizer's own ``orthogonalize`` so the Newton-Schulz + scaling - path is identical by construction; the test then only validates the - distributed momentum / gather / scatter / sharded-update logic. - """ - buf = state.get("buf") - if buf is None: - buf = torch.zeros_like(g) - state["buf"] = buf - buf.mul_(momentum).add_(g) - pre_ns = g + momentum * buf if nesterov else buf.clone() - orth = opt.orthogonalize(param, pre_ns.to(torch.float32)).to(w.dtype) - w_new = w.clone() - if weight_decay != 0.0: - w_new.mul_(1.0 - lr * weight_decay) - w_new.add_(orth, alpha=-lr) - return w_new - - -# ------------------------------------------------------------------ # -# Core correctness: distributed Muon == single-process reference -# ------------------------------------------------------------------ # - - -@pytest.mark.parametrize("strategy", STRATEGIES) -def test_dist_muon_matches_reference(strategy): - # Shapes chosen so params land at different buffer offsets -> different - # load-balanced roots and shards that span rank boundaries unevenly. - shapes = [(3, 8), (7, 8), (5, 8), (9, 8), (12, 8)] - hp = dict(lr=0.1, momentum=0.9, nesterov=True, weight_decay=0.01, num_ns_steps=5) - - pg, params, device = _build_group(strategy, shapes) - - # Reference starts from the actual initialized master weights. - w_full = [_read_full_weight(pg, i, s) for i, s in enumerate(shapes)] - ref_state = [dict() for _ in shapes] - - opt = FullyShardV2Muon( - pg.dist_params, - lr=hp["lr"], - momentum=hp["momentum"], - nesterov=hp["nesterov"], - weight_decay=hp["weight_decay"], - num_ns_steps=hp["num_ns_steps"], - ) - - for step in range(3): - full_grads = _full_grads(shapes, device, step) - _plant_grads(pg, full_grads) - opt.step() - - # Advance the reference in lockstep (reusing opt.orthogonalize). - for i in range(len(shapes)): - w_full[i] = _ref_muon_step( - opt, params[i], w_full[i], full_grads[i], ref_state[i], - lr=hp["lr"], momentum=hp["momentum"], nesterov=hp["nesterov"], - weight_decay=hp["weight_decay"], - ) - - for i, s in enumerate(shapes): - got = _read_full_weight(pg, i, s) - # Tolerance reflects fp32 Newton-Schulz matmul non-determinism on GPU - # (the distributed root and the single-process reference run the SAME - # NS on the SAME assembled grad, so any real distribution bug would - # diverge by O(update magnitude), not ~1e-5). - assert torch.allclose(got, w_full[i], rtol=1e-3, atol=1e-4), ( - f"[{strategy}] param {i} shape {s} step {step}: max abs diff " - f"{(got - w_full[i]).abs().max().item():.3e}" - ) - - torch.distributed.barrier() - - -# ------------------------------------------------------------------ # -# Double-update guard: Muon only touches 2D matrix params -# ------------------------------------------------------------------ # - - @pytest.mark.parametrize("strategy", STRATEGIES) -def test_dist_muon_skips_non_matrix_params(strategy): +@pytest.mark.parametrize("nesterov", [True, False]) +@pytest.mark.parametrize("shapes", SHAPE_SETS) +def test_dist_muon_matches_reference(strategy, nesterov, shapes): rank = torch.distributed.get_rank() device = torch.device(f"cuda:{rank % torch.cuda.device_count()}") - torch.manual_seed(7) - matrix = nn.Parameter(torch.randn(8, 8, dtype=torch.float32, device=device)) - bias = nn.Parameter(torch.randn(8, dtype=torch.float32, device=device)) # 1D -> not Muon - params = [matrix, bias] + # shapes (parametrized): 2D matrices Muon updates + a trailing 1D bias it skips + # (``len(s) == 2`` tells them apart). + lr, momentum, weight_decay, num_ns_steps = 0.1, 0.9, 0.01, 5 + # Identical init weights on every rank; fp32 master + fp32 grads -> fp32 update. + torch.manual_seed(1234) + params = [nn.Parameter(torch.randn(*s, dtype=torch.float32, device=device)) for s in shapes] pg = ParameterGroup( params=params, param_group_id=ParamGroupIdx(0, 0), - mp_policy=FullyShardMixedPrecisionPolicy( + mp_policy=MixedPrecisionPolicy( main_params_dtype=torch.float32, main_grads_dtype=torch.float32 ), mesh=None, sharding_strategy=strategy, ) + opt = FullyShardV2Muon( + pg.dist_params, + lr=lr, + momentum=momentum, + nesterov=nesterov, + weight_decay=weight_decay, + num_ns_steps=num_ns_steps, + ) - matrix_before = _read_full_weight(pg, 0, (8, 8)) - bias_before = _read_full_weight(pg, 1, (8,)) + # Reference, identical on every rank: a full master weight per param (the 1D + # bias's stays frozen at init -> matching it asserts Muon skipped it) plus a + # momentum buffer per param. Snapshot the inits to confirm matrices actually move. + ref_w = [_read_full_weight(pg, i, s) for i, s in enumerate(shapes)] + ref_buf = [torch.zeros(*s, dtype=torch.float32, device=device) for s in shapes] + init_w = [w.clone() for w in ref_w] - _plant_grads(pg, [torch.ones(8, 8, device=device), torch.ones(8, device=device)]) - FullyShardV2Muon(pg.dist_params, lr=0.5).step() + for step in range(3): + # 1. Construct a full grad per param, identical across ranks (CPU RNG -> GPU). + full_grads = [] + for i, s in enumerate(shapes): + gen = torch.Generator(device="cpu").manual_seed(100 * step + i) + full_grads.append(torch.randn(*s, generator=gen, dtype=torch.float32).to(device)) - matrix_after = _read_full_weight(pg, 0, (8, 8)) - bias_after = _read_full_weight(pg, 1, (8,)) + # 2. Plant the full grads into the sharded grad buffer (each rank keeps its slice). + for i, g in enumerate(full_grads): + pg.main_grad_buffer.set_item(i, g) - assert torch.equal(bias_before, bias_after), "Muon must not modify the 1D bias param" - assert not torch.allclose(matrix_after, matrix_before), "Muon should update the 2D matrix" + # 3. Reference: full-tensor Muon on every rank, 2D matrices only (the 1D bias + # is left frozen). pre_ns mirrors step()'s Phase 1: momentum + nesterov. + for i, s in enumerate(shapes): + if len(s) != 2: + continue + ref_buf[i].mul_(momentum).add_(full_grads[i]) + # Mirror step()'s exact ops so the comparison can be bit-exact: the + # nesterov look-ahead is a single alpha-add (g + m*buf), not g + (m*buf). + pre_ns = full_grads[i].add(ref_buf[i], alpha=momentum) if nesterov else ref_buf[i].clone() + orth = opt.orthogonalize(params[i], pre_ns) + if weight_decay != 0.0: + ref_w[i].mul_(1.0 - lr * weight_decay) + ref_w[i].add_(orth, alpha=-lr) + + # 4. Distributed Muon: gather -> NS on root -> scatter -> sharded update. + opt.step() + # 5. Compare every param exactly: 2D matrices match the updated reference, + # the 1D bias matches its untouched init. + for i, s in enumerate(shapes): + got = _read_full_weight(pg, i, s) + assert torch.equal(got, ref_w[i]), ( + f"[{strategy}] param {i} shape {s} step {step}: max abs diff " + f"{(got - ref_w[i]).abs().max().item():.3e}" + ) + torch.distributed.barrier() From a4d9231be0b7ddd06362bdb10751b947d1d6dd24 Mon Sep 17 00:00:00 2001 From: jianbinc Date: Fri, 5 Jun 2026 09:55:56 +0800 Subject: [PATCH 11/28] docs(mfsdp-v2): update README sharding strategies and MixedPrecisionPolicy - 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 --- .../fsdp/src/megatron_fsdp/v2/README.md | 48 ++++++++++++------- 1 file changed, 31 insertions(+), 17 deletions(-) diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/README.md b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/README.md index 3fb04b2f2f4..2d28f74bddd 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/README.md +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/README.md @@ -42,12 +42,6 @@ Wraps a module with FSDP sharding semantics: Controls parameter/gradient dtypes and communication precision. -| Policy | Notes | -|--------|-------| -| `MixedPrecisionPolicy` | Base policy; all fields default to ``None`` | -| `FullyShardFP8Policy` | MXFP8 rowwise/colwise quantized weights | -| `FullyShardNVFP4Policy` | NVFP4 primary weights | - **Key fields:** | Field | Default | Purpose | @@ -57,8 +51,18 @@ Controls parameter/gradient dtypes and communication precision. | `grad_comm_dtype` | ``None`` | Dtype for gradient reduce-scatter communication. ``None`` = use ``main_grads_dtype``. | | `use_decoupled_grad` | ``False`` | When ``False``, ``main_grads_dtype`` is inferred from ``main_params_dtype`` so the optimizer operates in a consistent precision context. | +**FP8 & NVFP4 recipes** + +`FullyShardFP8Policy` and `FullyShardNVFP4Policy` are recipe dataclasses +that configure quantized mixed-precision behavior within `MixedPrecisionPolicy`, +passed via the ``fp8`` and ``nvfp4`` fields respectively. They are not +standalone policies. + ```python -from megatron_fsdp.v2 import fully_shard, MixedPrecisionPolicy +from megatron_fsdp.v2 import ( + fully_shard, MixedPrecisionPolicy, + FullyShardFP8Policy, FullyShardNVFP4Policy, +) # No separate main buffer — optimizer mutates model params directly mp_policy = MixedPrecisionPolicy() @@ -68,11 +72,17 @@ fully_shard(model, mp_policy=mp_policy) mp_policy = MixedPrecisionPolicy(main_params_dtype=torch.float32) fully_shard(model, mp_policy=mp_policy) -# FP8 mixed precision — fp32 main weights auto-created by adapter -from megatron_fsdp.v2 import FullyShardFP8Policy +# FP8 mixed precision — fp32 main weights + MXFP8 rowwise/colwise quantized compute +mp_policy = MixedPrecisionPolicy( + main_params_dtype=torch.float32, + fp8=FullyShardFP8Policy(enabled=True), +) +fully_shard(model, mp_policy=mp_policy) + +# NVFP4 mixed precision — fp32 main weights + NVFP4 primary compute weights mp_policy = MixedPrecisionPolicy( main_params_dtype=torch.float32, - fp8=FullyShardFP8Policy(enabled=True) + nvfp4=FullyShardNVFP4Policy(enabled=True), ) fully_shard(model, mp_policy=mp_policy) ``` @@ -117,12 +127,15 @@ See the parent directory `..` for `uneven_dtensor.py` which provides: ## Sharding Strategies +All strategies except `no_shard` use `Shard(0)` DTensor placements. The +strategy controls which buffers and communication collectives are used. + | Strategy | Shard Weights | Shard Gradients | Status | Notes | |----------|---------------|-----------------|--------|-------| -| `optim_grads_params` | Yes | Yes | **Supported** | Like ZeRO-3: full parameter/gradient/optimizer sharding | -| `no_shard` | No | No | **Not yet supported** | Like DDP: no sharding | -| `optim` | No | No | **Not yet supported** | Like ZeRO-1: shard optimizer states only | -| `optim_grads` | No | Yes | **Not yet supported** | Like ZeRO-2: shard optimizer states + gradients | +| `optim_grads_params` | Yes | Yes | **Supported** | Like ZeRO-3: all-gather weights pre-forward, reduce-scatter grads during backward, sharded optimizer states | +| `optim_grads` | No | Yes | **Supported** | Like ZeRO-2: replicated weights, reduce-scatter grads during backward, sharded optimizer states. No param-gather overlap. | +| `optim` | No | No | **Supported** | Like ZeRO-1: replicated weights, grads accumulated in replicated buffer, single reduce-scatter at ``finish_grad_sync``, sharded optimizer states. No param-gather overlap. | +| `no_shard` | No | No | **Not yet supported** | Like DDP: no sharding. Raises ``NotImplementedError``. | ## Known Limitations @@ -137,9 +150,10 @@ See the parent directory `..` for `uneven_dtensor.py` which provides: ### Sharding Strategies -Only `optim_grads_params` (ZeRO-3 equivalent) is implemented. `no_shard` -(DDP-like), `optim` (ZeRO-1), and `optim_grads` (ZeRO-2) are planned but -not yet available. +All strategies except `no_shard` are supported. `optim_grads` (ZeRO-2) and +`optim` (ZeRO-1) do not support parameter gather overlap (prefetch/unshard +pipelining) — weights are replicated so no all-gather is needed. +`no_shard` (DDP-like) is not yet implemented and raises ``NotImplementedError``. ### `fully_shard()` API Parameters From 63d0f8c1d734856463c6dd7996daea3128bc1c3c Mon Sep 17 00:00:00 2001 From: jianbinc Date: Fri, 5 Jun 2026 10:12:31 +0800 Subject: [PATCH 12/28] feat(mfsdp-v2): add Qwen3-30B-A3B 2-node MXFP8 training example with W&B logging --- ...s128_mbs4_seq4096_n2_mfsdp2_mxfp8_wandb.sh | 151 ++++++++++++++++++ 1 file changed, 151 insertions(+) create mode 100644 examples/megatron_fsdp/qwen3-30b-a3b.gbs128_mbs4_seq4096_n2_mfsdp2_mxfp8_wandb.sh diff --git a/examples/megatron_fsdp/qwen3-30b-a3b.gbs128_mbs4_seq4096_n2_mfsdp2_mxfp8_wandb.sh b/examples/megatron_fsdp/qwen3-30b-a3b.gbs128_mbs4_seq4096_n2_mfsdp2_mxfp8_wandb.sh new file mode 100644 index 00000000000..98491366f21 --- /dev/null +++ b/examples/megatron_fsdp/qwen3-30b-a3b.gbs128_mbs4_seq4096_n2_mfsdp2_mxfp8_wandb.sh @@ -0,0 +1,151 @@ +#!/bin/bash + +#SBATCH --job-name=qwen3-30b-a3b.gbs128_mbs4_seq4096_n2_mfsdp2_mxfp8 +#SBATCH --nodes=2 +#SBATCH --ntasks-per-node=4 +#SBATCH --partition=gb200 +#SBATCH --account= +#SBATCH --time=2:00:00 +#SBATCH --exclusive +#SBATCH --mem=0 +#SBATCH --dependency=singleton + +export NCCL_IB_SL=1 +export NCCL_IB_TIMEOUT=19 +export NVTE_FWD_LAYERNORM_SM_MARGIN=16 +export NVTE_BWD_LAYERNORM_SM_MARGIN=16 +export NCCL_P2P_NET_CHUNKSIZE=2097152 +export TORCH_NCCL_AVOID_RECORD_STREAMS=1 +export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True +export CUDA_DEVICE_MAX_CONNECTIONS=8 +export WANDB_API_KEY=${WANDB_API_KEY:-""} +export WANDB_ENTITY=${WANDB_ENTITY:-"nvidia"} + +MEGATRON_PATH=${MEGATRON_PATH:-"/path/to/Megatron-LM"} +OUTPUT_PATH=${OUTPUT_PATH:-"${MEGATRON_PATH}/output"} +NAME=qwen3-30b-a3b.gbs128_mbs4_seq4096_n2_mfsdp2_mxfp8_wandb +mkdir -p "$OUTPUT_PATH/checkpoints/$NAME" +CHECKPOINT_DIR=${CHECKPOINT_DIR:-"${MEGATRON_PATH}/checkpoints"} +DATA_PATH=${DATA_PATH:-"/path/to/data/c4/en/c4-train.en_6_text_document"} +TOKENIZER_MODEL=${TOKENIZER_MODEL:-"/path/to/data/c4/en/tokenizer"} + +PRETRAIN_ARGS=( + --seq-length 4096 + --max-position-embeddings 4096 + --tensor-model-parallel-size 1 + --context-parallel-size 1 + --pipeline-model-parallel-size 1 + --expert-tensor-parallel-size 1 + --use-mcore-models + --use-flash-attn + --no-check-for-nan-in-loss-and-grad + --manual-gc + --manual-gc-interval 10 + --recompute-granularity selective + --recompute-modules moe_act + --transformer-impl transformer_engine + --apply-layernorm-1p + --tokenizer-type HuggingFaceTokenizer + --tokenizer-model ${TOKENIZER_MODEL} + --data-path ${DATA_PATH} + --data-cache-path $MEGATRON_PATH/data_cache + --num-layers 48 + --hidden-size 2048 + --ffn-hidden-size 6144 + --num-attention-heads 32 + --norm-epsilon 1e-06 + --normalization RMSNorm + --attention-dropout 0.0 + --hidden-dropout 0.0 + --position-embedding-type rope + --rotary-base 1000000 + --swiglu + --untie-embeddings-and-output-weights + --disable-bias-linear + --group-query-attention + --num-query-groups 4 + --kv-channels 128 + --make-vocab-size-divisible-by 1187 + --qk-layernorm + --attention-backend fused + --rotary-percent 1.0 + --rotary-seq-len-interpolation-factor 1 + --clip-grad 1.0 + --weight-decay 0.1 + --adam-beta1 0.9 + --adam-beta2 0.95 + --lr 0.00012 + --min-lr 1.2e-05 + --lr-decay-samples 255126953 + --lr-warmup-samples 162761 + --lr-decay-style cosine + --init-method-std 0.02 + --ckpt-format fsdp_dtensor + --train-samples 268554688 + --exit-duration-in-mins 230 + --split 99,1,0 + --eval-iters 32 + --eval-interval 100 + --log-interval 1 + --distributed-timeout-minutes 60 + --no-mmap-bin-files + --no-create-attention-mask-in-dataloader + --log-throughput + --bf16 + --num-workers 6 + --enable-experimental + --auto-detect-ckpt-format + --save $OUTPUT_PATH/checkpoints/$NAME + --save-interval 500 + --load $OUTPUT_PATH/checkpoints/$NAME + --dist-ckpt-strictness log_all + --fp8-recipe mxfp8 + --fp8-format e4m3 + --overlap-grad-reduce + --overlap-param-gather + --use-megatron-fsdp + --data-parallel-sharding-strategy optim_grads_params + --calculate-per-token-loss + --init-model-with-meta-device + --grad-reduce-in-bf16 + --use-distributed-optimizer + --num-experts 128 + --moe-ffn-hidden-size 768 + --moe-router-load-balancing-type aux_loss + --moe-router-topk 8 + --moe-aux-loss-coeff 0.001 + --moe-grouped-gemm + --moe-router-dtype fp32 + --moe-permute-fusion + --moe-router-fusion + --moe-router-force-load-balancing + --moe-token-dispatcher-type alltoall + --sequence-parallel + --cross-entropy-loss-fusion + --cross-entropy-fusion-impl te + --wandb-project megatron-fsdp + --wandb-exp-name ${NAME} + --wandb-save-dir $OUTPUT_PATH/wandb + --log-timers-to-tensorboard + --log-memory-to-tensorboard + --log-num-zeros-in-grad + --log-params-norm + --log-validation-ppl-to-tensorboard + --tensorboard-dir $OUTPUT_PATH/tensorboard + --micro-batch-size 4 + --fp8-param-gather + --expert-model-parallel-size 4 + --global-batch-size 128 + --use-megatron-fsdp-v2 +) + +run_cmd=" +cd $MEGATRON_PATH; +git rev-parse HEAD; +export PYTHONPATH=$MEGATRON_PATH:$PYTHONPATH; +python3 -u pretrain_gpt.py ${PRETRAIN_ARGS[@]}" + +srun --mpi=pmix -l \ +--container-image nvcr.io/nvidia/nemo:26.04 \ +--container-mounts "/home:/home,/lustre:/lustre" --no-container-mount-home \ +bash -x -c "${run_cmd}" From 13c954ca99d98832e5b94b21959e835d851633f7 Mon Sep 17 00:00:00 2001 From: tongliu Date: Sun, 7 Jun 2026 21:43:02 -0700 Subject: [PATCH 13/28] Refactor Muon toward DTensor-driven, ParameterGroup-free (WIP) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- megatron/core/optimizer/__init__.py | 27 +--- .../core/optimizer/fully_shard_v2_muon.py | 150 ++++++++++++++++-- 2 files changed, 139 insertions(+), 38 deletions(-) diff --git a/megatron/core/optimizer/__init__.py b/megatron/core/optimizer/__init__.py index 7a9d10fb2d3..dde8580efb9 100644 --- a/megatron/core/optimizer/__init__.py +++ b/megatron/core/optimizer/__init__.py @@ -882,7 +882,7 @@ def _build_megatron_fsdp_v2_muon_optimizer( checkpointing, and passing full FSDP buffers to an Adam built over only the non-matrix param subset). """ - from .fully_shard_v2_muon import FullyShardV2Muon, FullyShardV2MuonOptimizer + from .fully_shard_v2_muon import FullyShardV2MuonOptimizer # Tag params and route non-linear/embedding params to Adam (same mechanism # as the non-FSDP emerging path in _get_megatron_emerging_optimizer). @@ -911,8 +911,6 @@ def _build_megatron_fsdp_v2_muon_optimizer( scale_mode=getattr(config, 'muon_scale_mode', 'spectral'), extra_scale_factor=getattr(config, 'muon_extra_scale_factor', 1.0), fp32_matmul_prec=getattr(config, 'muon_fp32_matmul_prec', 'medium'), - # TODO(phase3): wire pg_collection/tp_mode for tensor-parallel NS (Phase 2). - pg_collection=None, ) adam_config = copy.copy(config) @@ -954,26 +952,13 @@ def _build_megatron_fsdp_v2_muon_optimizer( adam_part.config = config optimizers.append(adam_part) - # --- Muon for the 2D matrix params (single-root FSDP v2) --- - # Symmetric with Adam: pass the muon bucket of optimizer-facing params - # (dist_params). FullyShardV2Muon resolves each back to its FSDP - # ParameterGroup via the dist_param back-references. - muon_groups, _ = _get_param_groups_and_buffers( - chunk_list, - model_chunk_offset=model_chunk_offset, - config=config, - config_overrides=config_overrides, - filter_fn=lambda g: g.get('optimizer') != 'adam', - buffer_name='buffers', - ) - if muon_groups: - muon = FullyShardV2Muon(muon_groups, **muon_kwargs) - optimizers.append( - FullyShardV2MuonOptimizer(muon, config, model_chunks=chunk_list) - ) - model_chunk_offset += 1 + # One single-root Muon over every chunk's 2D matrix params on this rank. The + # wrapper pulls those dist_params + grad DTensors out of the FSDP modules itself + # (in layer order), so the factory needs no per-chunk muon param-group plumbing. + optimizers.append(FullyShardV2MuonOptimizer(config, model_chunks, **muon_kwargs)) + if len(optimizers) == 1: return optimizers[0] return ChainedOptimizer(optimizers) diff --git a/megatron/core/optimizer/fully_shard_v2_muon.py b/megatron/core/optimizer/fully_shard_v2_muon.py index 19db3fbcad2..d9c00d20973 100644 --- a/megatron/core/optimizer/fully_shard_v2_muon.py +++ b/megatron/core/optimizer/fully_shard_v2_muon.py @@ -13,9 +13,12 @@ # limitations under the License. +from collections import OrderedDict from typing import List import torch +from torch.distributed.tensor import DTensor +from torch.distributed.tensor.placement_types import Replicate from .emerging_optimizers import HAVE_EMERGING_OPTIMIZERS from .optimizer import MegatronOptimizer @@ -56,21 +59,22 @@ class FullyShardV2Muon(torch.optim.Optimizer): """Single-root distributed Muon for Megatron-FSDP v2 (a torch optimizer). Args: - params: optimizer-facing params (FSDP-v2 ``dist_param`` DTensors) or param - groups — the same form every torch/Megatron optimizer takes. Each - dist_param carries ``_fsdp_param_group`` / ``_fsdp_orig_param`` - back-references (set in ``ParameterGroup._init_dist_params``) used to - drive the per-param gather/scatter. + params: optimizer-facing params — all assumed to be FSDP-v2 ``dist_param`` + DTensors (asserted). A flat list or param groups, as any torch optimizer. + grads: the gradient DTensor for each param, aligned 1:1 with the flattened + params (asserted to be DTensors). This is the sole gradient source — + replaces the old ParameterGroup back-reference path. lr / momentum / nesterov / weight_decay: Muon update hyperparameters. num_ns_steps / coefficient_type / scale_mode / extra_scale_factor: NS + scaling config used by ``orthogonalize``. - fp32_matmul_prec / pg_collection / tp_mode: reserved for the Phase-2 - tensor-parallel NS path (currently TP size 1; accepted and ignored). + fp32_matmul_prec / tp_mode: reserved for the Phase-2 tensor-parallel NS path + (currently TP size 1; accepted and ignored). """ def __init__( self, params, + grads, *, lr: float, momentum: float = 0.95, @@ -81,11 +85,25 @@ def __init__( scale_mode: str = "spectral", extra_scale_factor: float = 1.0, fp32_matmul_prec: str = "medium", - pg_collection=None, tp_mode: str = "duplicated", ) -> None: super().__init__(params, dict(lr=lr, weight_decay=weight_decay)) + # Refactor contract: this optimizer is driven purely by DTensors. Assert the + # two assumptions the new design rests on — every param (FSDP-v2 dist_param) + # and every grad is a DTensor, and grads align 1:1 with the flattened params. + flat_params = [p for grp in self.param_groups for p in grp["params"]] + assert all( + isinstance(p, DTensor) for p in flat_params + ), "FullyShardV2Muon expects every param to be a DTensor (FSDP-v2 dist_param)." + assert all( + isinstance(g, DTensor) for g in grads + ), "FullyShardV2Muon expects every grad to be a DTensor." + assert len(grads) == len(flat_params), ( + f"grads ({len(grads)}) must align 1:1 with params ({len(flat_params)})." + ) + self._grads = list(grads) + # Muon update hyperparameters owned by this class. lr is NOT stored — it # lives in param_groups[0]["lr"] (where an LR scheduler writes), read in step. self._momentum_coef = momentum @@ -104,8 +122,12 @@ def __init__( # keeps the collective send/recv in phases 2/3a aligned across ranks. # Each entry: (ParameterGroup, idx, original param, momentum shard). self._managed = [] + self._managed_grads = [] # grad DTensor for each managed entry (aligned with _managed) + gi = 0 # running index into self._grads (== flattened params order) for group in self.param_groups: for dist_param in group["params"]: + grad_dtensor = self._grads[gi] + gi += 1 param_group = getattr(dist_param, "_fsdp_param_group", None) orig = getattr(dist_param, "_fsdp_orig_param", None) if param_group is None or orig is None: @@ -125,10 +147,81 @@ def __init__( grad_shard = param_group.main_grad_buffer.get_item(idx, as_shard=True) momentum_shard = torch.zeros_like(grad_shard) self._managed.append((param_group, idx, orig, momentum_shard)) + self._managed_grads.append(grad_dtensor) if not self._managed: raise ValueError("FullyShardV2Muon got no 2D matrix parameters to manage.") + # Reconstruct, purely from the grad DTensors, every DP rank's (flat_offset, + # size) slice of each managed param — the global sharding the gather/scatter + # needs. A DTensor only knows its OWN rank's shard, so one all_gather per + # dp_group recovers the rest. (Root assignment is computed in a later step.) + self._shard_ranges = self._gather_shard_ranges() + self._crosscheck_shard_ranges_against_param_group() + + def _gather_shard_ranges(self): + """Per managed param, every DP rank's (flat_offset, size) within the full + param, derived from the grad DTensor sharding + one all_gather per dp_group. + + FSDP-v2 shards 2D grads as Shard(0) over the flattened buffer, aligned so each + rank owns whole rows (chunk_size_factor = lcm(cols)); to_local() is therefore a + contiguous block of full rows. A DTensor only knows its OWN rank's row count, so + we all_gather row counts within each grad's dp_group and prefix-sum them. + Replicated grads (ZeRO-1) live fully on every rank -> (0, numel) for all ranks. + """ + import torch.distributed as dist + + by_group = OrderedDict() # dp process group -> [managed indices] + for mi, g in enumerate(self._managed_grads): + by_group.setdefault(g.device_mesh.get_group(), []).append(mi) + + ranges = [None] * len(self._managed_grads) + for grp, mis in by_group.items(): + ws = dist.get_world_size(grp) + my_rows = [] # this rank's row count per param (-1 sentinel = replicated) + for mi in mis: + g = self._managed_grads[mi] + my_rows.append(-1 if isinstance(g.placements[0], Replicate) else g.to_local().shape[0]) + gathered = [None] * ws + dist.all_gather_object(gathered, my_rows, group=grp) + for j, mi in enumerate(mis): + g = self._managed_grads[mi] + row_numel = 1 + for d in g.shape[1:]: + row_numel *= d + rows_per_rank = [gathered[r][j] for r in range(ws)] + if rows_per_rank[0] == -1: # replicated: full param on every rank + ranges[mi] = [(0, g.shape.numel())] * ws + else: + rr, off = [], 0 + for r in range(ws): + sz = rows_per_rank[r] * row_numel + rr.append((off, sz)) + off += sz + ranges[mi] = rr + return ranges + + def _crosscheck_shard_ranges_against_param_group(self): + """TEMP (drop once step() is rewired): assert the DTensor-derived shard ranges + match the ParameterGroup's old ``_grad_gather_plans`` dst ranges + holders, + proving the pure-DTensor reconstruction is correct. Skips replicated (ZeRO-1).""" + for mi, (param_group, idx, _orig, _mom) in enumerate(self._managed): + if isinstance(self._managed_grads[mi].placements[0], Replicate): + continue + _old_root, _old_src, old_dst = param_group._grad_gather_plans[idx] + for r, (off, sz) in enumerate(self._shard_ranges[mi]): + old_lo, old_hi = old_dst[r] + if sz > 0: + assert (off, off + sz) == (old_lo, old_hi), ( + f"[muon-refactor] param {mi} rank {r}: DTensor range " + f"{(off, off + sz)} != old plan {(old_lo, old_hi)}" + ) + else: + assert old_hi <= old_lo, ( + f"[muon-refactor] param {mi} rank {r}: DTensor says empty " + f"but old plan has {(old_lo, old_hi)}" + ) + @torch.no_grad() def orthogonalize(self, param, grad: torch.Tensor) -> torch.Tensor: """Newton-Schulz orthogonalize + Muon spectral scale of a 2D gradient. @@ -221,13 +314,6 @@ def step(self, closure=None): weight_shard.mul_(1.0 - lr * self._weight_decay) weight_shard.add_(orth_shard.to(weight_shard.dtype), alpha=-lr) - refreshed = set() - for param_group, _, _, _ in items: - if id(param_group) in refreshed or param_group.main_weight_buffer is None: - continue - param_group.copy_main_weights_to_model_weights() - refreshed.add(id(param_group)) - class FullyShardV2MuonOptimizer(MegatronOptimizer): """MegatronOptimizer adapter wrapping FullyShardV2Muon for ChainedOptimizer. @@ -238,9 +324,29 @@ class FullyShardV2MuonOptimizer(MegatronOptimizer): return ``[]``) — orthogonalized updates are already self-normalized. """ - def __init__(self, muon: "FullyShardV2Muon", config, model_chunks=None): - super().__init__(muon, config) + def __init__(self, config, model_chunks, **muon_hyperparams): self.model_chunks = list(model_chunks) if model_chunks else [] + # Pull every chunk's 2D matrix dist_param + its grad DTensor straight from the + # FSDP v2 ParameterGroups, in module-declaration (== layer) order, and build the + # inner torch optimizer purely from those DTensors (no factory plumbing needed). + from megatron.core.distributed.fsdp.src.megatron_fsdp.v2 import FSDPModule + + params, grads = [], [] + for chunk in self.model_chunks: + root = chunk if isinstance(chunk, FSDPModule) else chunk.module + for m in root.modules(): + if not isinstance(m, FSDPModule): + continue + for pg in m._fsdp_param_groups: + for dp, dg in zip(pg.dist_params, pg.dist_grads): + if ( + dg is not None + and dp.dim() == 2 + and not getattr(dp, "is_embedding_or_output_parameter", False) + ): + params.append(dp) + grads.append(dg) + super().__init__(FullyShardV2Muon(params, grads, **muon_hyperparams), config) self.is_stub_optimizer = False # --- Excluded from the chained grad-norm / clip / zero-count machinery. --- @@ -259,6 +365,16 @@ def prepare_grads(self) -> bool: @torch.no_grad() def step_with_ready_grads(self) -> bool: self.optimizer.step() + # Cast the updated fp32 master weights back into the model (bf16) weight buffers + # via the Megatron-FSDP v2 FSDPModule API. The inner optimizer only updates + # master weights; this master->model cast (and all ZeRO-1/2/3 / quantization + # handling) lives in FSDP. _copy_main_weights_to_model_weights recurses over + # child FSDPModules, so one call on each chunk's root FSDPModule suffices. + from megatron.core.distributed.fsdp.src.megatron_fsdp.v2 import FSDPModule + + for model_chunk in self.model_chunks: + fsdp_module = model_chunk if isinstance(model_chunk, FSDPModule) else model_chunk.module + fsdp_module._copy_main_weights_to_model_weights() return True def step(self): From fda44d8894cd82ca880de1836d1f2ded3e2ac3dc Mon Sep 17 00:00:00 2001 From: tongliu Date: Mon, 8 Jun 2026 02:20:33 -0700 Subject: [PATCH 14/28] Make distributed Muon copy-free via batched isend/irecv 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 --- .../core/optimizer/fully_shard_v2_muon.py | 559 ++++++++++++------ .../megatron_fsdp/v2/test_dist_muon.py | 35 +- 2 files changed, 413 insertions(+), 181 deletions(-) diff --git a/megatron/core/optimizer/fully_shard_v2_muon.py b/megatron/core/optimizer/fully_shard_v2_muon.py index d9c00d20973..a4aa95af5d4 100644 --- a/megatron/core/optimizer/fully_shard_v2_muon.py +++ b/megatron/core/optimizer/fully_shard_v2_muon.py @@ -17,8 +17,8 @@ from typing import List import torch +import torch.distributed as dist from torch.distributed.tensor import DTensor -from torch.distributed.tensor.placement_types import Replicate from .emerging_optimizers import HAVE_EMERGING_OPTIMIZERS from .optimizer import MegatronOptimizer @@ -69,6 +69,10 @@ class FullyShardV2Muon(torch.optim.Optimizer): scaling config used by ``orthogonalize``. fp32_matmul_prec / tp_mode: reserved for the Phase-2 tensor-parallel NS path (currently TP size 1; accepted and ignored). + packages: list of param-index lists grouping params for batched P2P (one + batch_isend_irecv per package; packages pipeline so a package's gather + overlaps the previous package's NS). Every package must share one dp_group. + None -> one package per dp_group. """ def __init__( @@ -86,26 +90,25 @@ def __init__( extra_scale_factor: float = 1.0, fp32_matmul_prec: str = "medium", tp_mode: str = "duplicated", + packages=None, ) -> None: super().__init__(params, dict(lr=lr, weight_decay=weight_decay)) - # Refactor contract: this optimizer is driven purely by DTensors. Assert the - # two assumptions the new design rests on — every param (FSDP-v2 dist_param) - # and every grad is a DTensor, and grads align 1:1 with the flattened params. - flat_params = [p for grp in self.param_groups for p in grp["params"]] + flat_params = [param for group in self.param_groups for param in group["params"]] assert all( isinstance(p, DTensor) for p in flat_params ), "FullyShardV2Muon expects every param to be a DTensor (FSDP-v2 dist_param)." + # grads align 1:1 with params; each a Shard(0) DTensor, or None if our shard is empty. assert all( - isinstance(g, DTensor) for g in grads - ), "FullyShardV2Muon expects every grad to be a DTensor." + g is None or isinstance(g, DTensor) for g in grads + ), "FullyShardV2Muon expects every grad to be a DTensor or None." assert len(grads) == len(flat_params), ( f"grads ({len(grads)}) must align 1:1 with params ({len(flat_params)})." ) - self._grads = list(grads) + self._main_grads = list(grads) - # Muon update hyperparameters owned by this class. lr is NOT stored — it - # lives in param_groups[0]["lr"] (where an LR scheduler writes), read in step. + # Muon hyperparameters. lr is NOT stored — it lives in param_groups[0]["lr"] + # (where the LR scheduler writes), read in step(). self._momentum_coef = momentum self._nesterov = nesterov self._weight_decay = weight_decay @@ -115,112 +118,304 @@ def __init__( self._scale_mode = scale_mode self._extra_scale_factor = extra_scale_factor - # Resolve the managed 2D matrix params to their FSDP ParameterGroup via - # the dist_param back-references, and PRE-ALLOCATE the (sharded) momentum - # buffer for each one here at init (keeps step() free of allocation / - # lazy-init on the hot path). Iterating self.param_groups in a fixed order - # keeps the collective send/recv in phases 2/3a aligned across ranks. - # Each entry: (ParameterGroup, idx, original param, momentum shard). + if not flat_params: + raise ValueError("FullyShardV2Muon got no parameters to manage.") + self._main_weights = list(flat_params) + + # Per param: (dp_group, this rank's rank within it, that group's world size). + self._comm = [] + for main_weight in self._main_weights: + group = main_weight.device_mesh.get_group() + self._comm.append((group, dist.get_rank(group), dist.get_world_size(group))) + + # Per-rank (flat_offset, size) of each param + a load-balanced NS root + # (layout-only, so identical on every rank). + self._shard_ranges = self._compute_shard_ranges() + self._roots = self._assign_roots() + + # Persistent full-grad buffer per SPLIT param rooted here (the root assembles its + # shards for NS). Single-holder params need none — their grad shard already IS the + # full grad. + self._full_grads = {} + for i, main_weight in enumerate(self._main_weights): + _group, this_rank, _world_size = self._comm[i] + if self._roots[i] != this_rank: + continue + holders = [r for r, (_offset, size) in enumerate(self._shard_ranges[i]) if size > 0] + if len(holders) <= 1: + continue + grad = self._main_grads[i] + dtype = grad.dtype if grad is not None else main_weight.dtype + self._full_grads[i] = torch.empty( + main_weight.shape.numel(), dtype=dtype, device=main_weight.device + ) + + # Each entry: (master weight, grad or None, momentum, post_momentum_shard). + # post_momentum_shard is THIS rank's slice of the momentum-adjusted gradient — + # both where Phase 1 puts it AND where gather reads it (isend source / a single + # root's full grad), so there is no copy: + # nesterov: grad + coef*m (a fresh value) -> a split root's full-grad slice, else + # the grad shard written in place; momentum is a separate buffer. + # non-nesterov: == m -> it IS the momentum buffer (which for a split root already + # lives in the full-grad slice); Phase 1 just updates it, nothing else. self._managed = [] - self._managed_grads = [] # grad DTensor for each managed entry (aligned with _managed) - gi = 0 # running index into self._grads (== flattened params order) - for group in self.param_groups: - for dist_param in group["params"]: - grad_dtensor = self._grads[gi] - gi += 1 - param_group = getattr(dist_param, "_fsdp_param_group", None) - orig = getattr(dist_param, "_fsdp_orig_param", None) - if param_group is None or orig is None: - continue # not an FSDP-v2 param - if param_group.sharding_strategy == "no_shard": - raise NotImplementedError( - "FullyShardV2Muon does not support the 'no_shard' strategy." - ) - if param_group.main_grad_buffer is None: - continue - if ( - orig.requires_grad - and orig.dim() == 2 - and not getattr(orig, "is_embedding_or_output_parameter", False) - ): - idx = param_group.param_idx[orig] - grad_shard = param_group.main_grad_buffer.get_item(idx, as_shard=True) - momentum_shard = torch.zeros_like(grad_shard) - self._managed.append((param_group, idx, orig, momentum_shard)) - self._managed_grads.append(grad_dtensor) - - if not self._managed: - raise ValueError("FullyShardV2Muon got no 2D matrix parameters to manage.") - - # Reconstruct, purely from the grad DTensors, every DP rank's (flat_offset, - # size) slice of each managed param — the global sharding the gather/scatter - # needs. A DTensor only knows its OWN rank's shard, so one all_gather per - # dp_group recovers the rest. (Root assignment is computed in a later step.) - self._shard_ranges = self._gather_shard_ranges() - self._crosscheck_shard_ranges_against_param_group() - - def _gather_shard_ranges(self): - """Per managed param, every DP rank's (flat_offset, size) within the full - param, derived from the grad DTensor sharding + one all_gather per dp_group. - - FSDP-v2 shards 2D grads as Shard(0) over the flattened buffer, aligned so each - rank owns whole rows (chunk_size_factor = lcm(cols)); to_local() is therefore a - contiguous block of full rows. A DTensor only knows its OWN rank's row count, so - we all_gather row counts within each grad's dp_group and prefix-sum them. - Replicated grads (ZeRO-1) live fully on every rank -> (0, numel) for all ranks. + for i in range(len(self._main_weights)): + main_weight = self._main_weights[i] + grad = self._main_grads[i] + _group, this_rank, _world_size = self._comm[i] + this_offset, this_size = self._shard_ranges[i][this_rank] + dtype = grad.dtype if grad is not None else main_weight.dtype + # this rank's slice inside the preallocated full grad (split root only, else None) + grad_slice = ( + self._full_grads[i].reshape(-1)[this_offset : this_offset + this_size] + if i in self._full_grads + else None + ) + if self._nesterov: + momentum = torch.zeros(this_size, dtype=dtype, device=main_weight.device) + # nesterov's post-momentum (grad + coef*m) is a fresh value, written into a + # split root's full-grad slice, else the grad shard in place. + post_momentum_shard = grad_slice if grad_slice is not None else ( + grad.to_local().reshape(-1) if grad is not None else None + ) + else: + # non-nesterov post-momentum == m, so momentum IS the post-momentum; put it + # where gather reads it (the full-grad slice for a split root, else a + # standalone buffer) and let post_momentum_shard alias it. + if grad_slice is not None: + momentum = grad_slice + momentum.zero_() + else: + momentum = torch.zeros(this_size, dtype=dtype, device=main_weight.device) + post_momentum_shard = momentum + self._managed.append((main_weight, grad, momentum, post_momentum_shard)) + + # `packages`: list of param-index lists (into self._main_weights). Each package + # does ONE batched isend/irecv; packages pipeline (package i+1's gather overlaps + # package i's NS). Every package must share one dp_group. The wrapper passes one + # package per FSDPModule (== layer), so layers pipeline: + # packages = [[0, 1, 2], [3, 4, 5]] # layer 0 | layer 1 + # Default (None): one package per dp_group, e.g. dense vs expert params: + # packages = [[0, 1, 3], [2, 4]] # dp-group | edp-group; no intra-group pipeline + if packages is None: + by_group = OrderedDict() + for param_idx in range(len(self._main_weights)): + by_group.setdefault(self._comm[param_idx][0], []).append(param_idx) + packages = list(by_group.values()) + self._packages = [list(p) for p in packages] + + # Precompute the static per-package isend/irecv plan; step() only fires the P2Ps. + self._gather_plans, self._scatter_plans = self._build_comm_plans() + + def _compute_shard_ranges(self): + """Compute, per param, each DP rank's (flat_offset, size) within the full param. + + Steps: + 1. group params by dp_group (one all_gather per group); + 2. all_gather this rank's per-param row counts -> all_ranks_rows[rank][j]; + 3. per param: size = rows * row_numel; prefix-sum the sizes into (offset, size). + + Example (world_size=4, param shape (10, 8) so row_numel=8, rows [3, 3, 2, 2]): + ranges = [(0, 24), (24, 24), (48, 16), (64, 16)] + # rank0 rank1 rank2 rank3 (size == 0 => no shard) """ - import torch.distributed as dist - - by_group = OrderedDict() # dp process group -> [managed indices] - for mi, g in enumerate(self._managed_grads): - by_group.setdefault(g.device_mesh.get_group(), []).append(mi) - - ranges = [None] * len(self._managed_grads) - for grp, mis in by_group.items(): - ws = dist.get_world_size(grp) - my_rows = [] # this rank's row count per param (-1 sentinel = replicated) - for mi in mis: - g = self._managed_grads[mi] - my_rows.append(-1 if isinstance(g.placements[0], Replicate) else g.to_local().shape[0]) - gathered = [None] * ws - dist.all_gather_object(gathered, my_rows, group=grp) - for j, mi in enumerate(mis): - g = self._managed_grads[mi] + params_by_group = OrderedDict() + for param_idx in range(len(self._main_weights)): + params_by_group.setdefault(self._comm[param_idx][0], []).append(param_idx) + ranges = [None] * len(self._main_weights) + for group, param_indices in params_by_group.items(): + world_size = dist.get_world_size(group) + this_rank_rows = [self._main_weights[p].to_local().shape[0] for p in param_indices] + all_ranks_rows = [None] * world_size # all_ranks_rows[rank][j] = rank's rows for param j + dist.all_gather_object(all_ranks_rows, this_rank_rows, group=group) + for j, param_idx in enumerate(param_indices): + main_weight = self._main_weights[param_idx] row_numel = 1 - for d in g.shape[1:]: - row_numel *= d - rows_per_rank = [gathered[r][j] for r in range(ws)] - if rows_per_rank[0] == -1: # replicated: full param on every rank - ranges[mi] = [(0, g.shape.numel())] * ws - else: - rr, off = [], 0 - for r in range(ws): - sz = rows_per_rank[r] * row_numel - rr.append((off, sz)) - off += sz - ranges[mi] = rr + for dim in main_weight.shape[1:]: + row_numel *= dim + rank_ranges, offset = [], 0 + for rank in range(world_size): + size = all_ranks_rows[rank][j] * row_numel + rank_ranges.append((offset, size)) + offset += size + ranges[param_idx] = rank_ranges return ranges - def _crosscheck_shard_ranges_against_param_group(self): - """TEMP (drop once step() is rewired): assert the DTensor-derived shard ranges - match the ParameterGroup's old ``_grad_gather_plans`` dst ranges + holders, - proving the pure-DTensor reconstruction is correct. Skips replicated (ZeRO-1).""" - for mi, (param_group, idx, _orig, _mom) in enumerate(self._managed): - if isinstance(self._managed_grads[mi].placements[0], Replicate): + def _assign_roots(self): + """Assign each param an NS root rank, load-balancing NS work per dp_group. + + Steps: + 1. group params by dp_group; track each rank's NS load (cost ~ numel*min(dim)); + 2. pin single-holder params to their only holder (zero gather/scatter) and add + their cost to that rank's load; + 3. assign split params heaviest-first (LPT) to the least-loaded holder (ties: + the holder with more of this param's data, then lowest rank). + + Deterministic + identical on every rank (depends only on shard layout). + + Example (world_size=2): + param A: single holder {0}, cost 100 -> root 0 (load [100, 0]) + param B: split {0,1}, cost 300 -> root 1 (idler) (load [100, 300]) + param C: split {0,1}, cost 200 -> root 0 (load [300, 300]) + """ + params_by_group = OrderedDict() + for param_idx in range(len(self._main_weights)): + params_by_group.setdefault(self._comm[param_idx][0], []).append(param_idx) + roots = [0] * len(self._main_weights) + for _group, param_indices in params_by_group.items(): + world_size = self._comm[param_indices[0]][2] + load = [0] * world_size + split = [] + for param_idx in param_indices: + holders = [ + rank + for rank, (_offset, size) in enumerate(self._shard_ranges[param_idx]) + if size > 0 + ] + shape = self._main_weights[param_idx].shape + ns_cost = shape.numel() * min(shape) if len(shape) == 2 else 0 + if len(holders) == 1: + roots[param_idx] = holders[0] + load[holders[0]] += ns_cost + elif len(holders) > 1: + split.append((ns_cost, param_idx, holders)) + for ns_cost, param_idx, holders in sorted(split, key=lambda e: (-e[0], e[1])): + ranges = self._shard_ranges[param_idx] + best = min(holders, key=lambda rank: (load[rank], -ranges[rank][1], rank)) + roots[param_idx] = best + load[best] += ns_cost + return roots + + def _build_comm_plans(self): + """Precompute the STATIC isend/irecv plan per package; step() just fires the P2Ps, + which read/write the grad / full-grad / orth tensors directly (no pack, no alloc — + split-root full grads are preallocated in __init__). + + Both ends walk params in ``package`` order so same-(sender, receiver) P2Ps pair + up, and each plan is the symmetric {group, send, recv}. + 1. gather — route each rank's grad shard to that param's root: + - root == this rank, SPLIT: irecv the OTHER holders' shards into the + preallocated full grad (Phase 1 wrote our own segment). + - root == this rank, SINGLE holder: no P2P — the grad shard already IS the + full grad. + - else we hold a shard: isend it to the root. + 2. scatter — the inverse: isend each OTHER holder's shard out of the orth, irecv + our shard into the grad. Our own shard is read straight from the orth by + Phase 3b, never touching the grad. + + Returns (gather_plans, scatter_plans), parallel to self._packages; each a dict: + gather: group, + send = [(param_idx, root)], # isend grad.to_local() -> root + recv = [(param_idx, src, offset, size)], # irecv src shard -> full_grad[offset:] + scatter: group, + send = [(param_idx, dst, offset, size)], # isend orth[offset:] -> dst + recv = [(param_idx, root)], # irecv -> grad.to_local() + (_finish_gather derives the rooted params on the fly from self._roots / + self._full_grads.) + """ + gather_plans, scatter_plans = [], [] + for package in self._packages: + group, this_rank, world_size = self._comm[package[0]] + + # ---- gather: each rank's grad shard -> that param's root ---- + send, recv = [], [] + for param_idx in package: + root = self._roots[param_idx] + ranges = self._shard_ranges[param_idx] + if root == this_rank: + if param_idx in self._full_grads: # split: assemble from other holders + for src in range(world_size): + offset, size = ranges[src] + if size == 0 or src == this_rank: # own segment came from Phase 1 + continue + recv.append((param_idx, src, offset, size)) + # single holder: grad shard already is the full grad -> no P2P + elif ranges[this_rank][1] > 0: + send.append((param_idx, root)) + gather_plans.append({"group": group, "send": send, "recv": recv}) + + # ---- scatter: each root's orth shard -> its OTHER holders (inverse of gather); + # the root's own shard is read straight from the orth by Phase 3b ---- + send, recv = [], [] + for param_idx in package: + root = self._roots[param_idx] + ranges = self._shard_ranges[param_idx] + if root == this_rank: + for dst in range(world_size): + offset, size = ranges[dst] + if size == 0 or dst == this_rank: # own shard: used directly, not sent + continue + send.append((param_idx, dst, offset, size)) + elif ranges[this_rank][1] > 0: + recv.append((param_idx, root)) + scatter_plans.append({"group": group, "send": send, "recv": recv}) + return gather_plans, scatter_plans + + @torch.no_grad() + def _issue_gather(self, p): + """Fire the package's gather P2Ps: isend our shards to their roots, irecv remote + shards straight into the preallocated full grad. Returns the in-flight Works.""" + plan = self._gather_plans[p] + group = plan["group"] + ops = [] + for param_idx, root in plan["send"]: + post_momentum_shard = self._managed[param_idx][3] + ops.append(torch.distributed.P2POp( + torch.distributed.isend, post_momentum_shard, root, group=group, + )) + for param_idx, src, offset, size in plan["recv"]: + ops.append(torch.distributed.P2POp( + torch.distributed.irecv, + self._full_grads[param_idx].reshape(-1)[offset : offset + size], src, group=group, + )) + return torch.distributed.batch_isend_irecv(ops) if ops else [] + + @torch.no_grad() + def _finish_gather(self, p, reqs): + """Wait the gather P2Ps and return {param_idx: full_grad (2D)} for params rooted + here — split from the preallocated buffer, single from the grad view.""" + for req in reqs: + req.wait() + package = self._packages[p] + this_rank = self._comm[package[0]][1] + full_grads = {} + for param_idx in package: + if self._roots[param_idx] != this_rank: continue - _old_root, _old_src, old_dst = param_group._grad_gather_plans[idx] - for r, (off, sz) in enumerate(self._shard_ranges[mi]): - old_lo, old_hi = old_dst[r] - if sz > 0: - assert (off, off + sz) == (old_lo, old_hi), ( - f"[muon-refactor] param {mi} rank {r}: DTensor range " - f"{(off, off + sz)} != old plan {(old_lo, old_hi)}" - ) - else: - assert old_hi <= old_lo, ( - f"[muon-refactor] param {mi} rank {r}: DTensor says empty " - f"but old plan has {(old_lo, old_hi)}" - ) + shape = self._main_weights[param_idx].shape + if param_idx in self._full_grads: # split: assembled in the preallocated buffer + full_grads[param_idx] = self._full_grads[param_idx].view(shape) + else: # single holder: this rank's post-momentum shard already is the full grad + post_momentum_shard = self._managed[param_idx][3] + full_grads[param_idx] = post_momentum_shard.view(shape) + return full_grads + + @torch.no_grad() + def _issue_scatter(self, p, orths): + """Fire the package's scatter P2Ps: isend each remote holder's shard out of the + orth, irecv our shards into the grad. (The root's own shard is read straight from + the orth by Phase 3b, never round-tripping the grad.) Returns the in-flight Works. + The isend source ``orths`` is kept alive by the caller's ``owned_orths``.""" + plan = self._scatter_plans[p] + group = plan["group"] + ops = [] + for param_idx, dst, offset, size in plan["send"]: + ops.append(torch.distributed.P2POp( + torch.distributed.isend, + orths[param_idx].reshape(-1)[offset : offset + size], dst, group=group, + )) + for param_idx, root in plan["recv"]: + ops.append(torch.distributed.P2POp( + torch.distributed.irecv, + self._main_grads[param_idx].to_local().reshape(-1), root, group=group, + )) + return torch.distributed.batch_isend_irecv(ops) if ops else [] + + @torch.no_grad() + def _finish_scatter(self, p, reqs): + """Wait the scatter P2Ps (the orth shards have landed in this rank's grad).""" + for req in reqs: + req.wait() @torch.no_grad() def orthogonalize(self, param, grad: torch.Tensor) -> torch.Tensor: @@ -250,66 +445,76 @@ def orthogonalize(self, param, grad: torch.Tensor) -> torch.Tensor: def step(self, closure=None): """Run one Muon update over all managed 2D matrix params. - Per param, in order: - 1. momentum - update this rank's grad shard locally (stays sharded); - 2. gather - send the full grad to the param's root rank; - 3. NS - root orthogonalizes the full grad (Newton-Schulz); - 4. scatter - root sends each orthogonalized shard back to its holder; - 5. update - each rank applies a local decoupled-WD step to its master shard. + Per param: (1) momentum on this rank's grad shard; (2) gather the full grad to + its root; (3) root orthogonalizes (NS); (4) scatter the orthogonalized shards + back; (5) decoupled-WD update of this rank's master shard. - Gather/scatter are async: all gathers (step 2) fire up front so the later ones - overlap the NS (step 3) of earlier params. + Gather/scatter are batched per package (one async batch_isend_irecv each, into/out + of the grad and full-grad tensors directly) and pipelined across packages: package + i+1's gather (NCCL stream) is fired before package i's NS (compute stream) so the + two overlap. """ assert closure is None, "FullyShardV2Muon does not support a closure." lr = self.param_groups[0]["lr"] # set by the LR scheduler items = self._managed - # Phase 1: in-place momentum on each rank's grad shard, so gathers see post-momentum grad. - for param_group, idx, _, momentum in items: - grad_shard = param_group.main_grad_buffer.get_item(idx, as_shard=True) - if grad_shard.numel() > 0: - momentum.mul_(self._momentum_coef).add_(grad_shard) - if self._nesterov: - grad_shard.add_(momentum, alpha=self._momentum_coef) - else: - grad_shard.copy_(momentum) - - # Phase 2: fire ALL gathers up front so later ones overlap Phase-3 NS. - gather_work = [] # (param_group, param, full_grad_or_None, reqs) - for param_group, idx, param, _ in items: - full_grad, reqs = param_group.unshard_grad_to_root_async(param) - gather_work.append((param_group, param, full_grad, reqs)) - - # Phase 3: root waits its own gather, orthogonalizes, fires scatter; non-root fires recv. - scatter_work = [] # (reqs, orth_keepalive) - for param_group, param, full_grad, reqs in gather_work: - if full_grad is not None: # this rank is the param's root - for req in reqs: - req.wait() - orth = self.orthogonalize(param, full_grad).to(full_grad.dtype) - scatter_work.append((param_group.scatter_grad_from_root_async(param, orth), orth)) - else: # non-root holder / uninvolved - scatter_work.append((param_group.scatter_grad_from_root_async(param, None), None)) - - # Drain: non-root gather isends, then all scatters (recvs wrote each grad shard in place). - for _param_group, _param, _full_grad, reqs in gather_work: - for req in reqs: - req.wait() - for reqs, _orth in scatter_work: - for req in reqs: - req.wait() - - # Phase 3b: decoupled-WD master-shard update from the orthogonalized grad, then refresh weights. - for param_group, idx, _, _ in items: - wbuf = ( - param_group.main_weight_buffer - if param_group.main_weight_buffer is not None - else param_group.model_weight_buffer - ) - weight_shard = wbuf.get_item(idx, as_shard=True) + # Phase 1: momentum on each rank's grad shard, producing post_momentum_shard with + # no copy. nesterov writes the look-ahead grad + coef*m into it (a split root's + # full-grad slice / the grad shard). non-nesterov's value == m and + # post_momentum_shard already aliases momentum, so updating momentum is all we do. + coef = self._momentum_coef + for _main_weight, grad, momentum, post_momentum_shard in items: + if grad is None: + continue + grad_shard = grad.to_local().reshape(-1) + if grad_shard.numel() == 0: + continue + momentum.mul_(coef).add_(grad_shard) + if self._nesterov: + torch.add(grad_shard, momentum, alpha=coef, out=post_momentum_shard) + + # Phase 2/3 (per-package batched P2P, pipelined): one batch_isend_irecv per + # package gathers grads to roots; NS runs there; one more scatters orth shards + # back. Each package's gather is fired before the previous package's NS so the + # P2Ps (NCCL stream) overlap the NS GEMMs (compute stream). + num_packages = len(self._packages) + scatter_reqs = [] # (package index, in-flight scatter Works) + owned_orths = {} # param_idx -> full orth this rank rooted (read by Phase 3b) + gather_reqs = self._issue_gather(0) if num_packages else None + for i in range(num_packages): + next_gather_reqs = self._issue_gather(i + 1) if i + 1 < num_packages else None + full_grads = self._finish_gather(i, gather_reqs) + orths = { + param_idx: self.orthogonalize(self._main_weights[param_idx], full_grad).to( + full_grad.dtype + ) + for param_idx, full_grad in full_grads.items() + } + owned_orths.update(orths) + scatter_reqs.append((i, self._issue_scatter(i, orths))) + gather_reqs = next_gather_reqs + + # Drain scatters: each non-root holder now has its orth shard irecv'd into its grad. + for i, reqs in scatter_reqs: + self._finish_scatter(i, reqs) + + # Phase 3b: decoupled-WD update of each rank's master shard (fp32) from its orth + # shard. The root reads its own segment straight out of the orth it computed (no + # grad round-trip); every other holder reads the orth shard scatter delivered into + # its grad. + for param_idx, (main_weight, grad, _momentum, _post_momentum_shard) in enumerate(items): + if grad is None: + continue + weight_shard = main_weight.to_local().reshape(-1) if weight_shard.numel() == 0: continue - orth_shard = param_group.main_grad_buffer.get_item(idx, as_shard=True) + orth = owned_orths.get(param_idx) + if orth is not None: # this rank is the root: read its own orth segment directly + this_rank = self._comm[param_idx][1] + own_offset, own_size = self._shard_ranges[param_idx][this_rank] + orth_shard = orth.reshape(-1)[own_offset : own_offset + own_size] + else: # non-root holder: scatter delivered the orth shard into the grad + orth_shard = grad.to_local().reshape(-1) if self._weight_decay != 0.0: weight_shard.mul_(1.0 - lr * self._weight_decay) weight_shard.add_(orth_shard.to(weight_shard.dtype), alpha=-lr) @@ -331,22 +536,32 @@ def __init__(self, config, model_chunks, **muon_hyperparams): # inner torch optimizer purely from those DTensors (no factory plumbing needed). from megatron.core.distributed.fsdp.src.megatron_fsdp.v2 import FSDPModule - params, grads = [], [] + # One package per FSDPModule (== fsdp unit / transformer layer), so each package's + # gather can pipeline with the next layer's NS. A single FSDPModule's params share + # one dp_group (the per-package batch_isend_irecv requirement). + params, grads, packages = [], [], [] for chunk in self.model_chunks: root = chunk if isinstance(chunk, FSDPModule) else chunk.module for m in root.modules(): if not isinstance(m, FSDPModule): continue - for pg in m._fsdp_param_groups: - for dp, dg in zip(pg.dist_params, pg.dist_grads): - if ( - dg is not None - and dp.dim() == 2 - and not getattr(dp, "is_embedding_or_output_parameter", False) + package = [] + for param_group in m._fsdp_param_groups: + for dist_param, dist_grad in zip(param_group.dist_params, param_group.dist_grads): + # Filter by the PARAM's global attrs only (rank-consistent); the + # grad may be None where this rank's shard is empty — keep it + # aligned so every rank's managed set matches (collectives align). + if dist_param.dim() == 2 and not getattr( + dist_param, "is_embedding_or_output_parameter", False ): - params.append(dp) - grads.append(dg) - super().__init__(FullyShardV2Muon(params, grads, **muon_hyperparams), config) + package.append(len(params)) + params.append(dist_param) + grads.append(dist_grad) + if package: + packages.append(package) + super().__init__( + FullyShardV2Muon(params, grads, packages=packages, **muon_hyperparams), config + ) self.is_stub_optimizer = False # --- Excluded from the chained grad-norm / clip / zero-count machinery. --- diff --git a/tests/unit_tests/distributed/megatron_fsdp/v2/test_dist_muon.py b/tests/unit_tests/distributed/megatron_fsdp/v2/test_dist_muon.py index bd56b130a6a..a18068b247f 100644 --- a/tests/unit_tests/distributed/megatron_fsdp/v2/test_dist_muon.py +++ b/tests/unit_tests/distributed/megatron_fsdp/v2/test_dist_muon.py @@ -56,15 +56,17 @@ STRATEGIES = ["optim", "optim_grads", "optim_grads_params"] -# Each set is a list of param shapes at a different scale: a wide matrix -# (rows < cols), a tall one (rows > cols -> NS transposes), and a trailing 1D bias -# Muon must skip. Sizes are picked so that under grad sharding (ZeRO-2/3) each -# matrix straddles a shard boundary and is split across ranks, so the gather/scatter -# is actually exercised. +# Each set is a list of param shapes at a different scale: four 2D matrices (mixing +# wide rowscols -> NS transposes) + a trailing 1D bias Muon must +# skip. Four matrices let the test split them into 2 multi-param packages, exercising +# both the batched all_to_all packing and the cross-package pipeline. Sizes are picked +# so each matrix straddles a shard boundary (ZeRO-2/3), so gather/scatter is real. SHAPE_SETS = [ - pytest.param([(5, 8), (9, 8), (8,)], id="small"), # dims < 10 - pytest.param([(200, 256), (400, 256), (256,)], id="medium"), # dims in the hundreds - pytest.param([(1500, 2048), (3000, 2048), (2048,)], id="large"), # dims in the thousands + pytest.param([(5, 8), (9, 8), (8, 6), (7, 8), (8,)], id="small"), # dims < 10 + pytest.param([(200, 256), (400, 256), (256, 300), (256, 128), (256,)], id="medium"), + pytest.param( + [(1500, 2048), (3000, 2048), (2048, 1200), (2048, 2600), (2048,)], id="large" + ), ] @@ -113,8 +115,23 @@ def test_dist_muon_matches_reference(strategy, nesterov, shapes): mesh=None, sharding_strategy=strategy, ) + # Pass only the 2D-matrix dist_params + their grad DTensors (in production the + # FullyShardV2MuonOptimizer wrapper does this filtering). The 1D bias is excluded, + # so the optimizer never touches it. grad DTensors may be None on ranks whose shard + # of a param is empty. + muon_params, muon_grads = [], [] + for dp, dg in zip(pg.dist_params, pg.dist_grads): + if dp.dim() == 2: + muon_params.append(dp) + muon_grads.append(dg) + # Split the matrices into two multi-param packages to exercise both the batched + # all_to_all packing and the cross-package pipeline. + half = len(muon_params) // 2 + packages = [list(range(half)), list(range(half, len(muon_params)))] opt = FullyShardV2Muon( - pg.dist_params, + muon_params, + muon_grads, + packages=packages, lr=lr, momentum=momentum, nesterov=nesterov, From 811cac7fdcbabfc1aa53a4b52859d10d706f24fd Mon Sep 17 00:00:00 2001 From: tongliu Date: Mon, 8 Jun 2026 02:59:03 -0700 Subject: [PATCH 15/28] Clean up Muon: trim comments + drop dead single-root plumbing MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 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 --- .../fsdp/src/megatron_fsdp/v2/dp_buffer.py | 73 ---------- .../fsdp/src/megatron_fsdp/v2/param_group.py | 136 ------------------ .../core/optimizer/fully_shard_v2_muon.py | 73 ++++------ 3 files changed, 26 insertions(+), 256 deletions(-) diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/dp_buffer.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/dp_buffer.py index c118051666f..9771267566f 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/dp_buffer.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/dp_buffer.py @@ -61,79 +61,6 @@ def __init__( self.shard_meta = self._build_shard_meta( self.bucket_meta, is_distributed, dp_world_size, dp_rank ) - self.item_root_ranks = self._build_item_root_ranks() - - def _build_item_root_ranks(self) -> Dict[int, int]: - """Assign each item a root rank, greedily balancing Newton-Schulz load. - - Only 2D matrices run NS (on their root rank), and that cost dominates the - single-root Muon step. NS flops scale ~ ``numel * min(dim)`` per matrix - (constant #steps dropped). We use Longest-Processing-Time greedy: assign - items heaviest-first to the currently least-loaded *holder* rank. - - Properties: - * Items with a single holder (not split across a shard boundary) are - forced onto their owner -> zero gather/scatter, comm unchanged. - * Only split items have a real choice; among holders we pick the - least-loaded, tie-broken toward the largest holder so an indifferent - balance still minimises the gather/scatter volume (old behaviour). - * Non-2D items get cost 0 (their root is never used by Muon). - - Must be deterministic and identical on every rank (gather/scatter must - agree on the root): inputs (param_shapes, shard_size, dp_world_size) are - rank-independent, and the sort/min tie-breaks are fully ordered by item_id. - """ - shard_size = self.bucket_meta.size // self.dp_world_size - load = [0] * self.dp_world_size - root_ranks: Dict[int, int] = {} - split_items = [] - - # Pass 1: a param living entirely on one rank is rooted there -- no choice, - # zero gather/scatter. Account its NS cost first so the split params in - # pass 2 balance on top of the complete fixed (non-split) load per rank. - for item_id, idx in self.item_index_map.items(): - item_start = idx.global_data_index - item_end = item_start + idx.size - overlaps = [ - max(0, min(item_end, (r + 1) * shard_size) - max(item_start, r * shard_size)) - for r in range(self.dp_world_size) - ] - holders = [r for r, o in enumerate(overlaps) if o > 0] - ns_cost = idx.size * min(idx.shape) if len(idx.shape) == 2 else 0 - if len(holders) == 1: - root_ranks[item_id] = holders[0] - load[holders[0]] += ns_cost - else: - split_items.append((ns_cost, item_id, holders, overlaps)) - - # Pass 2: only split params have a choice. LPT (heaviest first; ties by - # item_id for cross-rank determinism) -> assign each to its least-loaded - # holder (ties -> largest holder = min gather/scatter volume -> lowest rank). - for ns_cost, item_id, holders, overlaps in sorted(split_items, key=lambda e: (-e[0], e[1])): - best = min(holders, key=lambda r: (load[r], -overlaps[r], r)) - root_ranks[item_id] = best - load[best] += ns_cost - - # TODO(muon-debug): temporary one-time dump of the NS root assignment so we - # can see whether split-param roots actually land on different ranks. Drop - # once the load-balance behaviour is confirmed. - if self.dp_rank == 0: - twod = [iid for iid, idx in self.item_index_map.items() if len(idx.shape) == 2] - if twod: - roots_per_rank = [0] * self.dp_world_size - for iid in twod: - roots_per_rank[root_ranks[iid]] += 1 - split_dump = [ - (iid, tuple(self.item_index_map[iid].shape), holders, root_ranks[iid]) - for _, iid, holders, _ in sorted(split_items, key=lambda e: (-e[0], e[1])) - ] - logger.warning( - "[muon-root] pg=%s ws=%d: %d 2D items (%d split). " - "NS-load/rank=%s 2D-roots/rank=%s | split(item,shape,holders,root)=%s", - self.param_group_id, self.dp_world_size, len(twod), len(split_items), - load, roots_per_rank, split_dump, - ) - return root_ranks # ------------------------------------------------------------------ # # Layout construction (global, rank-independent) diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/param_group.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/param_group.py index eb39c469131..ce86674abea 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/param_group.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/param_group.py @@ -279,90 +279,6 @@ def release_grad_buffer(self): del param.main_grad self.main_grad_buffer.reshard() - @torch.no_grad() - def unshard_grad_to_root_async(self, param): - """Gather this param's grad to its assigned root rank (async, non-blocking). - - All ranks must call this collectively (after reduce_grad), in the same param - order. The cross-rank gather is issued as non-blocking P2P - (``batch_isend_irecv``), which runs on the NCCL backend's own stream and so - overlaps the caller's compute; the root's own-piece copy is a local memcpy on - the current stream. Returns ``(full_grad_or_None, reqs)`` — the full grad on - root (valid only after ``reqs`` complete, so the caller waits them before NS), - None on others; the returned tensors must stay alive until ``reqs`` complete. - """ - idx = self.param_idx[param] - root, src, dst = self._grad_gather_plans[idx] - gbuf = self.main_grad_buffer - ops = [] - if self._dp_rank == root: - full_grad = torch.empty(param.shape, dtype=gbuf.dtype, device=gbuf.device).flatten() - for r in range(self._dp_world_size): - src_lo, src_hi = src[r] - dst_lo, dst_hi = dst[r] - if dst_hi <= dst_lo: - continue # rank r holds no part of this param - if r == root: - full_grad[dst_lo:dst_hi].copy_(gbuf.data[src_lo:src_hi]) # local - else: - ops.append( - torch.distributed.P2POp( - torch.distributed.irecv, full_grad[dst_lo:dst_hi], r, self.dp_group - ) - ) - reqs = torch.distributed.batch_isend_irecv(ops) if ops else [] - return full_grad.view(param.shape), reqs - else: - src_lo, src_hi = src[self._dp_rank] - if src_hi <= src_lo: - return None, [] - ops.append( - torch.distributed.P2POp( - torch.distributed.isend, gbuf.data[src_lo:src_hi], root, self.dp_group - ) - ) - return None, torch.distributed.batch_isend_irecv(ops) - - @torch.no_grad() - def scatter_grad_from_root_async(self, param, data=None): - """Scatter this param's transformed grad from root back (async, non-blocking). - - All ranks must call this collectively, in the same param order. P2P runs on - the NCCL backend's stream (overlaps compute). Root copies its own piece - locally and isends the rest; each non-root holder irecvs its piece straight - into its grad-buffer shard (safe: the matching gather isend that read that - shard was enqueued earlier on the same NCCL stream, so it completes before - this recv writes). Returns the in-flight P2P Works (empty if uninvolved). - """ - idx = self.param_idx[param] - root, src, dst = self._grad_gather_plans[idx] - gbuf = self.main_grad_buffer - ops = [] - if self._dp_rank == root: - full = data.to(dtype=gbuf.dtype).flatten() - for r in range(self._dp_world_size): - src_lo, src_hi = src[r] - dst_lo, dst_hi = dst[r] - if dst_hi <= dst_lo: - continue # rank r holds no part of this param - if r == root: - gbuf.data[src_lo:src_hi].copy_(full[dst_lo:dst_hi]) # local - else: - ops.append( - torch.distributed.P2POp( - torch.distributed.isend, full[dst_lo:dst_hi], r, self.dp_group - ) - ) - else: - src_lo, src_hi = src[self._dp_rank] - if src_hi > src_lo: - ops.append( - torch.distributed.P2POp( - torch.distributed.irecv, gbuf.data[src_lo:src_hi], root, self.dp_group - ) - ) - return torch.distributed.batch_isend_irecv(ops) if ops else [] - def _init_dist_params(self): """ Initialize distributed parameter views (DTensors) into the buffers. @@ -404,11 +320,6 @@ def _init_dist_params(self): # Mark as FSDP parameter for special handling setattr(param, "__fsdp_param__", True) setattr(dist_param, "__fsdp_param__", True) - # Back-references so an optimizer that only sees the optimizer-facing - # dist_param (e.g. FullyShardV2Muon) can reach this ParameterGroup and - # the original param it indexes (for per-param gather/scatter). - dist_param._fsdp_param_group = self - dist_param._fsdp_orig_param = param self.dist_params.append(dist_param) # Update dist_param chunk metadata for checkpointing and debugging. @@ -436,50 +347,3 @@ def _init_dist_params(self): ) else: self.dist_grads.append(None) - - # Per-param send/recv plans from BufferIndex (identical on all ranks, so - # no metadata comm). Entry is (root, src, dst): - # root — DP rank that assembles the full grad. - # src[r] — (lo, hi) of rank r's piece in its LOCAL grad buffer. - # dst[r] — (lo, hi) of that piece in the ASSEMBLED full grad on root. - # Same elements live at gbuf.data[src[r]] (rank r) and full_grad[dst[r]] - # (root); unshard moves src->dst, scatter moves dst->src. Empty range - # (lo==hi) => rank r owns nothing (runtime skips via lo (lo-shard_start, hi-shard_start) (gbuf.data is the shard); - # replicated buffer (ZeRO-1) -> (lo, hi) (gbuf.data is the full buffer). - # Only src[self._dp_rank] ever indexes gbuf.data. - # - # Example (distributed): 4 ranks, shard_size=100, param [150,380). - # rank 0: shard [0,100) ∩ param = ∅ -> src=(0,0) dst=(0,0) - # rank 1: shard [100,200) ∩ [150,200) -> src=(50,100) dst=(0,50) - # rank 2: shard [200,300) ∩ [200,300) -> src=(0,100) dst=(50,150) - # rank 3: shard [300,400) ∩ [300,380) -> src=(0,80) dst=(150,230) - # (rank 1's src starts at 50 because the param begins 50 into - # its shard; its dst starts at 0 because it is the param's head.) - bi = self.main_grad_buffer.buffer_index - is_distributed = self.main_grad_buffer.is_distributed - shard_size = bi.bucket_meta.size // self._dp_world_size - self._grad_gather_plans: List[tuple] = [] - for idx in range(len(self.params)): - item = bi.item_index_map[idx] - root = bi.item_root_ranks[idx] - param_start = item.global_data_index - param_end = param_start + item.size - - src, dst = [], [] - for r in range(self._dp_world_size): - shard_start = r * shard_size - shard_end = shard_start + shard_size - lo = max(param_start, shard_start) - hi = min(param_end, shard_end) - if hi <= lo: - src.append((0, 0)) - dst.append((0, 0)) - else: - if is_distributed: - src.append((lo - shard_start, hi - shard_start)) - else: - src.append((lo, hi)) - dst.append((lo - param_start, hi - param_start)) - self._grad_gather_plans.append((root, src, dst)) diff --git a/megatron/core/optimizer/fully_shard_v2_muon.py b/megatron/core/optimizer/fully_shard_v2_muon.py index a4aa95af5d4..ab4f2c49699 100644 --- a/megatron/core/optimizer/fully_shard_v2_muon.py +++ b/megatron/core/optimizer/fully_shard_v2_muon.py @@ -151,13 +151,11 @@ def __init__( ) # Each entry: (master weight, grad or None, momentum, post_momentum_shard). - # post_momentum_shard is THIS rank's slice of the momentum-adjusted gradient — - # both where Phase 1 puts it AND where gather reads it (isend source / a single - # root's full grad), so there is no copy: - # nesterov: grad + coef*m (a fresh value) -> a split root's full-grad slice, else - # the grad shard written in place; momentum is a separate buffer. - # non-nesterov: == m -> it IS the momentum buffer (which for a split root already - # lives in the full-grad slice); Phase 1 just updates it, nothing else. + # post_momentum_shard is this rank's momentum-adjusted gradient slice, written by + # Phase 1 and read by gather in its FINAL location (no copy). nesterov: a fresh + # grad+coef*m, into the split root's full-grad slice or the grad shard (momentum + # separate). non-nesterov: == m, so it aliases momentum (which for a split root + # lives in the full-grad slice). self._managed = [] for i in range(len(self._main_weights)): main_weight = self._main_weights[i] @@ -173,20 +171,15 @@ def __init__( ) if self._nesterov: momentum = torch.zeros(this_size, dtype=dtype, device=main_weight.device) - # nesterov's post-momentum (grad + coef*m) is a fresh value, written into a - # split root's full-grad slice, else the grad shard in place. post_momentum_shard = grad_slice if grad_slice is not None else ( grad.to_local().reshape(-1) if grad is not None else None ) + elif grad_slice is not None: # split root: momentum lives in the full grad + momentum = grad_slice + momentum.zero_() + post_momentum_shard = momentum else: - # non-nesterov post-momentum == m, so momentum IS the post-momentum; put it - # where gather reads it (the full-grad slice for a split root, else a - # standalone buffer) and let post_momentum_shard alias it. - if grad_slice is not None: - momentum = grad_slice - momentum.zero_() - else: - momentum = torch.zeros(this_size, dtype=dtype, device=main_weight.device) + momentum = torch.zeros(this_size, dtype=dtype, device=main_weight.device) post_momentum_shard = momentum self._managed.append((main_weight, grad, momentum, post_momentum_shard)) @@ -449,19 +442,16 @@ def step(self, closure=None): its root; (3) root orthogonalizes (NS); (4) scatter the orthogonalized shards back; (5) decoupled-WD update of this rank's master shard. - Gather/scatter are batched per package (one async batch_isend_irecv each, into/out - of the grad and full-grad tensors directly) and pipelined across packages: package - i+1's gather (NCCL stream) is fired before package i's NS (compute stream) so the - two overlap. + Gather/scatter are batched per package (one async batch_isend_irecv each, directly + into/out of the grad and full-grad tensors) and pipelined: package i+1's gather is + fired before package i's NS, so the P2Ps overlap the NS GEMMs. """ assert closure is None, "FullyShardV2Muon does not support a closure." lr = self.param_groups[0]["lr"] # set by the LR scheduler items = self._managed - # Phase 1: momentum on each rank's grad shard, producing post_momentum_shard with - # no copy. nesterov writes the look-ahead grad + coef*m into it (a split root's - # full-grad slice / the grad shard). non-nesterov's value == m and - # post_momentum_shard already aliases momentum, so updating momentum is all we do. + # Phase 1: momentum on each rank's grad shard. nesterov also writes the look-ahead + # grad+coef*m into post_momentum_shard; non-nesterov's == m (already aliased). coef = self._momentum_coef for _main_weight, grad, momentum, post_momentum_shard in items: if grad is None: @@ -473,10 +463,7 @@ def step(self, closure=None): if self._nesterov: torch.add(grad_shard, momentum, alpha=coef, out=post_momentum_shard) - # Phase 2/3 (per-package batched P2P, pipelined): one batch_isend_irecv per - # package gathers grads to roots; NS runs there; one more scatters orth shards - # back. Each package's gather is fired before the previous package's NS so the - # P2Ps (NCCL stream) overlap the NS GEMMs (compute stream). + # Phase 2/3: per-package gather -> NS -> scatter, pipelined (gather i+1 before NS i). num_packages = len(self._packages) scatter_reqs = [] # (package index, in-flight scatter Works) owned_orths = {} # param_idx -> full orth this rank rooted (read by Phase 3b) @@ -498,10 +485,8 @@ def step(self, closure=None): for i, reqs in scatter_reqs: self._finish_scatter(i, reqs) - # Phase 3b: decoupled-WD update of each rank's master shard (fp32) from its orth - # shard. The root reads its own segment straight out of the orth it computed (no - # grad round-trip); every other holder reads the orth shard scatter delivered into - # its grad. + # Phase 3b: decoupled-WD master update from each rank's orth shard — the root reads + # its own segment from the orth it computed, other holders read scatter's grad. for param_idx, (main_weight, grad, _momentum, _post_momentum_shard) in enumerate(items): if grad is None: continue @@ -531,14 +516,11 @@ class FullyShardV2MuonOptimizer(MegatronOptimizer): def __init__(self, config, model_chunks, **muon_hyperparams): self.model_chunks = list(model_chunks) if model_chunks else [] - # Pull every chunk's 2D matrix dist_param + its grad DTensor straight from the - # FSDP v2 ParameterGroups, in module-declaration (== layer) order, and build the - # inner torch optimizer purely from those DTensors (no factory plumbing needed). from megatron.core.distributed.fsdp.src.megatron_fsdp.v2 import FSDPModule - # One package per FSDPModule (== fsdp unit / transformer layer), so each package's - # gather can pipeline with the next layer's NS. A single FSDPModule's params share - # one dp_group (the per-package batch_isend_irecv requirement). + # Build the inner optimizer from the FSDP-v2 2D-matrix dist_params + their grad + # DTensors, one package per FSDPModule (== layer) so each layer's gather pipelines + # with the next layer's NS (a package shares one dp_group, as the P2P needs). params, grads, packages = [], [], [] for chunk in self.model_chunks: root = chunk if isinstance(chunk, FSDPModule) else chunk.module @@ -548,9 +530,8 @@ def __init__(self, config, model_chunks, **muon_hyperparams): package = [] for param_group in m._fsdp_param_groups: for dist_param, dist_grad in zip(param_group.dist_params, param_group.dist_grads): - # Filter by the PARAM's global attrs only (rank-consistent); the - # grad may be None where this rank's shard is empty — keep it - # aligned so every rank's managed set matches (collectives align). + # filter by the param's global attrs only (rank-consistent); keep + # grad aligned even when None (empty shard) so collectives match if dist_param.dim() == 2 and not getattr( dist_param, "is_embedding_or_output_parameter", False ): @@ -580,11 +561,9 @@ def prepare_grads(self) -> bool: @torch.no_grad() def step_with_ready_grads(self) -> bool: self.optimizer.step() - # Cast the updated fp32 master weights back into the model (bf16) weight buffers - # via the Megatron-FSDP v2 FSDPModule API. The inner optimizer only updates - # master weights; this master->model cast (and all ZeRO-1/2/3 / quantization - # handling) lives in FSDP. _copy_main_weights_to_model_weights recurses over - # child FSDPModules, so one call on each chunk's root FSDPModule suffices. + # The inner optimizer only updated the fp32 master weights; cast them back into the + # model (bf16) buffers via the v2 FSDPModule API (it recurses over child + # FSDPModules, so one call per chunk root suffices). from megatron.core.distributed.fsdp.src.megatron_fsdp.v2 import FSDPModule for model_chunk in self.model_chunks: From 651aadb74315b0a8bd8cc17c07e6c086a103a34b Mon Sep 17 00:00:00 2001 From: tongliu Date: Mon, 8 Jun 2026 21:42:28 -0700 Subject: [PATCH 16/28] Refine FSDP v2 Muon gather flow --- .../core/optimizer/fully_shard_v2_muon.py | 66 +++++++------------ 1 file changed, 25 insertions(+), 41 deletions(-) diff --git a/megatron/core/optimizer/fully_shard_v2_muon.py b/megatron/core/optimizer/fully_shard_v2_muon.py index ab4f2c49699..7e606ed96bb 100644 --- a/megatron/core/optimizer/fully_shard_v2_muon.py +++ b/megatron/core/optimizer/fully_shard_v2_muon.py @@ -298,24 +298,26 @@ def _build_comm_plans(self): Returns (gather_plans, scatter_plans), parallel to self._packages; each a dict: gather: group, + rooted = [param_idx], # params rooted on this rank send = [(param_idx, root)], # isend grad.to_local() -> root recv = [(param_idx, src, offset, size)], # irecv src shard -> full_grad[offset:] scatter: group, send = [(param_idx, dst, offset, size)], # isend orth[offset:] -> dst recv = [(param_idx, root)], # irecv -> grad.to_local() - (_finish_gather derives the rooted params on the fly from self._roots / - self._full_grads.) + _issue_gather returns the rooted full-grad views along with the in-flight + P2P works; the views become valid after the works complete. """ gather_plans, scatter_plans = [], [] for package in self._packages: group, this_rank, world_size = self._comm[package[0]] # ---- gather: each rank's grad shard -> that param's root ---- - send, recv = [], [] + rooted, send, recv = [], [], [] for param_idx in package: root = self._roots[param_idx] ranges = self._shard_ranges[param_idx] if root == this_rank: + rooted.append(param_idx) if param_idx in self._full_grads: # split: assemble from other holders for src in range(world_size): offset, size = ranges[src] @@ -325,7 +327,7 @@ def _build_comm_plans(self): # single holder: grad shard already is the full grad -> no P2P elif ranges[this_rank][1] > 0: send.append((param_idx, root)) - gather_plans.append({"group": group, "send": send, "recv": recv}) + gather_plans.append({"group": group, "rooted": rooted, "send": send, "recv": recv}) # ---- scatter: each root's orth shard -> its OTHER holders (inverse of gather); # the root's own shard is read straight from the orth by Phase 3b ---- @@ -346,10 +348,15 @@ def _build_comm_plans(self): @torch.no_grad() def _issue_gather(self, p): - """Fire the package's gather P2Ps: isend our shards to their roots, irecv remote - shards straight into the preallocated full grad. Returns the in-flight Works.""" + """Fire the package's gather P2Ps and return rooted full-grad views + Works.""" plan = self._gather_plans[p] group = plan["group"] + full_grads = {} + for param_idx in plan["rooted"]: + shape = self._main_weights[param_idx].shape + full_grad = self._full_grads.get(param_idx, self._managed[param_idx][3]) + full_grads[param_idx] = full_grad.view(shape) + ops = [] for param_idx, root in plan["send"]: post_momentum_shard = self._managed[param_idx][3] @@ -361,27 +368,7 @@ def _issue_gather(self, p): torch.distributed.irecv, self._full_grads[param_idx].reshape(-1)[offset : offset + size], src, group=group, )) - return torch.distributed.batch_isend_irecv(ops) if ops else [] - - @torch.no_grad() - def _finish_gather(self, p, reqs): - """Wait the gather P2Ps and return {param_idx: full_grad (2D)} for params rooted - here — split from the preallocated buffer, single from the grad view.""" - for req in reqs: - req.wait() - package = self._packages[p] - this_rank = self._comm[package[0]][1] - full_grads = {} - for param_idx in package: - if self._roots[param_idx] != this_rank: - continue - shape = self._main_weights[param_idx].shape - if param_idx in self._full_grads: # split: assembled in the preallocated buffer - full_grads[param_idx] = self._full_grads[param_idx].view(shape) - else: # single holder: this rank's post-momentum shard already is the full grad - post_momentum_shard = self._managed[param_idx][3] - full_grads[param_idx] = post_momentum_shard.view(shape) - return full_grads + return full_grads, torch.distributed.batch_isend_irecv(ops) if ops else [] @torch.no_grad() def _issue_scatter(self, p, orths): @@ -404,12 +391,6 @@ def _issue_scatter(self, p, orths): )) return torch.distributed.batch_isend_irecv(ops) if ops else [] - @torch.no_grad() - def _finish_scatter(self, p, reqs): - """Wait the scatter P2Ps (the orth shards have landed in this rank's grad).""" - for req in reqs: - req.wait() - @torch.no_grad() def orthogonalize(self, param, grad: torch.Tensor) -> torch.Tensor: """Newton-Schulz orthogonalize + Muon spectral scale of a 2D gradient. @@ -465,12 +446,14 @@ def step(self, closure=None): # Phase 2/3: per-package gather -> NS -> scatter, pipelined (gather i+1 before NS i). num_packages = len(self._packages) - scatter_reqs = [] # (package index, in-flight scatter Works) + scatter_reqs = [] # in-flight scatter Works per package owned_orths = {} # param_idx -> full orth this rank rooted (read by Phase 3b) - gather_reqs = self._issue_gather(0) if num_packages else None + gather_work = self._issue_gather(0) if num_packages else None for i in range(num_packages): - next_gather_reqs = self._issue_gather(i + 1) if i + 1 < num_packages else None - full_grads = self._finish_gather(i, gather_reqs) + next_gather_work = self._issue_gather(i + 1) if i + 1 < num_packages else None + full_grads, gather_reqs = gather_work + for req in gather_reqs: + req.wait() orths = { param_idx: self.orthogonalize(self._main_weights[param_idx], full_grad).to( full_grad.dtype @@ -478,12 +461,13 @@ def step(self, closure=None): for param_idx, full_grad in full_grads.items() } owned_orths.update(orths) - scatter_reqs.append((i, self._issue_scatter(i, orths))) - gather_reqs = next_gather_reqs + scatter_reqs.append(self._issue_scatter(i, orths)) + gather_work = next_gather_work # Drain scatters: each non-root holder now has its orth shard irecv'd into its grad. - for i, reqs in scatter_reqs: - self._finish_scatter(i, reqs) + for reqs in scatter_reqs: + for req in reqs: + req.wait() # Phase 3b: decoupled-WD master update from each rank's orth shard — the root reads # its own segment from the orth it computed, other holders read scatter's grad. From d0d8e70a5edffd5c2fbda4399853b68b11bb8615 Mon Sep 17 00:00:00 2001 From: jianbinc Date: Tue, 9 Jun 2026 13:28:02 +0800 Subject: [PATCH 17/28] Megatron FSDP v2 built-in CUDA graph, allocator, checkpoint 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 --- examples/megatron_fsdp/README.md | 52 ++ examples/megatron_fsdp/fsdp_toy.py | 22 +- .../distributed_data_parallel_config.py | 11 +- megatron/core/distributed/fsdp/checkpoint.py | 256 +++++++- .../distributed/fsdp/mcore_fsdp_adapter.py | 45 +- .../fsdp/src/megatron_fsdp/uneven_dtensor.py | 30 +- .../fsdp/src/megatron_fsdp/v2/README.md | 71 ++- .../fsdp/src/megatron_fsdp/v2/allocator.py | 586 ++++++++++-------- .../src/megatron_fsdp/v2/allocator_design.md | 538 ++++++++++++++++ .../src/megatron_fsdp/v2/cuda_graph_design.md | 350 +++++++++++ .../src/megatron_fsdp/v2/cuda_graph_runner.py | 373 +++++++++++ .../fsdp/src/megatron_fsdp/v2/design.md | 87 +-- .../fsdp/src/megatron_fsdp/v2/dp_buffer.py | 93 ++- .../fsdp/src/megatron_fsdp/v2/fsdp_module.py | 220 +++++-- .../fsdp/src/megatron_fsdp/v2/fully_shard.py | 23 +- .../fsdp/src/megatron_fsdp/v2/hooks.py | 247 ++++++-- .../v2/mcore_fsdp_checkpoint_design.md | 81 ++- .../src/megatron_fsdp/v2/mixed_precision.py | 74 +-- .../fsdp/src/megatron_fsdp/v2/param_group.py | 35 +- megatron/core/transformer/moe/router.py | 2 +- megatron/training/arguments.py | 17 + megatron/training/checkpointing.py | 24 +- megatron/training/initialize.py | 2 +- .../megatron_fsdp/v2/test_dist_muon.py | 7 +- .../megatron_fsdp/v2/test_param_group.py | 7 +- 25 files changed, 2649 insertions(+), 604 deletions(-) create mode 100644 megatron/core/distributed/fsdp/src/megatron_fsdp/v2/allocator_design.md create mode 100644 megatron/core/distributed/fsdp/src/megatron_fsdp/v2/cuda_graph_design.md create mode 100644 megatron/core/distributed/fsdp/src/megatron_fsdp/v2/cuda_graph_runner.py diff --git a/examples/megatron_fsdp/README.md b/examples/megatron_fsdp/README.md index eaf5eca1364..ad166e65ee6 100644 --- a/examples/megatron_fsdp/README.md +++ b/examples/megatron_fsdp/README.md @@ -4,6 +4,58 @@ Example scripts for training and checkpoint conversion using [Megatron-FSDP](../ ## Scripts +### `fsdp_toy.py` + +Standalone toy example (not Megatron-LM) demonstrating Megatron-FSDP v2 usage: + +- Basic model wrapping with `fully_shard()` +- CUDA graph capture (`--cuda-graph` / `--no-cuda-graph`) + > **Experimental** — CUDA graph support is experimental and may change. +- Activation checkpointing (`--activation-checkpoint`) +- Distributed checkpointing + +```bash +torchrun --nproc_per_node=2 examples/megatron_fsdp/fsdp_toy.py \ + --model-dim 512 --n-layers 2 --batch-size 4 \ + --use-megatron-fsdp + +# With CUDA graph +torchrun --nproc_per_node=2 examples/megatron_fsdp/fsdp_toy.py \ + --model-dim 512 --n-layers 2 --batch-size 4 \ + --use-megatron-fsdp --cuda-graph + +# Without CUDA graph +torchrun --nproc_per_node=2 examples/megatron_fsdp/fsdp_toy.py \ + --model-dim 512 --n-layers 2 --batch-size 4 \ + --use-megatron-fsdp --no-cuda-graph +``` + +| Flag | Default | Description | +|------|---------|-------------| +| `--model-dim` | `1024` | Hidden dimension size | +| `--n-layers` | `3` | Number of transformer layers | +| `--use-megatron-fsdp` | off | Use Megatron-FSDP v2 instead of PyTorch FSDP2 | +| `--cuda-graph` | on | Enable CUDA graph capture on transformer layers | +| `--no-cuda-graph` | — | Disable CUDA graph capture | +| `--activation-checkpoint` | off | Enable activation checkpointing | + +### `qwen3-30b-a3b.gbs128_mbs4_seq4096_n2_mfsdp2_mxfp8_wandb.sh` + +2-node SLURM training script for **Qwen3-30B-A3B** (MoE) using Megatron-FSDP v2 +with MXFP8 precision and Weights & Biases logging. Serves as a reference for +production-scale MoE training with FSDP v2. + +```bash +export MEGATRON_PATH=/path/to/Megatron-LM +export DATA_PATH=/path/to/data/c4/en/c4-train.en_6_text_document +export TOKENIZER_MODEL=/path/to/data/c4/en/tokenizer + +sbatch examples/megatron_fsdp/qwen3-30b-a3b.gbs128_mbs4_seq4096_n2_mfsdp2_mxfp8_wandb.sh +``` + +Update the `#SBATCH` directives and `--container-image` in the script to match +your cluster configuration before submitting. + ### `train_llama3_8b_fsdp_h100_fp8.sh` Single-node training script for **Llama 3 8B** using Megatron-FSDP with FP8 precision on H100 GPUs. Uses `torchrun` for local distributed training and supports both mock data (for benchmarking) and real datasets. diff --git a/examples/megatron_fsdp/fsdp_toy.py b/examples/megatron_fsdp/fsdp_toy.py index c519f224be8..3e224c8a9b6 100644 --- a/examples/megatron_fsdp/fsdp_toy.py +++ b/examples/megatron_fsdp/fsdp_toy.py @@ -84,6 +84,7 @@ def build_fsdp_model( n_layers: int, use_megatron_fsdp: bool, use_activation_checkpointing: bool = False, + enable_cuda_graph: bool = True, ) -> Tuple["FSDPModule", torch.distributed.device_mesh.DeviceMesh]: if use_megatron_fsdp: try: @@ -101,12 +102,18 @@ def build_fsdp_model( if use_activation_checkpointing: model.enable_activation_checkpointing() - # Example: per-layer sharding - for layer in model.layers: - fully_shard(layer, mesh=mesh) + if use_megatron_fsdp: + sublayer_kwargs = dict( + enable_cuda_graph=enable_cuda_graph, + ) + kwargs = {} + else: + sublayer_kwargs = {} + kwargs = {} - # Optionally shard the root as well - fully_shard(model, mesh=mesh) + for layer in model.layers: + fully_shard(layer, mesh=mesh, **sublayer_kwargs) + fully_shard(model, mesh=mesh, **kwargs) assert isinstance(model, ToyModel) assert isinstance(model, FSDPModule) @@ -227,7 +234,7 @@ def train( loss = y.sum() / (world_size * args.batch_size) loss.backward() optimizer.step() - optimizer.zero_grad(set_to_none=True) + optimizer.zero_grad() if step % args.log_interval == 0 and rank == 0: print(f"[rank0] epoch={epoch} step={step} loss={loss.item():.4f}") @@ -259,6 +266,8 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--log-interval", type=int, default=5) parser.add_argument("--use-megatron-fsdp", action="store_true", help="Use Megatron-FSDP instead of PyTorch FSDP2") parser.add_argument("--activation-checkpoint", action="store_true", help="Enable activation checkpointing on transformer layers") + parser.add_argument("--cuda-graph", action="store_true", default=True, help="Enable CUDA graph capture (default: on)") + parser.add_argument("--no-cuda-graph", action="store_false", dest="cuda_graph", help="Disable CUDA graph capture") return parser.parse_args() @@ -270,6 +279,7 @@ def main() -> None: n_layers=args.n_layers, use_megatron_fsdp=args.use_megatron_fsdp, use_activation_checkpointing=args.activation_checkpoint, + enable_cuda_graph=args.cuda_graph, ) optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr) diff --git a/megatron/core/distributed/distributed_data_parallel_config.py b/megatron/core/distributed/distributed_data_parallel_config.py index f4dbcddab2c..70445819e4f 100644 --- a/megatron/core/distributed/distributed_data_parallel_config.py +++ b/megatron/core/distributed/distributed_data_parallel_config.py @@ -1,6 +1,6 @@ # Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import Optional, Tuple import torch @@ -83,6 +83,15 @@ class DistributedDataParallelConfig: use_megatron_fsdp: bool = False """If true, use the FSDP code path for DDP.""" + use_megatron_fsdp_v2: bool = False + """If true, use the `fully_shard` API for FSDP sharding the model. + """ + + mfsdp_cuda_graph_modules: list = field(default_factory=list) + """If set and ``use_megatron_fsdp_v2`` is set, enable CUDA graph capture + on specific FSDP module types. Valid values: ``'mamba'`` (MambaLayer), + ``'transformer'`` (TransformerLayer). Example: ``['mamba', 'transformer']``. + Only leaf FSDP modules (without FSDP children) are eligible.""" use_custom_fsdp: bool = False """ NOTE: The flag `use_custom_fsdp` is deprecated and will be removed in future versions. diff --git a/megatron/core/distributed/fsdp/checkpoint.py b/megatron/core/distributed/fsdp/checkpoint.py index 98059f1bf12..06667c4e711 100644 --- a/megatron/core/distributed/fsdp/checkpoint.py +++ b/megatron/core/distributed/fsdp/checkpoint.py @@ -22,7 +22,8 @@ and applies MCore post-processing. - Post-processing functions for Megatron FSDP v2 state dicts: ``handle_swiglu_in_state_dict_v2``, ``handle_gdn_in_state_dict_v2``, - ``handle_experts_in_state_dict``, ``handle_fp8_extra_state_case``. + ``handle_mamba_in_state_dict_v2``, ``handle_experts_in_state_dict``, + ``handle_fp8_extra_state_case``. """ import logging @@ -39,6 +40,7 @@ import megatron.core.parallel_state as mpu from megatron.core import dist_checkpointing from megatron.core.distributed.fsdp.src.megatron_fsdp.uneven_dtensor import copy_chunk_metadata +from megatron.core.distributed.fsdp.src.megatron_fsdp.uneven_dtensor import get_chunk_meta_source from megatron.core.distributed.fsdp.src.megatron_fsdp.uneven_dtensor import ( get_state_dict as _get_state_dict, ) @@ -50,6 +52,11 @@ ) from megatron.core.transformer.transformer_layer import TransformerLayer +try: + from megatron.core.ssm.mamba_mixer import MambaMixer +except ImportError: + MambaMixer = None + logger = logging.getLogger(__name__) __all__ = [ @@ -63,6 +70,7 @@ "handle_experts_in_state_dict", "handle_swiglu_in_state_dict_v2", "handle_gdn_in_state_dict_v2", + "handle_mamba_in_state_dict_v2", ] _MODULE_PREFIX = "module." @@ -79,15 +87,9 @@ class MegatronFSDPStateful(Stateful): Implements DCP's ``Stateful`` protocol. ``state_dict()`` uses ``get_state_dict`` from ``uneven_dtensor`` (which preprocesses DTensors for uneven sharding and produces FQN-keyed optimizer states), then applies - MCore post-processing (SwiGLU, GDN, FP8, expert remapping — see individual - ``handle_*_v2`` functions). ``load_state_dict()`` uses PyTorch's + MCore post-processing (SwiGLU, GDN, MambaMixer, FP8, expert remapping — see + individual ``handle_*_v2`` functions). ``load_state_dict()`` uses PyTorch's ``set_state_dict``. - - Parameters: - model: FSDP v2 wrapped model. - optimizer: Optional optimizer to checkpoint alongside the model. - args: Optional MCore args namespace for post-processing config - (``swiglu``, ``num_experts``, ``gdn``). """ def __init__( @@ -159,6 +161,13 @@ def _should_keep(expert_index): def _replace(key, expert_index, sd): new_idx = expert_index + local_expert_start + if new_idx == expert_index and not _should_keep(expert_index): + logger.warning( + "Identity transform for non-local expert key '%s' (expert_index=%d, " + "local_expert_start=%d). This expert does not belong to EP rank %d " + "but survived in the state dict. Consider removing it instead.", + key, expert_index, local_expert_start, ep_rank, + ) # GroupedMLP: 'mlp.experts.linear_fc1.weight0', 'mlp.experts.linear_fc2.weight0' if 'mlp.experts.linear_fc1.weight' in key or 'mlp.experts.linear_fc2.weight' in key: if key.endswith('_w') or key.endswith('_v'): @@ -443,6 +452,105 @@ def handle_gdn_in_state_dict_v2( ) +# ------------------------------------------------------------------ +# MambaMixer key patterns and helpers +# ------------------------------------------------------------------ + +_MAMBA_MIXER_IN_PROJ_NAMES = ["z", "x", "B", "C", "dt"] +_MAMBA_MIXER_CONV1D_NAMES = ["x", "B", "C"] + +_MAMBA_MIXER_KEY_PATTERNS = [ + r"(.*)\.mixer\.in_proj\.weight$", + r"(.*)\.mixer\.conv1d\.weight$", + r"(.*)\.mixer\.conv1d\.bias$", +] + + +def _detect_mamba_mixers(model: nn.Module) -> dict: + """Return ``{layer_path: MambaMixer_module}`` for MambaMixer layers.""" + if MambaMixer is None: + return {} + _mixers = {} + for name, module in model.named_modules(): + if isinstance(module, MambaMixer): + _mixers[_strip_wrappers(name)] = module + return _mixers + + +def _mamba_mixer_detector(key, dtensor, model, mixer_map): + """Detector for MambaMixer fused parameters. + + Returns ``(sizes, names, dim)`` for keys matching ``mixer.in_proj.*`` + or ``mixer.conv1d.*``, using the TP-local dimensions from the owning + ``MambaMixer`` module. Returns ``None`` for non-mamba keys. + """ + if not _MAMBA_MIXER_KEY_PATTERNS or MambaMixer is None: + return None + dim = 0 + for pat in _MAMBA_MIXER_KEY_PATTERNS: + m = re.match(pat, key) + if not m: + continue + prefix = m.group(1) + mixer_module = mixer_map.get(prefix) + if mixer_module is None: + # Strip module. prefix and retry + alt = prefix + while alt.startswith(_MODULE_PREFIX): + alt = alt[len(_MODULE_PREFIX):] + mixer_module = mixer_map.get(alt) + if mixer_module is not None: + break + if mixer_module is None: + return None + if "in_proj.weight" in key: + sizes = [ + mixer_module.d_inner_local_tp, + mixer_module.d_inner_local_tp, + mixer_module.ngroups_local_tp * mixer_module.d_state, + mixer_module.ngroups_local_tp * mixer_module.d_state, + mixer_module.nheads_local_tp, + ] + return (sizes, _MAMBA_MIXER_IN_PROJ_NAMES, dim) + if "conv1d.weight" in key or "conv1d.bias" in key: + sizes = [ + mixer_module.d_inner_local_tp, + mixer_module.ngroups_local_tp * mixer_module.d_state, + mixer_module.ngroups_local_tp * mixer_module.d_state, + ] + return (sizes, _MAMBA_MIXER_CONV1D_NAMES, dim) + return None + + +def handle_mamba_in_state_dict_v2( + model: nn.Module, model_state_dict: dict, optimizer_state_dict: Optional[dict] +) -> Tuple[dict, Optional[dict]]: + """Split fused MambaMixer parameters into per-component DTensors. + + Splits ``mixer.in_proj.weight`` → ``.z`` / ``.x`` / ``.B`` / ``.C`` / ``.dt`` + and ``mixer.conv1d.{weight,bias}`` → ``.x`` / ``.B`` / ``.C``, using + TP-local dimensions read from each ``MambaMixer`` module. + Delegates to :func:`_split_fused_params_v2`. + """ + if MambaMixer is None: + return model_state_dict, optimizer_state_dict + mixer_map = _detect_mamba_mixers(model) + if not mixer_map: + return model_state_dict, optimizer_state_dict + + def detector(key, dtensor, _model): + return _mamba_mixer_detector(key, dtensor, _model, mixer_map) + + return _split_fused_params_v2( + model, + model_state_dict, + optimizer_state_dict, + detector, + lambda k, s: f"{k}.{s}", + "MambaMixer", + ) + + # ------------------------------------------------------------------ # Unified post-processing # ------------------------------------------------------------------ @@ -475,19 +583,37 @@ def _propagate_chunk_metadata_to_state_dict(model: nn.Module, state_dict: dict) if isinstance(param, DTensor) and hasattr(param._local_tensor, "__create_chunk_list__"): param_map[name] = param + missing_src_metadata = [] for key, value in state_dict.items(): if not isinstance(value, DTensor): continue param = _find_param_in_map(key, param_map) if param is not None: copy_chunk_metadata(param, value) + else: + has_meta = hasattr(value._local_tensor, "__create_chunk_list__") + missing_src_metadata.append((key, value.shape, value._local_tensor.shape, has_meta)) + + if missing_src_metadata: + if torch.distributed.get_rank() == 0: + logger.warning( + "[chunk_metadata_diag] _propagate_chunk_metadata_to_state_dict: " + f"{len(missing_src_metadata)} DTensor(s) in state_dict could NOT be matched to " + "a model parameter with __create_chunk_list__. Their metadata (if any) comes " + "from preprocess_state_dict_for_uneven_dtensor only." + ) + for key, global_shape, local_shape, has_meta in missing_src_metadata: + logger.warning( + f" key={key} global_shape={tuple(global_shape)} " + f"local_shape={tuple(local_shape)} has_create_chunk_list={has_meta}" + ) def _apply_mcore_postprocess(raw_state_dict, args, model): """Apply MCore-specific state dict post-processing. Copies *raw_state_dict*, wraps optimizer states as DTensors, then - applies FP8 cleanup, SwiGLU/GDN split, and expert key remapping. + applies FP8 cleanup, SwiGLU/GDN/MambaMixer split, and expert key remapping. The original *raw_state_dict* is not mutated. """ state_dict = raw_state_dict.copy() @@ -511,6 +637,12 @@ def _apply_mcore_postprocess(raw_state_dict, args, model): if new_opt is not None: state_dict["optimizer"] = new_opt + opt_sd = state_dict.get("optimizer") + model_sd, new_opt = handle_mamba_in_state_dict_v2(model, state_dict["model"], opt_sd) + state_dict["model"] = model_sd + if new_opt is not None: + state_dict["optimizer"] = new_opt + num_experts = getattr(args, "num_experts", None) if num_experts: state_dict["model"] = handle_experts_in_state_dict(state_dict["model"], num_experts) @@ -522,15 +654,65 @@ def _apply_mcore_postprocess(raw_state_dict, args, model): ) flattened_sd = flatten_state_dict_keys(state_dict) + _verify_chunk_metadata(flattened_sd) + + return state_dict + + +def _numel(shape: tuple) -> int: + """Return the product of all dimensions in *shape*.""" + n = 1 + for d in shape: + n *= d + return n + + +def _verify_chunk_metadata(flattened_sd: dict) -> None: + """Verify every DTensor has correct ``__create_chunk_list__`` metadata. + + Checks both existence AND consistency (chunks total numel must match + local tensor numel). On failure, prints diagnostic information + including the metadata source tag and shape details. + """ + failures = [] for key, value in flattened_sd.items(): - if isinstance(value, DTensor): - assert hasattr(value._local_tensor, "__create_chunk_list__"), ( - f"DTensor for key '{key}' is missing chunk metadata. This may cause issues" - f" with checkpointing. Ensure that the model parameters have chunk metadata" - f" and that it is properly propagated." + if not isinstance(value, DTensor): + continue + lt = value._local_tensor + if not hasattr(lt, "__create_chunk_list__"): + failures.append( + f"MISSING metadata: key={key} global_shape={tuple(value.shape)} " + f"local_shape={tuple(lt.shape)} source={get_chunk_meta_source(value)}" ) + continue - return state_dict + cl = lt.__create_chunk_list__() + cl_total = sum(_numel(c.sizes) for c in cl) + local_numel = lt.numel() + if cl_total != local_numel: + cl_detail = [(tuple(c.offsets), tuple(c.sizes)) for c in cl] + failures.append( + f"NUMEL MISMATCH: key={key} " + f"global_shape={tuple(value.shape)} " + f"local_shape={tuple(lt.shape)} " + f"local_numel={local_numel} " + f"chunks_total={cl_total} " + f"chunk_list={cl_detail} " + f"source={get_chunk_meta_source(value)} " + f"device_mesh={value.device_mesh}" + ) + + if failures: + logger.error( + "[chunk_metadata_verify] %d DTensor(s) have invalid chunk metadata:", + len(failures), + ) + for msg in failures: + logger.error(" %s", msg) + raise AssertionError( + f"{len(failures)} DTensor(s) have invalid chunk metadata. " + "See log above for details." + ) def flatten_state_dict_keys(state_dict, parent_key="", sep="."): @@ -663,7 +845,7 @@ def _maybe_wrap_as_uneven_dtensor(tensor, dist_param: DTensor): if isinstance(tensor, DTensor): dt = tensor else: - if isinstance(tensor, DTensor) or not isinstance(tensor, torch.Tensor): + if not isinstance(tensor, torch.Tensor): return tensor if tensor.shape != dist_param._local_tensor.shape: return tensor @@ -715,7 +897,7 @@ def _build_dtensor_optim_sd(raw_opt_state_dict: dict, model: nn.Module) -> dict: return opt_state_dict -def _preprocess_and_verify_v2_state_dict(v2_state_dict): +def _preprocess_and_verify_v2_state_dict(v2_state_dict, model=None): """Preprocess and verify the Megatron FSDP v2 state dict before DCP loading. DCP requires DTensors for distributed load, but ``optimizer.load_state_dict`` @@ -728,6 +910,11 @@ def _preprocess_and_verify_v2_state_dict(v2_state_dict): Also verifies that every model and shadow optimizer DTensor has ``__create_chunk_list__`` and ``__create_write_items__`` metadata. + Args: + v2_state_dict: The v2 state dict (from ``_build_megatron_fsdp_v2_state_dict``). + model: The FSDP model (for chunk metadata propagation). If ``None``, + metadata propagation is skipped. + Returns: v2_by_canonical: ``{canonical_name: model_DTensor}`` mapping. v2_optim_state: ``{canonical_name: {state_key: DTensor}}`` shadow dict. @@ -736,10 +923,7 @@ def _preprocess_and_verify_v2_state_dict(v2_state_dict): v2_optim_state_raw = v2_state_dict.get("optimizer", {}).get("state", {}) # ---- Propagate chunk metadata from model to state dict DTensors ---- - model = v2_state_dict.get("_model") if model is not None: - if isinstance(model, (list, tuple)): - model = model[0] _propagate_chunk_metadata_to_state_dict(model, v2_model) # ---- Strip ``module.`` prefix from optimizer state keys ---- @@ -1177,7 +1361,13 @@ def _canonicalize_td_key(td_key, *, strip_model_prefix=True): return key -def _load_torch_dist_into_megatron_fsdp_v2(args, checkpoint_name, v2_state_dict, strict=True): +def _load_torch_dist_into_megatron_fsdp_v2( + args, + checkpoint_name, + model, + v2_state_dict, + strict=True, +): """Load a torch_dist checkpoint into a Megatron FSDP v2 skeleton via DCP. This is the entry point for online checkpoint conversion from @@ -1188,6 +1378,10 @@ def _load_torch_dist_into_megatron_fsdp_v2(args, checkpoint_name, v2_state_dict, 1. **Preprocess & verify** — wrap optimizer states as uneven DTensors and verify ``__create_chunk_list__`` / ``__create_write_items__`` metadata. + When ``args.mamba`` is True, split fused MambaMixer params + (``in_proj.weight`` → ``.z`` / ``.x`` / ``.B`` / ``.C`` / ``.dt``, + ``conv1d.*`` → ``.x`` / ``.B`` / ``.C``) so they match the torch_dist + split keys. 2. **Build name mapping** — match torch_dist metadata keys to v2 DTensors, handling fused layers, GroupedMLP, and SequentialMLP expert formats. 3. **DCP load 1:1 entries** — load regular model weights and optimizer @@ -1204,7 +1398,21 @@ def _load_torch_dist_into_megatron_fsdp_v2(args, checkpoint_name, v2_state_dict, from torch.distributed.checkpoint.default_planner import DefaultLoadPlanner # ---- Phase 1: Preprocess & verify v2 state dict ---- - v2_by_canonical, v2_optim_state = _preprocess_and_verify_v2_state_dict(v2_state_dict) + + # Split fused MambaMixer params so that the sub-keys match the + # torch_dist checkpoint (which stores e.g. ``in_proj.weight.z``). + if model is not None: + if isinstance(model, (list, tuple)): + assert len(model) == 1 + model = model[0] + model_sd, new_opt = handle_mamba_in_state_dict_v2( + model, v2_state_dict["model"], v2_state_dict.get("optimizer") + ) + v2_state_dict["model"] = model_sd + if new_opt is not None: + v2_state_dict["optimizer"] = new_opt + + v2_by_canonical, v2_optim_state = _preprocess_and_verify_v2_state_dict(v2_state_dict, model) # ---- Phase 2: Read torch_dist metadata & build name mapping ---- reader = FileSystemReader(checkpoint_name) @@ -1309,7 +1517,7 @@ def _load_torch_dist_into_megatron_fsdp_v2(args, checkpoint_name, v2_state_dict, chunk = flat[layer_idx] if isinstance(v2_val, DTensor): lt = v2_val._local_tensor - if lt.numel() != 0: + if lt.numel() != 0 and hasattr(lt, "__create_chunk_list__"): local_off = 0 for c in lt.__create_chunk_list__(): off = c.offsets diff --git a/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py b/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py index aac55076c72..d58afaa09d5 100644 --- a/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py +++ b/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py @@ -42,6 +42,9 @@ from megatron.core.distributed.distributed_data_parallel_config import DistributedDataParallelConfig from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.transformer.transformer_config import TransformerConfig +from megatron.core.ssm.mamba_layer import MambaLayer +from megatron.core.transformer.moe.router import Router as MoERouter +from megatron.core.transformer.attention import Attention from megatron.core.transformer.transformer_layer import TransformerLayer from megatron.core.utils import is_te_min_version, log_single_rank @@ -234,7 +237,7 @@ def _init_with_fully_shard( fsdp_unit_modules is None and ddp_config.data_parallel_sharding_strategy == "optim_grads_params" ): - fsdp_unit_modules = [TransformerLayer] + fsdp_unit_modules = [TransformerLayer, MambaLayer, TEGroupedMLP, SequentialMLP] if pg_collection is None: pg_collection = ProcessGroupCollection.use_mpu_process_groups() @@ -284,19 +287,37 @@ def _init_with_fully_shard( gradient_scaling_factor = 1.0 / dp_world_size expert_gradient_scaling_factor = 1.0 / dp_world_size - for m in module.modules(): - if isinstance(m, (TEGroupedMLP, SequentialMLP)): - fully_shard( - m, - mesh=edp_mesh, - gradient_scaling_factor=expert_gradient_scaling_factor, - **kwargs, - ) if fsdp_unit_modules is not None: - for m in module.modules(): - if isinstance(m, tuple(fsdp_unit_modules)): + cuda_graph_on = set(ddp_config.mfsdp_cuda_graph_modules) + # Iterate modules post order to ensure that child modules are fully sharded + # before their parents, which is required for correct param group divide. + for name, m in reversed(list(module.named_modules())): + if isinstance(m, (TEGroupedMLP, SequentialMLP)): + grad_sf = expert_gradient_scaling_factor + mesh = edp_mesh + else: + grad_sf = gradient_scaling_factor + mesh = dp_mesh + + if any([ + isinstance(m, MambaLayer) and "mamba" in cuda_graph_on, + isinstance(m, Attention) and "attn" in cuda_graph_on, + isinstance(m, MoERouter) and "moe_router" in cuda_graph_on, + ]): + fully_shard( + m, + enable_cuda_graph=True, + mesh=mesh, + gradient_scaling_factor=grad_sf, + **kwargs + ) + elif isinstance(m, tuple(fsdp_unit_modules)): fully_shard( - m, mesh=dp_mesh, gradient_scaling_factor=gradient_scaling_factor, **kwargs + m, + mesh=mesh, + gradient_scaling_factor=grad_sf, + enable_cuda_graph=False, + **kwargs ) fully_shard(module, mesh=dp_mesh, gradient_scaling_factor=gradient_scaling_factor, **kwargs) diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/uneven_dtensor.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/uneven_dtensor.py index 1efdeec2376..e70d26ebf67 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/uneven_dtensor.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/uneven_dtensor.py @@ -96,7 +96,7 @@ def _update_offsets_and_cumulative_shape( return ChunkStorageMetadata(offsets=tuple(offsets), sizes=tuple(local_shape)) -def update_uneven_dtensor_chunk_metadata(dtensor: DTensor) -> dict: +def update_uneven_dtensor_chunk_metadata(dtensor: DTensor, source: str = "init") -> dict: """ Update the DTensor's chunk metadata to handle uneven sharding. This function modifies the DTensor in-place to include chunk metadata @@ -111,7 +111,7 @@ def update_uneven_dtensor_chunk_metadata(dtensor: DTensor) -> dict: uneven_chunk_meta = gather_and_compute_chunk_metadata(dtensor) # Set the chunk list and write items closure for the DTensor - _set_chunk_metadata(dtensor, uneven_chunk_meta.offsets, uneven_chunk_meta.sizes) + _set_chunk_metadata(dtensor, uneven_chunk_meta.offsets, uneven_chunk_meta.sizes, source=source) def validate_uneven_dtensor(dtensor: DTensor) -> None: @@ -227,7 +227,7 @@ def preprocess_state_dict_for_uneven_dtensor(state_dict: dict) -> dict: for key_chain in sorted(visit_dtensor): # Get the DTensor at the key chain dtensor = get_unflattened_state_dict(state_dict, key_chain) - update_uneven_dtensor_chunk_metadata(dtensor) + update_uneven_dtensor_chunk_metadata(dtensor, source="preprocess") return state_dict @@ -562,7 +562,7 @@ def split_dtensor( new_offsets[dim] = 0 new_sizes[dim] = 0 - _set_chunk_metadata(new_dtensor, tuple(new_offsets), tuple(new_sizes)) + _set_chunk_metadata(new_dtensor, tuple(new_offsets), tuple(new_sizes), source="split") yield new_dtensor @@ -613,7 +613,7 @@ def make_uneven_dtensor( # existing uneven DTensor instead of recomputing it. copy_chunk_metadata(copy_chunk_meta_from, dtensor) elif chunk_metadata is not None: - _set_chunk_metadata(dtensor, *chunk_metadata) + _set_chunk_metadata(dtensor, *chunk_metadata, source="make_uneven") return dtensor @@ -645,6 +645,16 @@ def copy_chunk_metadata(src: DTensor, dst: DTensor) -> None: """Copy ``__create_chunk_list__`` / ``__create_write_items__`` from *src* to *dst*.""" dst._local_tensor.__create_chunk_list__ = src._local_tensor.__create_chunk_list__ dst._local_tensor.__create_write_items__ = src._local_tensor.__create_write_items__ + src_source = getattr(src._local_tensor, "_chunk_meta_source", None) + if src_source is not None: + dst._local_tensor._chunk_meta_source = f"propagate:{src_source}" + else: + dst._local_tensor._chunk_meta_source = "propagate:unknown" + + +def get_chunk_meta_source(dtensor: DTensor) -> str: + """Return the source tag for *dtensor*'s chunk metadata, or ``"none"``.""" + return getattr(dtensor._local_tensor, "_chunk_meta_source", "none") def compute_split_offsets_and_sizes(dist_param, split_dim, comp_idx, total_split, comp_data): @@ -710,10 +720,17 @@ def get_fsdp_slice_from_uneven_dtensor(dist_param: DTensor) -> slice: return slice(start, start + local_numel) -def _set_chunk_metadata(dtensor: DTensor, offsets: tuple, sizes: tuple) -> None: +def _set_chunk_metadata( + dtensor: DTensor, offsets: tuple, sizes: tuple, source: str = "unknown" +) -> None: """Set ``__create_chunk_list__`` / ``__create_write_items__`` closures on *dtensor*. No collective ops — *offsets* and *sizes* are computed locally. + + Args: + source: A tag describing where the metadata was set, e.g. ``"init"``, + ``"preprocess"``, ``"split"``. Stored on the local tensor as + ``_chunk_meta_source`` for diagnostic tracing. """ chunk_meta = ChunkStorageMetadata(offsets=tuple(offsets), sizes=tuple(sizes)) @@ -734,3 +751,4 @@ def _write_items(fqn: str, tensor: DTensor) -> list: dtensor._local_tensor.__create_chunk_list__ = lambda: [chunk_meta] dtensor._local_tensor.__create_write_items__ = _write_items + dtensor._local_tensor._chunk_meta_source = source diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/README.md b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/README.md index 2d28f74bddd..870b0e64d6f 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/README.md +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/README.md @@ -97,6 +97,52 @@ Mixin class added to wrapped modules. Methods: | `reshard()` | Post-forward, post-backward | Release unsharded buffer | | `reduce_grad()` | Post-backward / grad sync | Reduce-scatter gradients into optimizer-facing shards | +### CUDA Graph Capture + +> **Experimental** — CUDA graph support in Megatron FSDP v2 is an experimental +> feature. The API and behaviour may change in future releases without notice. + +FSDP v2 supports transparent CUDA graph capture for individual FSDP modules. +Enable it per-module with ``enable_cuda_graph=True``: + +```python +for layer in model.layers: + fully_shard(layer, enable_cuda_graph=True) +fully_shard(model) # root without CUDA graph +``` + +Capture happens automatically on the first optimized forward pass +(after the trace pool has been planned). Subsequent forward passes replay +the captured graph, bypassing FSDP hooks entirely for that module. + +**How it works:** + +1. The forward pre-hook detects ``enable_cuda_graph=True`` and creates an + ``FSDPCudaGraphRunner`` for the module. +2. The runner pops FSDP hooks, warms up via ``make_graphed_callables``, + and captures the module's ``forward()`` (without FSDP collectives) into a + CUDA graph. +3. After capture, ``install()`` patches ``module.forward`` to a wrapper + that replays the graph on subsequent calls. +4. After capture, ``install()`` patches ``module.forward`` to a wrapper that + replays the graph on subsequent calls. During replay, the patched forward + calls the graphed callable directly, bypassing FSDP forward hooks entirely. + Backward hooks (unshard/reshard/reduce_grad) still fire normally. + +**Limitation — nesting:** A parent FSDP module that contains other FSDP +modules as children **cannot** use ``enable_cuda_graph=True``. Only leaf +FSDP modules (those without FSDP children) are eligible. Attempting to +enable CUDA graph on a module with FSDP children raises a ``RuntimeError``. + +```python +# OK — layers are leaf FSDP modules (no FSDP children inside them) +for layer in model.layers: + fully_shard(layer, enable_cuda_graph=True) + +# NOT OK — model contains FSDP layers as children +fully_shard(model, enable_cuda_graph=True) # raises RuntimeError +``` + ### DataParallelBuffer Flat buffer managing (a shard of) parameter/gradient data: @@ -155,6 +201,17 @@ All strategies except `no_shard` are supported. `optim_grads` (ZeRO-2) and pipelining) — weights are replicated so no all-gather is needed. `no_shard` (DDP-like) is not yet implemented and raises ``NotImplementedError``. +### CUDA Graph + +- **Experimental.** CUDA graph support is experimental. The API and behavior + may change in future releases without notice. Enable via + ``enable_cuda_graph=True`` on leaf FSDP modules only. +- **Nesting not supported.** A parent FSDP module that contains other FSDP + modules as children cannot use ``enable_cuda_graph=True``. Only leaf FSDP + modules (those without FSDP children) are eligible for CUDA graph capture. + Attempting to enable CUDA graph on a nested FSDP module raises a + ``RuntimeError``. + ### `fully_shard()` API Parameters The following parameters are accepted in the function signature but are **not @@ -231,6 +288,7 @@ without any Megatron-LM dependency once the package is released. See `examples/megatron_fsdp/fsdp_toy.py` for a standalone example showing: - Basic model wrapping with `fully_shard()` +- CUDA graph capture (`--cuda-graph` / `--no-cuda-graph`) - Training loop with gradient accumulation - Activation checkpointing (`--activation-checkpoint`) - Distributed checkpointing with `torch.distributed.checkpoint` @@ -243,12 +301,23 @@ torchrun --nproc_per_node=2 examples/megatron_fsdp/fsdp_toy.py \ torchrun --nproc_per_node=2 examples/megatron_fsdp/fsdp_toy.py \ --model-dim 512 --n-layers 2 --batch-size 4 \ --use-megatron-fsdp --activation-checkpoint + +# Disable CUDA graph capture +torchrun --nproc_per_node=2 examples/megatron_fsdp/fsdp_toy.py \ + --model-dim 512 --n-layers 2 --batch-size 4 \ + --use-megatron-fsdp --no-cuda-graph ``` ## Gotchas / Pitfalls - **Zero-numel gradient shards and fused optimizers.** When a parameter's local shard is empty on some DP ranks (e.g., small biases on high DP counts), creating a `DTensor` gradient with `numel() == 0` and passing it to fused multi-tensor optimizers (TE `FusedAdam`) can silently corrupt updates for neighboring non-empty parameters. This manifests only as convergence divergence with no error — see [design.md § Pitfall](design.md) for details and the fix in `param_group.py`. -- **Temporary communication buckets must record their CUDA stream.** `DataParallelBuffer.reshard()` and `release_grad_buffer()` can return temporary all-gather / reduce-scatter bucket storage to the allocator while side-stream CUDA work is still queued. A CUDA event wait orders streams, but it does not tell PyTorch's caching allocator that the freed storage is still used by `ag_stream` or `rs_stream`. If the block is recycled on another stream, the communication payload can be silently corrupted. Call `record_stream()` on any temporary CUDA tensor that participates in side-stream communication before it may be freed; a full synchronization only masks the bug and hurts overlap. +- **Temporary communication bucket lifecycle.** All temporary all-gather / + reduce-scatter buckets are allocated on the default CUDA stream and only + compute operations (all-gather, reduce-scatter) run on side streams + (`ag_stream`, `rs_stream`). CUDA events inserted at the boundary between + allocation and compute, and between compute and free, guarantee ordering + without ``record_stream``. Keeping allocations on the default stream avoids + non-deterministic caching-allocator behaviour and peak memory regressions. ## Unit Tests diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/allocator.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/allocator.py index ff7e7eaa89f..6569dabc0ed 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/allocator.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/allocator.py @@ -13,17 +13,20 @@ # limitations under the License. import dataclasses -from typing import Dict, Hashable, List, Optional, Tuple +import logging +from collections import defaultdict +from typing import Dict, Hashable, List, Optional, Set, Tuple import torch +logger = logging.getLogger(__name__) + AllocatorKey = Hashable def _resolve_key(key: Optional[AllocatorKey], param_group_id: Optional[AllocatorKey]): if key is not None: return key - # Backward-compatible alias for older callers that keyed only by param group. assert param_group_id is not None, "allocator key is required" return param_group_id @@ -51,7 +54,10 @@ def allocate( raise NotImplementedError def free( - self, key: Optional[AllocatorKey] = None, *, param_group_id: Optional[AllocatorKey] = None + self, + key: Optional[AllocatorKey] = None, + *, + param_group_id: Optional[AllocatorKey] = None, ) -> None: """Free the bucket associated with the given key.""" raise NotImplementedError @@ -80,11 +86,16 @@ def allocate( key = _resolve_key(key, param_group_id) assert dtype is not None and device is not None if key not in self.buckets: - self.buckets[key] = Bucket(data=torch.empty(size, dtype=dtype, device=device)) + self.buckets[key] = Bucket( + data=torch.empty(size, dtype=dtype, device=device) + ) return self.buckets[key] def free( - self, key: Optional[AllocatorKey] = None, *, param_group_id: Optional[AllocatorKey] = None + self, + key: Optional[AllocatorKey] = None, + *, + param_group_id: Optional[AllocatorKey] = None, ) -> None: key = _resolve_key(key, param_group_id) if key in self.buckets: @@ -93,11 +104,7 @@ def free( class StorageFreeingBucketAllocator(BucketAllocator): - """Manages temporary flat buffers keyed by caller-provided allocation key. - - Freeing releases the underlying storage without deleting the bucket entry, - so the same tensor object can be reused on the next allocation. - """ + """Manages temporary flat buffers keyed by caller-provided allocation key.""" def __init__(self): super().__init__() @@ -115,13 +122,18 @@ def allocate( key = _resolve_key(key, param_group_id) assert dtype is not None and device is not None if key not in self.buckets: - self.buckets[key] = Bucket(data=torch.empty(size, dtype=dtype, device=device)) + self.buckets[key] = Bucket( + data=torch.empty(size, dtype=dtype, device=device) + ) return self.buckets[key] _alloc_storage(self.buckets[key].data, torch.Size([size])) return self.buckets[key] def free( - self, key: Optional[AllocatorKey] = None, *, param_group_id: Optional[AllocatorKey] = None + self, + key: Optional[AllocatorKey] = None, + *, + param_group_id: Optional[AllocatorKey] = None, ) -> None: key = _resolve_key(key, param_group_id) if key in self.buckets: @@ -129,72 +141,53 @@ def free( class TracePoolAllocator(BucketAllocator): - """Two-phase bucket allocator that eliminates per-call ``torch.empty`` overhead. - - **Design** + """Two-phase bucket allocator for CUDA graph-compatible training. - The FSDP framework allocates and frees temporary flat buffers (for - all-gather input/output and gradient accumulation) in a deterministic, - repeatable order across micro-batches. ``TracePoolAllocator`` exploits - this by profiling one pass and then serving all subsequent passes from - a pre-allocated pool. + Profiles one micro-batch to record allocation patterns, then builds a + static key-to-address plan that is identical across all subsequent + micro-batches — essential for CUDA graph capture. - **Phase 1 — Trace** (``plan()`` not yet called) + **Phase 1 — Trace** (first micro-batch) - Behaves like ``TemporaryBucketAllocator``: ``allocate`` creates a - ``torch.empty`` bucket on first use, ``free`` releases its storage. - Additionally, every alloc/free call is recorded as a ``_TraceEvent`` - with a monotonic ``seq`` number, and metadata ``(size, dtype, device)`` - is stored per allocation key for later planning. + Records alloc/free calls with monotonic sequence numbers. Buckets are + created with ``torch.empty`` and freed via ``_free_storage`` so the same + tensor object can be resurrected on re-alloc (keeping outstanding views + alive, e.g. NVFP4 ``_rowwise_data`` references). **Phase 2 — Plan** (``plan()``) - The trace is replayed to extract *intervals*: for each alloc/free pair - an ``_Interval(alloc_seq, free_seq, size)`` is built. Intervals are - grouped by ``(dtype, device)`` and then colored with a greedy - left-edge algorithm: + Replays the trace to build per-key live intervals, then uses + **conflict-graph coloring** to assign slots: - 1. Sort intervals by ``alloc_seq``. - 2. For each interval, try to reuse a *slot* whose previous occupant - freed before this interval starts (``slot_free_seq < alloc_seq``). - 3. If no slot is free, allocate a new one. - 4. Grow the slot's capacity to ``max(size, current)``. - 5. Record the assignment: append the slot index to the per-key list - in ``_slot_map``. + * An interval-overlap graph is built: edges connect keys whose live + intervals overlap. + * Nodes are colored greedily (largest-size-first, best-fit bin packing) + so two keys share a slot iff they never overlap. + * Yields the **theoretical minimum** number of slots. - After coloring, slots are laid out contiguously and a single - ``torch.empty`` per ``(dtype, device)`` group is allocated. If the - trace is empty, planning is a no-op and later cursor resets also no-op. + Each slot is a **separate** ``torch.empty()`` tensor (per-slot allocation), + not a slice of a monolithic pool. The CUDA caching allocator can place + them independently, reducing fragmentation pressure from giant contiguous + blocks, while each key still resolves to a fixed memory address. **Phase 3 — Optimized** (after ``plan()``) - ``allocate`` returns a ``Bucket`` with a slice-view into the pool; - ``free`` marks the slot as unused but never releases storage. Because - the same allocation key can appear in multiple intervals (e.g., - forward unshard → free → backward unshard → free), ``_slot_map`` maps - each key to a **list** of slot indices in alloc order. A per-key - ``_slot_cursors`` counter tracks which index to consume next. Call - ``reset_cursor()`` at the start of each micro-batch to rewind all - cursors to 0. - - The trace pattern must be **repeatable** — the same alloc/free call - sequence is expected every micro-batch. + ``allocate`` / ``free`` are O(1) dict lookups that return pre-computed + tensor views. No allocations or storage resizes occur in this phase — + memory addresses are stable across all micro-batches. """ # -- Inner types ---------------------------------------------------- # - class _Slot: - """A contiguous slice of the pool tensor assigned to one or more - non-overlapping intervals.""" - - __slots__ = ("offset", "size", "dtype", "device", "in_use") + @dataclasses.dataclass + class _SlotInfo: + """Metadata for a physical slot (backed by its own tensor).""" - def __init__(self, offset: int, size: int, dtype: torch.dtype, device: torch.device): - self.offset = offset - self.size = size - self.dtype = dtype - self.device = device - self.in_use = False + tensor: torch.Tensor # The actual backing tensor for this slot + size: int # Capacity in elements + dtype: torch.dtype + device: torch.device + in_use: bool = False @dataclasses.dataclass class _TraceEvent: @@ -204,35 +197,26 @@ class _TraceEvent: op: str # "alloc" | "free" key: AllocatorKey - @dataclasses.dataclass - class _Interval: - """An allocation's lifetime: from alloc_seq to free_seq with a given size.""" - - key: AllocatorKey - size: int - alloc_seq: int - free_seq: int - # -- Init ----------------------------------------------------------- # def __init__(self) -> None: super().__init__() - # Phase bookkeeping self._phase: str = "trace" # "trace" | "optimized" - self._seq: int = 0 # monotonic alloc/free counter # Trace state + self._seq: int = 0 self._trace: List["TracePoolAllocator._TraceEvent"] = [] self._trace_meta: Dict[AllocatorKey, Tuple[int, torch.dtype, torch.device]] = {} - self._buckets: Dict[AllocatorKey, Bucket] = {} # only used in trace phase + self._buckets: Dict[AllocatorKey, Bucket] = {} + self._active_keys: Set[AllocatorKey] = set() # Pool state — populated by plan(), used in optimized phase - self._pools: Dict[Tuple[torch.dtype, torch.device], torch.Tensor] = {} - self._slot_map: Dict[AllocatorKey, List[int]] = {} # key -> [slot indices] - self._slot_cursors: Dict[AllocatorKey, int] = {} # key -> next index to use - self._slots: List["TracePoolAllocator._Slot"] = [] + self._slots: List["TracePoolAllocator._SlotInfo"] = [] + self._key_to_slot: Dict[AllocatorKey, int] = {} + # For each key, the view into its slot (pre-computed for O(1) access) + self._key_to_view: Dict[AllocatorKey, torch.Tensor] = {} - # -- Phase 1: trace -------------------------------------------------- # + # -- Public interface ------------------------------------------------ # def allocate( self, @@ -243,236 +227,281 @@ def allocate( *, param_group_id: Optional[AllocatorKey] = None, ) -> Bucket: - """Dispatch to trace or pool path depending on phase.""" key = _resolve_key(key, param_group_id) assert dtype is not None and device is not None if self._phase != "optimized": return self._trace_allocate(key, size, dtype, device) - return self._pool_allocate(key, size, dtype, device) + else: + return self._optimized_allocate(key, size, dtype, device) def free( - self, key: Optional[AllocatorKey] = None, *, param_group_id: Optional[AllocatorKey] = None + self, + key: Optional[AllocatorKey] = None, + *, + param_group_id: Optional[AllocatorKey] = None, ) -> None: - """Dispatch to trace or pool path depending on phase.""" key = _resolve_key(key, param_group_id) if self._phase != "optimized": self._trace_free(key) else: - self._pool_free(key) + self._optimized_free(key) + + # -- Phase 1: trace -------------------------------------------------- # def _trace_allocate( self, key: AllocatorKey, size: int, dtype: torch.dtype, device: torch.device ) -> Bucket: - """Trace-phase allocate: record the event and create a bucket on first use. + if key in self._active_keys: + return self._buckets[key] - Duplicate allocs (without an intervening free) do NOT generate - new trace events — they are no-ops that return the existing bucket. - """ if key not in self._buckets: self._trace.append(self._TraceEvent(seq=self._seq, op="alloc", key=key)) self._seq += 1 self._trace_meta[key] = (size, dtype, device) - self._buckets[key] = Bucket(data=torch.empty(size, dtype=dtype, device=device)) + self._buckets[key] = Bucket( + data=torch.empty(size, dtype=dtype, device=device) + ) + else: + self._trace.append(self._TraceEvent(seq=self._seq, op="alloc", key=key)) + self._seq += 1 + _alloc_storage(self._buckets[key].data, torch.Size([size])) + + self._active_keys.add(key) return self._buckets[key] def _trace_free(self, key: AllocatorKey) -> None: - """Trace-phase free: record the event and release the bucket storage.""" + if key not in self._active_keys: + return self._trace.append(self._TraceEvent(seq=self._seq, op="free", key=key)) self._seq += 1 if key in self._buckets: _free_storage(self._buckets[key].data) - del self._buckets[key] + self._active_keys.discard(key) # -- Phase 2: plan --------------------------------------------------- # def plan(self) -> int: - """Build the static pool from the recorded trace. + """Build the static key→slot plan from the recorded trace. - 1. Replay the trace to pair alloc/free events into ``_Interval`` objects. - 2. Group intervals by ``(dtype, device)``. - 3. Color each group with the greedy left-edge algorithm. - 4. Allocate one flat pool tensor per group. + Uses conflict-graph coloring to achieve optimal memory usage, and + allocates each slot as a separate tensor (per-slot allocation) to + minimize fragmentation pressure on the CUDA caching allocator. Returns: - Total pool size in **elements** (sum across all groups). - Multiply by ``element_size(dtype)`` for bytes. + Total pool size in elements (sum across all dtype/device groups). """ assert self._phase == "trace", "plan() can only be called in trace phase" if len(self._trace) == 0: self._phase = "optimized" return 0 - # ---- step 1: build intervals from alloc/free pairs ---- - alloc_stack: Dict[AllocatorKey, List[int]] = {} # key -> [alloc_seq, ...] - intervals: List["TracePoolAllocator._Interval"] = [] + # Step 1: Build per-key intervals from alloc/free pairs + alloc_stack: Dict[AllocatorKey, List[int]] = {} + intervals_per_key: Dict[AllocatorKey, List[Tuple[int, int]]] = defaultdict(list) for ev in self._trace: if ev.op == "alloc": alloc_stack.setdefault(ev.key, []).append(ev.seq) - else: # "free" + else: if ev.key in alloc_stack and alloc_stack[ev.key]: alloc_seq = alloc_stack[ev.key].pop(0) - meta = self._trace_meta.get(ev.key) - if meta is not None: - size, dtype, device = meta - intervals.append( - self._Interval( - key=ev.key, size=size, alloc_seq=alloc_seq, free_seq=ev.seq - ) - ) - - if len(intervals) == 0: - self._phase = "optimized" - return 0 + intervals_per_key[ev.key].append((alloc_seq, ev.seq)) - # ---- step 2 & 3: color and allocate ---- - return self._assign_pool(intervals) + # Keys allocated but never freed get a sentinel free_seq + _SENTINEL_FREE_SEQ = 1 << 60 + sentinel_seq = _SENTINEL_FREE_SEQ + for key, pending_allocs in alloc_stack.items(): + for alloc_seq in pending_allocs: + intervals_per_key[key].append((alloc_seq, sentinel_seq)) + sentinel_seq += 1 - def _assign_pool(self, intervals: List["TracePoolAllocator._Interval"]) -> int: - """Group intervals by (dtype, device), color each group, sum sizes.""" - groups: Dict[Tuple[torch.dtype, torch.device], List["TracePoolAllocator._Interval"]] = {} - for iv in intervals: - meta = self._trace_meta[iv.key] - dtype_device = (meta[1], meta[2]) - groups.setdefault(dtype_device, []).append(iv) + if not intervals_per_key: + self._phase = "optimized" + return 0 - # Clear any previous plan state before rebuilding - self._slot_map.clear() - self._slot_cursors.clear() + # Step 2: Compute per-key max size + key_max_size: Dict[AllocatorKey, int] = {} + for key in intervals_per_key: + meta = self._trace_meta.get(key) + if meta is not None: + key_max_size[key] = meta[0] + + # Step 3: Group keys by (dtype, device) + groups: Dict[ + Tuple[torch.dtype, torch.device], List[AllocatorKey] + ] = defaultdict(list) + for key in intervals_per_key: + meta = self._trace_meta.get(key) + if meta is not None: + groups[(meta[1], meta[2])].append(key) + + # Step 4: Color each group and allocate per-slot tensors self._slots.clear() - self._pools.clear() + self._key_to_slot.clear() + self._key_to_view.clear() total_elems = 0 - for (dtype, device), group in groups.items(): - total_elems += self._color_group(group, dtype, device) + for (dtype, device), keys in groups.items(): + total_elems += self._color_and_allocate_slots( + keys, intervals_per_key, key_max_size, dtype, device + ) + + # Free trace-phase resources + self._buckets.clear() + self._active_keys.clear() self._phase = "optimized" + + if torch.distributed.get_rank() == 0: + logger.debug( + f"TracePoolAllocator plan complete: {len(self._slots)} slots, " + f"{total_elems} total elements, " + f"{self.total_pool_bytes / 1024 / 1024:.1f} MB" + ) return total_elems - def _color_group( + def _color_and_allocate_slots( self, - intervals: List["TracePoolAllocator._Interval"], + keys: List[AllocatorKey], + intervals_per_key: Dict[AllocatorKey, List[Tuple[int, int]]], + key_max_size: Dict[AllocatorKey, int], dtype: torch.dtype, device: torch.device, ) -> int: - """Greedy left-edge interval coloring for one (dtype, device) group. + """Conflict-graph coloring + per-slot tensor allocation. - Algorithm: + Each color/slot gets its own ``torch.empty()`` tensor rather than + being a slice of a monolithic pool. This reduces CUDA caching + allocator fragmentation: slot tensors can be placed in gaps between + other allocations rather than requiring one massive contiguous block. - 1. Sort intervals by ``alloc_seq``. - 2. Maintain a *free-list* of ``(slot_index, free_seq)`` — slots that - become free at ``free_seq``. - 3. For each interval: - a. Scan the free-list for a slot whose ``free_seq < alloc_seq``. - If found, reuse it (grow its size if needed) and update its - free time to this interval's ``free_seq``. - b. If no slot is free, create a new one. - c. Append the assigned slot index to ``_slot_map[key]``. + Returns total elements allocated across all slots in this group. + """ + n = len(keys) + if n == 0: + return 0 + + # Build conflict graph + conflicts: Dict[AllocatorKey, Set[AllocatorKey]] = defaultdict(set) + for i in range(n): + key_a = keys[i] + ivs_a = intervals_per_key[key_a] + for j in range(i + 1, n): + key_b = keys[j] + ivs_b = intervals_per_key[key_b] + if _intervals_overlap(ivs_a, ivs_b): + conflicts[key_a].add(key_b) + conflicts[key_b].add(key_a) + + # Greedy graph coloring: largest-first, best-fit + keys_sorted = sorted(keys, key=lambda k: key_max_size.get(k, 0), reverse=True) + color_of: Dict[AllocatorKey, int] = {} + slot_sizes: List[int] = [] # color_idx -> capacity in elements + + for k in keys_sorted: + size_k = key_max_size.get(k, 0) + neighbor_colors: Set[int] = set() + for neighbor in conflicts[k]: + if neighbor in color_of: + neighbor_colors.add(color_of[neighbor]) + + # Best-fit: find smallest existing slot that fits and doesn't conflict + best_slot: Optional[int] = None + best_waste = -1 + + for slot_idx in range(len(slot_sizes)): + if slot_idx in neighbor_colors: + continue + new_capacity = max(slot_sizes[slot_idx], size_k) + waste = new_capacity - size_k + if best_slot is None or waste < best_waste: + best_waste = waste + best_slot = slot_idx + + if best_slot is not None: + color_of[k] = best_slot + slot_sizes[best_slot] = max(slot_sizes[best_slot], size_k) + else: + color_of[k] = len(slot_sizes) + slot_sizes.append(size_k) + + # Allocate each slot as a SEPARATE tensor + global_slot_offset = len(self._slots) + slot_tensors: List[torch.Tensor] = [] + + for slot_size in slot_sizes: + t = torch.empty(slot_size, dtype=dtype, device=device) + slot_tensors.append(t) + self._slots.append( + self._SlotInfo( + tensor=t, size=slot_size, dtype=dtype, device=device + ) + ) - After all intervals are colored, slots are laid out contiguously - and a single ``torch.empty`` is issued for the group. + # Map each key to its slot and pre-compute the view + for k in keys: + local_idx = color_of[k] + global_idx = global_slot_offset + local_idx + self._key_to_slot[k] = global_idx + size_k = key_max_size.get(k, 0) + # View into the slot tensor (first size_k elements) + self._key_to_view[k] = slot_tensors[local_idx][:size_k] - Complexity: O(n²) worst-case (linear scan per interval), negligible - for the typical tens-to-hundreds of param groups per dtype/device. - """ - intervals = sorted(intervals, key=lambda iv: iv.alloc_seq) - - # free_slots: list of (local_slot_index, free_seq) - free_slots: List[Tuple[int, int]] = [] - group_slots: List["TracePoolAllocator._Slot"] = [] # slots for this group - local_to_global: Dict[int, int] = {} # local index -> global index - - for iv in intervals: - assigned = False - # Try to reuse an existing slot whose previous occupant has already freed - for i, (slot_idx, slot_free_seq) in enumerate(free_slots): - if slot_free_seq < iv.alloc_seq: - slot = group_slots[slot_idx] - if iv.size > slot.size: - slot.size = iv.size # grow if needed - free_slots[i] = (slot_idx, iv.free_seq) # update free time - self._slot_map.setdefault(iv.key, []).append(local_to_global[slot_idx]) - assigned = True - break - - if not assigned: - # No reusable slot — allocate a new one - local_idx = len(group_slots) - global_idx = len(self._slots) - local_to_global[local_idx] = global_idx - slot = self._Slot(offset=0, size=iv.size, dtype=dtype, device=device) - group_slots.append(slot) - self._slots.append(slot) - free_slots.append((local_idx, iv.free_seq)) - self._slot_map.setdefault(iv.key, []).append(global_idx) - - # Lay out slots contiguously within the group pool - offset = 0 - for slot in group_slots: - slot.offset = offset - offset += slot.size - - if offset > 0: - self._pools[(dtype, device)] = torch.empty(offset, dtype=dtype, device=device) - return offset + return sum(slot_sizes) # -- Phase 3: optimized runtime ------------------------------------- # - # - # Each micro-batch replays the same alloc/free sequence. Because a - # single allocation key may appear multiple times (e.g., forward - # unshard → free → backward unshard → free), ``_slot_map[key]`` is - # a **list** of slot indices in alloc order. A per-key cursor in - # ``_slot_cursors`` tracks which index to consume next. Between - # micro-batches, ``reset_cursor()`` rewinds all cursors to 0. - - def _pool_allocate( + + def _optimized_allocate( self, key: AllocatorKey, size: int, dtype: torch.dtype, device: torch.device ) -> Bucket: - """Return a ``Bucket`` whose data is a slice of the pre-allocated pool. - - Advances the per-key slot cursor after consuming the slot. - """ - slot_list = self._slot_map[key] - cursor = self._slot_cursors.get(key, 0) - assert cursor < len(slot_list), ( - f"no slot available for key={key} " f"(cursor={cursor}, slots={slot_list})" - ) - slot_idx = slot_list[cursor] - self._slot_cursors[key] = cursor + 1 - + slot_idx = self._key_to_slot[key] slot = self._slots[slot_idx] - assert not slot.in_use, ( - f"slot {slot_idx} already in use (key={key}, " - f"cursor={cursor}, slot_list={slot_list})" + assert size <= slot.size, ( + f"requested {size} > slot capacity {slot.size} (key={key!r})" ) - assert size <= slot.size, f"requested {size} > slot capacity {slot.size} (key={key})" - pool = self._pools[(slot.dtype, slot.device)] slot.in_use = True - self._seq += 1 - return Bucket(data=pool[slot.offset : slot.offset + size]) - - def _pool_free(self, key: AllocatorKey) -> None: - """Mark the most recently allocated slot for this key as free. + self._active_keys.add(key) + # Return a view of the pre-allocated slot tensor + view = self._key_to_view[key] + return Bucket(data=view) - Double-frees are silently ignored (idempotent). - """ - # The last allocated slot for this key is at cursor - 1. - slot_idx = self._slot_map[key][self._slot_cursors.get(key, 1) - 1] - slot = self._slots[slot_idx] - self._seq += 1 - if not slot.in_use: + def _optimized_free(self, key: AllocatorKey) -> None: + if key not in self._active_keys: return - slot.in_use = False + self._slots[self._key_to_slot[key]].in_use = False + self._active_keys.discard(key) + + # -- Lifecycle ------------------------------------------------------- # + + def reset(self) -> None: + """Full teardown: discard pool, plan, and trace; return to trace phase.""" + self._phase = "trace" + self._seq = 0 + self._trace.clear() + self._trace_meta.clear() + self._buckets.clear() + self._active_keys.clear() + self._slots.clear() + self._key_to_slot.clear() + self._key_to_view.clear() + + @property + def phase(self) -> str: + return self._phase + + @property + def total_pool_bytes(self) -> int: + total = 0 + for slot in self._slots: + total += slot.size * slot.tensor.element_size() + return total # -- Debug ---------------------------------------------------------- # def dump_trace(self) -> str: - """Return a human-readable dump of the trace and pool plan. - - Useful for debugging slot-conflict errors (e.g., "slot already in use"). - """ + """Return a human-readable dump of the trace and pool plan.""" lines = [] - lines.append(f"=== TracePoolAllocator (phase={self._phase}, seq={self._seq}) ===") + lines.append(f"=== TracePoolAllocator (phase={self._phase}) ===") lines.append(f"trace events: {len(self._trace)}") for ev in self._trace: meta = self._trace_meta.get(ev.key) @@ -487,52 +516,69 @@ def dump_trace(self) -> str: if self._phase == "optimized": lines.append(f"\nslots: {len(self._slots)}") for i, slot in enumerate(self._slots): + keys_in_slot = [ + k for k, idx in self._key_to_slot.items() if idx == i + ] lines.append( - f" slot[{i}]: offset={slot.offset} size={slot.size} " - f"dtype={slot.dtype} device={slot.device}" + f" slot[{i}]: size={slot.size} " + f"dtype={slot.dtype} device={slot.device} " + f"addr=0x{slot.tensor.data_ptr():x} " + f"{'IN_USE' if slot.in_use else 'free'} " + f"keys={keys_in_slot}" ) - lines.append("\nslot_map (key -> [slot indices]):") - for key, slot_list in self._slot_map.items(): - cursor = self._slot_cursors.get(key, 0) - lines.append(f" {key} -> {slot_list} cursor={cursor}") + lines.append( + f"\ntotal pool: {len(self._slots)} slots, " + f"{self.total_pool_bytes} bytes " + f"({self.total_pool_bytes / 1024 / 1024:.1f} MB)" + ) return "\n".join(lines) - # -- Lifecycle ------------------------------------------------------- # - - def reset_cursor(self) -> None: - """Reset all slot cursors to 0 for the next iteration / micro-batch. - Must be called between micro-batches in the optimized phase so that - the alloc/free sequence replays from the start of the slot lists. - """ - self._slot_cursors.clear() - self._seq = 0 +def _intervals_overlap( + ivs_a: List[Tuple[int, int]], ivs_b: List[Tuple[int, int]] +) -> bool: + """Check if any interval in ivs_a overlaps with any interval in ivs_b. - def reset(self) -> None: - """Reset to trace phase, discarding the pool and all recorded state.""" - self._phase = "trace" - self._seq = 0 - self._trace.clear() - self._trace_meta.clear() - self._buckets.clear() - self._pools.clear() - self._slot_map.clear() - self._slot_cursors.clear() - self._slots.clear() - - @property - def phase(self) -> str: - """Current allocator phase: 'trace', 'plan', or 'optimized'.""" - return self._phase - - @property - def total_pool_bytes(self) -> int: - """Total pool size in bytes across all dtype/device groups.""" - total = 0 - for (dtype, _), pool in self._pools.items(): - total += pool.numel() * pool.element_size() - return total + Two intervals (a_start, a_end) and (b_start, b_end) overlap iff + a_start < b_end AND b_start < a_end. + """ + # For small lists (common case: 1-3 intervals per key), brute force + if len(ivs_a) * len(ivs_b) <= 16: + for a_start, a_end in ivs_a: + for b_start, b_end in ivs_b: + if a_start < b_end and b_start < a_end: + return True + return False + + # Sweep-line for larger sets + events: List[Tuple[int, int, int]] = [] + for start, end in ivs_a: + events.append((start, 0, 0)) + events.append((end, 1, 0)) + for start, end in ivs_b: + events.append((start, 0, 1)) + events.append((end, 1, 1)) + events.sort(key=lambda e: (e[0], -e[1])) + + active_a = 0 + active_b = 0 + for time, typ, group in events: + if typ == 0: + if group == 0: + active_a += 1 + if active_b > 0: + return True + else: + active_b += 1 + if active_a > 0: + return True + else: + if group == 0: + active_a -= 1 + else: + active_b -= 1 + return False def _free_storage(tensor: torch.Tensor) -> None: diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/allocator_design.md b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/allocator_design.md new file mode 100644 index 00000000000..e38129ee72f --- /dev/null +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/allocator_design.md @@ -0,0 +1,538 @@ +# TracePoolAllocator — v3 design (static key→slot plan) + +## Background: Megatron FSDP v2 buffer allocation + +Megatron FSDP v2 shards model parameters across GPUs. Before a layer computes +forward/backward, it must **all-gather** (unshard) the parameters — collecting +shards from all ranks into a temporary full-sized buffer. After compute, it +**reshards** — freeing that temporary buffer. The same pattern applies to +gradient reduction. + +These temporary buffers are managed by a `BucketAllocator` that provides +`allocate(key, size, dtype, device) → Bucket` and `free(key)`. Each FSDP +parameter group gets its own key — e.g. `(pg_id, "model_weight")` for +parameter all-gather and `(pg_id, "main_grad")` for gradient reduction. + +``` +Forward: allocate(model_weight) → [compute] → free(model_weight) +Backward: allocate(main_grad) → [reduce] → free(main_grad) +``` + +## Why CUDA graph needs stable addresses + +CUDA graphs capture a sequence of GPU kernel launches into a single replayable +object. During capture, PyTorch records every kernel launch and the exact GPU +memory addresses those kernels read/write. During replay, the graph replays +those exact operations — it cannot handle the addresses changing. + +So if a layer's forward reads from address `0x7f00`, that same layer's replay +must also read from `0x7f00`. Every buffer that a CUDA graph touches must have +a **fixed, deterministic memory address** across all micro-batches. + +## Why the current TracePoolAllocator breaks + +The current allocator uses a **seq-driven replay schedule**: + +1. **Trace** (micro-batch 0): Record every alloc/free call as `(seq, op, key)`. +2. **Plan** (`plan()`): Build intervals, color them, produce a `seq → (op, key, slot)` schedule (`_seq_ops`). +3. **Replay** (subsequent micro-batches): Walk `_seq_ops` linearly — each + `allocate`/`free` call must arrive at the exact seq position in the exact + order the trace produced. + +**The failure**: When CUDA graph is introduced, the seq-driven replay cannot +adapt. Hook calls during capture warmup, capture, and replay may arrive in +different orders or at different relative positions from the trace. The +`_seq` counter drifts from `_seq_ops`, and future `allocate`/`free` calls +hit `RuntimeError` because the seq-driven schedule expects a different call +at the current position. + +**The root cause**: The allocator assumes a fixed, repeatable alloc/free call +order and wraps it into a position-indexed schedule (`_seq_ops`). CUDA graph +capture and replay break this assumption — calls can be reordered, omitted, +or duplicated relative to the trace. + +## Design goal + +> Keep `TracePoolAllocator` as a single class. Replace the seq-driven replay +> schedule (`_seq_ops`) with a **static key→slot plan** (`_key_to_slot`). +> Each key maps to exactly one slot, giving one fixed memory address — +> derived once from the trace. Per-key intervals are merged into a +> single time range during ``plan()`` so every key gets exactly one slot. + +## Approach: `plan()` builds a static key→slot mapping + +Instead of using the trace as a replay script, use it as **input data for +memory planning**. The trace tells us which allocations overlap in time. +Per-key intervals are merged into a single time range, then colored with +greedy left-edge. Runtime is just a dict lookup: + +```python +def allocate(key, ...): + slot = _key_to_slot[key] # dict lookup, O(1) + return pool[slot.offset : slot.offset + size] # always same address +``` + +No `_seq` counter. No `_seq_ops` schedule. No dependency on call order. +Per-key cursors reset to 0 between micro-batches, guaranteeing forward-pass +address stability. + +## What stays, what changes, what's removed + +| Stays unchanged | Changes | Removed | +|---|---|---| +| Trace phase (`_trace_allocate`, `_trace_free`, `_TraceEvent`) | `plan()` builds `_key_to_slot` instead of `_seq_ops` | `_seq`, `_seq_ops` | +| Interval construction from alloc/free pairs | `allocate()` / `free()` dispatch to key→slot lookup instead of seq walk | `_pool_allocate`, `_pool_free` | +| Per-(dtype, device) pool tensors | `enable_flexible_mode` / `disable_flexible_mode` not needed | `_flexible`, `_flex_key_to_slot` | +| `Bucket` dataclass, `BucketAllocator` interface | `_phase` transitions: `"trace"` → `"optimized"` (no intermediate `"plan"` phase held) | `snapshot_slots`, `restore_slots` | +| `_key_to_slot: Dict[alloc_key, slot_idx]` | | `reset_cursor()` | + +### Why flexible mode is removed + +In the current design, `enable_flexible_mode` / `disable_flexible_mode` provide +key→slot lookup for auxiliary allocations between micro-batches (e.g. weight +quantisation). In v3, **every** `allocate()` is already a key→slot lookup — the +flexible-mode path is the **default** path. A separate toggle is unnecessary. + +Between micro-batches, the allocator is idle. An auxiliary `allocate(quant_key)` will +find the slot free → works exactly as before. + +## `plan()` — the core method + +```python +def plan(self) -> int: + """Build a static key→slot plan from the recorded trace. + + 1. Replay trace events to pair alloc↔free into intervals. + 2. Group intervals by (dtype, device), color each group with + greedy left-edge algorithm. + 3. Allocate one flat pool tensor per group. + 4. Build _key_to_slot: every alloc_key → exactly one slot_idx. + """ +``` + +### The coloring algorithm + +Per-key intervals are merged into a single time range before coloring, +so each key appears exactly once — plain left-edge, no same-key +enforcement needed: + +```python +def _color_group(intervals, dtype, device) -> int: + sorted_intervals = sorted(intervals, key=lambda iv: iv.alloc_seq) + + free_slots = [] # (local_slot_idx, free_seq) + group_slots = [] # Slot objects for this group + local_to_global = {} # local → global slot index + + for iv in sorted_intervals: + # ── left-edge: reuse an existing free slot, or create new ── + assigned_local = None + for local_idx, slot_free_seq in free_slots: + if slot_free_seq < iv.alloc_seq: + slot = group_slots[local_idx] + if iv.size > slot.size: + slot.size = iv.size + free_slots[...] = (local_idx, iv.free_seq) + assigned_local = local_idx + break + + if assigned_local is None: + assigned_local = len(group_slots) + global_idx = len(self._slots) + local_to_global[assigned_local] = global_idx + slot = Slot(offset=0, size=iv.size, dtype=dtype, device=device) + group_slots.append(slot) + self._slots.append(slot) + free_slots.append((assigned_local, iv.free_seq)) + else: + global_idx = local_to_global[assigned_local] + + self._key_to_slot[iv.key] = global_idx + + # Lay out slots contiguously with alignment + offset = 0 + alignment = _get_alignment(device, dtype) + for slot in group_slots: + offset = (offset + alignment - 1) // alignment * alignment + slot.offset = offset + offset += slot.size + + if offset > 0: + self._pools[(dtype, device)] = torch.empty(offset, dtype=dtype, device=device) + + return offset +``` + +### Memory alignment + +Each slot's offset is aligned to a device- and dtype-aware minimum. This is +critical for NVFP4 (sub-byte) types and CUDA kernel alignment requirements +(e.g. 256-byte base alignment from `cudaMalloc`): + +```python +def _get_alignment(device, dtype): + """Return the minimum alignment (in elements) for the given device/dtype.""" + element_bytes = torch.empty(0, dtype=dtype, device=device).element_size() + # At minimum, align to the element size itself + align_bytes = max( + element_bytes, + torch.cuda.get_device_properties(device).texture_alignment + if device.type == "cuda" else 1, + ) + return align_bytes // element_bytes +``` + +### Coloring algorithm — known limitations and future improvements + +The current greedy left-edge algorithm with per-key interval merging +meets the basic requirements but has room for improvement: + +1. **Per-group isolation**: Coloring runs independently per `(dtype, device)` + group. Groups on the same device could share a single pool with a unified + coloring pass, potentially reducing total memory. This requires handling + different element sizes in the same pool (offset arithmetic must account + for dtype-specific strides). + +2. **Merged range over-estimates active time**: Merging all of a key's + intervals into a single range (first alloc → last free) may prevent + other keys from reusing a slot during gaps between the key's intervals. + For FSDP forward/backward patterns, these gaps are typically small + (consecutive passes), so the overhead is minimal in practice. + +3. **Alternatives worth exploring**: + - **Global (cross-group) coloring**: Run one coloring pass across all + ``(dtype, device)`` groups that share a device, reducing total pool + bytes. + - **Slot merging pass**: After coloring, adjacent slots with compatible + dtypes could be merged if their intervals never overlap and no alignment + constraints are violated. + - **Size-class bucketing**: Group intervals by size class to reduce + internal fragmentation when a small interval forces a large slot. + +4. **Evaluation criteria**: + - Total pool bytes vs. the per-key `torch.empty` baseline. + - Slot count (fewer slots = less metadata overhead). + - Internal fragmentation (unused bytes within each slot, `slot.size - max_iv_size`). + - Alignment waste (padding bytes between slots). + +**Plan**: Ship with the current algorithm and instrument it with pool-size +and fragmentation metrics. Collect traces from real workloads (large models, +varied parallelism configs). Use the data to guide which improvements are +worth the complexity. + +## Runtime: allocate / free / reset + +State is driven entirely by ``_optimized_allocate`` and ``_optimized_free``. +No batched reset is needed — if every ``allocate`` is matched by a ``free`` +within the micro-batch, ``in_use`` flags and ``_active_keys`` are already +correct by the end. + +```python +def allocate(key, size, dtype, device): + slot_idx = _key_to_slot[key] # raises KeyError if key never traced + slot = _slots[slot_idx] + if slot.in_use and key not in _active_keys: + raise RuntimeError("Slot collision") + assert size <= slot.size + if key in _active_keys: + # Re-entrant: key already allocated this micro-batch — idempotent + return Bucket(data=pool[slot.offset : slot.offset + size]) + slot.in_use = True + _active_keys.add(key) + return Bucket(data=pool[slot.offset : slot.offset + size]) + +def free(key): + if key not in _active_keys: + return # double-free or never-allocated + _slots[_key_to_slot[key]].in_use = False + _active_keys.discard(key) + +def reset(): + """Full teardown: discard pool, plan, and trace.""" + self._phase = "trace" + self._seq = 0 + self._trace.clear() + self._trace_meta.clear() + self._buckets.clear() + self._active_keys.clear() + self._pools.clear() + self._key_to_slot.clear() + self._slots.clear() +``` + +### Error handling for unknown keys + +If `allocate(key)` is called with a key never seen during trace, `_key_to_slot` +raises `KeyError`. The caller must re-trace (call `reset()`, re-run micro-batch +0, then `plan()`) to pick up new allocation patterns. This guarantees that all +keys used during CUDA graph replay were planned, keeping addresses stable. + +No on-the-fly slot allocation is supported — it would change pool tensor sizes +and invalidate all captured CUDA graphs. + +## CUDA graph integration lifecycle + +The full lifecycle spanning trace, plan, capture, and replay: + +``` + ┌─────────────────────────────────────────┐ + │ Micro-batch 0 (trace) │ + │ Phase: "trace" │ + │ │ + │ forward_pre_hook (per-module): │ + │ if is_root: forward_phase=True │ + │ unshard → _trace_allocate │ + │ Backward: alloc/free → _trace_* │ + │ records │ + │ Post-backward: plan() builds │ + │ _key_to_slot, allocates │ + │ pool tensors │ + │ Phase → "optimized" │ + └──────────────────┬──────────────────────┘ + │ + ┌────────────────────────┼────────────────────────┐ + │ │ │ + ┌─────────▼──────────┐ ┌────────▼───────────┐ ┌───────▼───────────┐ + │ Micro-batch 1 │ │ Micro-batch 2 │ │ Micro-batch N │ + │ (capture) │ │ (replay) │ │ (replay) │ + │ Phase: "optimized" │ │ Phase: "optimized" │ │ Phase: "optimized"│ + │ │ │ │ │ │ + │ forward_pre_hook │ │ forward_pre_hook │ │ forward_pre_hook │ + │ (per-module): │ │ (per-module): │ │ (per-module): │ + │ assert not │ │ assert not │ │ assert not │ + │ cuda_graph_active│ │ cuda_graph_active │ │ cuda_graph_active│ + │ unshard→allocate │ │ unshard→allocate │ │ unshard→allocate│ + │ → key→slot │ │ → key→slot │ │ → key→slot │ + │ │ │ │ │ │ + │ if not captured: │ │ graphed() call │ │ graphed() call │ + │ capture_fwd(): │ │ (replay capture) │ │ (replay capture)│ + │ • pop hooks │ │ │ │ │ + │ (recursive) │ │ pool tensors │ │ pool tensors │ + │ • set cuda_ │ │ at SAME addr │ │ at SAME addr │ + │ graph_active │ │ as MB 1 │ │ as MB 1 │ + │ • warmup │ │ │ │ │ + │ • capture fwd │ │ │ │ │ + │ • clear flag │ │ │ │ │ + │ • restore hooks│ │ │ │ │ + │ • install() │ │ │ │ │ + │ │ │ │ │ │ + │ Post-bwd hooks │ │ Post-bwd hooks │ │ Post-bwd hooks │ + │ fire (eager) │ │ fire │ │ fire │ + │ │ │ │ │ │ + │ Post-backward: │ │ Post-backward: │ │ Post-backward: │ + │ (hooks clean up) │ │ (hooks clean up) │ │ (hooks clean up)│ + └────────────────────┘ └─────────────────────┘ └───────────────────┘ +``` + +### Per-module CUDA graph capture detail (MB 1) + +FSDP captures one CUDA graph per compatible leaf module (not the whole model). +The capture is triggered inside the unified per-module forward pre-hook +(``forward_pre_hook``, which also handles root-level phase init): + +1. **Pre-forward hook allocates param**: The hook calls + `allocate(pg_id, "model_weight")` → key→slot lookup → pool view at fixed + address. FSDP all-gathers shards into that view. The param buffer is now + resident at a deterministic address for the rest of the pool's lifetime. + +2. **main_grad is manually allocated**: Before capture, the gradient buffer + is pre-allocated via `allocate(pg_id, "main_grad")` to ensure both forward + and backward passes see fixed addresses. + +3. **Hooks popped recursively**: `_pop_hooks_recursive` removes all FSDP + hooks on the target module and every submodule. ``cuda_graph_active`` is + set to ``True`` so any stray hook that fires during this window asserts. + +4. **Warmup + capture**: `make_graphed_callables` runs 3 warmup iterations + (no hooks → raw forward only) + the actual graph capture. The graph + records GPU kernels reading/writing the pool views at their fixed addresses. + +5. **Cleanup**: ``cuda_graph_active`` is cleared, hooks are restored + recursively, the patched forward is installed. + +6. **Replay**: The patched forward calls `graphed(*flat)`. Since the graph + captured the fixed pool addresses, replay reads/writes the **same** + addresses every time — no re-allocation, no address change. + +### Key safety property + +During capture, ``cuda_graph_active`` is ``True`` and all hooks are popped +(recursively). The CUDA graph records kernel operations against the pool +tensor addresses. Because the pool tensors are allocated **once** in +``plan()`` and never resized, the same ``key`` will always resolve to the +same address — regardless of call ordering variations across micro-batches. + +### TE wgrad fusion under CUDA graph + +When Transformer Engine's gradient-accumulation fusion writes weight +gradients directly into ``param.main_grad``, it sets the Python-side +``grad_added_to_main_grad = True`` to tell FSDP not to overwrite +``main_grad``. Under CUDA graph replay, TE's GPU kernel still runs but the +``setattr`` is not part of the graph. + +**Fix**: Recording and consume live together in ``reduce_grad()``. On the +first call (trace micro-batch, eager backward), if TE set the flag, we +persist it as ``param._mfsdp_recorded_te_wgrad = True``. On all subsequent +calls, ``reduce_grad`` checks both the live flag (eager path) and the +recorded flag (CUDA graph replay path). ``pre_backward_hook`` restores +the recorded value before each backward pass. + +``` +reduce_grad(): + if grad_added_to_main_grad or _mfsdp_recorded_te_wgrad: + discard .grad # TE already populated main_grad + if grad_added and enable_cuda_graph: + _mfsdp_recorded_te_wgrad = True # persist for future replays +``` + +## Visual + +``` + ┌───────────────────────────────┐ + │ Micro-batch 0 (trace) │ + │ Phase: "trace" │ + │ Forward + backward runs │ + │ All alloc/free calls logged │ + └──────────────┬────────────────┘ + │ + ┌────────▼──────────┐ + │ plan() │ + │ Phase → "optimized" │ + │ │ + │ 1. Build intervals │ + │ from trace │ + │ 2. Merge per-key │ + │ intervals │ + │ 3. Interval coloring │ + │ → slots │ + │ 3. Align & allocate │ + │ pool tensors │ + │ 4. Build static │ + │ key → slot map │ + └─────────┬───────────┘ + │ + ┌───────────────────────┼───────────────────────┐ + │ │ │ + ┌─────────▼───────┐ ┌────────▼────────┐ ┌───────▼─────────┐ + │ Micro-batch 1 │ │ Micro-batch 2 │ │ Micro-batch N │ + │ (capture) │ │ (replay) │ │ (replay) │ + │ Phase: │ │ Phase: │ │ Phase: │ + │ "optimized" │ │ "optimized" │ │ "optimized" │ + │ │ │ │ │ │ + │ alloc(key) → │ │ alloc(key) → │ │ alloc(key) → │ + │ same address │ │ same address │ │ same address │ + └─────────────────┘ └──────────────────┘ └─────────────────┘ +``` + +## Hooks integration (updated for v3) + +`_register_forward_pre_hook` is a single hook that serves all FSDP modules +(root and non-root). Root-only logic (phase flags, CUDA graph stream init) +is gated by ``is_root``. + +``` +Micro-batch 0 (trace) +┌───────────────────────────────────────────────────────────────────────────┐ +│ forward_pre_hook (per-module) │ +│ if is_root: forward_phase = True │ +│ assert not cuda_graph_active │ +│ unshard params → _trace_allocate(); forward (trace) │ +│ │ +│ pre_backward_hook (per-module, reverse order) │ +│ if is_root: switch to backward phase, enqueue final callback │ +│ assert not cuda_graph_active │ +│ unshard (bwd) → _trace_allocate() │ +│ reset: post_backward_issued=False, grad_added_to_main_grad=False │ +│ TE wgrad groups: overwrite_main_grad=True │ +│ backward (trace, eager) → TE sets grad_added_to_main_grad │ +│ │ +│ post_backward (per-module): reshard, reduce_grad, post_backward_issued │ +│ │ +│ _post_backward_final_callback │ +│ handle modules with post_backward_issued=False (activation ckpt) │ +│ drain async reduce-grad events │ +│ plan() → "optimized" │ +│ │ +│ reduce_grad (per-module, inside post_backward): │ +│ if grad_added_to_main_grad (TE set it eagerly): │ +│ discard .grad; param._mfsdp_recorded_te_wgrad ← True │ +│ (recorded for CUDA graph replays where setattr doesn't fire) │ +└───────────────────────────────────────────────────────────────────────────┘ + +Micro-batch 1+ (optimized) +┌───────────────────────────────────────────────────────────────────────────┐ +│ forward_pre_hook (per-module) │ +│ if is_root: forward_phase = True │ +│ assert not cuda_graph_active │ +│ unshard params → allocate() → key→slot │ +│ → if enable_cuda_graph and not yet captured and not backward: │ +│ FSDPCudaGraphRunner.capture_forward() │ +│ (pops hooks recursively, sets cuda_graph_active, │ +│ runs warmup+capture, restores hooks, clears flag) │ +│ forward (optimized, graph replays if captured) │ +│ │ +│ pre_backward_hook (per-module, reverse order) │ +│ if is_root: switch to backward phase, enqueue final callback │ +│ assert not cuda_graph_active │ +│ unshard (bwd) → allocate() → key→slot │ +│ reset: post_backward_issued=False │ +│ grad_added_to_main_grad ← _mfsdp_recorded_te_wgrad (CUDA graph) │ +│ TE wgrad groups: overwrite_main_grad=True │ +│ backward (optimized, graph replays if captured) │ +│ │ +│ post_backward (per-module): reshard, reduce_grad, post_backward_issued │ +│ │ +│ _post_backward_final_callback │ +│ handle modules with post_backward_issued=False │ +│ drain async reduce-grad events │ +│ reset context state for next micro-batch │ +└───────────────────────────────────────────────────────────────────────────┘ +``` + +## Debug: `dump_trace()` + +Updated for v3 to show the full static key→slot mapping (not just active keys): + +```python +def dump_trace(self) -> str: + lines = [f"=== TracePoolAllocator (phase={self._phase}) ==="] + # ... trace events (unchanged) ... + + if self._phase == "optimized": + lines.append(f"\nslots: {len(self._slots)}") + for i, slot in enumerate(self._slots): + lines.append( + f" slot[{i}]: offset={slot.offset} size={slot.size} " + f"dtype={slot.dtype} device={slot.device} " + f"{'in_use' if slot.in_use else 'free'}" + ) + total_bytes = sum( + s.size * torch.empty(0, dtype=s.dtype).element_size() + for s in self._slots + ) + lines.append(f"\ntotal pool: {len(self._slots)} slots, {total_bytes} bytes") + lines.append(f"\nkey_to_slot ({len(self._key_to_slot)} entries):") + for key, slot_idx in sorted(self._key_to_slot.items(), key=lambda x: str(x[0])): + slot = self._slots[slot_idx] + lines.append( + f" {key} -> slot[{slot_idx}] " + f"(offset={slot.offset}, size={slot.size}, " + f"address=0x{pool[slot.offset].data_ptr():x})" + ) + lines.append(f"\nactive_keys ({len(self._active_keys)}):") + for key in self._active_keys: + lines.append(f" {key}") + return "\n".join(lines) +``` + +## Key properties + +| Property | How it's achieved | +|---|---| +| Fixed address per key | Pool tensors allocated once in `plan()`, never resized. Each key maps to one slot at a fixed, aligned offset. `pool[offset:offset+size]` returns identical view every time. | +| Memory efficiency | Left-edge interval coloring reuses slots for non-overlapping allocations. Per-key intervals are merged into a single time range before coloring, so each key gets exactly one slot — guaranteed stable address with minimal slot count. | +| CUDA graph compatible | One key → one slot → one address. Per-key intervals merged into a single time range during ``plan()`` ensures each key gets exactly one dedicated slot at a fixed offset. | +| No fragmentation within pool | Slots are laid out contiguously with alignment padding. No gaps between slots (only alignment-padding gaps). No dynamic allocation/deallocation — pool is a single `torch.empty`. | +| Simple implementation | `allocate()` is a dict lookup + guard. `free()` is a flag clear. State is self-driving — no batch reset, no cursors, no seq walking. | +| Same trace mechanism | Phase 1 (trace) is identical. Only the plan output and runtime dispatch change. | +| Backward compatible for non-CUDA-graph users | The key→slot runtime also works for regular eager execution — it's strictly simpler than the seq-driven approach. | diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/cuda_graph_design.md b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/cuda_graph_design.md new file mode 100644 index 00000000000..637655353fb --- /dev/null +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/cuda_graph_design.md @@ -0,0 +1,350 @@ +# Silent CUDA Graph inside Megatron FSDP v2 — Design + +> **Experimental** — CUDA graph support in Megatron FSDP v2 is an experimental +> feature. The API and behaviour may change in future releases without notice. + +## 1. Motivation + +mcore's CUDA graph system (`cuda_graph_impl="local"`, `cuda_graphs.py`) was +designed for DDP's memory model — each layer receives freshly-allocated tensor +inputs/outputs. FSDP v2 shares the same pool-backed buffers across layers, +and the FSDP hooks (unshard/reshard) are not captured in the graph in the +same way. + +This doc describes a CUDA graph system built INTO FSDP v2, using +`TracePoolAllocator` as the stable-memory foundation. The user enables it +with a single flag — everything else is automatic. + +## 2. One knob + +```python +fully_shard(module, enable_cuda_graph=True) +``` + +No `--cuda-graph-warmup-steps`, no `--cuda-graph-scope`, no coordination with +the pipeline schedule. The system automatically captures and replays on the +first optimized microbatch. + +## 3. Why TracePoolAllocator is the enabler + +CUDA graphs require **stable buffer addresses**. After `plan()` allocates +the pool tensor, every slot has a fixed offset. The returned views' addresses +are deterministic every micro-batch. + +During graph replay, the graphed CUDA kernels operate on the exact addresses +that were recorded during capture. The allocator is **not called** for +graphed modules during replay — the graph uses the addresses directly. Pool +slots are pre-allocated by `plan()`, so no dynamic allocation occurs inside +the graph region. + +## 4. Lifecycle + +``` +╔═══════════════════════════════════════════════════════════════╗ +║ Stage 1 — Microbatch 0 (trace) ║ +║ ║ +║ forward: _trace_allocate / _trace_free → trace recorded ║ +║ backward: _trace_allocate / _trace_free → trace continues ║ +║ plan(): pool built, _seq_ops populated, phase="optimized" ║ +║ snapshot_slots(): freeze slot state for replay ║ +╚═══════════════════════════════════════════════════════════════╝ + │ +╔═══════════════════════════════════════════════════════════════╗ +║ Stage 2 — Microbatch 1+ (capture + eager forward) ║ +║ ║ +║ root_forward_pre_hook: ║ +║ reset_cursor() ║ +║ restore_slots() (restore capture-time slot state) ║ +║ forward (per graphed FSDPModule): ║ +║ forward_pre_hook detects enable_cuda_graph=True: ║ +║ 1. Creates FSDPCudaGraphRunner(module) ║ +║ 2. Runner sets cuda_graph_active=True, pops hooks ║ +║ 3. Calls torch.cuda.make_graphed_callables(shim, ...) ║ +║ → warms up, captures forward+backward graph ║ +║ 4. Runner restores hooks, clears cuda_graph_active ║ +║ 5. install() patches module.forward to graph replay ║ +║ Module then runs _orig_forward eagerly (this call) ║ +║ backward: ║ +║ post_backward hooks fire normally ║ +║ reshard / reduce_grad run normally ║ +║ post-bwd: ║ +║ enable_flexible_mode() ║ +╚═══════════════════════════════════════════════════════════════╝ + │ +╔═══════════════════════════════════════════════════════════════╗ +║ Stage 3 — Microbatch 2+ (replay) ║ +║ ║ +║ root_forward_pre_hook: ║ +║ reset_cursor() ║ +║ restore_slots() ║ +║ forward (per graphed FSDPModule): ║ +║ forward_pre_hook: runner already exists → skip capture ║ +║ patched forward: graphed(*flat) → replays captured graph ║ +║ → forward+backward complete within graph callable ║ +║ → NO Python FSDP hooks fire during graph ║ +║ → allocator NOT called for graphed modules ║ +║ forward (per non-graphed FSDPModule): ║ +║ normal eager: hooks fire → alloc calls ║ +║ backward: ║ +║ post_backward hooks fire normally (after graph returns) ║ +║ reshard / reduce_grad run normally ║ +║ post-bwd: ║ +║ enable_flexible_mode() ║ +╚═══════════════════════════════════════════════════════════════╝ +``` + +## 5. FSDPCudaGraphRunner + +The runner encapsulates per-module CUDA graph capture and replay: + +```python +class FSDPCudaGraphRunner: + def capture_forward(self, *sample_args, **sample_kwargs): + # 1. Build a _ForwardShim that freezes non-tensor kwargs + # 2. Disable side-stream collectives on the root context + # 3. Pop FSDP hooks from the module (save/restore) + # 4. Call torch.cuda.make_graphed_callables(shim, ...) + # → handles warmup (3 iters) + capture internally + # 5. Restore hooks and side-stream settings + + def install(self): + # Patches module.forward to a wrapper that flattens + # args/kwargs → positional tensors and calls graphed(*flat) +``` + +**Why `make_graphed_callables`?** PyTorch's built-in utility handles: +- Warmup iterations to settle cuDNN/cuBLAS auto-tuning and TE FP8 scales +- Memory pool allocation for captured tensors +- AccumulateGrad node creation on the capture stream (avoids + `cudaErrorStreamCaptureImplicit` by keeping `.grad` tensors alive from + warmup through capture) + +**Why pop hooks?** The graph records the module's `forward()` only — no +FSDP all-gather or reduce-scatter collectives. Hooks are restored +immediately after capture, before the eager forward call. + +## 6. Per-FSDPModule selectivity + +```python +# Only specific leaf layers graphed +for layer in model.layers: + fully_shard(layer, enable_cuda_graph=True) +fully_shard(model, enable_cuda_graph=False) +``` + +Each `FSDPModule` carries a flag in `_FSDPState`: + +```python +class _FSDPState: + enable_cuda_graph: bool = False +``` + +All modules share the same `TracePoolAllocator` — slot assignments are fixed +by `plan()` regardless of which modules are graphed. + +### Nesting limitation + +A parent FSDP module that contains other FSDP modules as children **cannot** +use `enable_cuda_graph=True`. Only leaf FSDP modules (those without FSDP +children) are eligible. This is enforced in `_init_fsdp_state` with a +`RuntimeError`. + +The reason: CUDA graph capture runs the module's forward without FSDP hooks +(the hooks are popped). For a parent module, the forward calls child +FSDP-module forwards, which may themselves have CUDA graph logic or require +FSDP unshard/reshard. Limiting CUDA graph to leaf modules avoids this +complexity. + +## 7. Capture stream invariants + +During CUDA graph capture, `FSDPCudaGraphRunner.capture_forward` disables +side-stream collectives and sets `cuda_graph_active=True` so that hooks +defer reshard and skip post-backward cleanup while the graph is being built. + +The root context provides: + +```python +@dataclass +class _FSDPRootContext: + cuda_graph_active: bool = False + """True only during CUDA graph capture (inside FSDPCudaGraphRunner). + Suppresses side-stream vs default-stream mismatches and defers reshard.""" + + @property + def cuda_graph_compatible(self) -> bool: + """Return True when both side-stream collectives are disabled.""" + return (not self.enable_unshard_prefetch) and (not self.enable_async_reduce_grad) +``` + +``FSDPModule.cuda_graph_compatible`` adds an additional check that the +allocator is in ``"optimized"`` phase: + +```python +@property +def cuda_graph_compatible(self) -> bool: + ctx = self._fsdp_root_context + if not isinstance(ctx.bucket_allocator, TracePoolAllocator): + return False + if ctx.bucket_allocator.phase != "optimized": + return False + return ctx.cuda_graph_compatible +``` + +During replay, ``cuda_graph_active`` is **not** set — the graph callable +runs forward+backward atomically inside ``graphed(*flat)``, and the standard +hook lifecycle handles reshard/reduce_grad normally after the graph returns. + +## 8. Hook changes + +### 8.1 Forward pre-hook — capture trigger + +```python +# hooks.py — _register_forward_pre_hook +def unshard_param_groups(fsdp_module, args, kwargs): + # ... unshard as normal ... + if fsdp_module._fsdp_state.enable_cuda_graph and ( + not hasattr(fsdp_module, "_fsdp_cg_runner") + ): + cg_runner = FSDPCudaGraphRunner(fsdp_module) + cg_runner.capture_forward(*args, **kwargs) + cg_runner.install() + fsdp_module._fsdp_cg_runner = cg_runner +``` + +Capture happens exactly once — the first time the module is called in +an optimized forward pass. Subsequent calls skip because `_fsdp_cg_runner` +already exists. + +### 8.2 Defer post-forward reshard + +```python +# hooks.py — _register_forward_hook → reshard_param_groups +def reshard_param_groups(module, *unused): + ctx = module._fsdp_root_context + if ctx.backward_phase and id(module) == ctx.backward_module: + return + if ctx.cuda_graph_active: + return # deferred; cleanup runs after capture returns + module.reshard() +``` + +### 8.3 Suppress post_backward during capture + +```python +# hooks.py — _register_backward_hook → post_backward +def post_backward(module): + ctx = module._fsdp_root_context + if ctx.cuda_graph_active: + return # hooks are popped during capture; nothing to do + ctx.backward_done_modules.add(id(module)) + ctx._advance_backward_module() + module.reshard() + ... +``` + +### 8.4 Final callback — skips during capture + +```python +# hooks.py — _post_backward_final_callback +def _post_backward_final_callback(root_state, root_module): + ctx = root_module._fsdp_root_context + if ctx.cuda_graph_active: + return # capture manages cleanup; hooks are already popped + stream = ctx.rs_stream + for module in reversed(ctx.forward_order): + ... + ctx.cuda_graph_active = False # not reached during capture — clears + # any stale state at end of replay +``` + +``cuda_graph_active`` is **only needed during capture** (when +``FSDPCudaGraphRunner.capture_forward`` calls ``make_graphed_callables``). +The runner pops hooks before capture and restores them after, so normal +post-forward / post-backward cleanup runs correctly outside the graph region. + +During replay, the graph callable runs forward+backward atomically inside +``graphed(*flat)``. After it returns, the standard hook lifecycle handles +reshard and reduce_grad normally — no ``cuda_graph_active`` guard is required. + +### 8.5 Root forward pre-hook — restore slots + +```python +# hooks.py — _register_root_forward_pre_hook +def root_forward_pre_hook(_hook_module, args): + ctx = fsdp_module._fsdp_root_context + ... + if isinstance(ba, TracePoolAllocator) and ba.phase == "optimized": + ba.disable_flexible_mode() + ba.reset_cursor() + if any(getattr(m._fsdp_state, "enable_cuda_graph", False) + for m in ctx.forward_order): + ba.restore_slots() +``` + +## 9. Allocator interface + +The allocator requires three methods to support CUDA graph replay: + +```python +class TracePoolAllocator: + _captured_slot_state: List[bool] # snapshot of slot.in_use at capture end + + def snapshot_slots(self): + """Freeze current slot in_use state after all graphs are captured.""" + self._captured_slot_state = [s.in_use for s in self._slots] + + def restore_slots(self): + """Restore slot state to capture-time snapshot before each replay.""" + for i, in_use in enumerate(self._captured_slot_state): + self._slots[i].in_use = in_use +``` + +`snapshot_slots()` is called in the post-backward callback after `plan()` +(trace→optimized transition). `restore_slots()` is called in the root +forward pre-hook before each replay to ensure the same slot availability +as during capture. + +## 10. Slot state management + +After the trace phase and `plan()`, `snapshot_slots()` freezes the +allocator's slot `in_use` state. Before each forward pass during replay, +`restore_slots()` resets slots to that snapshot. This ensures: + +- Flexible-mode allocs between microbatches see correct slot availability. +- Non-graphed modules during replay advance `in_use` flags correctly. +- The next replay starts from the same clean state. + +## 11. Risks + +| Risk | Mitigation | +|------|-----------| +| Side-stream all-gather invisible to graph → corrupt param buffers | `cuda_graph_compatible` assertion + runner disables side streams during capture | +| Post-forward reshard frees buffers too early | `cuda_graph_active` guard in `reshard_param_groups` defers reshard during capture | +| AccumulateGrad nodes carry stale stream → `cudaErrorStreamCaptureImplicit` | `make_graphed_callables` handles warmup + capture on the same stream, keeping `.grad` tensors alive through warmup | +| Non-deterministic FP8/RNG across warmup iterations | `make_graphed_callables` runs 3 warmup iterations before capture — FP8 scales settle, RNG advances | +| Slot `in_use` flags stale after replay | `restore_slots()` resets to capture-end snapshot before each forward | +| Capture triggered for nested FSDP modules | `_init_fsdp_state` rejects `enable_cuda_graph=True` on modules with FSDP children | +| Graph capture OOM (pool too large) | Pool is pre-allocated by `plan()`; CUDA graph mempool uses same pattern | + +## 12. Files + +| File | Role | +|------|------| +| `cuda_graph_runner.py` | `FSDPCudaGraphRunner` — capture via `make_graphed_callables`, install/uninstall forward patch | +| `fsdp_module.py` | `_FSDPRootContext.cuda_graph_active` flag, `cuda_graph_compatible` property, nested FSDP check in `_init_fsdp_state` | +| `hooks.py` | Capture trigger in forward pre-hook, reshard deferral in forward/backward hooks, slot restore in root forward pre-hook, cleanup in final callback | +| `allocator.py` | `TracePoolAllocator.snapshot_slots()` / `restore_slots()` for stable slot state | +| `fully_shard.py` | Accept `enable_cuda_graph` kwarg; passes to `_enable_cuda_graph()` | + +## 13. No user-visible knobs + +| What happens | User action | +|---|---| +| Trace collection | Automatic (MB0 forward) | +| Pool planning | Automatic (MB0 post-backward) | +| Slot snapshot | Automatic (after plan) | +| Graph capture | Automatic (first optimized MB forward, via `FSDPCudaGraphRunner`) | +| Graph replay | Automatic (subsequent MBs, via patched forward) | +| Slot state management | Automatic (snapshot/restore around replay) | +| Side-stream suppression | Automatic (`cuda_graph_active` flag in hooks) | +| Flexible mode toggling | Automatic (existing hook lifecycle) | diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/cuda_graph_runner.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/cuda_graph_runner.py new file mode 100644 index 00000000000..ebcbb297580 --- /dev/null +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/cuda_graph_runner.py @@ -0,0 +1,373 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""CUDA graph capture / replay for individual FSDP modules.""" + +import inspect +import gc +import warnings +from collections import OrderedDict +from typing import Any, Dict, List, Optional, Tuple + +import torch + + +# ------------------------------------------------------------------ +# Helpers +# ------------------------------------------------------------------ + +# All known hook attributes across PyTorch versions (including 2.x additions) +_HOOK_ATTRS = [ + "_forward_pre_hooks", + "_forward_hooks", + "_forward_hooks_with_kwargs", + "_forward_pre_hooks_with_kwargs", + "_backward_hooks", + "_backward_pre_hooks", + "_state_dict_hooks", + "_load_state_dict_pre_hooks", + "_load_state_dict_post_hooks", +] + + +def _get_forward_param_names(module: torch.nn.Module) -> List[str]: + """Return the ordered parameter names of module.forward (excluding 'self').""" + sig = inspect.signature(module.forward) + return [ + name + for name, p in sig.parameters.items() + if name != "self" + and p.kind + in ( + inspect.Parameter.POSITIONAL_OR_KEYWORD, + inspect.Parameter.KEYWORD_ONLY, + inspect.Parameter.POSITIONAL_ONLY, + ) + ] + + +class _ForwardShim(torch.nn.Module): + """Wraps module.forward so that non-tensor kwargs are frozen at + capture time and tensor inputs are passed positionally in signature + order. + + Handles None outputs by filtering them out before returning to + make_graphed_callables, and records their positions so they can be + restored during replay. + """ + + def __init__( + self, + module: torch.nn.Module, + tensor_param_names: List[str], + frozen_kwargs: dict, + ): + super().__init__() + self.module = module + self.tensor_param_names = tensor_param_names + self.frozen_kwargs = frozen_kwargs + + # Populated on first forward: tracks which output positions are None + self._none_mask: Optional[List[bool]] = None + # Whether the original output was a tuple (vs single tensor) + self._output_is_tuple: bool = True + + def forward(self, *flat_tensor_args): + kwargs = dict(zip(self.tensor_param_names, flat_tensor_args)) + kwargs.update(self.frozen_kwargs) + outputs = self.module.forward(**kwargs) + + # Handle single tensor output + if not isinstance(outputs, tuple): + self._output_is_tuple = False + if outputs is None: + raise RuntimeError( + "Module returned None as its only output — cannot capture " + "with CUDA graphs. The module must return at least one tensor." + ) + return outputs + + # Handle tuple outputs — filter out None values + self._output_is_tuple = True + if self._none_mask is None: + self._none_mask = [o is None for o in outputs] + + filtered = tuple(o for o in outputs if o is not None) + if len(filtered) == 0: + raise RuntimeError( + "Module returned a tuple of all None values — cannot capture " + "with CUDA graphs. At least one output must be a tensor." + ) + if len(filtered) == 1: + return filtered[0] + return filtered + + def restore_none_positions(self, graph_output) -> Any: + """Re-insert None values at their original positions after graph replay.""" + if not self._output_is_tuple or self._none_mask is None: + return graph_output + + # If all non-None outputs were collapsed to a single tensor + if not isinstance(graph_output, tuple): + graph_output = (graph_output,) + + full: List[Any] = [] + tensor_iter = iter(graph_output) + for is_none in self._none_mask: + if is_none: + full.append(None) + else: + full.append(next(tensor_iter)) + return tuple(full) + + +def _pop_hooks(module: torch.nn.Module) -> Dict[str, Any]: + """Remove all hooks from *module* (non-recursive) and return a snapshot.""" + saved: Dict[str, Any] = {} + for attr in _HOOK_ATTRS: + if hasattr(module, attr): + saved[attr] = getattr(module, attr) + setattr(module, attr, OrderedDict()) + return saved + + +def _pop_hooks_recursive(module: torch.nn.Module) -> List[Tuple[torch.nn.Module, Dict[str, Any]]]: + """Remove all hooks from *module* and all its submodules recursively. + + Returns a list of (submodule, saved_hooks) tuples for restore. + Using direct module references instead of id() for safety. + """ + saved: List[Tuple[torch.nn.Module, Dict[str, Any]]] = [] + for submodule in module.modules(): + saved.append((submodule, _pop_hooks(submodule))) + return saved + + +def _restore_hooks(module: torch.nn.Module, saved: Dict[str, Any]) -> None: + """Put the hooks back exactly as they were.""" + for name, value in saved.items(): + if value is not None: + setattr(module, name, value) + + +def _restore_hooks_recursive( + module: torch.nn.Module, saved: List[Tuple[torch.nn.Module, Dict[str, Any]]] +) -> None: + """Restore hooks for all submodules saved by ``_pop_hooks_recursive``.""" + for submodule, sub_saved in saved: + _restore_hooks(submodule, sub_saved) + + +# ------------------------------------------------------------------ +# Runner +# ------------------------------------------------------------------ + + +class FSDPCudaGraphRunner: + """Captures a forward+backward CUDA graph for one FSDP module. + + During capture hooks are temporarily removed so the graph records + only the user's ``forward()``, not FSDP all-gather / reduce-scatter + collectives. FSDP side streams are disabled for the capture region. + + Parameters: + fsdp_module: The FSDP module to capture. + gc_freeze: If True (default), call ``gc.collect()`` and ``gc.freeze()`` + before capture to prevent Python GC from stalling replay. + graph_pool: Optional shared CUDA graph memory pool handle obtained + via ``torch.cuda.graph_pool_handle()``. When provided, the + ``CUDAGraph`` is created with this handle so multiple FSDP + modules share the same backing memory pool, reducing total + GPU memory consumption. + + Usage:: + + runner = FSDPCudaGraphRunner(my_fsdp_module) + runner.capture_forward(sample_input) + runner.install() # patches module.forward + output = my_fsdp_module(input_batch) # replays graph, no hooks + runner.uninstall() # restore original behaviour + """ + + def __init__( + self, + fsdp_module: torch.nn.Module, + gc_freeze: bool = True, + graph_pool: Optional[Any] = None, + ): + warnings.warn( + "FSDPCudaGraphRunner is an experimental feature. The API and " + "behaviour may change in future releases without notice.", + FutureWarning, + stacklevel=2, + ) + self._module: torch.nn.Module = fsdp_module + self._gc_freeze: bool = gc_freeze + self._graph_pool: Optional[int] = graph_pool + + # Will hold the callable returned by make_graphed_callables + self._graphed: Optional[Any] = None + + self._orig_fwd: Optional[Any] = None + self._use_cuda_graph: bool = False + self._captured: bool = False + + # Saved during capture for install() replay flattening + self._tensor_param_names: List[str] = [] + self._frozen_kwargs: Dict[str, Any] = {} + + # The shim used during capture (needed for None restoration) + self._shim: Optional[_ForwardShim] = None + + # ------------------------------------------------------------------ + # 1. Capture + # ------------------------------------------------------------------ + + def capture_forward( + self, + *sample_args, + **sample_kwargs, + ) -> None: + assert self._module.cuda_graph_compatible, ( + "CUDA graph capture requires TracePoolAllocator in optimized phase" + ) + + # Introspect the module's forward signature + param_names = _get_forward_param_names(self._module.__class__) + + # Separate tensor vs non-tensor inputs + bound = {} + for i, val in enumerate(sample_args): + if i < len(param_names): + bound[param_names[i]] = val + bound.update(sample_kwargs) + + tensor_names = [ + n for n in param_names if n in bound and isinstance(bound[n], torch.Tensor) + ] + frozen_kwargs = { + n: v for n, v in bound.items() if not isinstance(v, torch.Tensor) + } + flat_sample = tuple( + bound[n].clone().detach().requires_grad_(True) for n in tensor_names + ) + + # Build shim (handles None filtering) + shim = _ForwardShim(self._module, tensor_names, frozen_kwargs) + + for param in self._module.parameters(): + param.grad = None + + # For gradient accumulate fusion + self.unshard_main_grad_buffer() + + if self._gc_freeze: + gc.collect() + gc.freeze() + + # Disable side-stream collectives during capture so every CUDA + # operation lands on the default (capture) stream. + saved_hooks = _pop_hooks_recursive(self._module) + ctx = self._module._fsdp_root_context + ctx.cuda_graph_active = True + try: + torch.cuda.synchronize() + self._graphed = torch.cuda.make_graphed_callables( + shim, + sample_args=flat_sample, + num_warmup_iters=3, + allow_unused_input=True, + # pool=self._graph_pool, + ) + finally: + ctx.cuda_graph_active = False + _restore_hooks_recursive(self._module, saved_hooks) + self.reshard_main_grad_buffer() + + self._shim = shim + self._tensor_param_names = tensor_names + self._frozen_kwargs = frozen_kwargs + self._captured = True + + # ------------------------------------------------------------------ + # 2. Install / uninstall the patched forward + # ------------------------------------------------------------------ + def install(self) -> None: + if not self._captured: + raise RuntimeError("Call capture_forward() first") + if self._orig_fwd is not None: + return + + self._orig_fwd = self._module.forward + graphed = self._graphed + shim = self._shim + param_names = _get_forward_param_names(self._module.__class__) + tensor_names = self._tensor_param_names + + def _patched_fwd(*args, **kwargs): + if self._use_cuda_graph: + bound = {} + for i, val in enumerate(args): + if i < len(param_names): + bound[param_names[i]] = val + bound.update(kwargs) + flat = tuple(bound[n] for n in tensor_names) + result = graphed(*flat) + # Restore None positions that were filtered during capture + return shim.restore_none_positions(result) + return self._orig_fwd(*args, **kwargs) + + self._module.forward = _patched_fwd + self._use_cuda_graph = True + + def uninstall(self) -> None: + """Restore the original ``forward``.""" + if self._orig_fwd is None: + return + self._module.forward = self._orig_fwd + self._orig_fwd = None + self._use_cuda_graph = False + + # ------------------------------------------------------------------ + # 3. Properties + # ------------------------------------------------------------------ + + @property + def captured(self) -> bool: + """True if ``capture_forward`` has been called successfully.""" + return self._captured + + @property + def using_cuda_graph(self) -> bool: + """True if the patch is currently active (install() called).""" + return self._use_cuda_graph + + def reset(self) -> None: + """Uninstall the patch and allow a fresh capture later.""" + self.uninstall() + self._graphed = None + self._shim = None + self._captured = False + self._graph = None + + def unshard_main_grad_buffer(self): + """Unshard the main grad buffer for all param groups.""" + for group in self._module._fsdp_param_groups: + if hasattr(group, "main_grad_buffer"): + group.main_grad_buffer.fetch_buffer() + + def reshard_main_grad_buffer(self): + """Reshard the main grad buffer for all param groups.""" + for group in self._module._fsdp_param_groups: + group.release_grad_buffer() diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/design.md b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/design.md index b4c5f9b0345..ce5cc7bcc7e 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/design.md +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/design.md @@ -792,34 +792,42 @@ that eliminates allocation overhead and fragmentation. | Phase | Behaviour | |---|---| -| **Trace** (``plan()`` not yet called) | Records every ``allocate`` / ``free`` call as a ``(seq, op, key)`` event. Also stores ``(size, dtype, device)`` metadata per key. Buckets are allocated via ``torch.empty`` as usual. Duplicate allocs (without an intervening free) do not generate new trace events. | -| **Plan** (``plan()``) | Replays the trace to build intervals ``(alloc_seq, free_seq, size)`` for each matched alloc/free pair, groups them by ``(dtype, device)``, and runs a greedy left-edge interval-coloring algorithm per group. Each color is a **slot** in a contiguous flat pool tensor. Because the same key may appear in multiple intervals, ``_slot_map[key]`` is a **list** of slot indices in alloc order. | -| **Optimized** (after ``plan()``) | ``allocate`` returns a ``Bucket`` with a slice-view into the pool, advancing a per-key cursor through the slot list. ``free`` marks the most recently allocated slot as unused (idempotent). ``reset_cursor()`` rewinds all cursors between micro-batches so the same sequence replays. | - -**Slot lists and cursors.** A single allocation key can appear in multiple -intervals — e.g., forward unshard → free → backward unshard → free. The plan -may assign these intervals to *different* slots (if they overlap) or *reuse* -the same slot (if they don't). ``_slot_map`` therefore maps each key to a -**list** of slot indices in the exact alloc order. During optimized-phase -runtime a per-key **cursor** tracks which list entry to consume next; between -micro-batches ``reset_cursor()`` rewinds all cursors to 0. - -**Greedy left-edge coloring.** For each ``(dtype, device)`` group, intervals are -sorted by ``alloc_seq``. For each interval the algorithm tries to reuse a slot -whose previous occupant has already freed (``slot_free_seq < alloc_seq``). If no -slot is free a new one is allocated. The slot is sized to the maximum bucket -assigned to it. After coloring, slots are laid out contiguously and a single -``torch.empty`` is issued per group. +| **Trace** (``plan()`` not yet called) | Records alloc/free pairs via an ``_active_keys`` set: the first ``allocate`` per key records a trace event and marks the key active; duplicate allocs (key still active) are no-ops. The first ``free`` per key records a trace event and marks it inactive; double-frees and free-before-alloc are ignored. A key that is freed then re-allocated generates a new pair of events. Buckets are created with ``torch.empty`` and **never deleted** — on re-alloc the same tensor object is resurrected via ``_alloc_storage``, keeping outstanding views (NVFP4 ``_rowwise_data`` references) live. | +| **Plan** (``plan()``) | Replays the trace to extract per-key live intervals ``(alloc_seq, free_seq)``, groups keys by ``(dtype, device)``, and runs a **conflict-graph coloring** algorithm per group. Two keys are connected if any of their live intervals overlap; the graph is then colored greedily (largest-size-first, best-fit bin packing). Each color is a **slot** backed by its own ``torch.empty()`` tensor (per-slot allocation). Each key maps to exactly one fixed slot via ``_key_to_slot``, with views pre-computed in ``_key_to_view`` for O(1) access. | +| **Optimized** (after ``plan()``) | ``allocate`` returns a ``Bucket`` with a pre-computed view of the per-slot tensor. ``free`` marks the slot as unused. Both are O(1) dict lookups — no allocations, storage resizes, slot-lists, or cursor management. Addresses are fixed per-key across all micro-batches. | + +**Conflict-graph coloring vs. left-edge.** The previous design used greedy +left-edge interval coloring on a monolithic pool tensor, requiring per-key +slot *lists* and *cursors* because the same key could appear in multiple +intervals (pre-merged into a "super-interval" to guarantee one address per +key). The current design builds an explicit interval-overlap graph and +colors it with greedy graph coloring. Each key maps to exactly one slot +(``_key_to_slot``) backed by a separate tensor (per-slot allocation). This: + +* Eliminates the need for slot lists, cursors, and ``reset_cursor()``. +* Produces the **theoretical minimum** slot count (optimal coloring of the + interval-overlap graph). +* Reduces CUDA caching-allocator fragmentation by replacing one giant + contiguous block with many smaller, independently-placed tensors. +* Still guarantees fixed per-key addresses (CUDA graph safe). **Properties.** -- **Optimal slot count:** left-edge produces the minimum number of slots for - interval graphs — it is impossible to use fewer without causing a conflict. -- **Repeatable trace required:** the same allocate/free call sequence must - repeat across micro-batches. Call ``reset_cursor()`` between micro-batches; - call ``reset()`` to re-profile if the pattern changes. -- **Double-free safe:** ``free`` is idempotent (silently returns if the slot - is already free), matching ``TemporaryBucketAllocator``'s behavior. +- **Optimal slot count:** Conflict-graph coloring yields the theoretical + minimum (peak simultaneous live memory). +- **Repeatable trace required:** The FSDP framework must execute the same + alloc/free calls in the same order per micro-batch. +- **Fully idempotent:** ``allocate`` and ``free`` are always safe to call + multiple times. `allocate` → `allocate` is a no-op (returns existing). + `free` → `free` is a no-op (ignored). `free` without a prior `allocate` + is a no-op. Both trace and optimized phases share this guarantee. +- **Stable tensor objects:** Buckets are never deleted — the same Python + tensor object is reused across alloc/free cycles, preventing dangling + views (e.g., NVFP4 parameter ``_rowwise_data``). +- **Per-slot tensors:** Each slot is a separate ``torch.empty`` allocation, + not a slice of a monolithic pool. Slots are independently placed by the + CUDA caching allocator, avoiding fragmentation from giant contiguous + blocks. **API.** @@ -827,26 +835,25 @@ assigned to it. After coloring, slots are laid out contiguously and a single allocator = TracePoolAllocator() # … run one iteration (trace phase) … pool_elems = allocator.plan() # returns total element count -# … subsequent micro-batches use the pool … -allocator.reset_cursor() # between micro-batches -print(allocator.total_pool_bytes) # bytes across all groups +# … subsequent micro-batches use the pre-allocated slots … +print(allocator.total_pool_bytes) # bytes across all slots allocator.reset() # back to trace phase ``` **Lifecycle diagram for one allocation key across two micro-batches.** ``` -Trace phase Optimized phase ------------ --------------- -allocate(key) → torch.empty --. allocate(key) → pool slot 0 (cursor 0→1) -free(key) → _free_storage | plan free(key) → slot free -allocate(key) → torch.empty --' allocate(key) → pool slot 1 (cursor 1→2) -free(key) → _free_storage free(key) → slot free - -- reset_cursor() -- - allocate(key) → pool slot 0 (cursor 0→1) - free(key) → slot free - ... +Trace phase Optimized phase +----------- --------------- +allocate(key) → torch.empty --. allocate(key) → slot_tensor view (fixed addr) +free(key) → _free_storage | plan(A) free(key) → slot free +allocate(key) → torch.empty --' allocate(key) → slot_tensor view (same addr) +free(key) → _free_storage free(key) → slot free + -- next micro-batch -- + allocate(key) → slot_tensor view (same addr) + free(key) → slot free + ... ``` -No ``torch.empty`` or storage resizing occurs in the optimized phase — the pool -owns all memory, and buckets are lightweight views. +No ``torch.empty`` or storage resizing occurs in the optimized phase — each +slot owns its own tensor, and buckets are lightweight views into them. diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/dp_buffer.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/dp_buffer.py index 9771267566f..f210dfae063 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/dp_buffer.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/dp_buffer.py @@ -279,15 +279,15 @@ def compact(self, factor: float, compact_shapes: List[torch.Size]) -> None: # falls within the current rank's shard. # _get_item_local_range → (start, end) within self.data (the local # GPU buffer). Where to read/write bytes. - # _get_item_global_range → (global_data_index, size) in the full - # logical (unsharded) buffer, same on all + # _get_item_global_range → (start, end) in the full logical + # (unsharded) buffer, same on all # ranks. # ------------------------------------------------------------------ # def _get_item_global_range(self, item_id: int) -> Tuple[int, int]: - """Return (global_data_index, size) for the given item.""" + """Return (start, end) in the full unsharded buffer for the given item.""" idx = self.item_index_map[item_id] - return (idx.global_data_index, idx.size) + return (idx.global_data_index, idx.global_data_index + idx.size) def _get_item_self_range(self, item_id: int, *, as_shard: bool = True) -> Tuple[int, int]: """Return coordinates relative to the item's own start. @@ -410,8 +410,6 @@ def __init__( self.data: Optional[torch.Tensor] = None self._unsharded_buffer: Optional[torch.Tensor] = None - # Set when a replicated buffer only has this rank's updated shard. - self._dirty = False # ------------------------------------------------------------------ # # Public API @@ -538,14 +536,15 @@ def get_item(self, item_id: int, *, as_shard: bool = False) -> torch.Tensor: def is_unsharded(self) -> bool: """Return whether this buffer currently has a full unsharded view.""" - return not self._dirty and ( - not self.is_distributed or self._unsharded_buffer is not None - ) + full_tensor = self._unsharded_buffer if self.is_distributed else self.data + if full_tensor is not None and not getattr(full_tensor, "_dirty", True): + return True + return False @torch.no_grad() def unshard( self, - bind_params: bool = True, + bind_params: bool = False, ) -> torch.Tensor: """All-gather the full buffer from all shards and bind parameter storage. @@ -553,45 +552,25 @@ def unshard( self.data is returned directly. If a replicated buffer only has this rank's updated shard, the shard is all-gathered into self.data first. """ - assert self.data is not None + full_buffer = self.fetch_buffer(as_shard=False) - shard_buffer = None - if self._dirty: - assert not self.is_distributed, "dirty unshard requires a replicated buffer" - full_buffer = self.data - sm = self.buffer_index.shard_meta - shard_buffer = self.data[sm.local_data_index : sm.local_data_index + sm.size] - elif self.is_distributed: - if self._unsharded_buffer is None: - bucket = self.allocator.allocate( - key=self.alloc_key, - size=self.buffer_index.bucket_meta.size, - dtype=self.dtype, - device=self.device, - ) - self._unsharded_buffer = bucket.data - sm = self.buffer_index.shard_meta - shard_buffer = self.data[sm.local_data_index : sm.local_data_index + sm.size] - full_buffer = self._unsharded_buffer - else: - full_buffer = self.data - - if shard_buffer is not None: - torch.distributed.all_gather_into_tensor( - output_tensor=full_buffer, - input_tensor=shard_buffer, - group=self.dp_group, - ) - if self.is_distributed and full_buffer.is_cuda: - # Temporary all-gather buckets may be released from another stream before - # the collective finishes; record the producer stream for allocator safety. - full_buffer.record_stream(torch.cuda.current_stream()) - if self._dirty: - self._dirty = False + sm = self.buffer_index.shard_meta + shard_buffer = self.data[sm.local_data_index : sm.local_data_index + sm.size] + torch.distributed.all_gather_into_tensor( + output_tensor=full_buffer, + input_tensor=shard_buffer, + group=self.dp_group, + ) + if full_buffer.is_cuda: + # Temporary all-gather buckets may be released from another stream before + # the collective finishes; record the producer stream for allocator safety. + full_buffer.record_stream(torch.cuda.current_stream()) if bind_params: self._bind_buffer_to_params(full_buffer) + setattr(full_buffer, "_dirty", False) # mark the buffer as clean (unsharded) + return full_buffer def _bind_buffer_to_params(self, buffer: torch.Tensor) -> None: @@ -602,9 +581,9 @@ def _bind_buffer_to_params(self, buffer: torch.Tensor) -> None: ) for p in self.params: item_id = self.param_idx[p] - offset, size = self.buffer_index._get_item_global_range(item_id) + start, end = self.buffer_index._get_item_global_range(item_id) idx_shape = self.buffer_index.item_index_map[item_id].shape - param_data = buffer[offset : offset + size].view(idx_shape) + param_data = buffer[start:end].view(idx_shape) self.mp_policy.bind_unsharded_param(p, param_data, self.buffer_role) @torch.no_grad() @@ -623,6 +602,9 @@ def fetch_buffer(self, *, as_shard: bool = False) -> torch.Tensor: as_shard : bool If True, return only this rank's shard slice of the full buffer. Default (False) returns the full unsharded buffer. + + Memory allocation always occurs on the default stream for deterministic + caching-allocator behaviour. """ if self.is_distributed: if self._unsharded_buffer is None: @@ -635,6 +617,7 @@ def fetch_buffer(self, *, as_shard: bool = False) -> torch.Tensor: self._unsharded_buffer = bucket.data full = self._unsharded_buffer else: + assert self.data is not None, "DataParallelBuffer data not initialized" full = self.data if as_shard: @@ -643,7 +626,11 @@ def fetch_buffer(self, *, as_shard: bool = False) -> torch.Tensor: return full @torch.no_grad() - def reduce_grad(self, grad_comm_dtype: Optional[torch.dtype] = None): + def reduce_grad( + self, + grad_comm_dtype: Optional[torch.dtype] = None, + overwrite_grad: bool = False, + ): """Reduce gradients into the optimizer-facing local shard. For distributed buffers, this reduce-scatters a temporary full gradient @@ -687,7 +674,10 @@ def reduce_grad(self, grad_comm_dtype: Optional[torch.dtype] = None): # accumulated into ``local_grad_shard`` for gradient accumulation. input_buffer = self.fetch_buffer() output_offset = sm.bucket_data_index - accumulate_output = True + if overwrite_grad: + accumulate_output = False + else: + accumulate_output = True if input_buffer.is_cuda: # Keep temporary reduce-scatter buffers tied to the stream that uses them. input_buffer.record_stream(torch.cuda.current_stream()) @@ -699,6 +689,7 @@ def reduce_grad(self, grad_comm_dtype: Optional[torch.dtype] = None): input_buffer = self.data output_offset = sm.local_data_index accumulate_output = False + overwrite_grad = False comm_input = ( input_buffer if grad_comm_dtype == self.dtype else input_buffer.to(grad_comm_dtype) @@ -711,10 +702,12 @@ def reduce_grad(self, grad_comm_dtype: Optional[torch.dtype] = None): output=reduced_grad_shard, input=comm_input, group=self.dp_group, op=op ) - if accumulate_output: + if overwrite_grad: + local_grad_shard.copy_(reduced_grad_shard) + elif accumulate_output: local_grad_shard += reduced_grad_shard elif grad_comm_dtype != self.dtype: - local_grad_shard.copy_(reduced_grad_shard.to(self.dtype)) + local_grad_shard.copy_(reduced_grad_shard) def check_all_fsdp_buffers(module) -> bool: diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/fsdp_module.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/fsdp_module.py index 884c51d5275..ee241a448c9 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/fsdp_module.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/fsdp_module.py @@ -15,16 +15,16 @@ """FSDPModule implementation for Megatron-FSDP2.""" import logging -from contextlib import nullcontext +from contextlib import contextmanager, nullcontext from dataclasses import dataclass, field -from typing import Dict, List, Optional, Tuple +from typing import Any, Dict, List, Optional, Tuple import torch import torch.nn as nn from torch.distributed.device_mesh import DeviceMesh from torch.distributed.tensor import DTensor -from .allocator import BucketAllocator +from .allocator import BucketAllocator, TracePoolAllocator from .mixed_precision import MixedPrecisionPolicy from .param_group import ParameterGroup from .utils import ParamGroupIdx, _replace_module_parameter @@ -44,6 +44,7 @@ class _FSDPState: def __init__(self): self._is_root = True self._post_backward_callback_queued = False + self.enable_cuda_graph: bool = False @dataclass @@ -89,7 +90,9 @@ class _FSDPRootContext: # ------------------------------------------------------------------ # Unshard (all-gather) tracking # ------------------------------------------------------------------ - unshard_done_events: Dict[int, Optional[torch.cuda.Event]] = field(default_factory=dict) + unshard_done_events: Dict[int, Optional[torch.cuda.Event]] = field( + default_factory=dict + ) """ Maps module_id -> CUDA event signaling completion of parameter unshard. @@ -102,9 +105,9 @@ class _FSDPRootContext: # ------------------------------------------------------------------ # Reduce-scatter (gradient sync) tracking # ------------------------------------------------------------------ - reduce_grad_buckets: Dict[ - int, List[Tuple[torch.cuda.Event, "ParameterGroup"]] - ] = field(default_factory=dict) + reduce_grad_buckets: Dict[int, List[Tuple[torch.cuda.Event, "ParameterGroup"]]] = ( + field(default_factory=dict) + ) """ Maps module_id -> list of (event, parameter_group) tuples. @@ -125,7 +128,34 @@ class _FSDPRootContext: # Activation recompute / gradient checkpointing support # ------------------------------------------------------------------ backward_phase: bool = False - """True from the root backward pre-hook until the final callback.""" + """True from the root backward pre-hook until the final callback. + ``forward_phase`` is set to ``False`` when this becomes ``True``.""" + + forward_phase: bool = False + """True from the root forward pre-hook until the root backward pre-hook. + ``backward_phase`` is set to ``False`` when this becomes ``True``.""" + + enable_cuda_graph: bool = False + """Set by enable_cuda_graph() — tells hooks to manage the side stream.""" + + cuda_graph_stream: Optional[torch.cuda.Stream] = None + """Side stream for CUDA graph capture/replay. Created lazily on the + first forward pre-hook and shared across all FSDP modules.""" + + cuda_graph_active: bool = False + """True while ``make_graphed_callables`` is inside its capture + region. All hooks must assert ``not cuda_graph_active`` — hooks + are popped during capture, so a callback firing inside the capture + window indicates a bug.""" + + cuda_graph_pool: Optional[Any] = None + f"""Shared CUDA graph memory pool handle for CUDA graph capture. + + Obtained via ``torch.cuda.graph_pool_handle()``. Multiple + ``torch.cuda.CUDAGraph`` objects created with this handle share + the same backing memory pool, allowing the CUDA driver to reuse + graph memory across FSDP modules and reduce total GPU memory + consumption.""" backward_module: Optional[int] = None """``id(module)`` of the FSDP module whose backward is pending next. @@ -152,7 +182,9 @@ def get_prefetch_next_modules( self, module: "FSDPModule", bwd_pass: bool = False ) -> List["FSDPModule"]: """Return the next FSDP module to prefetch in forward or backward order.""" - module_order = list(reversed(self.forward_order)) if bwd_pass else self.forward_order + module_order = ( + list(reversed(self.forward_order)) if bwd_pass else self.forward_order + ) for module_index, candidate_module in enumerate(module_order): if candidate_module is module: @@ -167,7 +199,7 @@ def get_root_module(self): return self.forward_order[0] if self.forward_order else None -class FSDPModule(nn.Module): +class FSDPModule: """ Mixin class for FSDP-wrapped modules. @@ -178,12 +210,43 @@ class FSDPModule(nn.Module): - reduce_grad(): Reduce gradients after backward """ + @property + def cuda_graph_compatible(self) -> bool: + """Return True when the root context is configured for CUDA graph capture. + + Requires side-stream collectives to be disabled so every CUDA + operation lands on the default stream. Can be used as a guard + before entering a graph capture region:: + + assert module.cuda_graph_compatible + """ + ctx = self._fsdp_root_context + if not isinstance(ctx.bucket_allocator, TracePoolAllocator): + return False + if ctx.bucket_allocator.phase != "optimized": + return False + return True + + def enable_cuda_graph(self) -> None: + """Enable CUDA graph capture for this FSDP module. + + Must be called while the module is resharded (before the first + forward pass). Sets the ``enable_cuda_graph`` flag on both the + per-module FSDP state and the shared root context. The side + stream is created lazily on the first root forward pre-hook. + + Usage:: + + fully_shard(layer, mesh=mesh) + layer.enable_cuda_graph() # call before first forward + """ + self._fsdp_state.enable_cuda_graph = True + def _init_named_param_groups( self, mesh: Optional[DeviceMesh], ignored_params: Optional[set], mp_policy: MixedPrecisionPolicy, - bucket_allocator: BucketAllocator, gradient_scaling_factor: Optional[float] = None, sharding_strategy: str = "optim_grads_params", ): @@ -215,7 +278,6 @@ def _init_named_param_groups( mp_policy=mp_policy, mesh=mesh, ignored_params=ignored_params, - allocator=bucket_allocator, gradient_scaling_factor=gradient_scaling_factor, sharding_strategy=sharding_strategy, ) @@ -251,9 +313,9 @@ def main_grad_getter(p): assert gbuf_data.numel() > 0 # Get offset and size from buffer index - offset, size = gbuf.buffer_index._get_item_global_range(item_id) + start, end = gbuf.buffer_index._get_item_global_range(item_id) param_shape = gbuf.buffer_index.item_index_map[item_id].shape - grad_data = gbuf_data[offset : offset + size].view(param_shape) + grad_data = gbuf_data[start:end].view(param_shape) return grad_data @@ -287,11 +349,15 @@ def _materialize_meta_module( continue m._apply( - lambda t: torch.empty_like(t, device=materialization_device) if t.is_meta else t, + lambda t: torch.empty_like(t, device=materialization_device) + if t.is_meta + else t, recurse=False, ) init_context = ( - mp_policy.model_init_context(m) if mp_policy is not None else nullcontext() + mp_policy.model_init_context(m) + if mp_policy is not None + else nullcontext() ) with init_context: if hasattr(m, "reset_parameters"): @@ -315,7 +381,11 @@ def _materialize_meta_module( torch.distributed.broadcast(param.data, src=src_rank, group=dp_group) def _init_fsdp_state( - self, enable_unshard_prefetch, enable_async_reduce_grad, bucket_allocator: BucketAllocator + self, + enable_unshard_prefetch, + enable_async_reduce_grad, + bucket_allocator: BucketAllocator, + enable_cuda_graph: bool = False, ): """Initialize FSDP state and mark nested FSDP modules as non-root. @@ -325,7 +395,9 @@ def _init_fsdp_state( check below enforces that constraint. """ named_forward_modules = [ - (name, child) for name, child in self.named_modules() if isinstance(child, FSDPModule) + (name, child) + for name, child in self.named_modules() + if isinstance(child, FSDPModule) ] forward_order = [child for name, child in named_forward_modules] @@ -357,10 +429,14 @@ def _init_fsdp_state( root_context = _FSDPRootContext( ag_stream=( - torch.cuda.Stream() if enable_unshard_prefetch else torch.cuda.current_stream() + torch.cuda.Stream() + if enable_unshard_prefetch + else torch.cuda.current_stream() ), rs_stream=( - torch.cuda.Stream() if enable_async_reduce_grad else torch.cuda.current_stream() + torch.cuda.Stream() + if enable_async_reduce_grad + else torch.cuda.current_stream() ), forward_order=forward_order, reduce_grad_buckets={id(module): [] for module in forward_order}, @@ -386,6 +462,23 @@ def _init_fsdp_state( setattr(module, "_fsdp_module_name", name) module_idx += 1 + if enable_cuda_graph: + if len(forward_order) > 1: + child_names = [ + name for name, m in named_forward_modules if m is not self + ] + raise RuntimeError( + f"enable_cuda_graph=True is not supported for FSDP modules that contain " + f"other FSDP modules as children. " + f"Module '{self._fsdp_module_name}' (type={type(self).__name__}) " + f"has FSDP children: {child_names}. " + f"Only leaf FSDP modules (no FSDP children) can use CUDA graph capture." + ) + self.enable_cuda_graph() + + if any(module._fsdp_state.enable_cuda_graph for module in forward_order): + root_context.enable_cuda_graph = True + def unshard(self, async_op: bool = False, bwd_pass: bool = False): """ Unshard parameters by all-gathering from the sharded buffer. @@ -424,9 +517,9 @@ def unshard(self, async_op: bool = False, bwd_pass: bool = False): # Optional NaN checking for debugging if getattr(module, "_enable_nan_checks", False): for name, dist_param in zip(param_names, param_group.dist_params): - assert not torch.isnan( - dist_param._local_tensor - ).any(), f"NaN detected in dist param for parameter {name}" + assert not torch.isnan(dist_param._local_tensor).any(), ( + f"NaN detected in dist param for parameter {name}" + ) with torch.cuda.stream(stream): param_group.unshard(bwd_pass=bwd_pass) @@ -451,7 +544,9 @@ def unshard(self, async_op: bool = False, bwd_pass: bool = False): # Optional NaN checking for debugging if getattr(self, "_enable_nan_checks", False): for name, param in zip(param_names, param_group.params): - assert not torch.isnan(param).any(), f"NaN detected in parameter {name}" + assert not torch.isnan(param).any(), ( + f"NaN detected in parameter {name}" + ) torch.cuda.nvtx.range_pop() @@ -508,19 +603,34 @@ def reduce_grad(self, async_op: bool = False): if getattr(self, "_enable_nan_checks", False): for name, param in zip(param_names, param_group.params): if param.grad is not None: - assert not torch.isnan( - param.grad - ).any(), f"NaN in parameter grad for {name}" + assert not torch.isnan(param.grad).any(), ( + f"NaN in parameter grad for {name}" + ) # Copy .grad -> main grad buffer on main stream (fast memcpy). # When gradient_accumulation_fusion is active for FSDP params, the backward # kernel writes directly into main_grad (weight.main_grad = get_main_grad() in # layers.py) and sets grad_added_to_main_grad=True. In that case we must NOT # zero or overwrite main_grad; discard the dummy .grad tensor if present. + # + # Under CUDA graph replay, TE's pure-GPU backward kernel still runs, but + # the Python-side ``setattr(param, "grad_added_to_main_grad", True)`` that + # accompanies the eager backward is captured away. We record the per-param + # flag during the trace micro-batch and restore it here. for name, param in zip(param_names, param_group.params): - if getattr(param, "grad_added_to_main_grad", False): + grad_added = getattr(param, "grad_added_to_main_grad", False) + recorded = getattr(param, "_mfsdp_recorded_te_wgrad", False) + if grad_added or recorded: if param.grad is not None: del param.grad + # Record TE wgrad-fusion flags for CUDA graph restore. + # The trace backward ran eagerly, so TE set + # grad_added_to_main_grad on each param it wrote to. + # Under CUDA graph replay only the GPU kernel runs; + # we record the flags here and restore them in + # the CG replay backward. + if grad_added and self._fsdp_state.enable_cuda_graph: + setattr(param, "_mfsdp_recorded_te_wgrad", True) elif param.grad is None: main_grad = param.get_main_grad() param_main_grad = getattr(param, "main_grad", None) @@ -563,7 +673,9 @@ def reduce_grad(self, async_op: bool = False): else: assert ( dist_grad is None or dist_param.dtype == dist_grad.dtype - ), f"{name} Dist param dtype {dist_param.dtype} does not match dist grad dtype {dist_grad.dtype}" + ), ( + f"{name} Dist param dtype {dist_param.dtype} does not match dist grad dtype {dist_grad.dtype}" + ) setattr(dist_param, "grad", dist_grad) if hasattr(dist_param, "decoupled_grad"): dist_param.decoupled_grad = None @@ -576,9 +688,9 @@ def reduce_grad(self, async_op: bool = False): if getattr(self, "_enable_nan_checks", False): for name, dist_grad in zip(param_names, param_group.dist_grads): if dist_grad is not None: - assert not torch.isnan( - dist_grad._local_tensor - ).any(), f"NaN in dist grad for parameter {name}" + assert not torch.isnan(dist_grad._local_tensor).any(), ( + f"NaN in dist grad for parameter {name}" + ) torch.cuda.nvtx.range_pop() @@ -631,14 +743,7 @@ def _zero_grad_buffer(self): if not isinstance(child, FSDPModule): continue for param_group in child._fsdp_param_groups: - if param_group.main_grad_buffer is not None: - param_group.main_grad_buffer.data.zero_() - param_group.release_grad_buffer() - for dist_param in param_group.dist_params: - if dist_param.grad is not None: - del dist_param.grad - if hasattr(dist_param, "decoupled_grad"): - dist_param.decoupled_grad = None + param_group.zero_grad() def _copy_main_weights_to_model_weights(self): """Copy main weight buffer to model weight buffer.""" @@ -667,7 +772,9 @@ def _compute_per_param_norms(self) -> Dict[str, Dict[str, float]]: for param_name, dist_param, dist_grad in zip( param_names, param_group.dist_params, param_group.dist_grads ): - full_name = f"{module_name}.{param_name}" if module_name else param_name + full_name = ( + f"{module_name}.{param_name}" if module_name else param_name + ) results[full_name] = {"param_norm": 0.0, "grad_norm": 0.0} if dist_param._local_tensor.numel() > 0: results[full_name]["param_norm"] = ( @@ -779,15 +886,18 @@ def _mb(num_bytes: int | float) -> str: ): dist_idx = param_group.param_idx.get(param) offset_info = "" - if param_group.model_weight_buffer is not None and dist_idx is not None: - item_index = ( - param_group.model_weight_buffer.buffer_index.item_index_map.get( - dist_idx - ) + if ( + param_group.model_weight_buffer is not None + and dist_idx is not None + ): + item_index = param_group.model_weight_buffer.buffer_index.item_index_map.get( + dist_idx ) if item_index is not None: offset_info = f" @{item_index.global_data_index:,}+{item_index.size:,}" - lines.append(f" {param_name:50s} {str(tuple(param_shape)):24s}{offset_info}") + lines.append( + f" {param_name:50s} {str(tuple(param_shape)):24s}{offset_info}" + ) group_idx += 1 lines.append( @@ -809,17 +919,21 @@ def _set_nan_check(self, enable_nan_checks: bool): param_data = param.data._local_tensor else: param_data = param.data - assert not torch.isnan(param_data).any(), f"NaN detected in parameter {name}" + assert not torch.isnan(param_data).any(), ( + f"NaN detected in parameter {name}" + ) for child in self.modules(): if not isinstance(child, FSDPModule): continue for param_group in child._fsdp_param_groups: for param in param_group.params: wbuf = param_group.model_weight_buffer - param_data = wbuf.get_item(param_group.param_idx[param], as_shard=False) - assert not torch.isnan( - param_data - ).any(), "NaN detected in model weight buffer" + param_data = wbuf.get_item( + param_group.param_idx[param], as_shard=False + ) + assert not torch.isnan(param_data).any(), ( + "NaN detected in model weight buffer" + ) def get_root_module(self): """Return the root FSDP module associated with this module.""" @@ -832,7 +946,6 @@ def _sync_module_states_after_load(self): def _get_module_fsdp_param_groups( module: nn.Module, mp_policy: MixedPrecisionPolicy, - allocator: BucketAllocator, mesh: Optional[DeviceMesh] = None, ignored_params: Optional[set[nn.Parameter]] = None, gradient_scaling_factor: Optional[float] = None, @@ -868,7 +981,6 @@ def _get_module_fsdp_param_groups( param_group_id=ParamGroupIdx(id(module), i), mp_policy=mp_policy, gradient_scaling_factor=gradient_scaling_factor, - allocator=allocator, sharding_strategy=sharding_strategy, ) ) diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/fully_shard.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/fully_shard.py index ed7b59a7905..0f63e9b6fe8 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/fully_shard.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/fully_shard.py @@ -22,6 +22,7 @@ from typing import Callable, Optional +import torch import torch.nn as nn from torch.distributed.device_mesh import DeviceMesh from torch.distributed.tensor.placement_types import Shard @@ -31,7 +32,6 @@ from .hooks import ( _register_backward_hook, _register_backward_pre_hook, - _register_fine_grained_forward_pre_hooks, _register_forward_hook, _register_forward_pre_hook, ) @@ -58,6 +58,7 @@ def fully_shard( gradient_scaling_factor: Optional[float] = None, enable_trace_pool: bool = False, sharding_strategy: str = "optim_grads_params", + enable_cuda_graph: bool = False, ) -> nn.Module: """ Wrap a module with FSDP sharding semantics. @@ -83,7 +84,15 @@ def fully_shard( new_cls = type(f"FSDP{cls.__name__}", (FSDPModule, cls), {}) module.__class__ = new_cls - use_trace_pool = enable_trace_pool and sharding_strategy in ( + use_trace_pool = ( + enable_trace_pool + or enable_cuda_graph + or any( + getattr(m._fsdp_state, "enable_cuda_graph", False) + for m in module.modules() + if isinstance(m, FSDPModule) and m is not module + ) + ) and sharding_strategy in ( "optim", "optim_grads", "optim_grads_params", @@ -94,7 +103,6 @@ def fully_shard( mesh, ignored_params, mp_policy=mp_policy, - bucket_allocator=bucket_allocator, gradient_scaling_factor=gradient_scaling_factor, sharding_strategy=sharding_strategy, ) @@ -102,13 +110,14 @@ def fully_shard( enable_unshard_prefetch=enable_unshard_prefetch, enable_async_reduce_grad=enable_async_reduce_grad, bucket_allocator=bucket_allocator, + enable_cuda_graph=enable_cuda_graph, ) module._init_param_main_grad_func() - if mp_policy.fine_grained_forward_hooks_required(module._fsdp_param_groups): - _register_fine_grained_forward_pre_hooks(module) - else: - _register_forward_pre_hook(module) + _register_forward_pre_hook( + fsdp_module=module, + fine_grained=mp_policy.fine_grained_forward_hooks_required(module._fsdp_param_groups), + ) _register_forward_hook(module) _register_backward_pre_hook(module) _register_backward_hook(module) diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/hooks.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/hooks.py index 2ecbe30ab06..98325cb267b 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/hooks.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/hooks.py @@ -26,39 +26,76 @@ from .allocator import TracePoolAllocator from .fsdp_module import FSDPModule, _FSDPState from .utils import RegisterFSDPBackwardFunction +from .cuda_graph_runner import FSDPCudaGraphRunner logger = logging.getLogger(__name__) -def _register_forward_pre_hook(fsdp_module: FSDPModule, hook_module: nn.Module | None = None): - """Register a pre-forward hook to unshard parameters for this FSDP unit.""" +def _register_forward_pre_hook(fsdp_module: FSDPModule, fine_grained: bool = False): + """Register a pre-forward hook on *hook_module*. - def unshard_param_groups(_hook_module, *unused): + Called before every ``forward()`` of the FSDP module — both in the + forward and backward (activation recomputation) passes. Handles + root-level phase bookkeeping, parameter unshard, and CUDA graph + capture (once per compatible module). + """ + + def forward_pre_hook(hook_module, args, kwargs): ctx = fsdp_module._fsdp_root_context + assert not ctx.cuda_graph_active, ( + "hooks must not fire during CUDA graph capture" + ) + + # ---- root: forward-phase setup (once per micro-batch) -------------- + if hook_module is fsdp_module and fsdp_module._fsdp_state._is_root: + if ctx.enable_cuda_graph and ctx.cuda_graph_stream is None: + ctx.cuda_graph_stream = torch.cuda.Stream() + torch.cuda.set_stream(ctx.cuda_graph_stream) + ctx.cuda_graph_pool = torch.cuda.graph_pool_handle() + ctx.forward_phase = True + ctx.backward_phase = False + + # ---- unshard parameters for this module --------------------------- if ctx.backward_phase: fsdp_module.unshard(async_op=ctx.enable_unshard_prefetch, bwd_pass=True) fsdp_module.unshard(async_op=ctx.enable_unshard_prefetch, bwd_pass=False) - # In the normal path, the hook is installed on the FSDP module itself. - # Fine-grained recompute can replay a child forward without entering the - # FSDP module's forward, so the same owner hook may need to be installed - # on a child module while still unsharding the owning FSDP module above. - module_to_hook = hook_module if hook_module is not None else fsdp_module - return module_to_hook.register_forward_pre_hook(unshard_param_groups, prepend=True) - - -def _register_fine_grained_forward_pre_hooks(module: FSDPModule): - """Register pre-forward hooks on this FSDP unit and its owned submodules.""" - - nested_fsdp_modules = set() - for submodule in module.modules(): - if submodule is not module and isinstance(submodule, FSDPModule): - nested_fsdp_modules.update(submodule.modules()) - - for submodule in module.modules(): - if submodule in nested_fsdp_modules: - continue - _register_forward_pre_hook(module, hook_module=submodule) + # ---- CUDA graph capture (once per compatible module) -------------- + if ( + hook_module is fsdp_module + and fsdp_module._fsdp_state.enable_cuda_graph + and (not hasattr(fsdp_module, "_fsdp_cg_runner")) + and not ctx.backward_phase + and fsdp_module.cuda_graph_compatible + ): + if torch.distributed.get_rank() == 0: + logger.debug( + "Capturing CUDA graph for module %s (id=%s)", + fsdp_module._fsdp_module_name, + id(fsdp_module), + ) + cg_runner = FSDPCudaGraphRunner( + fsdp_module, graph_pool=ctx.cuda_graph_pool + ) + cg_runner.capture_forward(*args, **kwargs) + cg_runner.install() + fsdp_module._fsdp_cg_runner = cg_runner + if torch.distributed.get_rank() == 0: + logger.debug( + "Captured CUDA graph for module %s (id=%s)", + fsdp_module._fsdp_module_name, + id(fsdp_module), + ) + + if fine_grained: + for submodule in fsdp_module.modules(): + submodule.register_forward_pre_hook( + forward_pre_hook, prepend=True, with_kwargs=True, + ) + else: + fsdp_module.register_forward_pre_hook( + forward_pre_hook, prepend=True, with_kwargs=True + ) def _register_forward_hook(module: FSDPModule): @@ -66,6 +103,9 @@ def _register_forward_hook(module: FSDPModule): def reshard_param_groups(module, *unused): ctx = module._fsdp_root_context + assert not ctx.cuda_graph_active, ( + "hooks must not fire during CUDA graph capture" + ) if ctx.backward_phase and id(module) == ctx.backward_module: return module.reshard() @@ -75,59 +115,118 @@ def reshard_param_groups(module, *unused): def _register_backward_pre_hook(module: FSDPModule): """ - Register backward pre-hook to handle gradient computation. + Register backward pre-hook using multi-grad hooks on output tensors. - This uses a custom backward hook that attaches to output tensors - to trigger unshard at the right time during backward pass. + Attaches a ``register_multi_grad_hook`` to every tensor output of + ``module.forward()``. When autograd reaches this module during the + backward pass, the hook fires *before* the module's own backward, + giving FSDP a chance to unshard parameters for gradient computation. """ - def create_custom_backward_hook(module: FSDPModule, custom_backward_handler: Callable): - """Create a custom backward hook attached to output tensors.""" + def create_custom_backward_hook( + module: FSDPModule, custom_backward_handler: Callable + ): + """Wrap *module* so that ``custom_backward_handler`` fires as a + pre-backward hook via ``register_multi_grad_hook``.""" @torch.compiler.disable def forward_hook(_module, inputs, output): - # View-as to avoid output being the same tensor object - output = tree_map(lambda t: t.view_as(t) if torch.is_tensor(t) else t, output) + assert not _module._fsdp_root_context.cuda_graph_active, ( + "hooks must not fire during CUDA graph capture" + ) + # ``view_as`` ensures the autograd graph sees a distinct tensor + # object even when the module returns a view of an input. + output = tree_map( + lambda t: t.view_as(t) if torch.is_tensor(t) else t, output + ) - # Collect tensor outputs output_list = [] if isinstance(output, torch.Tensor): output_list = [output] elif isinstance(output, (tuple, list)): output_list = [t for t in output if isinstance(t, torch.Tensor)] - # Register pre-backward hook on output tensors. - # This triggers when gradients are computed. torch.autograd.graph.register_multi_grad_hook( - output_list, lambda grads: custom_backward_handler(_module, grads), mode="any" + output_list, + lambda grads: custom_backward_handler(_module, grads), + mode="any", ) return output return module.register_forward_hook(forward_hook) def pre_backward_hook(module: FSDPModule, grads): - """Hook called before backward pass for this module.""" + """Pre-backward callback for a single FSDP module. + + Invoked by ``register_multi_grad_hook`` when autograd reaches + this module during the backward pass — *before* the module's + own ``backward()`` runs. Executed in reverse forward order. + + Execution flow + -------------- + 1. **Assert** not inside CUDA graph capture. + 2. **Root setup** (first module in backward): + clear backward tracking, switch phase flags, advance the + backward module cursor, and enqueue the post-backward final + callback (once). + 3. **Unshard** parameters for this module (bwd pass). + 4. **Reset per-module state**: mark ``post_backward_issued`` + as not yet handled, and reset Transformer Engine gradient + accumulation flags. + 5. **TE wgrad fusion** (``optim_grads`` / ``optim_grads_params``): + set ``overwrite_main_grad`` so TE writes weight gradients + directly into ``param.main_grad`` (overwrite, not accumulate) + to avoid double-counting across micro-batches. + 6. **Unshard main_grad_buffer** if the parameter group owns one. + """ ctx = module._fsdp_root_context + assert not ctx.cuda_graph_active, ( + "hooks must not fire during CUDA graph capture" + ) + + # ---- root: backward-phase setup ----------------------------------- if module._fsdp_state._is_root: ctx.backward_done_modules.clear() + ctx.forward_phase = False ctx.backward_phase = True ctx._advance_backward_module() + if not module._fsdp_state._post_backward_callback_queued: + _register_post_backward_final_callback(module._fsdp_state, module) + + # ---- unshard params for backward compute -------------------------- + module.unshard(async_op=ctx.enable_unshard_prefetch, bwd_pass=True) + + # ---- reset per-module bookkeeping --------------------------------- + # ``post_backward_issued`` guards against the post-backward path + # (reshard + reduce_grad) being skipped. setattr(module, "post_backward_issued", False) + + # ---- Transformer Engine gradient-accumulation fusion --------------- + # + # In the eager (trace) backward, TE sets + # ``grad_added_to_main_grad = True`` on each param whose weight + # gradient it writes directly to ``main_grad``. Under CUDA graph + # replay the GPU kernel still runs, but the Python-side + # ``setattr`` is not part of the graph. We recorded the per-param + # flag after the trace backward (see ``_post_backward_final_callback``) + # and restore it here so that ``reduce_grad`` knows TE already + # populated ``main_grad``. for param_group in module._fsdp_param_groups: for param in param_group.params: setattr(param, "grad_added_to_main_grad", False) - # When gradient_accumulation_fusion is active, the TE backward - # kernel writes weight gradients directly into param.main_grad. - # By default TE accumulates (adds) into main_grad, which silently - # doubles gradients when the buffer isn't zeroed between - # micro-batches. Setting overwrite_main_grad tells TE to - # overwrite instead of accumulate, preventing NaN from corrupted - # gradient buffers. - if param_group.sharding_strategy in ["optim_grads_params", "optim_grads"]: + if param_group.sharding_strategy in ( + "optim_grads_params", + "optim_grads", + ): + # TE's backward kernel writes weight gradients directly + # into param.main_grad. By default TE *accumulates* + # (adds) into main_grad, which silently doubles + # gradients when the buffer isn't zeroed between + # micro-batches. ``overwrite_main_grad`` tells TE to + # overwrite instead. setattr(param, "overwrite_main_grad", True) - if module._fsdp_state._is_root and not module._fsdp_state._post_backward_callback_queued: - _register_post_backward_final_callback(module._fsdp_state, module) - module.unshard(async_op=ctx.enable_unshard_prefetch, bwd_pass=True) + if param_group.main_grad_buffer is not None: + param_group.main_grad_buffer.unshard() module._mfsdp_backward_pre_hook = create_custom_backward_hook( module, custom_backward_handler=pre_backward_hook @@ -145,6 +244,9 @@ def _register_backward_hook(module: FSDPModule): def post_backward(module: FSDPModule): """Hook called after backward pass for this module.""" ctx = module._fsdp_root_context + assert not ctx.cuda_graph_active, ( + "hooks must not fire during CUDA graph capture" + ) ctx.backward_done_modules.add(id(module)) ctx._advance_backward_module() module.reshard() @@ -169,6 +271,9 @@ def _register_post_backward_hook( input tensors in an autograd Function. The Function's backward calls the post_backward_hook after gradients are computed. """ + assert not module._fsdp_root_context.cuda_graph_active, ( + "hooks must not fire during CUDA graph capture" + ) if not torch.is_grad_enabled(): return args, kwargs @@ -209,22 +314,36 @@ def _register_post_backward_hook( ) -def _register_post_backward_final_callback(state: _FSDPState, module: nn.Module) -> None: +def _register_post_backward_final_callback( + state: _FSDPState, module: nn.Module +) -> None: """ - Register the final callback that runs after all backward passes complete. - - This is only registered by the root FSDP module to avoid duplicate - callbacks. It reshards all modules and reduces gradients at the end - of the backward pass. + Enqueue a *single* engine callback that fires after every module's + backward pass has completed. + + Registered once by the root FSDP module (avoids duplicates). + The callback: + - Reshards and reduces gradients for any module whose per-module + post-backward hook was silently skipped (e.g. activation + recomputation). + - Waits for all async reduce-grad operations to finish. + - Resets root / context state for the next micro-batch. + - On the first call (trace → optimized transition): builds the + pool plan and records TE wgrad-fusion flags for CUDA graph restore. """ assert state._is_root, "Only root FSDP should register post-backward callback" if state._post_backward_callback_queued: return def _post_backward_final_callback(root_state: _FSDPState, root_module: nn.Module): - """Final callback: reshard all modules and reduce gradients.""" + """Engine callback — the last thing autograd runs after the backward + pass of every micro-batch.""" ctx = root_module._fsdp_root_context - stream = ctx.rs_stream + assert not ctx.cuda_graph_active, ( + "hooks must not fire during CUDA graph capture" + ) + + # ---- handle modules whose per-module post-backward was skipped ---- for module in reversed(ctx.forward_order): if getattr(module, "post_backward_issued", False): continue @@ -234,29 +353,37 @@ def _post_backward_final_callback(root_state: _FSDPState, root_module: nn.Module for param_group in module._fsdp_param_groups ): module.reduce_grad(async_op=ctx.enable_async_reduce_grad) + + # ---- drain pending async reduce-grad events ----------------------- + stream = ctx.rs_stream for buckets in ctx.reduce_grad_buckets.values(): while len(buckets) > 0: event, param_group = buckets.pop() event.wait() param_group.release_grad_buffer() torch.cuda.current_stream().wait_stream(stream) + + # ---- reset root / context state for the next micro-batch ---------- root_state._post_backward_callback_queued = False ctx.backward_phase = False ctx.backward_module = None ctx.backward_done_modules.clear() - # After the first backward pass, we have the full trace of bucket allocations - # and releases. We can now plan the memory pool based on this trace. + # ---- trace → optimized transition (first micro-batch only) -------- if isinstance(ctx.bucket_allocator, TracePoolAllocator): bucket_alloc = ctx.bucket_allocator if bucket_alloc.phase == "trace": if torch.distributed.get_rank() == 0: logger.debug(bucket_alloc.dump_trace()) + for m in ctx.forward_order: + logger.debug( + f"module_id={id(m)}, module_name={m._fsdp_module_name}" + ) bucket_alloc.plan() - elif bucket_alloc.phase == "optimized": - bucket_alloc.reset_cursor() - else: - raise ValueError(f"Unexpected bucket allocator phase: {bucket_alloc.phase}") + elif bucket_alloc.phase != "optimized": + raise ValueError( + f"Unexpected bucket allocator phase: {bucket_alloc.phase}" + ) state._post_backward_callback_queued = True Variable._execution_engine.queue_callback( diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/mcore_fsdp_checkpoint_design.md b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/mcore_fsdp_checkpoint_design.md index ffbeb1fd245..7b1051bc880 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/mcore_fsdp_checkpoint_design.md +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/mcore_fsdp_checkpoint_design.md @@ -25,7 +25,7 @@ the checkpoint save/load architecture and how it integrates with MCore's | Source Format | Target Format | Model | Optimizer | Notes | |---------------|--------------|-------|-----------|-------| | MFSDP v2 | MFSDP v2 | ✓ | ✓ | Full round-trip and cross-setting (`optim_grads` ↔ `optim_grads_params`) | -| MFSDP v1 baseline | MFSDP v2 | ✓ | ✓ | Both `fsdp_dtensor` format, canonical key matching | +| MFSDP v1 baseline | MFSDP v2 | ✗ | ✗ | Currently skipped in tests (``pytest.skip("v1 checkpoint format not available")``). Both use ``fsdp_dtensor`` format but key names may differ. Planned for future support. | | ND-parallel (`torch_dist`) | MFSDP v2 | ✓ | ✓ (``fully_reshardable`` only) | Online conversion via `_load_torch_dist_into_megatron_fsdp_v2` in `checkpointing.py`. Expert weights are split from flattened multi-expert tensors. Optimizer states (`exp_avg`, `exp_avg_sq`) are loaded into V2's name-based format. Hi-precision FP32 optimizer param copies are used for model weights when available. ``dp_reshardable`` format is not supported (bucket-based layout incompatible with name-based V2 format). | | ND-parallel (`torch`) | MFSDP v2 | ✗ | ✗ | Not supported (different serialization) | @@ -53,10 +53,10 @@ format for both paths. It uses DCP directly, storing each parameter as a `DTenso | Module | Location | Responsibility | |--------|----------|----------------| | `uneven_dtensor.py` | `megatron/core/distributed/fsdp/src/megatron_fsdp/` | `get_state_dict`, `preprocess_state_dict_for_uneven_dtensor`, chunk metadata for uneven DTensors | -| `fsdp_dtensor_checkpoint.py` | `megatron/core/transformer/` | SWiGLU split, GDN split, expert key remapping, FP8 cleanup | +| `fsdp_dtensor_checkpoint.py` | `megatron/core/transformer/` | (v1 only) SWiGLU split, GDN split, expert key remapping, FP8 cleanup. v2 equivalents live in ``checkpoint.py``. | `distrib_optimizer.py` | `megatron/core/optimizer/` | `state_dict()`, `load_state_dict()`, `sharded_state_dict()`, `sharded_param_state_fsdp_dtensor()` | | `checkpointing.py` | `megatron/training/` | High-level save/load orchestration, `_build_megatron_fsdp_v2_state_dict`, `preprocess_fsdp_dtensor_state_dict()` | -| `checkpoint.py` | `megatron/core/distributed/fsdp/` | `MegatronFSDPStateful`, `_apply_mcore_postprocess`, `_split_fused_params_v2` (SwiGLU+GDN), `_build_dtensor_optim_sd`, `_preprocess_and_verify_v2_state_dict`, `_build_torch_dist_to_v2_map`, `load_torch_dist_into_fsdp_v2` | +| `checkpoint.py` | `megatron/core/distributed/fsdp/` | `MegatronFSDPStateful`, `_apply_mcore_postprocess`, `_verify_chunk_metadata`, `_propagate_chunk_metadata_to_state_dict`, `_split_fused_params_v2` (SwiGLU+GDN+MambaMixer), `_build_dtensor_optim_sd`, `_preprocess_and_verify_v2_state_dict`, `_build_torch_dist_to_v2_map`, `load_torch_dist_into_fsdp_v2` | | `mcore_fsdp_adapter.py` | `megatron/core/distributed/fsdp/` | Routes to v1 `MegatronFSDP` or Megatron FSDP v2 `fully_shard` | ### 3.1.1 `checkpoint.py` Functions (Megatron FSDP v2) @@ -77,7 +77,11 @@ format for both paths. It uses DCP directly, storing each parameter as a `DTenso | `handle_experts_in_state_dict(model_sd, num_experts)` | Rename expert parameter keys for expert-parallel sharding. | | `handle_swiglu_in_state_dict_v2(model, model_sd, opt_sd)` | Thin wrapper: precomputes ``layer_glu`` map, delegates to ``_split_fused_params_v2`` with SwiGLU detector and ``_w``/``_v`` suffix format. | | `handle_gdn_in_state_dict_v2(model, model_sd, opt_sd)` | Thin wrapper: delegates to ``_split_fused_params_v2`` with ``_match_gdn_key`` detector and ``.query``/``.key``/``.value`` suffix format. | -| `_split_fused_params_v2(model, model_sd, opt_sd, detector, key_fmt, tag)` | Unified fused-parameter splitting skeleton. Iterates model and optimizer state dicts, calls the ``detector(key, dtensor, model) -> (sizes, names, dim)`` callback to identify fused tensors, splits via ``split_dtensor``, and renames keys via ``key_fmt(key, sub_name)``. Shared by SwiGLU and GDN. | +| `handle_mamba_in_state_dict_v2(model, model_sd, opt_sd)` | Thin wrapper: delegates to ``_split_fused_params_v2`` with ``_mamba_mixer_detector`` and MambaMixer suffix format. | +| `_split_fused_params_v2(...)` | Unified fused-parameter splitting skeleton. Iterates model and optimizer state dicts, calls the ``detector(key, dtensor, model) -> (sizes, names, dim)`` callback to identify fused tensors, splits via ``split_dtensor``, and renames keys via ``key_fmt(key, sub_name)``. Shared by SwiGLU, GDN, and MambaMixer. | +| `_verify_chunk_metadata(flattened_sd)` | Final verification: checks every DTensor has ``__create_chunk_list__`` AND that ``chunks_total_numel == local_numel``. On failure, logs key, shape, chunk details, source tag, and device mesh before raising ``AssertionError``. | +| `_propagate_chunk_metadata_to_state_dict(model, state_dict)` | Copy ``__create_chunk_list__`` / ``__create_write_items__`` from model parameters (which have them from init) to state dict DTensors (fresh from ``model.state_dict()``). Matches by key name. Logs unmatched DTensors at rank 0 for diagnostic tracing. | +| `get_chunk_meta_source(dtensor)` | Return the source tag (e.g. ``"init"``, ``"preprocess"``, ``"split"``, ``"propagate:init"``) stored on ``dtensor._local_tensor._chunk_meta_source``. | | `_is_swiglu_key(key)` | Check whether a key matches any SwiGLU fc1 pattern. | | `_match_gdn_key(key, dtensor)` | GDN detector: returns ``(sizes, names, dim)`` for fused QKV projections, or ``None``. | | `_detect_glu_layers(model)` | Return ``{layer_path: gated_linear_unit}`` for all TransformerLayers. | @@ -398,11 +402,15 @@ SwiGLU split, expert key remapping) can be applied as a separate layer. │ ┌──────────────────────────▼───────────────────────────────────┐ │ Layer 2: MCore Post-Processing │ -│ preprocess_fsdp_dtensor_state_dict() │ -│ → handle_fp8_extra_state_case() # strip FP8 artifacts │ -│ → handle_swiglu_in_state_dict() # split fc1 into _w/_v │ -│ → handle_experts_in_state_dict() # EP key remapping │ -│ (uneven DTensor is already handled by Layer 1) │ +│ _apply_mcore_postprocess(state_dict, args, model) │ +│ → _propagate_chunk_metadata_to_state_dict (copy metadata) │ +│ → _build_dtensor_optim_sd (wrap as DTensor)│ +│ → handle_fp8_extra_state_case (FP8 cleanup) │ +│ → handle_swiglu_in_state_dict_v2 (split fc1) │ +│ → handle_gdn_in_state_dict_v2 (split QKV) │ +│ → handle_mamba_in_state_dict_v2 (split Mamba) │ +│ → handle_experts_in_state_dict (EP key remap) │ +│ → _verify_chunk_metadata (consistency ✓) │ └──────────────────────────┬───────────────────────────────────┘ │ ┌──────────────────────────▼───────────────────────────────────┐ @@ -511,8 +519,7 @@ def load_checkpoint_fsdp_v2(model_chunks, optimizer, ...): `uneven_dtensor.get_state_dict` adds uneven DTensor chunk metadata in the same call. With the raw PyTorch version, `preprocess_state_dict_for_uneven_dtensor` must be called -separately, and the caller must ensure it is not double-called (once by the wrapper, -once by `preprocess_fsdp_dtensor_state_dict`). Using `uneven_dtensor.get_state_dict` +once by `preprocess_state_dict_for_uneven_dtensor`). Using `uneven_dtensor.get_state_dict` consolidates this into a single well-tested entry point. **Why separate the layers?** @@ -802,6 +809,37 @@ This function: This is called on **both** model and optimizer state dicts. +### 9.3 Chunk Metadata Source Tag + +Each DTensor's local tensor carries a `_chunk_meta_source` attribute that records +where its `__create_chunk_list__` metadata was set. This enables tracing the +provenance of chunk metadata when debugging size mismatches. + +| Source Tag | Set By | When | +|---|---|---| +| `"init"` | `update_uneven_dtensor_chunk_metadata` (default) | FSDP v2 construction (`param_group.py`) | +| `"preprocess"` | `preprocess_state_dict_for_uneven_dtensor` | Save Phase 1 (re-computes via all_gather) | +| `"propagate:"` | `copy_chunk_metadata` | Save Phase 2 (copied from model param; `` is the original tag) | +| `"split"` | `split_dtensor` | Save Phase 3 (locally derived from parent's metadata) | +| `"make_uneven"` | `make_uneven_dtensor(chunk_metadata=...)` | Optimizer state wrapping as DTensors | + +Use `get_chunk_meta_source(dtensor)` to read the tag (returns `"none"` if unset). + +### 9.4 Final Consistency Check (`_verify_chunk_metadata`) + +At the end of `_apply_mcore_postprocess`, `_verify_chunk_metadata` validates +every DTensor in the flattened state dict: + +1. **Existence check** — every DTensor must have `__create_chunk_list__`. +2. **Numel consistency check** — `sum(prod(chunk.sizes)) == local_tensor.numel()`. + Each chunk's sizes may be multi-dimensional (e.g., `(256, 512)`), so the + element count is the product of all dimensions. A mismatch (e.g., 40960 vs + 2560) indicates chunk metadata was computed with the wrong device mesh + (DP instead of EDP) or against stale state. +3. **On failure** — logs the key, global/local shapes, chunk list with offsets + and sizes, source tag, and device mesh before raising `AssertionError`. + This pinpoints which phase produced incorrect metadata and what mesh was used. + --- ## 10. DTensor Attribute Propagation @@ -891,8 +929,8 @@ ensures model and optimizer state dict keys are consistent in the checkpoint. - [x] Skip `preprocess_fsdp_dtensor_state_dict` for v2 in both save and load paths (post-processing already handled by ``_apply_mcore_postprocess`` inside ``MegatronFSDPStateful.state_dict()``) -- [ ] Handle PP: iterate model chunks, build per-chunk state dicts -- [ ] Handle multi-optimizer (ChainedOptimizer: expert + non-expert optimizers) +- [ ] Handle PP: iterate model chunks, build per-chunk state dicts (not yet required — all current tests use PP=1) +- [ ] Handle multi-optimizer (ChainedOptimizer: expert + non-expert optimizers) (not yet required — all current tests use single optimizer) ### Phase 2: Path A — Standalone Save/Load (`checkpoint.py`) @@ -991,7 +1029,22 @@ ensures model and optimizer state dict keys are consistent in the checkpoint. 5. **Optimizer state restored?** — Check that `optimizer.state` is non-empty and contains expected keys (`exp_avg`, `exp_avg_sq`, `step`). -6. **NaN after load?** — Common causes: +6. **Chunk metadata present?** — After save, verify `__create_chunk_list__` exists + on all DTensors and the state dict passed `_verify_chunk_metadata` without error. + Check the log for `[chunk_metadata_diag]` warnings about unmatched DTensors. + +7. **Chunk metadata consistent?** — If DCP reports `invalid fill tensor-volume`, + inspect the `[chunk_metadata_verify]` error log. It shows global/local shapes, + chunk offsets/sizes, the source tag (which phase set the metadata), and the + device mesh. Common causes: + - `source=preprocess` + wrong device mesh → `all_gather_object` computed + offsets against DP mesh when parameter is sharded across EDP mesh. + - `source=propagate:init` + key mismatch → metadata was copied from a + different parameter (or not copied at all, falling back to Phase 1 metadata). + - `source=split` + wrong split → fused param split derived wrong local + offsets from parent's chunk metadata. + +8. **NaN after load?** — Common causes: - Missing `allreduce` attribute propagation (expert params misclassified) - Missing `overwrite_main_grad=True` for wgrad fusion (gradient doubling) - Wrong `gradient_accumulation_fusion` setting diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/mixed_precision.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/mixed_precision.py index 254ceb3d2dd..782cc466a1e 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/mixed_precision.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/mixed_precision.py @@ -553,17 +553,11 @@ def copy_main_weights_to_model_weights( assert model_weight_buffer is not None, "main weights require a model-weight buffer" - dirty_replicated_buffer = ( - main_weight_buffer.is_distributed and not model_weight_buffer.is_distributed - ) - if self.is_nvfp4_param(params[0]): quantize_main_weights_to_nvfp4( params, param_idx, data_parallel_group, model_weight_buffer, main_weight_buffer ) - return - - if not self.is_fp8_param(params[0]): + elif not self.is_fp8_param(params[0]): if model_weight_buffer.is_distributed and not main_weight_buffer.is_distributed: raise RuntimeError( "Unsupported FSDP main/model weight buffer layout: " @@ -583,41 +577,41 @@ def copy_main_weights_to_model_weights( + main_shard_meta.size ] ) - if dirty_replicated_buffer: - model_weight_buffer._dirty = True - return - - fp8_params = [] - main_params = [] - start_offsets = [] - model_param_shards = [] - for param in params: - item_id = param_idx[param] - model_shard = model_weight_buffer.get_item(item_id, as_shard=True) - if model_shard.numel() == 0: + else: + fp8_params = [] + main_params = [] + start_offsets = [] + model_param_shards = [] + for param in params: + item_id = param_idx[param] + model_shard = model_weight_buffer.get_item(item_id, as_shard=True) + if model_shard.numel() == 0: + fp8_params.append(param) + main_params.append(None) + start_offsets.append(None) + model_param_shards.append((None, None)) + continue + + transpose_shard = None + if transpose_weight_buffer is not None: + transpose_shard = transpose_weight_buffer.get_item(item_id, as_shard=True) + main_weight = main_weight_buffer.get_item(item_id, as_shard=True) + start_offset, _ = model_weight_buffer.buffer_index._get_item_self_range(item_id) fp8_params.append(param) - main_params.append(None) - start_offsets.append(None) - model_param_shards.append((None, None)) - continue + main_params.append(main_weight) + start_offsets.append(start_offset) + model_param_shards.append((model_shard, transpose_shard)) - transpose_shard = None - if transpose_weight_buffer is not None: - transpose_shard = transpose_weight_buffer.get_item(item_id, as_shard=True) - main_weight = main_weight_buffer.get_item(item_id, as_shard=True) - start_offset, _ = model_weight_buffer.buffer_index._get_item_self_range(item_id) - fp8_params.append(param) - main_params.append(main_weight) - start_offsets.append(start_offset) - model_param_shards.append((model_shard, transpose_shard)) - - quantize_main_weights_to_fp8( - fp8_params, main_params, start_offsets, data_parallel_group, model_param_shards - ) - if dirty_replicated_buffer: - model_weight_buffer._dirty = True - if transpose_weight_buffer is not None: - transpose_weight_buffer._dirty = True + quantize_main_weights_to_fp8( + fp8_params, main_params, start_offsets, data_parallel_group, model_param_shards + ) + + # Mark the model weights dirty so FSDP knows to all-gather them before + # the forward or backward pass. + if not model_weight_buffer.is_distributed: + model_weight_buffer.data._dirty = True + if transpose_weight_buffer is not None and not transpose_weight_buffer.is_distributed: + transpose_weight_buffer.data._dirty = True def is_fp8_param(tensor: torch.Tensor) -> bool: diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/param_group.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/param_group.py index ce86674abea..0fffb8ddf1e 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/param_group.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/param_group.py @@ -209,7 +209,10 @@ def _init_buffers(self) -> None: # Create distributed parameter views self._init_dist_params() - def unshard(self, bwd_pass: bool = False): + # Initially, gradients are zeroed. + self.is_zero_grad = True + + def unshard(self, bwd_pass: bool = False, bind_params: bool = True): """ Unshard model weights by all-gathering from sharded buffer. @@ -220,7 +223,7 @@ def unshard(self, bwd_pass: bool = False): self.model_weight_buffer, self.transpose_weight_buffer, bwd_pass=bwd_pass ): if weight_buffer is not None: - weight_buffer.unshard(bind_params=True) + weight_buffer.unshard(bind_params=bind_params) self.mp_policy.post_unshard(self.params, bwd_pass=bwd_pass) @@ -231,10 +234,6 @@ def has_unsharded_weight_buffers(self, bwd_pass: bool = False) -> bool: ): if weight_buffer is None: continue - if not weight_buffer.is_distributed: - # Replicated buffers still need DataParallelBuffer.unshard() to - # rebind original parameter storage before the module uses it. - return False if not weight_buffer.is_unsharded(): return False return True @@ -266,7 +265,16 @@ def reduce_grad(self): ZeRO-1 keeps grads replicated during backward and reduce-scatters the replicated buffer once when the optimizer syncs. """ - self.main_grad_buffer.reduce_grad(grad_comm_dtype=self.mp_policy.grad_comm_dtype) + # FIXME: When optimizer.zero_grad(set_to_none=True) is used, dist_param.grad + # becomes None, but the underlying grad buffer may still contain stale data + # from previous iterations. If overwrite_grad is not set to True, reduce_grad + # will accumulate into this stale buffer, potentially introducing NaNs and + # destabilizing training. + self.main_grad_buffer.reduce_grad( + grad_comm_dtype=self.mp_policy.grad_comm_dtype, + overwrite_grad=False, + ) + self.is_zero_grad = False def release_grad_buffer(self): """Release the main gradient buffer to free memory.""" @@ -347,3 +355,16 @@ def _init_dist_params(self): ) else: self.dist_grads.append(None) + + def zero_grad(self, set_to_none: bool = True): + """Zero the main gradient buffer and mark grads as zeroed.""" + self.release_grad_buffer() + if self.main_grad_buffer is not None: + self.main_grad_buffer.data.zero_() + if set_to_none: + for dist_param in self.dist_params: + if dist_param.grad is not None: + del dist_param.grad + if hasattr(dist_param, "decoupled_grad"): + dist_param.decoupled_grad = None + self.is_zero_grad = True diff --git a/megatron/core/transformer/moe/router.py b/megatron/core/transformer/moe/router.py index 0b139e70443..4dacd1e07b0 100644 --- a/megatron/core/transformer/moe/router.py +++ b/megatron/core/transformer/moe/router.py @@ -26,7 +26,7 @@ from megatron.core.transformer.transformer_config import TransformerConfig -class Router(ABC, MegatronModule): +class Router(MegatronModule): """Base Router class""" def __init__( diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index 0b39188c930..89b89f49a77 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -1214,6 +1214,11 @@ def validate_args(args, defaults={}): if getattr(args, 'use_megatron_fsdp_v2', False): args.use_megatron_fsdp = True + if getattr(args, 'mfsdp_cuda_graph_modules', None): + assert getattr(args, 'use_megatron_fsdp_v2', False), ( + "--mfsdp-cuda-graph requires --use-megatron-fsdp-v2" + ) + if args.inference_dynamic_batching_sampling_backend == 'flashinfer': try: import flashinfer # noqa: F401 @@ -3907,6 +3912,18 @@ def _add_distributed_args(parser): help='Use PyTorch fully shard API for FSDP implementation. ' 'This option is only effective when use-megatron-fsdp is set. ', ) + group.add_argument( + '--mfsdp-cuda-graph', + nargs='+', + default=[], + choices=['mamba', 'transformer', 'moe_router', 'attn'], + dest='mfsdp_cuda_graph_modules', + help='Enable CUDA graph capture on specific FSDP module types ' + 'when using Megatron FSDP v2 (--use-megatron-fsdp-v2). ' + 'Choices: mamba (MambaLayer), transformer (TransformerLayer). ' + 'Can be combined, e.g. --mfsdp-cuda-graph mamba transformer. ' + 'Only non-nested (leaf) FSDP modules are eligible.', + ) group.add_argument( '--create-all-gather-group', action='store_true', diff --git a/megatron/training/checkpointing.py b/megatron/training/checkpointing.py index 1324782e17e..be0aedc5bdc 100644 --- a/megatron/training/checkpointing.py +++ b/megatron/training/checkpointing.py @@ -1454,7 +1454,7 @@ def _load_non_persistent_base_checkpoint( def _load_global_dist_base_checkpoint( - load_dir, args, rank0, sharded_state_dict, iteration, release, checkpointing_context=None + load_dir, args, rank0, sharded_state_dict, iteration, release, checkpointing_context=None, model=None ): """Load the base state_dict from the given directory containing the global distributed checkpoint""" if rank0: @@ -1462,10 +1462,15 @@ def _load_global_dist_base_checkpoint( state_dict = dist_checkpointing.load_common_state_dict(checkpoint_name) return state_dict, checkpoint_name, release, CheckpointType.GLOBAL + # Load torch_dist checkpoint into Megatron FSDP v2 if args.use_megatron_fsdp_v2: checkpoint_name = find_checkpoint_rank_0(load_dir, iteration, release) state_dict = _load_torch_dist_into_megatron_fsdp_v2( - args, checkpoint_name, sharded_state_dict, strict=True + args, + checkpoint_name, + model=model, + v2_state_dict=sharded_state_dict, + strict=True, ) return state_dict, checkpoint_name, release, CheckpointType.FSDP_DTENSOR @@ -1534,7 +1539,12 @@ def _get_checkpoint_format(checkpoint_name, args): def _load_base_checkpoint( - load_dir, args, rank0=False, sharded_state_dict=None, checkpointing_context=None + load_dir, + args, + rank0=False, + sharded_state_dict=None, + checkpointing_context=None, + model=None, ): """Load the base state_dict from the given directory @@ -1617,6 +1627,7 @@ def _load_base_checkpoint( iteration, release, checkpointing_context=checkpointing_context, + model=model, ) elif ckpt_format == "torch": ckpt_type = CheckpointType.LEGACY @@ -2169,7 +2180,12 @@ def load_checkpoint( load_kwargs["sharded_state_dict"] = state_dict state_dict, checkpoint_name, release, ckpt_type = _load_base_checkpoint( - load_dir, args, rank0=False, checkpointing_context=checkpointing_context, **load_kwargs + load_dir, + args, + rank0=False, + checkpointing_context=checkpointing_context, + model=model, + **load_kwargs ) # Checkpoint not loaded. diff --git a/megatron/training/initialize.py b/megatron/training/initialize.py index 5707b4732ae..2ae3c8b2466 100644 --- a/megatron/training/initialize.py +++ b/megatron/training/initialize.py @@ -102,7 +102,7 @@ def finish_mpu_init(): args.data_parallel_random_init, args.te_rng_tracker, args.inference_rng_tracker, - use_cudagraphable_rng=args.cuda_graph_impl != "none", + use_cudagraphable_rng=args.cuda_graph_impl != "none" or args.mfsdp_cuda_graph_modules, ) # Setup MoE aux loss scale value. diff --git a/tests/unit_tests/distributed/megatron_fsdp/v2/test_dist_muon.py b/tests/unit_tests/distributed/megatron_fsdp/v2/test_dist_muon.py index a18068b247f..ec89ecfb4e7 100644 --- a/tests/unit_tests/distributed/megatron_fsdp/v2/test_dist_muon.py +++ b/tests/unit_tests/distributed/megatron_fsdp/v2/test_dist_muon.py @@ -79,15 +79,16 @@ def dist_env(): device = torch.device(f"cuda:{rank % torch.cuda.device_count()}") torch.cuda.set_device(device) yield - torch.distributed.destroy_process_group() + if torch.distributed.is_initialized(): + torch.distributed.destroy_process_group() def _read_full_weight(pg, idx, shape): """All-gather this param's full master weight and return it (reshaped, cloned).""" mbuf = pg.main_weight_buffer full = mbuf.unshard(bind_params=False) - off, sz = mbuf.buffer_index._get_item_global_range(idx) - out = full[off : off + sz].clone().view(shape) + start, end = mbuf.buffer_index._get_item_global_range(idx) + out = full[start:end].clone().view(shape) mbuf.reshard() return out diff --git a/tests/unit_tests/distributed/megatron_fsdp/v2/test_param_group.py b/tests/unit_tests/distributed/megatron_fsdp/v2/test_param_group.py index dfd0372413b..e4ce8f47fe4 100644 --- a/tests/unit_tests/distributed/megatron_fsdp/v2/test_param_group.py +++ b/tests/unit_tests/distributed/megatron_fsdp/v2/test_param_group.py @@ -50,7 +50,8 @@ def dist_env(): device = torch.device(f"cuda:{rank % torch.cuda.device_count()}") torch.cuda.set_device(device) yield - torch.distributed.destroy_process_group() + if torch.distributed.is_initialized(): + torch.distributed.destroy_process_group() # ------------------------------------------------------------------ # @@ -237,8 +238,8 @@ def test_unshard_reshard(strategy): # Distributed: after all-gather, every param should be fully # recoverable from the unsharded buffer at its global offset for i, p in enumerate(orig): - off, sz = wbuf.buffer_index._get_item_global_range(i) - recovered = unsharded[off : off + sz] + start, end = wbuf.buffer_index._get_item_global_range(i) + recovered = unsharded[start:end] assert torch.equal(recovered, p.flatten()) # Reshard: release temporary buffer, persistent shard must be intact From 2b7efd58cb7d4fa7c3a98b49ed7c55be50c7c218 Mon Sep 17 00:00:00 2001 From: tongliu Date: Tue, 9 Jun 2026 03:02:41 -0700 Subject: [PATCH 18/28] Support no-shard in FSDP v2 --- .../fsdp/src/megatron_fsdp/v2/README.md | 19 +++++------ .../fsdp/src/megatron_fsdp/v2/__init__.py | 5 ++- .../fsdp/src/megatron_fsdp/v2/design.md | 9 +++++ .../fsdp/src/megatron_fsdp/v2/fsdp_module.py | 16 ++++++--- .../src/megatron_fsdp/v2/mixed_precision.py | 14 +++++--- .../fsdp/src/megatron_fsdp/v2/param_group.py | 33 ++++++++----------- .../megatron_fsdp/v2/test_fully_shard.py | 25 ++++++++++++++ .../megatron_fsdp/v2/test_param_group.py | 30 ++++++++--------- 8 files changed, 95 insertions(+), 56 deletions(-) diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/README.md b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/README.md index 3fb04b2f2f4..403c0df51ab 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/README.md +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/README.md @@ -85,7 +85,7 @@ Mixin class added to wrapped modules. Methods: |--------|-------------|---------| | `unshard()` | Pre-forward | All-gather params from sharded buffer | | `reshard()` | Post-forward, post-backward | Release unsharded buffer | -| `reduce_grad()` | Post-backward / grad sync | Reduce-scatter gradients into optimizer-facing shards | +| `reduce_grad()` | Post-backward / grad sync | All-reduce no-shard grads or reduce-scatter ZeRO grads | ### DataParallelBuffer @@ -93,14 +93,14 @@ Flat buffer managing (a shard of) parameter/gradient data: - `unshard()` — all-gather to full tensor - `reshard()` — free temporary buffer -- `reduce_grad()` — reduce-scatter gradients into optimizer-facing shards +- `reduce_grad()` — all-reduce no-shard grads or reduce-scatter ZeRO grads - Uses `BufferIndex` to track parameter layout within the buffer ### ParameterGroup Groups parameters sharing the same (device, dtype, requires_grad): -- `model_weight_buffer` — stores compute weights; replicated for ZeRO-1/2 and sharded for ZeRO-3 +- `model_weight_buffer` — stores compute weights; replicated for no-shard/ZeRO-1/2 and sharded for ZeRO-3 - `main_weight_buffer` — optional high-precision optimizer copy; sharded when optimizer state is sharded - `main_grad_buffer` — accumulates gradients before reduce - `dist_params` — DTensor views into the buffer @@ -120,9 +120,9 @@ See the parent directory `..` for `uneven_dtensor.py` which provides: | Strategy | Shard Weights | Shard Gradients | Status | Notes | |----------|---------------|-----------------|--------|-------| | `optim_grads_params` | Yes | Yes | **Supported** | Like ZeRO-3: full parameter/gradient/optimizer sharding | -| `no_shard` | No | No | **Not yet supported** | Like DDP: no sharding | -| `optim` | No | No | **Not yet supported** | Like ZeRO-1: shard optimizer states only | -| `optim_grads` | No | Yes | **Not yet supported** | Like ZeRO-2: shard optimizer states + gradients | +| `no_shard` | No | No | **Supported** | Like DDP: replicated weights and full-gradient all-reduce | +| `optim` | No | No | **Supported** | Like ZeRO-1: shard optimizer states only | +| `optim_grads` | No | Yes | **Supported** | Like ZeRO-2: shard optimizer states + gradients | ## Known Limitations @@ -137,9 +137,8 @@ See the parent directory `..` for `uneven_dtensor.py` which provides: ### Sharding Strategies -Only `optim_grads_params` (ZeRO-3 equivalent) is implemented. `no_shard` -(DDP-like), `optim` (ZeRO-1), and `optim_grads` (ZeRO-2) are planned but -not yet available. +`no_shard` (DDP-like), `optim` (ZeRO-1), `optim_grads` (ZeRO-2), and +`optim_grads_params` (ZeRO-3 equivalent) are implemented. ### `fully_shard()` API Parameters @@ -255,4 +254,4 @@ torchrun --nproc_per_node=2 -m pytest \ # Single-GPU tests pytest tests/unit_tests/distributed/megatron_fsdp/v2/ -v \ -k "test_double_shard_rejected or test_no_params_module or test_get_state_dict_strict" -``` \ No newline at end of file +``` diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/__init__.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/__init__.py index 0e78fc0e52c..222409d204a 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/__init__.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/__init__.py @@ -12,9 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. -# FSDP v2 currently supports ZeRO-1/2/3 strategies (`optim`, -# `optim_grads`, and `optim_grads_params`). `no_shard` is not implemented in -# this path yet. +# FSDP v2 currently supports `no_shard` plus ZeRO-1/2/3 strategies (`optim`, +# `optim_grads`, and `optim_grads_params`). from ..uneven_dtensor import ( gather_and_compute_chunk_metadata, diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/design.md b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/design.md index b4c5f9b0345..cfa1cee7748 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/design.md +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/design.md @@ -368,6 +368,15 @@ def _post_backward_final_callback(root_state, root_module): ## ZeRO-1 and ZeRO-2 Workflow +### No-Shard (`no_shard`) + +1. Forward and backward read replicated `model_weight_buffer`; no parameter + all-gather is needed beyond the normal buffer rebind. +2. Backward accumulates local full gradients across micro-batches. Post-backward + reduction skips `no_shard`, matching ZeRO-1's delayed-sync behavior. +3. `finish_grad_sync()` performs one delayed full-buffer all-reduce for each + `no_shard` grad buffer. The optimizer consumes replicated DTensor grads. + `optim` and `optim_grads` keep compute weights replicated but still expose optimizer-facing DTensor shards through `dist_params`. diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/fsdp_module.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/fsdp_module.py index 884c51d5275..6d2efb05977 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/fsdp_module.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/fsdp_module.py @@ -590,11 +590,12 @@ def finish_grad_sync(self, force_all_reduce: Optional[bool] = False): if not isinstance(child, FSDPModule): continue if any( - param_group.sharding_strategy == "optim" + param_group.sharding_strategy in ("no_shard", "optim") for param_group in child._fsdp_param_groups ): - # ZeRO-1 keeps gradients replicated during backward and performs - # exactly one reduce-scatter at the iteration grad-sync boundary. + # no_shard and ZeRO-1 keep gradients replicated during backward. + # Sync them once at the iteration grad-sync boundary: no_shard + # all-reduces full grads, ZeRO-1 reduce-scatters virtual shards. child.reduce_grad(async_op=False) for param_group in child._fsdp_param_groups: for param, dist_grad in zip(param_group.params, param_group.dist_grads): @@ -653,9 +654,11 @@ def _compute_per_param_norms(self) -> Dict[str, Dict[str, float]]: Compute per-parameter L2 norms for params and grads. Returns {param_name: {"param_norm": float, "grad_norm": float}}. - Local squared norms are all-reduced across the DP group. + Sharded DTensor squared norms are all-reduced across the DP group; + replicated no-shard norms are computed from the local full tensor. """ results = {} + reduce_norms = {} dp_group = None for module_name, child in self.named_modules(): @@ -669,6 +672,7 @@ def _compute_per_param_norms(self) -> Dict[str, Dict[str, float]]: ): full_name = f"{module_name}.{param_name}" if module_name else param_name results[full_name] = {"param_norm": 0.0, "grad_norm": 0.0} + reduce_norms[full_name] = param_group.sharding_strategy != "no_shard" if dist_param._local_tensor.numel() > 0: results[full_name]["param_norm"] = ( dist_param._local_tensor.float().norm(p=2).item() ** 2 @@ -680,6 +684,10 @@ def _compute_per_param_norms(self) -> Dict[str, Dict[str, float]]: if dp_group is not None: for param_name in results: + if not reduce_norms[param_name]: + for key in ("param_norm", "grad_norm"): + results[param_name][key] = results[param_name][key] ** 0.5 + continue for key in ("param_norm", "grad_norm"): value = torch.tensor([results[param_name][key]], device="cuda") torch.distributed.all_reduce(value, group=dp_group) diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/mixed_precision.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/mixed_precision.py index 254ceb3d2dd..cbe1195cb41 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/mixed_precision.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/mixed_precision.py @@ -591,9 +591,10 @@ def copy_main_weights_to_model_weights( main_params = [] start_offsets = [] model_param_shards = [] + no_shard = model_weight_buffer.sharding_strategy == "no_shard" for param in params: item_id = param_idx[param] - model_shard = model_weight_buffer.get_item(item_id, as_shard=True) + model_shard = model_weight_buffer.get_item(item_id, as_shard=not no_shard) if model_shard.numel() == 0: fp8_params.append(param) main_params.append(None) @@ -603,9 +604,14 @@ def copy_main_weights_to_model_weights( transpose_shard = None if transpose_weight_buffer is not None: - transpose_shard = transpose_weight_buffer.get_item(item_id, as_shard=True) - main_weight = main_weight_buffer.get_item(item_id, as_shard=True) - start_offset, _ = model_weight_buffer.buffer_index._get_item_self_range(item_id) + transpose_shard = transpose_weight_buffer.get_item( + item_id, as_shard=not no_shard + ) + main_weight = main_weight_buffer.get_item(item_id, as_shard=not no_shard) + if no_shard: + start_offset = 0 + else: + start_offset, _ = model_weight_buffer.buffer_index._get_item_self_range(item_id) fp8_params.append(param) main_params.append(main_weight) start_offsets.append(start_offset) diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/param_group.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/param_group.py index ce86674abea..5e0facd4b78 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/param_group.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/param_group.py @@ -85,11 +85,7 @@ def __init__( self._dp_rank = torch.distributed.get_rank(self.dp_group) self._dp_world_size = torch.distributed.get_world_size(self.dp_group) - if sharding_strategy == "no_shard": - raise NotImplementedError( - "Sharding strategy 'no_shard' is not yet supported in FSDP v2." - ) - if sharding_strategy not in ("optim", "optim_grads", "optim_grads_params"): + if sharding_strategy not in ("no_shard", "optim", "optim_grads", "optim_grads_params"): raise ValueError(f"Unsupported sharding strategy: {sharding_strategy}") self.sharding_strategy = sharding_strategy self.param_group_id = param_group_id @@ -151,7 +147,7 @@ def _init_buffers(self) -> None: Initialize all buffers based on sharding strategy. Buffer creation logic: - - model_weight_buffer: always created unless "no_shard" + - model_weight_buffer: always created; replicated unless "optim_grads_params" - main_weight_buffer: created if mp_policy.main_params_dtype is specified - main_grad_buffer: created if requires_grad """ @@ -162,20 +158,19 @@ def _init_buffers(self) -> None: # Create model weight buffers. The policy owns dtype-sensitive storage # choices and exposes the tensor view that should be packed. - if s != "no_shard": - model_weight_dtype = self.mp_policy.model_weight_buffer_dtype(self.params[0]) - wbuf = self._create_buffer(model_weight_dtype, shard_weights, "model_weight") - wbuf.init_data(torch.empty(wbuf.data_size, dtype=wbuf.dtype, device=self.device)) + model_weight_dtype = self.mp_policy.model_weight_buffer_dtype(self.params[0]) + wbuf = self._create_buffer(model_weight_dtype, shard_weights, "model_weight") + wbuf.init_data(torch.empty(wbuf.data_size, dtype=wbuf.dtype, device=self.device)) + for i, p in enumerate(self.params): + wbuf.set_item(i, self.mp_policy.get_param_data(p)) + self.model_weight_buffer = wbuf + + if self.mp_policy.needs_transpose_weight_buffer(self.params[0]): + tbuf = self._create_buffer(torch.uint8, shard_weights, "transpose_weight") + tbuf.init_data(torch.empty(tbuf.data_size, dtype=tbuf.dtype, device=self.device)) for i, p in enumerate(self.params): - wbuf.set_item(i, self.mp_policy.get_param_data(p)) - self.model_weight_buffer = wbuf - - if self.mp_policy.needs_transpose_weight_buffer(self.params[0]): - tbuf = self._create_buffer(torch.uint8, shard_weights, "transpose_weight") - tbuf.init_data(torch.empty(tbuf.data_size, dtype=tbuf.dtype, device=self.device)) - for i, p in enumerate(self.params): - tbuf.set_item(i, self.mp_policy.get_param_data(p, transpose=True)) - self.transpose_weight_buffer = tbuf + tbuf.set_item(i, self.mp_policy.get_param_data(p, transpose=True)) + self.transpose_weight_buffer = tbuf # Create main weight buffer for mixed precision main_params_dtype = self.mp_policy.main_params_dtype_for_param(self.params[0]) diff --git a/tests/unit_tests/distributed/megatron_fsdp/v2/test_fully_shard.py b/tests/unit_tests/distributed/megatron_fsdp/v2/test_fully_shard.py index 336d810a4dd..82984f3563b 100644 --- a/tests/unit_tests/distributed/megatron_fsdp/v2/test_fully_shard.py +++ b/tests/unit_tests/distributed/megatron_fsdp/v2/test_fully_shard.py @@ -256,6 +256,31 @@ def test_forward_backward_no_nan(self): assert not torch.isnan(torch.tensor(loss)), "Loss is NaN" assert not torch.isinf(torch.tensor(loss)), "Loss is Inf" + def test_no_shard_forward_backward_finish_grad_sync(self): + """no_shard keeps full replicated buffers and all-reduces at grad sync.""" + torch.manual_seed(42) + model = SimpleMLP(64).to(_device()) + fully_shard(model, sharding_strategy="no_shard", enable_async_reduce_grad=False) + + x = torch.randn(2, 64, device=_device()) + loss = _forward_backward(model, x) + assert not torch.isnan(torch.tensor(loss)), "Loss is NaN" + model.finish_grad_sync() + + for param_group in model._fsdp_param_groups: + assert param_group.model_weight_buffer is not None + assert not param_group.model_weight_buffer.is_distributed + assert param_group.main_grad_buffer is not None + assert not param_group.main_grad_buffer.is_distributed + for dist_grad in param_group.dist_grads: + if dist_grad is None: + continue + local_grad = dist_grad._local_tensor + gathered = [torch.empty_like(local_grad) for _ in range(_world_size())] + torch.distributed.all_gather(gathered, local_grad) + for replica in gathered: + torch.testing.assert_close(replica, local_grad) + @pytest.mark.parametrize( "enable_unshard_prefetch,enable_async_reduce_grad", [(False, False), (False, True), (True, False), (True, True)], diff --git a/tests/unit_tests/distributed/megatron_fsdp/v2/test_param_group.py b/tests/unit_tests/distributed/megatron_fsdp/v2/test_param_group.py index dfd0372413b..5390db49a18 100644 --- a/tests/unit_tests/distributed/megatron_fsdp/v2/test_param_group.py +++ b/tests/unit_tests/distributed/megatron_fsdp/v2/test_param_group.py @@ -50,7 +50,8 @@ def dist_env(): device = torch.device(f"cuda:{rank % torch.cuda.device_count()}") torch.cuda.set_device(device) yield - torch.distributed.destroy_process_group() + if torch.distributed.is_initialized(): + torch.distributed.destroy_process_group() # ------------------------------------------------------------------ # @@ -119,13 +120,13 @@ def _build_groups(strategy): def _flags(s): """Return (has_model_weight_buf, has_grad_buf, weight_distributed, grad_distributed). - - has_model_weight_buf: whether model_weight_buffer is created (False for no_shard) + - has_model_weight_buf: model_weight_buffer is created for all supported strategies - has_grad_buf: always True when requires_grad, across all strategies - weight_distributed: only True for optim_grads_params (full FSDP) - grad_distributed: True for optim_grads and optim_grads_params """ return { - "no_shard": (False, True, False, False), + "no_shard": (True, True, False, False), "optim": (True, True, False, False), "optim_grads": (True, True, False, True), "optim_grads_params": (True, True, True, True), @@ -166,9 +167,10 @@ def all_reduce(t, group): @pytest.mark.parametrize("strategy", ["no_shard", "optim", "optim_grads", "optim_grads_params"]) def test_init_buffers(strategy): - if strategy != "optim_grads_params": + if strategy not in ("no_shard", "optim_grads_params"): pytest.skip( - f"Only optim_grads_params strategy creates distributed buffers, skipping {strategy} test." + "This test currently covers no_shard and optim_grads_params, " + f"skipping {strategy}." ) groups, originals, dp_group, rank, ws, device = _build_groups(strategy) @@ -192,10 +194,6 @@ def test_init_buffers(strategy): else: expected = p.flatten() assert torch.equal(item, expected) - else: - # no_shard: model_weight_buffer should not exist - assert pg.model_weight_buffer is None - # -- main_grad_buffer -- if pg.requires_grad: assert pg.main_grad_buffer is not None @@ -213,9 +211,10 @@ def test_init_buffers(strategy): @pytest.mark.parametrize("strategy", ["no_shard", "optim", "optim_grads", "optim_grads_params"]) def test_unshard_reshard(strategy): - if strategy != "optim_grads_params": + if strategy not in ("no_shard", "optim_grads_params"): pytest.skip( - f"Only optim_grads_params strategy creates distributed buffers, skipping {strategy} test." + "This test currently covers no_shard and optim_grads_params, " + f"skipping {strategy}." ) groups, originals, dp_group, rank, ws, device = _build_groups(strategy) @@ -223,9 +222,7 @@ def test_unshard_reshard(strategy): for pg, orig in zip(groups, originals): wbuf = pg.model_weight_buffer - if wbuf is None: - # no_shard: nothing to unshard - continue + assert wbuf is not None shard_before = wbuf.data.clone() unsharded = wbuf.unshard() @@ -257,9 +254,10 @@ def test_unshard_reshard(strategy): @pytest.mark.parametrize("strategy", ["no_shard", "optim", "optim_grads", "optim_grads_params"]) def test_reduce_grad(strategy): - if strategy != "optim_grads_params": + if strategy not in ("no_shard", "optim_grads_params"): pytest.skip( - f"Only optim_grads_params strategy creates distributed buffers, skipping {strategy} test." + "This test currently covers no_shard and optim_grads_params, " + f"skipping {strategy}." ) groups, _, dp_group, rank, ws, device = _build_groups(strategy) From 276cd35ce774cc18dd806fdf9ff166d6a224ee1e Mon Sep 17 00:00:00 2001 From: tongliu Date: Tue, 9 Jun 2026 20:12:04 -0700 Subject: [PATCH 19/28] Improve Muon NS and P2P overlap --- .../core/optimizer/fully_shard_v2_muon.py | 104 ++++++++++-------- 1 file changed, 58 insertions(+), 46 deletions(-) diff --git a/megatron/core/optimizer/fully_shard_v2_muon.py b/megatron/core/optimizer/fully_shard_v2_muon.py index 7e606ed96bb..c66331e7dcb 100644 --- a/megatron/core/optimizer/fully_shard_v2_muon.py +++ b/megatron/core/optimizer/fully_shard_v2_muon.py @@ -197,8 +197,8 @@ def __init__( packages = list(by_group.values()) self._packages = [list(p) for p in packages] - # Precompute the static per-package isend/irecv plan; step() only fires the P2Ps. - self._gather_plans, self._scatter_plans = self._build_comm_plans() + # Precompute the static per-package NS and P2P plans; step() only fires the P2Ps. + self._ns_plans, self._gather_plans, self._scatter_plans = self._build_plans() def _compute_shard_ranges(self): """Compute, per param, each DP rank's (flat_offset, size) within the full param. @@ -279,35 +279,24 @@ def _assign_roots(self): load[best] += ns_cost return roots - def _build_comm_plans(self): - """Precompute the STATIC isend/irecv plan per package; step() just fires the P2Ps, - which read/write the grad / full-grad / orth tensors directly (no pack, no alloc — - split-root full grads are preallocated in __init__). - - Both ends walk params in ``package`` order so same-(sender, receiver) P2Ps pair - up, and each plan is the symmetric {group, send, recv}. - 1. gather — route each rank's grad shard to that param's root: - - root == this rank, SPLIT: irecv the OTHER holders' shards into the - preallocated full grad (Phase 1 wrote our own segment). - - root == this rank, SINGLE holder: no P2P — the grad shard already IS the - full grad. - - else we hold a shard: isend it to the root. - 2. scatter — the inverse: isend each OTHER holder's shard out of the orth, irecv - our shard into the grad. Our own shard is read straight from the orth by - Phase 3b, never touching the grad. - - Returns (gather_plans, scatter_plans), parallel to self._packages; each a dict: - gather: group, - rooted = [param_idx], # params rooted on this rank - send = [(param_idx, root)], # isend grad.to_local() -> root - recv = [(param_idx, src, offset, size)], # irecv src shard -> full_grad[offset:] - scatter: group, - send = [(param_idx, dst, offset, size)], # isend orth[offset:] -> dst - recv = [(param_idx, root)], # irecv -> grad.to_local() - _issue_gather returns the rooted full-grad views along with the in-flight - P2P works; the views become valid after the works complete. + def _build_plans(self): + """Precompute per-package NS order and P2P send/recv lists. + + NS plan splits rooted grads so comm-free NS covers exposed comm: + + issue gather(i + 1) + NS(pre_comm) # overlaps gather(i); single-holder grads + wait gather(i) + NS(wait_comm) # depends on gather(i); split-root grads + issue scatter(i) + NS(post_comm) # overlaps scatter(i); single-holder grads + + Gather routes local post-momentum shards to each param's root, e.g. + ``send=(param_idx, root)`` or ``recv=(param_idx, src, offset, size)``. + Scatter is the reverse: roots send orth shards to non-root holders, while + the root reads its own orth shard locally. """ - gather_plans, scatter_plans = [], [] + ns_plans, gather_plans, scatter_plans = [], [], [] for package in self._packages: group, this_rank, world_size = self._comm[package[0]] @@ -327,7 +316,25 @@ def _build_comm_plans(self): # single holder: grad shard already is the full grad -> no P2P elif ranges[this_rank][1] > 0: send.append((param_idx, root)) - gather_plans.append({"group": group, "rooted": rooted, "send": send, "recv": recv}) + comm_free_rooted, wait_comm_rooted = [], [] + for param_idx in rooted: + shape = self._main_weights[param_idx].shape + full_grad = self._full_grads.get(param_idx, self._managed[param_idx][3]) + full_grad = full_grad.view(shape) + if param_idx in self._full_grads: + wait_comm_rooted.append((param_idx, full_grad)) + else: + comm_free_rooted.append((param_idx, full_grad)) + + split = (len(comm_free_rooted) + 1) // 2 + ns_plans.append( + { + "pre_comm": comm_free_rooted[:split], + "wait_comm": wait_comm_rooted, + "post_comm": comm_free_rooted[split:], + } + ) + gather_plans.append({"group": group, "send": send, "recv": recv}) # ---- scatter: each root's orth shard -> its OTHER holders (inverse of gather); # the root's own shard is read straight from the orth by Phase 3b ---- @@ -344,18 +351,13 @@ def _build_comm_plans(self): elif ranges[this_rank][1] > 0: recv.append((param_idx, root)) scatter_plans.append({"group": group, "send": send, "recv": recv}) - return gather_plans, scatter_plans + return ns_plans, gather_plans, scatter_plans @torch.no_grad() def _issue_gather(self, p): - """Fire the package's gather P2Ps and return rooted full-grad views + Works.""" + """Fire the package's gather P2Ps and return Works.""" plan = self._gather_plans[p] group = plan["group"] - full_grads = {} - for param_idx in plan["rooted"]: - shape = self._main_weights[param_idx].shape - full_grad = self._full_grads.get(param_idx, self._managed[param_idx][3]) - full_grads[param_idx] = full_grad.view(shape) ops = [] for param_idx, root in plan["send"]: @@ -368,7 +370,7 @@ def _issue_gather(self, p): torch.distributed.irecv, self._full_grads[param_idx].reshape(-1)[offset : offset + size], src, group=group, )) - return full_grads, torch.distributed.batch_isend_irecv(ops) if ops else [] + return torch.distributed.batch_isend_irecv(ops) if ops else [] @torch.no_grad() def _issue_scatter(self, p, orths): @@ -448,21 +450,31 @@ def step(self, closure=None): num_packages = len(self._packages) scatter_reqs = [] # in-flight scatter Works per package owned_orths = {} # param_idx -> full orth this rank rooted (read by Phase 3b) - gather_work = self._issue_gather(0) if num_packages else None + gather_reqs_by_package = [[] for _ in range(num_packages)] + if num_packages: + gather_reqs_by_package[0] = self._issue_gather(0) for i in range(num_packages): - next_gather_work = self._issue_gather(i + 1) if i + 1 < num_packages else None - full_grads, gather_reqs = gather_work - for req in gather_reqs: - req.wait() + if i + 1 < num_packages: + gather_reqs_by_package[i + 1] = self._issue_gather(i + 1) + ns_plan = self._ns_plans[i] orths = { param_idx: self.orthogonalize(self._main_weights[param_idx], full_grad).to( full_grad.dtype ) - for param_idx, full_grad in full_grads.items() + for param_idx, full_grad in ns_plan["pre_comm"] } + for req in gather_reqs_by_package[i]: + req.wait() + for param_idx, full_grad in ns_plan["wait_comm"]: + orths[param_idx] = self.orthogonalize( + self._main_weights[param_idx], full_grad + ).to(full_grad.dtype) owned_orths.update(orths) scatter_reqs.append(self._issue_scatter(i, orths)) - gather_work = next_gather_work + for param_idx, full_grad in ns_plan["post_comm"]: + owned_orths[param_idx] = self.orthogonalize( + self._main_weights[param_idx], full_grad + ).to(full_grad.dtype) # Drain scatters: each non-root holder now has its orth shard irecv'd into its grad. for reqs in scatter_reqs: From b70e14d42fafb45d5c959ffd945a6d9ab5e03ee7 Mon Sep 17 00:00:00 2001 From: tongliu Date: Wed, 10 Jun 2026 04:56:02 -0700 Subject: [PATCH 20/28] Fix Muon subgroup P2P peers --- .../core/optimizer/fully_shard_v2_muon.py | 606 +++++++++--------- .../megatron_fsdp/v2/test_dist_muon.py | 106 ++- 2 files changed, 390 insertions(+), 322 deletions(-) diff --git a/megatron/core/optimizer/fully_shard_v2_muon.py b/megatron/core/optimizer/fully_shard_v2_muon.py index c66331e7dcb..176ff92846e 100644 --- a/megatron/core/optimizer/fully_shard_v2_muon.py +++ b/megatron/core/optimizer/fully_shard_v2_muon.py @@ -14,6 +14,7 @@ from collections import OrderedDict +from dataclasses import dataclass from typing import List import torch @@ -55,9 +56,181 @@ def _vanilla_muon_scale(rows: int, cols: int) -> float: return max(1.0, rows / cols) ** 0.5 +@dataclass +class MuonGradPackage: + """Static Muon package over aligned main-weight / grad / momentum indices.""" + + param_ids: List[int] + has_comm: bool + + @torch.no_grad() + def update_momentum(self, opt, coef): + for param_idx in self.param_ids: + grad = opt._main_grads[param_idx] + momentum = opt._momentum_buffers[param_idx] + if grad is None: + continue + grad_shard = grad.to_local().reshape(-1) + if grad_shard.numel() == 0: + continue + momentum.mul_(coef).add_(grad_shard) + if opt._nesterov: + torch.add(grad_shard, momentum, alpha=coef, out=grad_shard) + + @torch.no_grad() + def issue_gather(self, opt, full_grads): + """Gather this package's NS input shards to each param root.""" + if not self.has_comm: + return [] + + group, this_rank, world_size = opt._comm[self.param_ids[0]] + ops = [] + for param_idx in self.param_ids: + root = opt._roots[param_idx] + ranges = opt._shard_ranges[param_idx] + this_offset, this_size = ranges[this_rank] + ns_input_shard = None + if this_size > 0: + if opt._nesterov: + ns_input_shard = opt._main_grads[param_idx].to_local().reshape(-1) + else: + ns_input_shard = opt._momentum_buffers[param_idx] + + if root == this_rank: + main_weight = opt._main_weights[param_idx] + momentum = opt._momentum_buffers[param_idx] + full_grad = torch.empty( + main_weight.shape.numel(), + dtype=momentum.dtype, + device=main_weight.device, + ) + full_grads[param_idx] = full_grad + if this_size > 0: + full_grad[this_offset : this_offset + this_size].copy_(ns_input_shard) + for src in range(world_size): + offset, size = ranges[src] + if size == 0 or src == this_rank: + continue + ops.append( + dist.P2POp( + dist.irecv, + full_grad[offset : offset + size], + dist.get_global_rank(group, src), + group=group, + ) + ) + elif this_size > 0: + ops.append( + dist.P2POp( + dist.isend, + ns_input_shard, + dist.get_global_rank(group, root), + group=group, + ) + ) + return dist.batch_isend_irecv(ops) if ops else [] + + @torch.no_grad() + def finish_gather(self, requests): + for request in requests: + request.wait() + + @torch.no_grad() + def orthogonalize(self, opt, requests, full_grads): + self.finish_gather(requests) + orths = {} + for param_idx in self.param_ids: + if opt._roots[param_idx] != opt._comm[param_idx][1]: + continue + if param_idx in full_grads: + full_grad = full_grads.pop(param_idx) + elif opt._nesterov: + full_grad = opt._main_grads[param_idx].to_local().reshape(-1) + else: + full_grad = opt._momentum_buffers[param_idx] + full_grad = full_grad.view(opt._main_weights[param_idx].shape) + + grad = full_grad.to(torch.float32) + if HAVE_EMERGING_OPTIMIZERS: + orth = newton_schulz_tp( + grad, + steps=opt._num_ns_steps, + coefficient_type=opt._coefficient_type, + tp_group=None, + partition_dim=None, + tp_mode="duplicated", + ) + scale = get_muon_scale_factor( + grad.size(-2), grad.size(-1), mode=opt._scale_mode + ) + else: + orth = _vanilla_newton_schulz(grad, steps=opt._num_ns_steps) + scale = _vanilla_muon_scale(grad.size(-2), grad.size(-1)) + orths[param_idx] = (orth * (scale * opt._extra_scale_factor)).to(full_grad.dtype) + return orths + + @torch.no_grad() + def issue_scatter(self, opt, orths): + """Scatter this package's orthogonalized shards from each param root.""" + if not self.has_comm: + return [] + + group, this_rank, world_size = opt._comm[self.param_ids[0]] + ops = [] + for param_idx in self.param_ids: + root = opt._roots[param_idx] + ranges = opt._shard_ranges[param_idx] + if root == this_rank: + orth = orths[param_idx].reshape(-1) + for dst in range(world_size): + offset, size = ranges[dst] + if size == 0 or dst == this_rank: + continue + ops.append( + dist.P2POp( + dist.isend, + orth[offset : offset + size], + dist.get_global_rank(group, dst), + group=group, + ) + ) + elif ranges[this_rank][1] > 0: + ops.append( + dist.P2POp( + dist.irecv, + opt._main_grads[param_idx].to_local().reshape(-1), + dist.get_global_rank(group, root), + group=group, + ) + ) + return dist.batch_isend_irecv(ops) if ops else [] + + @torch.no_grad() + def finish_scatter(self, requests): + for request in requests: + request.wait() + + class FullyShardV2Muon(torch.optim.Optimizer): """Single-root distributed Muon for Megatron-FSDP v2 (a torch optimizer). + Initialization: + 1. Layout: compute each rank's flat shard range for every parameter. + 2. Root selection: assign each parameter's Newton-Schulz work to one rank. + The greedy policy balances estimated NS cost while preferring ranks that + already hold the shard. + 3. Package construction: group communicating params by DP group and sort by + NS cost. No-comm params are split around communication packages to cover + the first gather and final scatter. + + Step: + 1. Momentum: update local momentum shards; Nesterov writes the NS input into + the local grad shard. + 2. Gather: send NS input shards to each parameter's root rank. + 3. Orthogonalize: root ranks run Newton-Schulz and keep full orth updates. + 4. Scatter: roots send orthogonalized shards back to shard owners. + 5. Apply: each rank updates its local optimizer-weight shard. + Args: params: optimizer-facing params — all assumed to be FSDP-v2 ``dist_param`` DTensors (asserted). A flat list or param groups, as any torch optimizer. @@ -67,12 +240,9 @@ class FullyShardV2Muon(torch.optim.Optimizer): lr / momentum / nesterov / weight_decay: Muon update hyperparameters. num_ns_steps / coefficient_type / scale_mode / extra_scale_factor: NS + scaling config used by ``orthogonalize``. - fp32_matmul_prec / tp_mode: reserved for the Phase-2 tensor-parallel NS path - (currently TP size 1; accepted and ignored). - packages: list of param-index lists grouping params for batched P2P (one - batch_isend_irecv per package; packages pipeline so a package's gather - overlaps the previous package's NS). Every package must share one dp_group. - None -> one package per dp_group. + fp32_matmul_prec / tp_mode: reserved for the future tensor-parallel NS path + (currently TP size 1; accepted and ignored). Params are grouped into + communication and no-communication packages after root assignment. """ def __init__( @@ -90,115 +260,55 @@ def __init__( extra_scale_factor: float = 1.0, fp32_matmul_prec: str = "medium", tp_mode: str = "duplicated", - packages=None, ) -> None: super().__init__(params, dict(lr=lr, weight_decay=weight_decay)) - flat_params = [param for group in self.param_groups for param in group["params"]] + param_list = [param for group in self.param_groups for param in group["params"]] + grad_list = list(grads) + if not param_list: + raise ValueError("FullyShardV2Muon got no parameters to manage.") assert all( - isinstance(p, DTensor) for p in flat_params + isinstance(param, DTensor) for param in param_list ), "FullyShardV2Muon expects every param to be a DTensor (FSDP-v2 dist_param)." - # grads align 1:1 with params; each a Shard(0) DTensor, or None if our shard is empty. assert all( - g is None or isinstance(g, DTensor) for g in grads + grad is None or isinstance(grad, DTensor) for grad in grad_list ), "FullyShardV2Muon expects every grad to be a DTensor or None." - assert len(grads) == len(flat_params), ( - f"grads ({len(grads)}) must align 1:1 with params ({len(flat_params)})." - ) - self._main_grads = list(grads) - # Muon hyperparameters. lr is NOT stored — it lives in param_groups[0]["lr"] - # (where the LR scheduler writes), read in step(). + self._main_weights = param_list + self._main_grads = grad_list + + # Hyperparameters read by step() and MuonGradPackage methods. self._momentum_coef = momentum - self._nesterov = nesterov self._weight_decay = weight_decay - # NS + scaling config for orthogonalize(). + self._nesterov = nesterov self._num_ns_steps = num_ns_steps self._coefficient_type = coefficient_type self._scale_mode = scale_mode self._extra_scale_factor = extra_scale_factor - if not flat_params: - raise ValueError("FullyShardV2Muon got no parameters to manage.") - self._main_weights = list(flat_params) - # Per param: (dp_group, this rank's rank within it, that group's world size). self._comm = [] for main_weight in self._main_weights: group = main_weight.device_mesh.get_group() self._comm.append((group, dist.get_rank(group), dist.get_world_size(group))) - # Per-rank (flat_offset, size) of each param + a load-balanced NS root - # (layout-only, so identical on every rank). self._shard_ranges = self._compute_shard_ranges() self._roots = self._assign_roots() - # Persistent full-grad buffer per SPLIT param rooted here (the root assembles its - # shards for NS). Single-holder params need none — their grad shard already IS the - # full grad. - self._full_grads = {} - for i, main_weight in enumerate(self._main_weights): - _group, this_rank, _world_size = self._comm[i] - if self._roots[i] != this_rank: - continue - holders = [r for r, (_offset, size) in enumerate(self._shard_ranges[i]) if size > 0] - if len(holders) <= 1: - continue - grad = self._main_grads[i] + self._momentum_buffers = [] + for param_idx, (main_weight, grad) in enumerate(zip(self._main_weights, self._main_grads)): + _offset, local_size = self._shard_ranges[param_idx][self._comm[param_idx][1]] dtype = grad.dtype if grad is not None else main_weight.dtype - self._full_grads[i] = torch.empty( - main_weight.shape.numel(), dtype=dtype, device=main_weight.device - ) + momentum_buffer = torch.zeros(local_size, dtype=dtype, device=main_weight.device) + self._momentum_buffers.append(momentum_buffer) - # Each entry: (master weight, grad or None, momentum, post_momentum_shard). - # post_momentum_shard is this rank's momentum-adjusted gradient slice, written by - # Phase 1 and read by gather in its FINAL location (no copy). nesterov: a fresh - # grad+coef*m, into the split root's full-grad slice or the grad shard (momentum - # separate). non-nesterov: == m, so it aliases momentum (which for a split root - # lives in the full-grad slice). - self._managed = [] - for i in range(len(self._main_weights)): - main_weight = self._main_weights[i] - grad = self._main_grads[i] - _group, this_rank, _world_size = self._comm[i] - this_offset, this_size = self._shard_ranges[i][this_rank] - dtype = grad.dtype if grad is not None else main_weight.dtype - # this rank's slice inside the preallocated full grad (split root only, else None) - grad_slice = ( - self._full_grads[i].reshape(-1)[this_offset : this_offset + this_size] - if i in self._full_grads - else None - ) - if self._nesterov: - momentum = torch.zeros(this_size, dtype=dtype, device=main_weight.device) - post_momentum_shard = grad_slice if grad_slice is not None else ( - grad.to_local().reshape(-1) if grad is not None else None - ) - elif grad_slice is not None: # split root: momentum lives in the full grad - momentum = grad_slice - momentum.zero_() - post_momentum_shard = momentum - else: - momentum = torch.zeros(this_size, dtype=dtype, device=main_weight.device) - post_momentum_shard = momentum - self._managed.append((main_weight, grad, momentum, post_momentum_shard)) - - # `packages`: list of param-index lists (into self._main_weights). Each package - # does ONE batched isend/irecv; packages pipeline (package i+1's gather overlaps - # package i's NS). Every package must share one dp_group. The wrapper passes one - # package per FSDPModule (== layer), so layers pipeline: - # packages = [[0, 1, 2], [3, 4, 5]] # layer 0 | layer 1 - # Default (None): one package per dp_group, e.g. dense vs expert params: - # packages = [[0, 1, 3], [2, 4]] # dp-group | edp-group; no intra-group pipeline - if packages is None: - by_group = OrderedDict() - for param_idx in range(len(self._main_weights)): - by_group.setdefault(self._comm[param_idx][0], []).append(param_idx) - packages = list(by_group.values()) - self._packages = [list(p) for p in packages] - - # Precompute the static per-package NS and P2P plans; step() only fires the P2Ps. - self._ns_plans, self._gather_plans, self._scatter_plans = self._build_plans() + self._ns_costs = [ + main_weight.shape.numel() * min(main_weight.shape) + if len(main_weight.shape) == 2 + else 0 + for main_weight in self._main_weights + ] + self._packages = self._build_packages() def _compute_shard_ranges(self): """Compute, per param, each DP rank's (flat_offset, size) within the full param. @@ -238,18 +348,18 @@ def _assign_roots(self): """Assign each param an NS root rank, load-balancing NS work per dp_group. Steps: - 1. group params by dp_group; track each rank's NS load (cost ~ numel*min(dim)); - 2. pin single-holder params to their only holder (zero gather/scatter) and add - their cost to that rank's load; - 3. assign split params heaviest-first (LPT) to the least-loaded holder (ties: - the holder with more of this param's data, then lowest rank). + 1. group params by dp_group and estimate each param's NS cost; + 2. compute the average NS target load for the group; + 3. assign params heaviest-first, preferring no-comm roots until their + ranks reach the target, then allowing remote roots for balance. Deterministic + identical on every rank (depends only on shard layout). Example (world_size=2): - param A: single holder {0}, cost 100 -> root 0 (load [100, 0]) - param B: split {0,1}, cost 300 -> root 1 (idler) (load [100, 300]) - param C: split {0,1}, cost 200 -> root 0 (load [300, 300]) + target load = 400 + param A: single holder {0}, cost 300 -> root 0 (load [300, 0]) + param B: single holder {0}, cost 300 -> root 0 (load [600, 0]) + param C: split {0,1}, cost 200 -> root 1 (load [600, 200]) """ params_by_group = OrderedDict() for param_idx in range(len(self._main_weights)): @@ -258,164 +368,82 @@ def _assign_roots(self): for _group, param_indices in params_by_group.items(): world_size = self._comm[param_indices[0]][2] load = [0] * world_size - split = [] + work = [] for param_idx in param_indices: - holders = [ - rank - for rank, (_offset, size) in enumerate(self._shard_ranges[param_idx]) - if size > 0 - ] shape = self._main_weights[param_idx].shape ns_cost = shape.numel() * min(shape) if len(shape) == 2 else 0 - if len(holders) == 1: - roots[param_idx] = holders[0] - load[holders[0]] += ns_cost - elif len(holders) > 1: - split.append((ns_cost, param_idx, holders)) - for ns_cost, param_idx, holders in sorted(split, key=lambda e: (-e[0], e[1])): + work.append((ns_cost, param_idx)) + target_load = sum(ns_cost for ns_cost, _param_idx in work) / world_size + for ns_cost, param_idx in sorted(work, key=lambda e: (-e[0], e[1])): ranges = self._shard_ranges[param_idx] - best = min(holders, key=lambda rank: (load[rank], -ranges[rank][1], rank)) + param_numel = self._main_weights[param_idx].shape.numel() + full_holders = [ + rank for rank, (_offset, size) in enumerate(ranges) if size == param_numel + ] + holders = [rank for rank, (_offset, size) in enumerate(ranges) if size > 0] + + # Prefer no-comm full holders until they reach the average target. + # For split params, prefer shard holders under target before using + # a non-holder; larger local shards win ties to limit traffic. + candidates = [rank for rank in full_holders if load[rank] < target_load] + if not candidates: + candidates = [rank for rank in holders if load[rank] < target_load] + if not candidates: + candidates = range(world_size) + + best = min(candidates, key=lambda rank: (load[rank], -ranges[rank][1], rank)) roots[param_idx] = best load[best] += ns_cost return roots - def _build_plans(self): - """Precompute per-package NS order and P2P send/recv lists. - - NS plan splits rooted grads so comm-free NS covers exposed comm: - - issue gather(i + 1) - NS(pre_comm) # overlaps gather(i); single-holder grads - wait gather(i) - NS(wait_comm) # depends on gather(i); split-root grads - issue scatter(i) - NS(post_comm) # overlaps scatter(i); single-holder grads - - Gather routes local post-momentum shards to each param's root, e.g. - ``send=(param_idx, root)`` or ``recv=(param_idx, src, offset, size)``. - Scatter is the reverse: roots send orth shards to non-root holders, while - the root reads its own orth shard locally. - """ - ns_plans, gather_plans, scatter_plans = [], [], [] - for package in self._packages: - group, this_rank, world_size = self._comm[package[0]] - - # ---- gather: each rank's grad shard -> that param's root ---- - rooted, send, recv = [], [], [] - for param_idx in package: - root = self._roots[param_idx] - ranges = self._shard_ranges[param_idx] - if root == this_rank: - rooted.append(param_idx) - if param_idx in self._full_grads: # split: assemble from other holders - for src in range(world_size): - offset, size = ranges[src] - if size == 0 or src == this_rank: # own segment came from Phase 1 - continue - recv.append((param_idx, src, offset, size)) - # single holder: grad shard already is the full grad -> no P2P - elif ranges[this_rank][1] > 0: - send.append((param_idx, root)) - comm_free_rooted, wait_comm_rooted = [], [] - for param_idx in rooted: - shape = self._main_weights[param_idx].shape - full_grad = self._full_grads.get(param_idx, self._managed[param_idx][3]) - full_grad = full_grad.view(shape) - if param_idx in self._full_grads: - wait_comm_rooted.append((param_idx, full_grad)) - else: - comm_free_rooted.append((param_idx, full_grad)) - - split = (len(comm_free_rooted) + 1) // 2 - ns_plans.append( - { - "pre_comm": comm_free_rooted[:split], - "wait_comm": wait_comm_rooted, - "post_comm": comm_free_rooted[split:], - } + def _build_packages(self, package_size=4): + """Build static Muon packages from the assigned roots.""" + ns_costs = self._ns_costs + comm_params, comm_free_params = [], [] + for param_idx in range(len(self._main_weights)): + root = self._roots[param_idx] + has_comm = any( + rank != root + for rank, (_offset, size) in enumerate(self._shard_ranges[param_idx]) + if size > 0 ) - gather_plans.append({"group": group, "send": send, "recv": recv}) - - # ---- scatter: each root's orth shard -> its OTHER holders (inverse of gather); - # the root's own shard is read straight from the orth by Phase 3b ---- - send, recv = [], [] - for param_idx in package: - root = self._roots[param_idx] - ranges = self._shard_ranges[param_idx] - if root == this_rank: - for dst in range(world_size): - offset, size = ranges[dst] - if size == 0 or dst == this_rank: # own shard: used directly, not sent - continue - send.append((param_idx, dst, offset, size)) - elif ranges[this_rank][1] > 0: - recv.append((param_idx, root)) - scatter_plans.append({"group": group, "send": send, "recv": recv}) - return ns_plans, gather_plans, scatter_plans - - @torch.no_grad() - def _issue_gather(self, p): - """Fire the package's gather P2Ps and return Works.""" - plan = self._gather_plans[p] - group = plan["group"] - - ops = [] - for param_idx, root in plan["send"]: - post_momentum_shard = self._managed[param_idx][3] - ops.append(torch.distributed.P2POp( - torch.distributed.isend, post_momentum_shard, root, group=group, - )) - for param_idx, src, offset, size in plan["recv"]: - ops.append(torch.distributed.P2POp( - torch.distributed.irecv, - self._full_grads[param_idx].reshape(-1)[offset : offset + size], src, group=group, - )) - return torch.distributed.batch_isend_irecv(ops) if ops else [] - - @torch.no_grad() - def _issue_scatter(self, p, orths): - """Fire the package's scatter P2Ps: isend each remote holder's shard out of the - orth, irecv our shards into the grad. (The root's own shard is read straight from - the orth by Phase 3b, never round-tripping the grad.) Returns the in-flight Works. - The isend source ``orths`` is kept alive by the caller's ``owned_orths``.""" - plan = self._scatter_plans[p] - group = plan["group"] - ops = [] - for param_idx, dst, offset, size in plan["send"]: - ops.append(torch.distributed.P2POp( - torch.distributed.isend, - orths[param_idx].reshape(-1)[offset : offset + size], dst, group=group, - )) - for param_idx, root in plan["recv"]: - ops.append(torch.distributed.P2POp( - torch.distributed.irecv, - self._main_grads[param_idx].to_local().reshape(-1), root, group=group, - )) - return torch.distributed.batch_isend_irecv(ops) if ops else [] - - @torch.no_grad() - def orthogonalize(self, param, grad: torch.Tensor) -> torch.Tensor: - """Newton-Schulz orthogonalize + Muon spectral scale of a 2D gradient. + (comm_params if has_comm else comm_free_params).append(param_idx) + + comm_params_by_group = OrderedDict() + for param_idx in comm_params: + comm_params_by_group.setdefault(self._comm[param_idx][0], []).append(param_idx) + + comm_packages = [] + for params_in_group in comm_params_by_group.values(): + params_in_group.sort(key=lambda param_idx: (-ns_costs[param_idx], param_idx)) + for start in range(0, len(params_in_group), package_size): + comm_packages.append( + MuonGradPackage( + param_ids=params_in_group[start : start + package_size], + has_comm=True, + ) + ) - Uses emerging_optimizers' production NS when available, else the built-in - vanilla NS fallback. (Phase 1 runs the non-TP path; ``param`` is accepted - for parity with the eventual TP / QKV-split path.) - """ - g = grad.to(torch.float32) - if HAVE_EMERGING_OPTIMIZERS: - orth = newton_schulz_tp( - g, - steps=self._num_ns_steps, - coefficient_type=self._coefficient_type, - tp_group=None, - partition_dim=None, - tp_mode="duplicated", + no_comm_packages = [ + MuonGradPackage( + param_ids=comm_free_params[start : start + package_size], + has_comm=False, ) - scale = get_muon_scale_factor(g.size(-2), g.size(-1), mode=self._scale_mode) - else: - orth = _vanilla_newton_schulz(g, steps=self._num_ns_steps) - scale = _vanilla_muon_scale(g.size(-2), g.size(-1)) - return orth * (scale * self._extra_scale_factor) + for start in range(0, len(comm_free_params), package_size) + ] + + # Put no-comm NS on both sides of the communication packages so it can + # cover the first gather and the final scatter. + prefix_target = sum(ns_costs[param_idx] for param_idx in comm_free_params) / 2 + split = 0 + prefix_cost = 0 + while split < len(no_comm_packages) and (split == 0 or prefix_cost < prefix_target): + prefix_cost += sum( + ns_costs[param_idx] for param_idx in no_comm_packages[split].param_ids + ) + split += 1 + + return no_comm_packages[:split] + comm_packages + no_comm_packages[split:] @torch.no_grad() def step(self, closure=None): @@ -425,65 +453,40 @@ def step(self, closure=None): its root; (3) root orthogonalizes (NS); (4) scatter the orthogonalized shards back; (5) decoupled-WD update of this rank's master shard. - Gather/scatter are batched per package (one async batch_isend_irecv each, directly - into/out of the grad and full-grad tensors) and pipelined: package i+1's gather is - fired before package i's NS, so the P2Ps overlap the NS GEMMs. + Gather/scatter use one batched P2P call per communication package. No-comm + packages are placed before and after the communicating packages so their NS + can cover the first gather and final scatter. """ assert closure is None, "FullyShardV2Muon does not support a closure." lr = self.param_groups[0]["lr"] # set by the LR scheduler - items = self._managed + gather_reqs = [[] for _ in range(len(self._packages))] + full_grads = {} - # Phase 1: momentum on each rank's grad shard. nesterov also writes the look-ahead - # grad+coef*m into post_momentum_shard; non-nesterov's == m (already aliased). + # Phase 1: update local momentum. Communication packages launch their + # gather immediately after their shards are ready, overlapping the rest of + # momentum work and the no-comm NS prefix. coef = self._momentum_coef - for _main_weight, grad, momentum, post_momentum_shard in items: - if grad is None: - continue - grad_shard = grad.to_local().reshape(-1) - if grad_shard.numel() == 0: - continue - momentum.mul_(coef).add_(grad_shard) - if self._nesterov: - torch.add(grad_shard, momentum, alpha=coef, out=post_momentum_shard) + for package_idx, package in enumerate(self._packages): + package.update_momentum(self, coef) + gather_reqs[package_idx] = package.issue_gather(self, full_grads) - # Phase 2/3: per-package gather -> NS -> scatter, pipelined (gather i+1 before NS i). - num_packages = len(self._packages) - scatter_reqs = [] # in-flight scatter Works per package + # Phase 2: consume packages in NS order, run root NS, then launch package + # scatter for the orthogonalized shards. + scatter_reqs = [] # in-flight scatter requests per communication package owned_orths = {} # param_idx -> full orth this rank rooted (read by Phase 3b) - gather_reqs_by_package = [[] for _ in range(num_packages)] - if num_packages: - gather_reqs_by_package[0] = self._issue_gather(0) - for i in range(num_packages): - if i + 1 < num_packages: - gather_reqs_by_package[i + 1] = self._issue_gather(i + 1) - ns_plan = self._ns_plans[i] - orths = { - param_idx: self.orthogonalize(self._main_weights[param_idx], full_grad).to( - full_grad.dtype - ) - for param_idx, full_grad in ns_plan["pre_comm"] - } - for req in gather_reqs_by_package[i]: - req.wait() - for param_idx, full_grad in ns_plan["wait_comm"]: - orths[param_idx] = self.orthogonalize( - self._main_weights[param_idx], full_grad - ).to(full_grad.dtype) + + for package_idx, package in enumerate(self._packages): + orths = package.orthogonalize(self, gather_reqs[package_idx], full_grads) owned_orths.update(orths) - scatter_reqs.append(self._issue_scatter(i, orths)) - for param_idx, full_grad in ns_plan["post_comm"]: - owned_orths[param_idx] = self.orthogonalize( - self._main_weights[param_idx], full_grad - ).to(full_grad.dtype) - - # Drain scatters: each non-root holder now has its orth shard irecv'd into its grad. - for reqs in scatter_reqs: - for req in reqs: - req.wait() - - # Phase 3b: decoupled-WD master update from each rank's orth shard — the root reads - # its own segment from the orth it computed, other holders read scatter's grad. - for param_idx, (main_weight, grad, _momentum, _post_momentum_shard) in enumerate(items): + if package.has_comm: + scatter_reqs.append((package, package.issue_scatter(self, orths))) + + # Phase 3a: drain scatters before reading remote orth shards from grad. + for package, requests in scatter_reqs: + package.finish_scatter(requests) + + # Phase 3b: apply decoupled-WD update from each rank's orth shard. + for param_idx, (main_weight, grad) in enumerate(zip(self._main_weights, self._main_grads)): if grad is None: continue weight_shard = main_weight.to_local().reshape(-1) @@ -515,15 +518,13 @@ def __init__(self, config, model_chunks, **muon_hyperparams): from megatron.core.distributed.fsdp.src.megatron_fsdp.v2 import FSDPModule # Build the inner optimizer from the FSDP-v2 2D-matrix dist_params + their grad - # DTensors, one package per FSDPModule (== layer) so each layer's gather pipelines - # with the next layer's NS (a package shares one dp_group, as the P2P needs). - params, grads, packages = [], [], [] + # DTensors. The inner optimizer orders params by communication need. + params, grads = [], [] for chunk in self.model_chunks: root = chunk if isinstance(chunk, FSDPModule) else chunk.module for m in root.modules(): if not isinstance(m, FSDPModule): continue - package = [] for param_group in m._fsdp_param_groups: for dist_param, dist_grad in zip(param_group.dist_params, param_group.dist_grads): # filter by the param's global attrs only (rank-consistent); keep @@ -531,14 +532,9 @@ def __init__(self, config, model_chunks, **muon_hyperparams): if dist_param.dim() == 2 and not getattr( dist_param, "is_embedding_or_output_parameter", False ): - package.append(len(params)) params.append(dist_param) grads.append(dist_grad) - if package: - packages.append(package) - super().__init__( - FullyShardV2Muon(params, grads, packages=packages, **muon_hyperparams), config - ) + super().__init__(FullyShardV2Muon(params, grads, **muon_hyperparams), config) self.is_stub_optimizer = False # --- Excluded from the chained grad-norm / clip / zero-count machinery. --- diff --git a/tests/unit_tests/distributed/megatron_fsdp/v2/test_dist_muon.py b/tests/unit_tests/distributed/megatron_fsdp/v2/test_dist_muon.py index ec89ecfb4e7..769d5026c94 100644 --- a/tests/unit_tests/distributed/megatron_fsdp/v2/test_dist_muon.py +++ b/tests/unit_tests/distributed/megatron_fsdp/v2/test_dist_muon.py @@ -19,8 +19,8 @@ 1. construct a full gradient, identical on every rank; 2. plant it into the sharded grad buffer (each rank keeps only its own slice); - 3. reference: run a full-tensor Muon update on every rank, reusing the - optimizer's own ``orthogonalize`` so the Newton-Schulz path is shared; + 3. reference: run a full-tensor Muon update on every rank, mirroring the + optimizer's Newton-Schulz path; 4. distributed: ``opt.step()`` gathers each grad to its root rank, runs NS there, scatters the result back, and updates each rank's shard; 5. compare the two results bit for bit (``torch.equal``). @@ -45,6 +45,7 @@ import pytest import torch import torch.nn as nn +from torch.distributed.tensor import DeviceMesh sys.path.insert(0, str(Path(__file__).parents[2])) from megatron.core.distributed.fsdp.src.megatron_fsdp.v2.mixed_precision import ( @@ -52,15 +53,31 @@ ) from megatron.core.distributed.fsdp.src.megatron_fsdp.v2.param_group import ParameterGroup from megatron.core.distributed.fsdp.src.megatron_fsdp.v2.utils import ParamGroupIdx -from megatron.core.optimizer.fully_shard_v2_muon import FullyShardV2Muon +from megatron.core.optimizer.fully_shard_v2_muon import ( + HAVE_EMERGING_OPTIMIZERS, + FullyShardV2Muon, + _vanilla_muon_scale, + _vanilla_newton_schulz, +) + +if HAVE_EMERGING_OPTIMIZERS: + from megatron.core.optimizer.fully_shard_v2_muon import ( + get_muon_scale_factor, + newton_schulz_tp, + ) STRATEGIES = ["optim", "optim_grads", "optim_grads_params"] +MESH_CASES = [ + pytest.param("world", id="world-mesh"), + pytest.param("rank_pairs", id="rank-pair-mesh"), +] + +_RANK_PAIR_MESH = None # Each set is a list of param shapes at a different scale: four 2D matrices (mixing # wide rowscols -> NS transposes) + a trailing 1D bias Muon must -# skip. Four matrices let the test split them into 2 multi-param packages, exercising -# both the batched all_to_all packing and the cross-package pipeline. Sizes are picked -# so each matrix straddles a shard boundary (ZeRO-2/3), so gather/scatter is real. +# skip. Sizes are picked so each matrix straddles a shard boundary (ZeRO-2/3), so +# gather/scatter is real. SHAPE_SETS = [ pytest.param([(5, 8), (9, 8), (8, 6), (7, 8), (8,)], id="small"), # dims < 10 pytest.param([(200, 256), (400, 256), (256, 300), (256, 128), (256,)], id="medium"), @@ -93,12 +110,68 @@ def _read_full_weight(pg, idx, shape): return out +def _make_mesh(mesh_case, device): + """Return either the default world mesh or a rank-pair subgroup mesh.""" + global _RANK_PAIR_MESH + + if mesh_case == "world": + return None + + world_size = torch.distributed.get_world_size() + if world_size < 4 or world_size % 2 != 0: + pytest.skip("rank-pair mesh coverage requires an even world size >= 4") + + if _RANK_PAIR_MESH is not None: + return _RANK_PAIR_MESH + + rank = torch.distributed.get_rank() + selected_group, selected_ranks = None, None + for start in range(0, world_size, 2): + ranks = list(range(start, start + 2)) + group = torch.distributed.new_group(ranks=ranks) + if rank in ranks: + selected_group = group + selected_ranks = ranks + + assert selected_group is not None + # Ranks 2/3 use group ranks 0/1 in a group whose global ranks are 2/3. This + # catches P2P code that accidentally passes group ranks as global peers. + _RANK_PAIR_MESH = DeviceMesh.from_group( + [selected_group], + device_type=device.type, + mesh=selected_ranks, + mesh_dim_names=("dp",), + ) + return _RANK_PAIR_MESH + + +def _reference_orthogonalize(opt, grad): + """Mirror FullyShardV2Muon's root-side Newton-Schulz path.""" + grad = grad.to(torch.float32) + if HAVE_EMERGING_OPTIMIZERS: + orth = newton_schulz_tp( + grad, + steps=opt._num_ns_steps, + coefficient_type=opt._coefficient_type, + tp_group=None, + partition_dim=None, + tp_mode="duplicated", + ) + scale = get_muon_scale_factor(grad.size(-2), grad.size(-1), mode=opt._scale_mode) + else: + orth = _vanilla_newton_schulz(grad, steps=opt._num_ns_steps) + scale = _vanilla_muon_scale(grad.size(-2), grad.size(-1)) + return orth * (scale * opt._extra_scale_factor) + + @pytest.mark.parametrize("strategy", STRATEGIES) @pytest.mark.parametrize("nesterov", [True, False]) +@pytest.mark.parametrize("mesh_case", MESH_CASES) @pytest.mark.parametrize("shapes", SHAPE_SETS) -def test_dist_muon_matches_reference(strategy, nesterov, shapes): +def test_dist_muon_matches_reference(strategy, nesterov, mesh_case, shapes): rank = torch.distributed.get_rank() device = torch.device(f"cuda:{rank % torch.cuda.device_count()}") + mesh = _make_mesh(mesh_case, device) # shapes (parametrized): 2D matrices Muon updates + a trailing 1D bias it skips # (``len(s) == 2`` tells them apart). @@ -113,7 +186,7 @@ def test_dist_muon_matches_reference(strategy, nesterov, shapes): mp_policy=MixedPrecisionPolicy( main_params_dtype=torch.float32, main_grads_dtype=torch.float32 ), - mesh=None, + mesh=mesh, sharding_strategy=strategy, ) # Pass only the 2D-matrix dist_params + their grad DTensors (in production the @@ -125,14 +198,9 @@ def test_dist_muon_matches_reference(strategy, nesterov, shapes): if dp.dim() == 2: muon_params.append(dp) muon_grads.append(dg) - # Split the matrices into two multi-param packages to exercise both the batched - # all_to_all packing and the cross-package pipeline. - half = len(muon_params) // 2 - packages = [list(range(half)), list(range(half, len(muon_params)))] opt = FullyShardV2Muon( muon_params, muon_grads, - packages=packages, lr=lr, momentum=momentum, nesterov=nesterov, @@ -166,8 +234,12 @@ def test_dist_muon_matches_reference(strategy, nesterov, shapes): ref_buf[i].mul_(momentum).add_(full_grads[i]) # Mirror step()'s exact ops so the comparison can be bit-exact: the # nesterov look-ahead is a single alpha-add (g + m*buf), not g + (m*buf). - pre_ns = full_grads[i].add(ref_buf[i], alpha=momentum) if nesterov else ref_buf[i].clone() - orth = opt.orthogonalize(params[i], pre_ns) + pre_ns = ( + full_grads[i].add(ref_buf[i], alpha=momentum) + if nesterov + else ref_buf[i].clone() + ) + orth = _reference_orthogonalize(opt, pre_ns) if weight_decay != 0.0: ref_w[i].mul_(1.0 - lr * weight_decay) ref_w[i].add_(orth, alpha=-lr) @@ -180,8 +252,8 @@ def test_dist_muon_matches_reference(strategy, nesterov, shapes): for i, s in enumerate(shapes): got = _read_full_weight(pg, i, s) assert torch.equal(got, ref_w[i]), ( - f"[{strategy}] param {i} shape {s} step {step}: max abs diff " + f"[{mesh_case}/{strategy}] param {i} shape {s} step {step}: max abs diff " f"{(got - ref_w[i]).abs().max().item():.3e}" ) - + torch.distributed.barrier() From 1ff3e08f8c2e6f08458e56696f9fd0f3c91d79d8 Mon Sep 17 00:00:00 2001 From: tongliu Date: Wed, 10 Jun 2026 05:39:01 -0700 Subject: [PATCH 21/28] Add Muon momentum checkpoint support --- .../core/optimizer/fully_shard_v2_muon.py | 185 ++++++++++++++++-- .../megatron_fsdp/v2/test_dist_muon.py | 92 ++++++++- 2 files changed, 254 insertions(+), 23 deletions(-) diff --git a/megatron/core/optimizer/fully_shard_v2_muon.py b/megatron/core/optimizer/fully_shard_v2_muon.py index 176ff92846e..04472cc225c 100644 --- a/megatron/core/optimizer/fully_shard_v2_muon.py +++ b/megatron/core/optimizer/fully_shard_v2_muon.py @@ -21,6 +21,8 @@ import torch.distributed as dist from torch.distributed.tensor import DTensor +from megatron.core.distributed.fsdp.src.megatron_fsdp.uneven_dtensor import copy_chunk_metadata + from .emerging_optimizers import HAVE_EMERGING_OPTIMIZERS from .optimizer import MegatronOptimizer @@ -67,12 +69,12 @@ class MuonGradPackage: def update_momentum(self, opt, coef): for param_idx in self.param_ids: grad = opt._main_grads[param_idx] - momentum = opt._momentum_buffers[param_idx] if grad is None: continue grad_shard = grad.to_local().reshape(-1) if grad_shard.numel() == 0: continue + momentum = opt._momentum_buffers[param_idx] momentum.mul_(coef).add_(grad_shard) if opt._nesterov: torch.add(grad_shard, momentum, alpha=coef, out=grad_shard) @@ -102,7 +104,7 @@ def issue_gather(self, opt, full_grads): full_grad = torch.empty( main_weight.shape.numel(), dtype=momentum.dtype, - device=main_weight.device, + device=momentum.device, ) full_grads[param_idx] = full_grad if this_size > 0: @@ -237,6 +239,8 @@ class FullyShardV2Muon(torch.optim.Optimizer): grads: the gradient DTensor for each param, aligned 1:1 with the flattened params (asserted to be DTensors). This is the sole gradient source — replaces the old ParameterGroup back-reference path. + momentum_buffers: flat local momentum tensors owned by the outer Megatron + optimizer adapter and updated in-place by this inner optimizer. lr / momentum / nesterov / weight_decay: Muon update hyperparameters. num_ns_steps / coefficient_type / scale_mode / extra_scale_factor: NS + scaling config used by ``orthogonalize``. @@ -249,6 +253,7 @@ def __init__( self, params, grads, + momentum_buffers, *, lr: float, momentum: float = 0.95, @@ -261,7 +266,19 @@ def __init__( fp32_matmul_prec: str = "medium", tp_mode: str = "duplicated", ) -> None: - super().__init__(params, dict(lr=lr, weight_decay=weight_decay)) + super().__init__( + params, + dict( + lr=lr, + momentum=momentum, + nesterov=nesterov, + weight_decay=weight_decay, + num_ns_steps=num_ns_steps, + coefficient_type=coefficient_type, + scale_mode=scale_mode, + extra_scale_factor=extra_scale_factor, + ), + ) param_list = [param for group in self.param_groups for param in group["params"]] grad_list = list(grads) @@ -273,9 +290,20 @@ def __init__( assert all( grad is None or isinstance(grad, DTensor) for grad in grad_list ), "FullyShardV2Muon expects every grad to be a DTensor or None." + momentum_list = list(momentum_buffers) + assert len(momentum_list) == len(param_list), ( + f"FullyShardV2Muon got {len(momentum_list)} momentum buffers for " + f"{len(param_list)} params." + ) + for param_idx, (param, momentum_buffer) in enumerate(zip(param_list, momentum_list)): + assert momentum_buffer.numel() == param.to_local().numel(), ( + f"Muon momentum buffer for param {param_idx} has {momentum_buffer.numel()} " + f"elements, expected {param.to_local().numel()}." + ) self._main_weights = param_list self._main_grads = grad_list + self._momentum_buffers = momentum_list # Hyperparameters read by step() and MuonGradPackage methods. self._momentum_coef = momentum @@ -295,13 +323,6 @@ def __init__( self._shard_ranges = self._compute_shard_ranges() self._roots = self._assign_roots() - self._momentum_buffers = [] - for param_idx, (main_weight, grad) in enumerate(zip(self._main_weights, self._main_grads)): - _offset, local_size = self._shard_ranges[param_idx][self._comm[param_idx][1]] - dtype = grad.dtype if grad is not None else main_weight.dtype - momentum_buffer = torch.zeros(local_size, dtype=dtype, device=main_weight.device) - self._momentum_buffers.append(momentum_buffer) - self._ns_costs = [ main_weight.shape.numel() * min(main_weight.shape) if len(main_weight.shape) == 2 @@ -519,14 +540,18 @@ def __init__(self, config, model_chunks, **muon_hyperparams): # Build the inner optimizer from the FSDP-v2 2D-matrix dist_params + their grad # DTensors. The inner optimizer orders params by communication need. - params, grads = [], [] + params, grads, names = [], [], [] + momentum_buffers, momentum_dtensors = [], [] + seen_names = set() for chunk in self.model_chunks: root = chunk if isinstance(chunk, FSDPModule) else chunk.module - for m in root.modules(): + for module_name, m in root.named_modules(): if not isinstance(m, FSDPModule): continue - for param_group in m._fsdp_param_groups: - for dist_param, dist_grad in zip(param_group.dist_params, param_group.dist_grads): + for param_names, param_group in m._named_param_groups: + for param_name, dist_param, dist_grad in zip( + param_names, param_group.dist_params, param_group.dist_grads + ): # filter by the param's global attrs only (rank-consistent); keep # grad aligned even when None (empty shard) so collectives match if dist_param.dim() == 2 and not getattr( @@ -534,7 +559,37 @@ def __init__(self, config, model_chunks, **muon_hyperparams): ): params.append(dist_param) grads.append(dist_grad) - super().__init__(FullyShardV2Muon(params, grads, **muon_hyperparams), config) + full_name = ( + f"{module_name}.{param_name}" if module_name else param_name + ) + if full_name in seen_names: + raise ValueError( + f"Duplicate Muon parameter name for checkpointing: {full_name}" + ) + dtype = dist_grad.dtype if dist_grad is not None else dist_param.dtype + local_param = dist_param.to_local() + local_momentum = torch.zeros( + local_param.numel(), dtype=dtype, device=local_param.device + ) + momentum_dtensor = DTensor.from_local( + local_tensor=local_momentum.view(local_param.shape), + device_mesh=dist_param.device_mesh, + placements=dist_param.placements, + run_check=False, + shape=dist_param.shape, + stride=dist_param.stride(), + ) + copy_chunk_metadata(dist_param, momentum_dtensor) + names.append(full_name) + momentum_buffers.append(local_momentum) + momentum_dtensors.append(momentum_dtensor) + seen_names.add(full_name) + super().__init__( + FullyShardV2Muon(params, grads, momentum_buffers, **muon_hyperparams), config + ) + self._param_names = names + self._momentum_buffers = momentum_buffers + self._momentum_dtensors = momentum_dtensors self.is_stub_optimizer = False # --- Excluded from the chained grad-norm / clip / zero-count machinery. --- @@ -579,14 +634,102 @@ def get_loss_scale(self) -> torch.Tensor: def reload_model_params(self, state_dict=None): pass + def _param_name(self, param: torch.Tensor) -> str: + params = [p for group in self.optimizer.param_groups for p in group["params"]] + param_to_name = dict(zip(params, self._param_names)) + assert param in param_to_name, f"Muon parameter {param} is not named." + return param_to_name[param] + + def _param_groups_to_param2group_meta(self) -> dict[str, dict]: + param_to_group_meta = {} + for group in self.optimizer.param_groups: + group_meta = group.copy() + del group_meta["params"] + for param in group["params"]: + param_to_group_meta[self._param_name(param)] = group_meta + return param_to_group_meta + + def _sync_hyperparams_from_state_dict(self, state_dict: dict): + param_groups = state_dict.get("param_groups") + if param_groups is None and "param_to_group_meta" in state_dict: + param_to_group_meta = state_dict["param_to_group_meta"] + param_groups = [] + for group in self.optimizer.param_groups: + group_meta = None + for param in group["params"]: + name = self._param_name(param) + if name in param_to_group_meta: + group_meta = param_to_group_meta[name] + break + if group_meta is not None: + param_groups.append(group_meta) + if not param_groups: + return + + group = param_groups[0] + group = {key: value for key, value in group.items() if key != "params"} + self.optimizer.param_groups[0].update(group) + self.optimizer._momentum_coef = group.get("momentum", self.optimizer._momentum_coef) + self.optimizer._weight_decay = group.get("weight_decay", self.optimizer._weight_decay) + self.optimizer._nesterov = group.get("nesterov", self.optimizer._nesterov) + self.optimizer._num_ns_steps = group.get("num_ns_steps", self.optimizer._num_ns_steps) + self.optimizer._coefficient_type = group.get( + "coefficient_type", self.optimizer._coefficient_type + ) + self.optimizer._scale_mode = group.get("scale_mode", self.optimizer._scale_mode) + self.optimizer._extra_scale_factor = group.get( + "extra_scale_factor", self.optimizer._extra_scale_factor + ) + + def _load_momentum_state(self, state: dict): + named = not all(isinstance(key, int) for key in state.keys()) + for idx, momentum_buffer in enumerate(self._momentum_buffers): + key = self._param_names[idx] if named else idx + if key not in state: + momentum_buffer.zero_() + continue + loaded = state[key].get("momentum_buffer") + if loaded is None: + momentum_buffer.zero_() + continue + loaded_local = loaded.to_local() if isinstance(loaded, DTensor) else loaded + assert loaded_local.numel() == momentum_buffer.numel(), ( + f"Loaded Muon momentum for param {idx} has {loaded_local.numel()} " + f"elements, expected {momentum_buffer.numel()}." + ) + momentum_buffer.copy_( + loaded_local.reshape(-1).to( + device=momentum_buffer.device, dtype=momentum_buffer.dtype + ) + ) + def state_dict(self): - # TODO: checkpoint the sharded Muon momentum buffers. - return {} + param_groups = [] + next_param_id = 0 + for group in self.optimizer.param_groups: + group_state = {key: value for key, value in group.items() if key != "params"} + param_count = len(group["params"]) + group_state["params"] = list(range(next_param_id, next_param_id + param_count)) + next_param_id += param_count + param_groups.append(group_state) + return { + "state": { + idx: {"momentum_buffer": momentum} + for idx, momentum in enumerate(self._momentum_dtensors) + }, + "param_groups": param_groups, + } def load_state_dict(self, state_dict): - # TODO: restore the sharded Muon momentum buffers. - pass + self._sync_hyperparams_from_state_dict(state_dict) + self._load_momentum_state(state_dict["state"]) def sharded_state_dict(self, model_sharded_state_dict, is_loading: bool = False, **kwargs): - # TODO: sharded checkpoint support for the Muon momentum. - return {} + named_state = { + param_name: {"momentum_buffer": momentum} + for param_name, momentum in zip(self._param_names, self._momentum_dtensors) + } + return { + "state": named_state, + "param_to_group_meta": self._param_groups_to_param2group_meta(), + } diff --git a/tests/unit_tests/distributed/megatron_fsdp/v2/test_dist_muon.py b/tests/unit_tests/distributed/megatron_fsdp/v2/test_dist_muon.py index 769d5026c94..40aae5e8246 100644 --- a/tests/unit_tests/distributed/megatron_fsdp/v2/test_dist_muon.py +++ b/tests/unit_tests/distributed/megatron_fsdp/v2/test_dist_muon.py @@ -39,15 +39,18 @@ torchrun --nproc_per_node=4 -m pytest tests/unit_tests/distributed/megatron_fsdp/v2/test_dist_muon.py -v """ +import shutil import sys from pathlib import Path import pytest import torch +import torch.distributed.checkpoint as dcp import torch.nn as nn -from torch.distributed.tensor import DeviceMesh +from torch.distributed.tensor import DTensor, DeviceMesh sys.path.insert(0, str(Path(__file__).parents[2])) +from megatron.core.distributed.fsdp.src.megatron_fsdp.v2.fully_shard import fully_shard from megatron.core.distributed.fsdp.src.megatron_fsdp.v2.mixed_precision import ( MixedPrecisionPolicy, ) @@ -56,9 +59,11 @@ from megatron.core.optimizer.fully_shard_v2_muon import ( HAVE_EMERGING_OPTIMIZERS, FullyShardV2Muon, + FullyShardV2MuonOptimizer, _vanilla_muon_scale, _vanilla_newton_schulz, ) +from megatron.core.optimizer.optimizer_config import OptimizerConfig if HAVE_EMERGING_OPTIMIZERS: from megatron.core.optimizer.fully_shard_v2_muon import ( @@ -73,6 +78,7 @@ ] _RANK_PAIR_MESH = None +SHARED_TMP_DIR = "/tmp/pytest-shared-tmp" # Each set is a list of param shapes at a different scale: four 2D matrices (mixing # wide rowscols -> NS transposes) + a trailing 1D bias Muon must @@ -87,6 +93,16 @@ ] +class _MuonCheckpointModel(nn.Module): + def __init__(self): + super().__init__() + self.proj = nn.Linear(32, 32) + self.mlp = nn.Linear(32, 24, bias=False) + + def forward(self, x): + return self.mlp(self.proj(x)) + + @pytest.fixture(scope="session", autouse=True) def dist_env(): """Initialize NCCL process group once and tear down at session end.""" @@ -164,6 +180,73 @@ def _reference_orthogonalize(opt, grad): return orth * (scale * opt._extra_scale_factor) +def _build_muon_checkpoint_optimizer(device): + torch.manual_seed(4321) + model = _MuonCheckpointModel().to(device) + fully_shard( + model, + mp_policy=MixedPrecisionPolicy( + main_params_dtype=torch.float32, main_grads_dtype=torch.float32 + ), + sharding_strategy="optim_grads_params", + enable_async_reduce_grad=False, + ) + opt = FullyShardV2MuonOptimizer( + OptimizerConfig(lr=0.05, weight_decay=0.01), + [model], + lr=0.05, + momentum=0.9, + nesterov=False, + weight_decay=0.01, + num_ns_steps=5, + ) + assert opt._momentum_buffers, "checkpoint test must manage at least one Muon param" + return model, opt + + +def test_dist_muon_checkpoint_save_load_bitwise(): + rank = torch.distributed.get_rank() + device = torch.device(f"cuda:{rank % torch.cuda.device_count()}") + + _, save_opt = _build_muon_checkpoint_optimizer(device) + for param_idx, momentum_buffer in enumerate(save_opt._momentum_buffers): + values = torch.arange(momentum_buffer.numel(), device=device, dtype=torch.float32) + values = values.add_(rank * 1000 + param_idx * 17) + momentum_buffer.copy_(values.to(momentum_buffer.dtype)) + expected = [momentum_buffer.clone() for momentum_buffer in save_opt._momentum_buffers] + + save_state = save_opt.sharded_state_dict(model_sharded_state_dict={}) + for param_name, param_state in save_state["state"].items(): + assert isinstance(param_state["momentum_buffer"], DTensor), param_name + + ckpt_dir = Path(SHARED_TMP_DIR) / "test_dist_muon_checkpoint_save_load_bitwise" + if rank == 0: + shutil.rmtree(ckpt_dir, ignore_errors=True) + ckpt_dir.mkdir(parents=True, exist_ok=True) + torch.distributed.barrier() + + dcp.save({"optimizer": save_state}, checkpoint_id=str(ckpt_dir)) + torch.distributed.barrier() + + _, load_opt = _build_muon_checkpoint_optimizer(device) + for momentum_buffer in load_opt._momentum_buffers: + momentum_buffer.fill_(-1) + load_state = load_opt.sharded_state_dict(model_sharded_state_dict={}, is_loading=True) + dcp.load({"optimizer": load_state}, checkpoint_id=str(ckpt_dir)) + load_opt.load_state_dict(load_state) + + for param_idx, (actual, ref) in enumerate(zip(load_opt._momentum_buffers, expected)): + assert torch.equal(actual, ref), ( + f"Muon momentum checkpoint mismatch for param {param_idx}: " + f"max abs diff {(actual - ref).abs().max().item() if actual.numel() else 0.0}" + ) + + torch.distributed.barrier() + if rank == 0: + shutil.rmtree(ckpt_dir, ignore_errors=True) + torch.distributed.barrier() + + @pytest.mark.parametrize("strategy", STRATEGIES) @pytest.mark.parametrize("nesterov", [True, False]) @pytest.mark.parametrize("mesh_case", MESH_CASES) @@ -193,14 +276,19 @@ def test_dist_muon_matches_reference(strategy, nesterov, mesh_case, shapes): # FullyShardV2MuonOptimizer wrapper does this filtering). The 1D bias is excluded, # so the optimizer never touches it. grad DTensors may be None on ranks whose shard # of a param is empty. - muon_params, muon_grads = [], [] + muon_params, muon_grads, muon_momentum_buffers = [], [], [] for dp, dg in zip(pg.dist_params, pg.dist_grads): if dp.dim() == 2: muon_params.append(dp) muon_grads.append(dg) + dtype = dg.dtype if dg is not None else dp.dtype + muon_momentum_buffers.append( + torch.zeros(dp.to_local().numel(), dtype=dtype, device=device) + ) opt = FullyShardV2Muon( muon_params, muon_grads, + muon_momentum_buffers, lr=lr, momentum=momentum, nesterov=nesterov, From 8c3bea7b171f5ca215628d3b52f8de9823a3f3f2 Mon Sep 17 00:00:00 2001 From: tongliu Date: Thu, 11 Jun 2026 02:28:03 -0700 Subject: [PATCH 22/28] Skip no-op unshard for clean replicated buffers --- .../fsdp/src/megatron_fsdp/v2/dp_buffer.py | 5 +++++ .../fsdp/src/megatron_fsdp/v2/mixed_precision.py | 14 ++++++++------ 2 files changed, 13 insertions(+), 6 deletions(-) diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/dp_buffer.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/dp_buffer.py index f210dfae063..9f4bced8d11 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/dp_buffer.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/dp_buffer.py @@ -554,6 +554,11 @@ def unshard( """ full_buffer = self.fetch_buffer(as_shard=False) + if not self.is_distributed and not getattr(full_buffer, "_dirty", False): + if bind_params: + self._bind_buffer_to_params(full_buffer) + return full_buffer + sm = self.buffer_index.shard_meta shard_buffer = self.data[sm.local_data_index : sm.local_data_index + sm.size] torch.distributed.all_gather_into_tensor( diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/mixed_precision.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/mixed_precision.py index 51dd0bebd4c..022c3efbbb1 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/mixed_precision.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/mixed_precision.py @@ -614,12 +614,14 @@ def copy_main_weights_to_model_weights( fp8_params, main_params, start_offsets, data_parallel_group, model_param_shards ) - # Mark the model weights dirty so FSDP knows to all-gather them before - # the forward or backward pass. - if not model_weight_buffer.is_distributed: - model_weight_buffer.data._dirty = True - if transpose_weight_buffer is not None and not transpose_weight_buffer.is_distributed: - transpose_weight_buffer.data._dirty = True + # ZeRO-1/2 refresh only this rank's slice; gather before next compute. + def mark_dirty(buffer): + if buffer is not None and not buffer.is_distributed: + buffer.data._dirty = True + + if model_weight_buffer.sharding_strategy != "no_shard": + mark_dirty(model_weight_buffer) + mark_dirty(transpose_weight_buffer) def is_fp8_param(tensor: torch.Tensor) -> bool: From 421f50fcfde28b1331a13b9798c396dff8e85b6b Mon Sep 17 00:00:00 2001 From: tongliu Date: Thu, 11 Jun 2026 08:52:15 -0700 Subject: [PATCH 23/28] Initialize clean replicated FSDP weight buffers --- .../core/distributed/fsdp/src/megatron_fsdp/v2/dp_buffer.py | 2 ++ .../core/distributed/fsdp/src/megatron_fsdp/v2/param_group.py | 4 ++++ 2 files changed, 6 insertions(+) diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/dp_buffer.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/dp_buffer.py index 9f4bced8d11..31a1ded1cdd 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/dp_buffer.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/dp_buffer.py @@ -420,6 +420,8 @@ def init_data(self, data: torch.Tensor) -> None: assert data.dtype == self.dtype, f"dtype mismatch: {data.dtype} vs {self.dtype}" assert data.numel() == self.data_size, f"size mismatch: {data.numel()} vs {self.data_size}" self.data = data + if self.buffer_role in ("model_weight", "transpose_weight") and not self.is_distributed: + self.data._dirty = False def check_no_local_overlap(self, label: str = "") -> bool: """ diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/param_group.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/param_group.py index 7ae6be1f53d..fccfba8ce93 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/param_group.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/param_group.py @@ -194,6 +194,10 @@ def _init_buffers(self) -> None: ): _free_storage(tensor) + for weight_buffer in (self.model_weight_buffer, self.transpose_weight_buffer): + if weight_buffer is not None and not weight_buffer.is_distributed: + weight_buffer._bind_buffer_to_params(weight_buffer.data) + # Create gradient buffer if self.requires_grad: main_grads_dtype = self.mp_policy.main_grads_dtype_for_param(self.params[0]) From 8502038275c49d7a847ee67ba2841dc8f329c024 Mon Sep 17 00:00:00 2001 From: jianbinc Date: Fri, 12 Jun 2026 19:06:35 +0800 Subject: [PATCH 24/28] feat(fsdp-v2): lazy grad buffer, CPU offload infra, allocator release/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. --- examples/megatron_fsdp/fsdp_toy.py | 122 ++++++-- .../distributed_data_parallel_config.py | 9 + .../distributed/fsdp/mcore_fsdp_adapter.py | 2 +- .../fsdp/src/megatron_fsdp/uneven_dtensor.py | 35 ++- .../fsdp/src/megatron_fsdp/v2/allocator.py | 121 +++++++- .../fsdp/src/megatron_fsdp/v2/dp_buffer.py | 58 +++- .../fsdp/src/megatron_fsdp/v2/fsdp_module.py | 165 +++++++++- .../fsdp/src/megatron_fsdp/v2/hooks.py | 6 +- .../v2/lazy_grad_buffer_design.md | 283 ++++++++++++++++++ .../fsdp/src/megatron_fsdp/v2/param_group.py | 137 +++++++-- megatron/training/arguments.py | 10 + .../megatron_fsdp/v2/test_dist_muon.py | 1 + .../megatron_fsdp/v2/test_fully_shard.py | 4 - .../megatron_fsdp/v2/test_mcore_checkpoint.py | 4 +- .../v2/test_mcore_nd_parallel.py | 1 + .../megatron_fsdp/v2/test_param_group.py | 4 +- 16 files changed, 862 insertions(+), 100 deletions(-) create mode 100644 megatron/core/distributed/fsdp/src/megatron_fsdp/v2/lazy_grad_buffer_design.md diff --git a/examples/megatron_fsdp/fsdp_toy.py b/examples/megatron_fsdp/fsdp_toy.py index 3e224c8a9b6..946f06b84d0 100644 --- a/examples/megatron_fsdp/fsdp_toy.py +++ b/examples/megatron_fsdp/fsdp_toy.py @@ -14,16 +14,15 @@ import argparse import os -import sys +import time from pathlib import Path from typing import Tuple -import warnings import torch import torch.distributed as dist import torch.distributed.checkpoint as dcp import torch.nn as nn -from torch.distributed.checkpoint.state_dict import get_state_dict, set_state_dict +from torch.distributed.checkpoint.state_dict import set_state_dict from torch.distributed.checkpoint.stateful import Stateful from torch.distributed.device_mesh import init_device_mesh from torch.distributed.tensor import DTensor @@ -33,10 +32,15 @@ # ----------------------- class ToyBlock(nn.Module): - def __init__(self, dim: int): + """MLP-style block with expand→project, similar to transformer FFN.""" + + def __init__(self, dim: int, expansion: int = 4, dropout: float = 0.0): super().__init__() - self.linear1 = nn.Linear(dim, dim) - self.linear2 = nn.Linear(dim, dim) + hidden_dim = dim * expansion + self.gate = nn.Linear(dim, hidden_dim, bias=False) + self.up = nn.Linear(dim, hidden_dim, bias=False) + self.down = nn.Linear(hidden_dim, dim, bias=False) + self.dropout = nn.Dropout(dropout) if dropout else nn.Identity() self._use_activation_checkpointing = False def forward(self, x): @@ -45,7 +49,7 @@ def forward(self, x): return self._forward_impl(x) def _forward_impl(self, x): - return self.linear2(torch.relu(self.linear1(x))) + return self.dropout(self.down(torch.nn.functional.gelu(self.gate(x)) * self.up(x))) class ToyModel(nn.Module): @@ -85,19 +89,18 @@ def build_fsdp_model( use_megatron_fsdp: bool, use_activation_checkpointing: bool = False, enable_cuda_graph: bool = True, + enable_trace_pool: bool = False, ) -> Tuple["FSDPModule", torch.distributed.device_mesh.DeviceMesh]: if use_megatron_fsdp: try: - from megatron_fsdp.uneven_dtensor import get_state_dict from megatron_fsdp.v2 import FSDPModule, fully_shard - except ImportError: - from megatron.core.distributed.fsdp.src.megatron_fsdp.uneven_dtensor import get_state_dict + except (ImportError, ModuleNotFoundError) as err: from megatron.core.distributed.fsdp.src.megatron_fsdp.v2 import FSDPModule, fully_shard else: from torch.distributed.fsdp import FSDPModule, fully_shard mesh = init_distributed() - model = ToyModel(dim=dim, n_layers=n_layers).to("cuda") + model = ToyModel(dim=dim, n_layers=n_layers).to(device="cuda", dtype=torch.bfloat16) if use_activation_checkpointing: model.enable_activation_checkpointing() @@ -105,8 +108,11 @@ def build_fsdp_model( if use_megatron_fsdp: sublayer_kwargs = dict( enable_cuda_graph=enable_cuda_graph, + enable_trace_pool=enable_trace_pool, + ) + kwargs = dict( + enable_trace_pool=enable_trace_pool, ) - kwargs = {} else: sublayer_kwargs = {} kwargs = {} @@ -118,6 +124,9 @@ def build_fsdp_model( assert isinstance(model, ToyModel) assert isinstance(model, FSDPModule) + if use_megatron_fsdp: + model._log_parameter_groups() + p = next(model.parameters()) assert isinstance(p, DTensor) return model, mesh @@ -144,24 +153,22 @@ def state_dict(self): try: from megatron_fsdp.uneven_dtensor import get_state_dict except ImportError: - from megatron.core.distributed.fsdp.src.megatron_fsdp.uneven_dtensor import get_state_dict - except ImportError: - warnings.warn("Could not import get_state_dict from megatron_fsdp.uneven_dtensor.") + from megatron.core.distributed.fsdp.src.megatron_fsdp.uneven_dtensor import ( + get_state_dict, + ) - # this line automatically manages FSDP FQN's, as well as sets the default state dict type to FSDP.SHARDED_STATE_DICT model_state_dict, optimizer_state_dict = get_state_dict(self.model, self.optimizer) return { "model": model_state_dict, - "optim": optimizer_state_dict + "optim": optimizer_state_dict, } def load_state_dict(self, state_dict): - # sets our state dicts on the model and optimizer, now that we've loaded set_state_dict( self.model, self.optimizer, model_state_dict=state_dict["model"], - optim_state_dict=state_dict["optim"] + optim_state_dict=state_dict["optim"], ) @@ -172,6 +179,10 @@ def save_checkpoint( ckpt_dir: str, ) -> None: rank = dist.get_rank() + if rank == 0: + print(f"[rank0] Saving checkpoint step={step} ...") + t0 = time.time() + Path(ckpt_dir).mkdir(parents=True, exist_ok=True) ckpt_step_dir = os.path.join(ckpt_dir, f"step_{step:06d}") @@ -180,7 +191,7 @@ def save_checkpoint( dcp.save(state_dict=state, checkpoint_id=ckpt_step_dir) if rank == 0: - print(f"[rank0] Saved checkpoint to {ckpt_dir}") + print(f"[rank0] Saved checkpoint to {ckpt_dir} ({time.time() - t0:.1f}s)") def load_checkpoint_if_available( @@ -203,10 +214,18 @@ def load_checkpoint_if_available( if last_ckpt is None: return 0 + rank = dist.get_rank() + if rank == 0: + print(f"[rank0] Loading checkpoint from {last_ckpt} ...") + t0 = time.time() + step = torch.zeros([1]) state = {"app": AppState(model, optimizer), "step": step} dcp.load(state_dict=state, checkpoint_id=last_ckpt) + if rank == 0: + print(f"[rank0] Loaded checkpoint ({time.time() - t0:.1f}s)") + return int(step.item()) + 1 @@ -214,6 +233,14 @@ def load_checkpoint_if_available( # Training loop # ----------------------- +def _fmt_bytes(n: int) -> str: + for unit in ("B", "KB", "MB", "GB"): + if n < 1024: + return f"{n:.1f} {unit}" + n /= 1024 + return f"{n:.1f} TB" + + def train( args: argparse.Namespace, model: nn.Module, @@ -221,23 +248,36 @@ def train( start_step: int = 0, ) -> None: rank = dist.get_rank() - world_size = dist.get_world_size() model.train() step = start_step + t_start = time.time() for epoch in range(args.epochs): for _ in range(args.steps_per_epoch): # Dummy data - x = torch.randn(args.batch_size, args.model_dim, device="cuda") + x = torch.randn(args.batch_size, args.seq_len, args.model_dim, + device="cuda", dtype=torch.bfloat16) y = model(x) - loss = y.sum() / (world_size * args.batch_size) + loss = y.sum() / (args.batch_size * args.seq_len) loss.backward() + if args.use_megatron_fsdp and args.release_memory_pool: + model.release_memory_pool() optimizer.step() optimizer.zero_grad() + model.zero_grad() if step % args.log_interval == 0 and rank == 0: - print(f"[rank0] epoch={epoch} step={step} loss={loss.item():.4f}") + t_now = time.time() + elapsed = t_now - t_start + it_s = elapsed / max(step - start_step + 1, 1) + alloc = _fmt_bytes(torch.cuda.memory_allocated()) + max_reserved = _fmt_bytes(torch.cuda.max_memory_reserved()) + print( + f"[rank0] epoch={epoch} step={step} loss={loss.item():.4f} " + f"alloc={alloc} max_reserved={max_reserved} " + f"it={elapsed:.1f}s ({it_s * 1000:.0f}ms/it)" + ) if args.ckpt_dir and step % args.ckpt_interval == 0 and step > 0: save_checkpoint(model, optimizer, step, args.ckpt_dir) @@ -258,28 +298,46 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--model-dim", type=int, default=1024) parser.add_argument("--n-layers", type=int, default=3) parser.add_argument("--batch-size", type=int, default=8) + parser.add_argument("--seq-len", type=int, default=128, + help="Sequence length (larger = more compute vs communication)") parser.add_argument("--epochs", type=int, default=2) parser.add_argument("--steps-per-epoch", type=int, default=10) parser.add_argument("--lr", type=float, default=1e-3) - parser.add_argument("--ckpt-dir", type=str, default="checkpoints") + parser.add_argument("--ckpt-dir", type=str, default=None, + help="Directory for DCP checkpoints (omit to skip)") parser.add_argument("--ckpt-interval", type=int, default=20) parser.add_argument("--log-interval", type=int, default=5) parser.add_argument("--use-megatron-fsdp", action="store_true", help="Use Megatron-FSDP instead of PyTorch FSDP2") parser.add_argument("--activation-checkpoint", action="store_true", help="Enable activation checkpointing on transformer layers") - parser.add_argument("--cuda-graph", action="store_true", default=True, help="Enable CUDA graph capture (default: on)") - parser.add_argument("--no-cuda-graph", action="store_false", dest="cuda_graph", help="Disable CUDA graph capture") + parser.add_argument("--no-cuda-graph", action="store_false", dest="cuda_graph", + default=True, help="Disable CUDA graph capture") + parser.add_argument("--use-trace-pool", action="store_true", default=False, + help="Use TracePoolAllocator for stable buffer addresses") + parser.add_argument("--release-memory-pool", action="store_true", default=False, + help="Call FSDPModule.release_memory_pool() after each backward " + "to release allocator slot tensors and CUDA graphs. " + "Only takes effect with --use-megatron-fsdp.") + parser.add_argument("--record-memory-history", type=str, default=None, metavar="DIR", + help="Enable CUDA memory recording, dump snapshot to this directory") return parser.parse_args() def main() -> None: args = parse_args() + if args.record_memory_history: + torch.cuda.memory._record_memory_history( + max_entries=100000, + stacks="all", + ) + model, _ = build_fsdp_model( dim=args.model_dim, n_layers=args.n_layers, use_megatron_fsdp=args.use_megatron_fsdp, use_activation_checkpointing=args.activation_checkpoint, enable_cuda_graph=args.cuda_graph, + enable_trace_pool=args.use_trace_pool, ) optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr) @@ -289,6 +347,16 @@ def main() -> None: train(args, model, optimizer, start_step=start_step) + if args.record_memory_history: + rank = dist.get_rank() + out_dir = args.record_memory_history + Path(out_dir).mkdir(parents=True, exist_ok=True) + snapshot_path = os.path.join(out_dir, f"memory_snapshot_rank{rank}.pickle") + torch.cuda.memory._dump_snapshot(snapshot_path) + if rank == 0: + print(f"[rank0] Memory snapshot dumped to {snapshot_path}") + torch.cuda.memory._record_memory_history(enabled=None) + dist.barrier() dist.destroy_process_group() diff --git a/megatron/core/distributed/distributed_data_parallel_config.py b/megatron/core/distributed/distributed_data_parallel_config.py index 70445819e4f..4c56ddb84b0 100644 --- a/megatron/core/distributed/distributed_data_parallel_config.py +++ b/megatron/core/distributed/distributed_data_parallel_config.py @@ -142,6 +142,15 @@ class DistributedDataParallelConfig: This option will be automatically set to True when nccl_ub=True. """ + fsdp_trace_pool: bool = False + """If true, use TracePoolAllocator for Megatron FSDP v2 communication buffers + instead of the default StorageFreeingBucketAllocator. The TracePoolAllocator + traces the first micro-batch to build a static key-to-address plan, providing + stable buffer addresses required for CUDA graph capture. This flag only takes + effect in the Megatron FSDP v2 path. In the v1 path, use ``fsdp_double_buffer`` + for FixedPoolAllocator-based double buffering. + """ + fsdp_db_use_persist_buf_on_alloc_fail: bool = False """Whether to fall back to persistent buffer when a bucket does not fit FSDP double buffer size. If true, FSDP will use the persistently diff --git a/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py b/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py index d58afaa09d5..51e852dd15d 100644 --- a/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py +++ b/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py @@ -271,7 +271,7 @@ def _init_with_fully_shard( "mp_policy": fully_shard_mp_policy, "enable_unshard_prefetch": ddp_config.overlap_param_gather, "enable_async_reduce_grad": ddp_config.overlap_grad_reduce, - "enable_trace_pool": ddp_config.fsdp_double_buffer, + "enable_trace_pool": ddp_config.fsdp_double_buffer or ddp_config.fsdp_trace_pool, "sharding_strategy": ddp_config.data_parallel_sharding_strategy, } if config.calculate_per_token_loss: diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/uneven_dtensor.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/uneven_dtensor.py index e70d26ebf67..d1ae1c97101 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/uneven_dtensor.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/uneven_dtensor.py @@ -624,15 +624,44 @@ def get_state_dict( submodules: Optional[set[nn.Module]] = None, options: Optional["StateDictOptions"] = None, ) -> tuple[dict[str, "ValueType"], "OptimizerStateType"]: - """Produce model and optimizer state dicts with uneven DTensor preprocessing.""" + """Produce model and optimizer state dicts with uneven DTensor preprocessing. + + PyTorch's ``get_state_dict`` clones every DTensor, so the returned + tensors lack ``__create_chunk_list__`` / ``__create_write_items__``. + Instead of recomputing metadata via ``all_gather_object``, we copy it + from the model's ``dist_params`` (which carry metadata from FSDP init) + to the corresponding state-dict entries — a zero-collective operation. + """ for param in model.parameters(): assert isinstance(param, DTensor), "Expected all parameters to be DTensors" model_state_dict, optimizer_state_dict = _get_state_dict( model=model, optimizers=optimizers, submodules=submodules, options=options ) - preprocess_state_dict_for_uneven_dtensor(model_state_dict) - preprocess_state_dict_for_uneven_dtensor(optimizer_state_dict) + + # Build FQN → model DTensor mapping (dist_params carry chunk metadata). + param_by_fqn: dict[str, DTensor] = {} + for fqn, p in model.named_parameters(): + if isinstance(p, DTensor) and hasattr(p._local_tensor, "__create_chunk_list__"): + param_by_fqn[fqn] = p + + # Copy chunk metadata into model state-dict entries. + for fqn, dt in model_state_dict.items(): + if isinstance(dt, DTensor) and not hasattr(dt._local_tensor, "__create_chunk_list__"): + src = param_by_fqn.get(fqn) + if src is not None: + copy_chunk_metadata(src, dt) + + # Copy chunk metadata into optimizer state-dict entries. + optim_state = optimizer_state_dict.get("state", {}) + for fqn, state_tensors in optim_state.items(): + src = param_by_fqn.get(fqn) + if src is None: + continue + for key, dt in state_tensors.items(): + if isinstance(dt, DTensor) and not hasattr(dt._local_tensor, "__create_chunk_list__"): + copy_chunk_metadata(src, dt) + return model_state_dict, optimizer_state_dict diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/allocator.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/allocator.py index 6569dabc0ed..d5505bfe955 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/allocator.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/allocator.py @@ -229,6 +229,8 @@ def allocate( ) -> Bucket: key = _resolve_key(key, param_group_id) assert dtype is not None and device is not None + if self._phase == "released": + self._auto_resume() if self._phase != "optimized": return self._trace_allocate(key, size, dtype, device) else: @@ -241,6 +243,8 @@ def free( param_group_id: Optional[AllocatorKey] = None, ) -> None: key = _resolve_key(key, param_group_id) + if self._phase == "released": + self._auto_resume() if self._phase != "optimized": self._trace_free(key) else: @@ -352,12 +356,13 @@ def plan(self) -> int: self._phase = "optimized" - if torch.distributed.get_rank() == 0: - logger.debug( - f"TracePoolAllocator plan complete: {len(self._slots)} slots, " - f"{total_elems} total elements, " - f"{self.total_pool_bytes / 1024 / 1024:.1f} MB" - ) + if torch.distributed.is_available() and torch.distributed.is_initialized(): + if torch.distributed.get_rank() == 0: + logger.debug( + f"TracePoolAllocator plan complete: {len(self._slots)} slots, " + f"{total_elems} total elements, " + f"{self.total_pool_bytes / 1024 / 1024:.1f} MB" + ) return total_elems def _color_and_allocate_slots( @@ -485,6 +490,72 @@ def reset(self) -> None: self._key_to_slot.clear() self._key_to_view.clear() + def release(self) -> None: + """Release all slot tensor memory while preserving the slot plan. + + In ``"trace"`` phase this is equivalent to ``reset()`` — discards all + trace data and returns to a clean trace state. + + In ``"optimized"`` phase this drops every ``_SlotInfo.tensor`` reference + (freeing GPU memory) but retains the plan metadata: ``_slots`` + (size/dtype/device), ``_key_to_slot``, ``_key_to_view``, and the trace + history. The allocator transitions to ``"released"``. + + On the next ``allocate()`` or ``free()`` call the allocator + **automatically re-allocates** all slots and returns to ``"optimized"`` + — no explicit ``resume()`` call is needed. + """ + if self._phase == "released": + return + + if self._phase == "trace": + self.reset() + return + + assert self._phase == "optimized", ( + f"release() requires 'optimized' or 'trace' phase, got '{self._phase}'" + ) + for slot in self._slots: + _free_storage(slot.tensor) + slot.tensor = torch.empty(0, dtype=slot.dtype, device=slot.device) + slot.in_use = False + self._active_keys.clear() + self._phase = "released" + + def _auto_resume(self) -> None: + """Lazily re-allocate all slot tensors when the first ``allocate``/``free`` + arrives in the ``"released"`` phase. + + Internal — called automatically from ``allocate`` and ``free``. + """ + if self._phase != "released": + return + + for slot_idx, slot in enumerate(self._slots): + new_tensor = torch.empty(slot.size, dtype=slot.dtype, device=slot.device) + slot.tensor = new_tensor + + for key, slot_idx in self._key_to_slot.items(): + slot = self._slots[slot_idx] + meta = self._trace_meta.get(key) + if meta is not None: + size_k, _, _ = meta + else: + size_k = slot.size + self._key_to_view[key] = slot.tensor[: min(size_k, slot.size)] + + self._active_keys.clear() + self._phase = "optimized" + + def resume(self) -> None: + """Explicitly re-allocate slots and return to ``"optimized"`` phase. + + Normally you do not need to call this — the first ``allocate`` or + ``free`` after ``release()`` will auto-resume. Use this only when + you need to restore the pool before any alloc/free call. + """ + self._auto_resume() + @property def phase(self) -> str: return self._phase @@ -513,24 +584,34 @@ def dump_trace(self) -> str: f"{size_str} {dtype_str} {device_str}" ) - if self._phase == "optimized": - lines.append(f"\nslots: {len(self._slots)}") + if self._phase in ("optimized", "released"): + lines.append(f"\nslots: {len(self._slots)} ({self._phase})") for i, slot in enumerate(self._slots): keys_in_slot = [ k for k, idx in self._key_to_slot.items() if idx == i ] + if self._phase == "optimized": + addr_str = f"addr=0x{slot.tensor.data_ptr():x}" + else: + addr_str = "addr=" lines.append( f" slot[{i}]: size={slot.size} " f"dtype={slot.dtype} device={slot.device} " - f"addr=0x{slot.tensor.data_ptr():x} " + f"{addr_str} " f"{'IN_USE' if slot.in_use else 'free'} " f"keys={keys_in_slot}" ) - lines.append( - f"\ntotal pool: {len(self._slots)} slots, " - f"{self.total_pool_bytes} bytes " - f"({self.total_pool_bytes / 1024 / 1024:.1f} MB)" - ) + if self._phase == "optimized": + lines.append( + f"\ntotal pool: {len(self._slots)} slots, " + f"{self.total_pool_bytes} bytes " + f"({self.total_pool_bytes / 1024 / 1024:.1f} MB)" + ) + else: + lines.append( + f"\ntotal pool: {len(self._slots)} slots, " + f"memory released (call resume() to restore)" + ) return "\n".join(lines) @@ -581,10 +662,18 @@ def _intervals_overlap( return False +def _is_torchdynamo_compiling() -> bool: + """Check whether torchdynamo is compiling — safe across PyTorch versions.""" + try: + return torch.distributed._functional_collectives.is_torchdynamo_compiling() + except (AttributeError, RuntimeError): + return False + + def _free_storage(tensor: torch.Tensor) -> None: """Free the underlying storage of ``tensor`` by resizing it to 0.""" with torch.no_grad(): - if not torch.distributed._functional_collectives.is_torchdynamo_compiling(): + if not _is_torchdynamo_compiling(): already_freed = tensor._typed_storage()._size() == 0 if not already_freed: assert tensor.storage_offset() == 0, ( @@ -604,7 +693,7 @@ def _alloc_storage(tensor: torch.Tensor, size: torch.Size) -> None: existing shape. """ with torch.no_grad(): - if not torch.distributed._functional_collectives.is_torchdynamo_compiling(): + if not _is_torchdynamo_compiling(): already_allocated = tensor._typed_storage()._size() == size.numel() if not already_allocated: tensor_storage_size = tensor._typed_storage()._size() diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/dp_buffer.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/dp_buffer.py index 31a1ded1cdd..6d2d2bd226e 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/dp_buffer.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/dp_buffer.py @@ -19,7 +19,7 @@ import torch -from .allocator import BucketAllocator, TemporaryBucketAllocator +from .allocator import BucketAllocator, TemporaryBucketAllocator, _free_storage from .mixed_precision import MixedPrecisionPolicy from .utils import ParamGroupIdx @@ -423,6 +423,45 @@ def init_data(self, data: torch.Tensor) -> None: if self.buffer_role in ("model_weight", "transpose_weight") and not self.is_distributed: self.data._dirty = False + # ------------------------------------------------------------------ # + # CPU offload + # ------------------------------------------------------------------ # + + def _is_on_cpu(self) -> bool: + """True if ``self.data`` is resident on CPU.""" + return self.data is not None and self.data.device.type == "cpu" + + def _ensure_data_on_gpu(self) -> bool: + """Move ``self.data`` to GPU if currently on CPU. + + Returns True if a move happened (caller must rebuild dist views). + """ + if not self._is_on_cpu(): + return False + self.data = self.data.to(self.device, non_blocking=True) + return True + + def _move_data_to( + self, + target_device: torch.device, + pin_memory: bool = False, + non_blocking: bool = True, + ) -> None: + """Move ``self.data`` to *target_device*, optionally using pinned memory. + + Caller must call ``ParameterGroup._rebuild_dist_views()`` afterwards + because ``dist_params._local_tensor`` views share ``self.data`` Storage. + """ + if self.data is None or self.data.device == target_device: + return + if target_device.type == "cpu" and pin_memory: + cpu_data = torch.empty(self.data.shape, dtype=self.data.dtype, pin_memory=True) + cpu_data.copy_(self.data, non_blocking=non_blocking) + _free_storage(self.data) + self.data = cpu_data + else: + self.data = self.data.to(target_device, non_blocking=non_blocking) + def check_no_local_overlap(self, label: str = "") -> bool: """ Runtime check: verify no two items' local slices overlap within ``self.data``. @@ -648,6 +687,9 @@ def reduce_grad( If grad_comm_dtype differs from self.dtype, communicate with a temporary casted tensor and cast the reduced result back before accumulation. """ + if self.sharding_strategy in ("no_shard", "optim"): + overwrite_grad = True + grad_comm_dtype = grad_comm_dtype or self.dtype if self.gradient_scaling_factor in (None, 1.0): @@ -681,10 +723,6 @@ def reduce_grad( # accumulated into ``local_grad_shard`` for gradient accumulation. input_buffer = self.fetch_buffer() output_offset = sm.bucket_data_index - if overwrite_grad: - accumulate_output = False - else: - accumulate_output = True if input_buffer.is_cuda: # Keep temporary reduce-scatter buffers tied to the stream that uses them. input_buffer.record_stream(torch.cuda.current_stream()) @@ -695,8 +733,6 @@ def reduce_grad( # slice instead of accumulating into a separate shard buffer. input_buffer = self.data output_offset = sm.local_data_index - accumulate_output = False - overwrite_grad = False comm_input = ( input_buffer if grad_comm_dtype == self.dtype else input_buffer.to(grad_comm_dtype) @@ -709,12 +745,14 @@ def reduce_grad( output=reduced_grad_shard, input=comm_input, group=self.dp_group, op=op ) + # If the reduced shard is already in the local grad buffer, skip copy/accumulation. + if local_grad_shard.data_ptr() == reduced_grad_shard.data_ptr(): + return + if overwrite_grad: local_grad_shard.copy_(reduced_grad_shard) - elif accumulate_output: + else: local_grad_shard += reduced_grad_shard - elif grad_comm_dtype != self.dtype: - local_grad_shard.copy_(reduced_grad_shard) def check_all_fsdp_buffers(module) -> bool: diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/fsdp_module.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/fsdp_module.py index 807bca11bc3..704e95a62c6 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/fsdp_module.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/fsdp_module.py @@ -242,6 +242,156 @@ def enable_cuda_graph(self) -> None: """ self._fsdp_state.enable_cuda_graph = True + def release_memory_pool(self) -> None: + """Release all persistent communication-buffer memory and any CUDA graphs. + + Tears down captured CUDA graphs across all FSDP modules, clears graph + sentinels so they auto-recapture on the next forward pass, and releases + the ``TracePoolAllocator`` slot tensors. + + On the next ``allocate`` / ``free`` call the allocator **automatically** + re-allocates slots, so no explicit "resume" call is needed. CUDA graphs + are re-captured by the hooks on the next forward pass. + + Typical use: temporarily free GPU memory (e.g. for checkpoint I/O). + """ + ctx = self._fsdp_root_context + allocator = ctx.bucket_allocator + for module in self._get_fsdp_modules(recursive=True): + for pg in module._fsdp_param_groups: + pg.release_grad_buffer() + + if not isinstance(allocator, TracePoolAllocator): + return + + self._release_cuda_graphs(ctx) + self._clear_cuda_graph_sentinels(ctx) + allocator.release() + + # ---------------------------------------------------------------- + # Internal: CUDA graph teardown / sentinel helpers + # ---------------------------------------------------------------- + + @staticmethod + def _release_cuda_graphs(ctx: "_FSDPRootContext") -> None: + """Tear down all captured CUDA graphs on every FSDP module. + + Supports both the per-module runner path (``_fsdp_cg_runner``) and + the batch helper path (``_fsdp_cuda_graphs``, ``_fsdp_cg_runner`` sentinel). + Restores original ``forward`` methods before deleting graph objects. + """ + if not ctx.enable_cuda_graph: + return + + for module in ctx.forward_order: + if hasattr(module, "_fsdp_cg_runner"): + runner = module._fsdp_cg_runner + if hasattr(runner, "reset"): + runner.reset() + delattr(module, "_fsdp_cg_runner") + + ctx.cuda_graph_active = False + ctx.cuda_graph_stream = None + ctx.cuda_graph_pool = None + + @staticmethod + def _clear_cuda_graph_sentinels(ctx: "_FSDPRootContext") -> None: + """Clear CUDA graph sentinels so hooks will re-capture on next forward.""" + for module in ctx.forward_order: + if hasattr(module, "_fsdp_cg_runner"): + delattr(module, "_fsdp_cg_runner") + + # ---------------------------------------------------------------- + # CPU offload + # ---------------------------------------------------------------- + + def _get_fsdp_modules(self, recursive: bool = True) -> List["FSDPModule"]: + """Return ``[self]`` plus optionally all child ``FSDPModules``.""" + if not recursive: + return [self] + result = [self] + for _, child in self.named_modules(): + if child is not self and isinstance(child, FSDPModule): + result.append(child) + return result + + def offload_to_cpu( + self, + recursive: bool = True, + pin_memory: bool = False, + max_cpu_bytes: Optional[int] = None, + ) -> Dict[str, int]: + """Offload FSDP-held GPU memory to CPU within an optional budget. + + Moves every DataParallelBuffer.data to CPU and releases + TracePoolAllocator slot tensors. All buffers auto-reload on + next access — no explicit reload call needed. + + Buffers are offloaded in descending size order so the largest + buffers (most GPU savings) are prioritized when the budget is + tight. + + Args: + recursive: If True (default), also offloads child FSDPModules. + pin_memory: If True, allocate pinned CPU memory for faster + CPU↔GPU transfers (~12 GB/s via DMA vs ~3-6 GB/s pageable). + max_cpu_bytes: Maximum CPU memory (in bytes) to consume. + Buffers beyond this limit are left on GPU. ``None`` + means no limit (offload everything). The allocator + slots are always released and do not count against + this budget. + + Returns: + Dict with ``"offloaded_bytes"`` and ``"skipped_bytes"``. + """ + ctx = self._fsdp_root_context + + # Collect (buffer, nbytes) pairs, largest first + entries: List[Tuple[Any, int]] = [] + for module in self._get_fsdp_modules(recursive): + for pg in module._fsdp_param_groups: + for buf in (pg.model_weight_buffer, pg.transpose_weight_buffer, + pg.main_weight_buffer, pg.main_grad_buffer): + if buf is not None and buf.data is not None and not buf._is_on_cpu(): + entries.append((buf, buf.data.nbytes)) + entries.sort(key=lambda x: x[1], reverse=True) + + offloaded_bytes = 0 + skipped_bytes = 0 + for buf, nbytes in entries: + if max_cpu_bytes is not None and offloaded_bytes + nbytes > max_cpu_bytes: + skipped_bytes += nbytes + continue + buf._move_data_to(torch.device("cpu"), pin_memory=pin_memory) + offloaded_bytes += nbytes + + # Rebuild views after all moves + for module in self._get_fsdp_modules(recursive): + for pg in module._fsdp_param_groups: + pg._rebuild_dist_views() + + # Release allocator slots (always — no CPU cost) + if isinstance(ctx.bucket_allocator, TracePoolAllocator): + ctx.bucket_allocator.release() + + return {"offloaded_bytes": offloaded_bytes, "skipped_bytes": skipped_bytes} + + def reload_to_gpu(self, recursive: bool = True) -> None: + """Explicitly move all buffers back to GPU and rebuild views. + + Normally not needed — every access path auto-reloads. + Useful to hide first-touch CPU→GPU copy latency. + """ + for module in self._get_fsdp_modules(recursive): + for pg in module._fsdp_param_groups: + for buf in (pg.model_weight_buffer, pg.transpose_weight_buffer, + pg.main_weight_buffer, pg.main_grad_buffer): + if buf is not None: + buf._move_data_to( + torch.device(f"cuda:{torch.cuda.current_device()}") + ) + pg._rebuild_dist_views() + def _init_named_param_groups( self, mesh: Optional[DeviceMesh], @@ -599,6 +749,9 @@ def reduce_grad(self, async_op: bool = False): if not param_group.requires_grad: continue + # Initialize main gradient buffer and param -> main_grad mapping if not already done. + param_group._init_dist_grads() + # NaN check before reduction if getattr(self, "_enable_nan_checks", False): for name, param in zip(param_names, param_group.params): @@ -738,13 +891,15 @@ def _scale_gradients(self, scaling_factor: float): else: grad.mul_(scaling_factor) + def zero_grad(self, set_to_none: bool = True): + """Zero gradients for all parameter groups.""" + for child in self._get_fsdp_modules(recursive=True): + for param_group in child._fsdp_param_groups: + param_group.zero_grad(set_to_none=set_to_none) + def _zero_grad_buffer(self): """Zero the gradient buffer for all parameter groups.""" - for _, child in self.named_modules(): - if not isinstance(child, FSDPModule): - continue - for param_group in child._fsdp_param_groups: - param_group.zero_grad() + self.zero_grad() def _copy_main_weights_to_model_weights(self): """Copy main weight buffer to model weight buffer.""" diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/hooks.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/hooks.py index 98325cb267b..738f6820526 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/hooks.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/hooks.py @@ -60,6 +60,9 @@ def forward_pre_hook(hook_module, args, kwargs): fsdp_module.unshard(async_op=ctx.enable_unshard_prefetch, bwd_pass=True) fsdp_module.unshard(async_op=ctx.enable_unshard_prefetch, bwd_pass=False) + for param_group in fsdp_module._fsdp_param_groups: + param_group._maybe_free_grad_data() + # ---- CUDA graph capture (once per compatible module) -------------- if ( hook_module is fsdp_module @@ -226,7 +229,8 @@ def pre_backward_hook(module: FSDPModule, grads): # overwrite instead. setattr(param, "overwrite_main_grad", True) if param_group.main_grad_buffer is not None: - param_group.main_grad_buffer.unshard() + param_group._init_dist_grads() + param_group.main_grad_buffer.fetch_buffer() module._mfsdp_backward_pre_hook = create_custom_backward_hook( module, custom_backward_handler=pre_backward_hook diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/lazy_grad_buffer_design.md b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/lazy_grad_buffer_design.md new file mode 100644 index 00000000000..d6cd56928a1 --- /dev/null +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/lazy_grad_buffer_design.md @@ -0,0 +1,283 @@ +# Lazy main_grad_buffer Design for Megatron FSDP v2 + +## 1. Motivation + +In the original Megatron FSDP v2, `main_grad_buffer.data` was allocated **eagerly** +during `_init_buffers()` for every param group with `requires_grad=True`, and was +**never freed** for the lifetime of training. The buffer size is `numel × 4 bytes +(FP32) ÷ dp_size` per param group — collectively a significant fraction of GPU memory. + +For a 70B BF16 model with `dp=8`, this is **35 GB of permanently-resident GPU memory** +per rank that is only actually needed during the `reduce_grad()` → `optimizer.step()` +window. + +Additionally, for large models with frozen sublayers or expert params that +never receive gradients in a given micro-batch, this memory is wasted +permanently. + +**Goal**: defer allocation to the first backward pass, and free the buffer between +steps — matching the dynamic memory behavior of PyTorch FSDP2. + +### 1.1 Why PyTorch FSDP2 uses less memory + +PyTorch's `fully_shard` uses a **dynamic** +strategy for gradient buffers: + +- `_lazy_init()` defers allocation until the first backward pass +- After reduce-scatter, the full (unsharded) gradient buffer is immediately freed +- Only the local shard persists — and can be offloaded to CPU between steps + +This gives PyTorch FSDP2 a memory advantage during forward, checkpoint I/O, +and model export. + +--- + +## 2. Design + +### 2.1 Key Methods + +| Method | Location | Role | +|--------|----------|------| +| `_init_dist_grads()` | `param_group.py` | Lazy-allocate `main_grad_buffer.data` and rebuild `dist_grads` DTensors on first use | +| `_maybe_free_grad_data()` | `param_group.py` | Free `main_grad_buffer.data` when all params are zero-graded | +| `zero_grad()` | `param_group.py` | Called by `optim.zero_grad()` → triggers `_maybe_free_grad_data()` | +| `_rebuild_dist_views()` | `param_group.py` | Rebuild `_local_tensor` views after buffer changes device | +| `_ensure_buffers_on_gpu()` | `param_group.py` | Auto-reload any buffer from CPU back to GPU | + +### 2.2 `_init_dist_grads()` — Lazy Allocation + +```python +def _init_dist_grads(self) -> None: + gbuf = self.main_grad_buffer + if gbuf is None or not self.requires_grad: + return + if gbuf.data is not None: + return # already initialised + + gbuf.init_data(torch.zeros(gbuf.data_size, dtype=gbuf.dtype, device=self.device)) + + s = self.sharding_strategy + is_grad_shard = s in ("optim", "optim_grads", "optim_grads_params") + placements = [Shard(dim=0)] if is_grad_shard else [Replicate()] + + self.dist_grads = [] + for p, dist_param in zip(self.params, self.dist_params): + grad_data = gbuf.get_item(self.param_idx[p], as_shard=is_grad_shard) + if p.requires_grad and grad_data.numel() > 0: + self.dist_grads.append( + make_uneven_dtensor(grad_data, p.shape, self.mesh, placements) + ) + else: + self.dist_grads.append(None) + + self._grad_buffer_is_fresh = True +``` + +Key properties: +- **Idempotent**: if `gbuf.data` is already allocated, returns immediately +- **Called from multiple safe points**: backward pre-hook, `reduce_grad()`, post-backward callback +- **Uses `torch.zeros`** (not `torch.empty`): the first reduce-scatter accumulates + (`overwrite_grad=False`) into the buffer, so it must start from zero — see §5 + +### 2.3 `_maybe_free_grad_data()` — Memory Refresh + +```python +def _maybe_free_grad_data(self) -> None: + if self.main_grad_buffer is None or self.main_grad_buffer.data is None: + return + if any(getattr(p, "grad", None) is not None for p in self.dist_params): + return + self.main_grad_buffer.data = None + self.dist_grads = [None for _ in self.params] +``` + +Guarded: only frees when no `dist_param` still holds a `.grad` reference. + +### 2.4 Call Sites + +``` +┌──────────────────────────────────────────────────────────────┐ +│ _init_buffers() │ +│ → main_grad_buffer created but data = None (layout only) │ +│ → dist_grads = [None] │ +│ │ +│ zero_grad() ← optim.zero_grad() │ +│ → dist_param.grad = None │ +│ → _maybe_free_grad_data() ← frees buffer │ +│ → _grad_buffer_is_fresh = True │ +│ │ +│ backward_pre_hook │ +│ → _init_dist_grads() ← re-allocates buffer │ +│ → fetch_buffer() ← allocates full buf │ +│ │ +│ reduce_grad() │ +│ → _ensure_buffers_on_gpu() ← auto-reload safety │ +│ → _init_dist_grads() ← no-op (already init) │ +│ → main_grad_buffer.reduce_grad() ← reduce-scatter │ +│ │ +│ finish_grad_sync() │ +│ → param.main_grad = dist_grad │ +│ │ +│ optim.step() │ +│ → copies main_grad → optimizer param │ +│ → updates weights → copies back to model │ +└──────────────────────────────────────────────────────────────┘ +``` + +--- + +## 3. Memory Lifecycle + +``` +Memory + ^ + │ ┌──────────────────────┐ + │ │ _init_dist_grads() │ ┌──────────────────────┐ + │ ┌──────────────────────┤ allocate grad buffer ├──┤ _init_dist_grads() │ + │ │ _maybe_free_grad_ │ (torch.zeros) │ │ re-allocate │ + │ │ data() frees buffer │ │ │ │ + │ ├──────────────────────┘ │ ├──────────────────────┘ + │ │ ┌──────────┐ ┌──────────┐ │ │ ┌──────────┐ + │ │ │ model │ │ model │ │ │ │ model │ + │ │ │ weight │ │ weight │ │ │ │ weight │ + │ │ │ buffer │ │ buffer │ │ │ │ buffer │ + └──┴──┴──────────┴─────────┴──────────┴──────────┴──┴──┴──────────┴──────► time + init fwd bwd optim zero_grad fwd bwd optim + + │←── step N ──→│ │←── step N+1 ──→│ + + Gradient buffer: ALLOCATED only during backward → optimizer window + FREED during forward and between steps +``` + +### 3.1 Savings Table + +| Phase | Before (eager) | After (lazy) | Delta | +|-------|---------------|--------------|-------| +| Init / construction | gradient shard on GPU | 0 bytes | `−model_params × 4 / dp` | +| Forward | gradient shard on GPU | FREE | `−model_params × 4 / dp` | +| Backward | gradient shard on GPU | gradient shard on GPU | 0 | +| Optimizer step | gradient shard on GPU | gradient shard on GPU | 0 | +| Between steps | gradient shard on GPU | FREE | `−model_params × 4 / dp` | +| Checkpoint I/O | gradient shard on GPU | FREE | freed for I/O buffers | +| Model export | gradient shard on GPU | FREE | freed for export | + +**Example** (70B BF16 model, `dp=8`): +- Gradient shard: `70B × 4 bytes ÷ 8 = 35 GB` (FP32) +- **35 GB freed** during forward, between steps, and checkpoint I/O + +--- + +## 4. Edge Cases + +### Frozen params (`requires_grad=False`) + +`main_grad_buffer` is never created — the existing `if self.requires_grad` guard +in `_init_buffers` already handles this. No change. + +### Param groups with no gradient flow + +Some param groups (e.g., expert params on ranks that don't own the expert) may +never see a backward pass. With lazy init, the buffer is never allocated — +saves memory permanently. + +### Multiple `reduce_grad()` calls (gradient accumulation) + +After the first call, `_init_dist_grads()` is a no-op. Subsequent calls just +do the reduce as before. `overwrite_grad=False` ensures proper accumulation +across micro-batches. + +### Activation recomputation + +During recomputation, the forward pre-hook fires again. `_maybe_free_grad_data()` +is a no-op (data already None or grads are live). The backward pre-hook re-allocates +via `_init_dist_grads()` before backward compute starts — no change in behavior. + +### `offload_to_cpu()` + +If called before the first `reduce_grad()`, `main_grad_buffer.data` is `None` — +skip. If called after, the buffer exists on GPU and is offloaded normally. +`_ensure_buffers_on_gpu()` auto-reloads on next access. + +--- + +## 5. The `torch.empty` Bug — A Bitter Lesson + +### 5.1 What Went Wrong + +The initial implementation used `torch.empty()` instead of `torch.zeros()`: + +```python +# BROKEN — uninitialised GPU memory +gbuf.init_data(torch.empty(gbuf.data_size, dtype=gbuf.dtype, device=self.device)) +``` + +`torch.empty()` returns a tensor backed by **uninitialised GPU memory** — +arbitrary bit patterns left by prior allocations. `reduce_grad()` then +accumulates into this garbage: + +```python +# DataParallelBuffer.reduce_grad() +if overwrite_grad: + local_grad_shard.copy_(reduced_grad_shard) +elif accumulate_output: + local_grad_shard += reduced_grad_shard # ← ADDS to garbage → NaN +``` + +### 5.2 Why It Hung + +1. Micro-batch 1: `+= garbage → NaN` +2. Micro-batches 2–8: `+= NaN → NaN` (NaN propagates) +3. `optimizer.step()` produces NaN weights +4. `_copy_main_params_to_model_params()` copies NaN to model +5. **Hang**: subsequent CUDA operations on NaN tensors stall the GPU or + trigger NCCL edge-case hangs + +### 5.3 The Fix + +```python +# CORRECT — zero-initialised +gbuf.init_data(torch.zeros(gbuf.data_size, dtype=gbuf.dtype, device=self.device)) +``` + +`torch.zeros` ensures the first accumulate has a clean slate. Paired with +`_grad_buffer_is_fresh` → `overwrite_grad` for defense-in-depth. + +--- + +## 6. Compatibility with CPU Offload + +- `_ensure_buffers_on_gpu()` auto-reloads CPU-offloaded buffers before use +- `_rebuild_dist_views()` re-slices `_local_tensor` views after device move +- `_maybe_free_grad_data()` frees GPU allocation regardless of CPU copy state + +Combined: the gradient buffer is **fully elastic** — on GPU only when needed, +on CPU (or nowhere) the rest of the time. + +--- + +## 7. Impact on Callers + +| Call site | Impact | +|-----------|--------| +| `reduce_grad()` | Lazy init fires — adds one `torch.zeros` + DTensor rebuild on first call | +| `zero_grad()` | Already guarded — no-op if data is None; frees buffer via `_maybe_free_grad_data` | +| `_rebuild_dist_views()` | Already guarded — no-op for grads if buffer is None | +| `_copy_main_weights_to_model_weights()` | No impact (uses model_weight + main_weight only) | +| `_compute_per_param_norms()` | Reads `dist_grads` — skips if None | +| `finish_grad_sync()` | Only called after `reduce_grad()` → buffer already created | +| Checkpoint `state_dict()` | Optimizer state dict references `dist_grads` → None for lazy groups (correct) | + +--- + +## 8. Related Files + +| File | Relevant Changes | +|------|-----------------| +| `param_group.py` | `_init_dist_grads()`, `_maybe_free_grad_data()`, `_rebuild_dist_views()`, `_ensure_buffers_on_gpu()`, lazy `_init_buffers()` | +| `dp_buffer.py` | `_is_on_cpu()`, `_ensure_data_on_gpu()`, `_move_data_to()` | +| `allocator.py` | `release()`, `_auto_resume()`, `resume()` — pool lifecycle | +| `fsdp_module.py` | `zero_grad()` override, `_get_fsdp_modules()`, `offload_to_cpu()`, `reload_to_gpu()` | +| `hooks.py` | `_init_dist_grads()` + `fetch_buffer()` call sites in backward pre-hook | +| `design.md` | Memory-pool release/resume documentation | +| `cpu_offload_design.md` | Full CPU offload architecture | diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/param_group.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/param_group.py index fccfba8ce93..66e37e74dbb 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/param_group.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/param_group.py @@ -202,15 +202,11 @@ def _init_buffers(self) -> None: if self.requires_grad: main_grads_dtype = self.mp_policy.main_grads_dtype_for_param(self.params[0]) gbuf = self._create_buffer(main_grads_dtype, shard_grads, "main_grad") - gbuf.init_data(torch.zeros(gbuf.data_size, dtype=gbuf.dtype, device=self.device)) self.main_grad_buffer = gbuf # Create distributed parameter views self._init_dist_params() - # Initially, gradients are zeroed. - self.is_zero_grad = True - def unshard(self, bwd_pass: bool = False, bind_params: bool = True): """ Unshard model weights by all-gathering from sharded buffer. @@ -218,6 +214,8 @@ def unshard(self, bwd_pass: bool = False, bind_params: bool = True): After unshard, parameters point to full unsharded storage. FP8 parameters rebind their TE raw payload instead of ``param.data``. """ + self._ensure_buffers_on_gpu() + for weight_buffer in self.mp_policy.weight_buffers_for_unshard( self.model_weight_buffer, self.transpose_weight_buffer, bwd_pass=bwd_pass ): @@ -247,6 +245,7 @@ def reshard(self): @torch.no_grad() def copy_main_weights_to_model_weights(self): """Install optimized main weights into model compute weights.""" + self._ensure_buffers_on_gpu() self.mp_policy.copy_main_weights_to_model_weights( self.params, self.param_idx, @@ -264,16 +263,15 @@ def reduce_grad(self): ZeRO-1 keeps grads replicated during backward and reduce-scatters the replicated buffer once when the optimizer syncs. """ - # FIXME: When optimizer.zero_grad(set_to_none=True) is used, dist_param.grad - # becomes None, but the underlying grad buffer may still contain stale data - # from previous iterations. If overwrite_grad is not set to True, reduce_grad - # will accumulate into this stale buffer, potentially introducing NaNs and - # destabilizing training. + self._ensure_buffers_on_gpu() + # _grad_buffer_is_fresh is True after zero_grad() or lazy buffer init, + # so the first reduce_grad after either event overwrites instead of + # accumulating — no stale data from uninitialised or zeroed buffers. self.main_grad_buffer.reduce_grad( grad_comm_dtype=self.mp_policy.grad_comm_dtype, - overwrite_grad=False, + overwrite_grad=self._grad_buffer_is_fresh, ) - self.is_zero_grad = False + self._grad_buffer_is_fresh = False def release_grad_buffer(self): """Release the main gradient buffer to free memory.""" @@ -286,6 +284,24 @@ def release_grad_buffer(self): del param.main_grad self.main_grad_buffer.reshard() + def _maybe_free_grad_data(self) -> None: + """Drop ``main_grad_buffer.data`` if all params are zero-graded. + + After ``zero_grad()`` (or before the first backward), all + ``dist_param.grad`` are ``None``, so the gradient buffer holds no + meaningful data. Free the backing tensor — ``_init_dist_grads`` + will re-allocate on the next ``reduce_grad``. + """ + if self.main_grad_buffer is None or self.main_grad_buffer.data is None: + return + if any( + [getattr(p, "grad", None) is not None for p in self.dist_params] + + [getattr(p, "decoupled_grad", None) is not None for p in self.dist_params] + ): + return + self.main_grad_buffer.data = None + self.dist_grads = [None for _ in self.params] + def _init_dist_params(self): """ Initialize distributed parameter views (DTensors) into the buffers. @@ -297,7 +313,7 @@ def _init_dist_params(self): - "no_shard": replicated, no sharding (DDP-equivalent) """ self.dist_params = [] - self.dist_grads = [] + self.dist_grads = [] # placeholder, populated in _init_dist_grads s = self.sharding_strategy # Determine placement based on sharding strategy @@ -328,26 +344,39 @@ def _init_dist_params(self): setattr(param, "__fsdp_param__", True) setattr(dist_param, "__fsdp_param__", True) self.dist_params.append(dist_param) + self.dist_grads.append(None) # placeholder, will be set in _init_dist_grads # Update dist_param chunk metadata for checkpointing and debugging. for dist_param in self.dist_params: update_uneven_dtensor_chunk_metadata(dist_param) - # Create gradient DTensor views. Some groups, e.g. uint8 FP8 model - # payloads, do not require grads and therefore have no grad buffer. - if self.main_grad_buffer is None: - self.dist_grads = [None for _ in self.params] + def _init_dist_grads(self) -> None: + """Lazily allocate ``main_grad_buffer.data`` and rebuild ``dist_grads``. + + The buffer layout (``BufferIndex``, offsets, shard) was created in + ``_init_buffers``; only the backing tensor is deferred. Called from + ``reduce_grad()`` on first use. Uses ``torch.empty`` to avoid the + zero-init cost; ``_grad_buffer_is_fresh`` is ``True`` after allocation + so the first reduce-scatter *overwrites* (``local_grad_shard.copy_``) + rather than accumulating — the uninitialized data is never read. + Subsequent calls are no-ops. + """ + gbuf = self.main_grad_buffer + if gbuf is None or not self.requires_grad: return + if gbuf.data is not None: + return # already initialised - is_grad_shard = is_param_shard - for p in self.params: - gbuf = self.main_grad_buffer + gbuf.init_data(torch.empty(gbuf.data_size, dtype=gbuf.dtype, device=self.device)) + + # Rebuild dist_grads views — dist_params are unchanged + s = self.sharding_strategy + is_grad_shard = s in ("optim", "optim_grads", "optim_grads_params") + placements = [Shard(dim=0)] if is_grad_shard else [Replicate()] + + self.dist_grads = [] + for p, dist_param in zip(self.params, self.dist_params): grad_data = gbuf.get_item(self.param_idx[p], as_shard=is_grad_shard) - # NOTE: Do not remove the grad_data.numel() > 0 check. - # Empty local grad shards are semantically no-ops, but materializing - # them as DTensor grads can pass zero-numel tensors into fused - # multi-tensor optimizers such as TE FusedAdam. That can break - # updates for neighboring non-empty shards. if p.requires_grad and grad_data.numel() > 0: self.dist_grads.append( make_uneven_dtensor(grad_data, p.shape, self.mesh, placements) @@ -355,15 +384,65 @@ def _init_dist_params(self): else: self.dist_grads.append(None) + self._grad_buffer_is_fresh = True + + def _rebuild_dist_views(self) -> None: + """In-place update ``dist_params._local_tensor`` / ``dist_grad._local_tensor``. + + Called after any buffer's ``self.data`` changes device (offload_to_cpu / + auto-reload). Updates the ``_local_tensor`` attribute inside existing + DTensor objects so optimizer references remain valid. + """ + s = self.sharding_strategy + is_param_shard = s in ("optim", "optim_grads", "optim_grads_params") + + for i, param in enumerate(self.params): + dist_param = self.dist_params[i] + if dist_param is not None: + if self.main_weight_buffer is not None: + data = self.main_weight_buffer.get_item(self.param_idx[param], + as_shard=is_param_shard) + elif self.model_weight_buffer is not None: + data = self.model_weight_buffer.get_item(self.param_idx[param], + as_shard=is_param_shard) + else: + continue + object.__setattr__(dist_param._local_tensor, 'data', data) + + if self.main_grad_buffer is not None and self.main_grad_buffer.data is not None: + is_grad_shard = is_param_shard + for i, param in enumerate(self.params): + dist_grad = self.dist_grads[i] + if dist_grad is not None: + grad_data = self.main_grad_buffer.get_item( + self.param_idx[param], as_shard=is_grad_shard + ) + object.__setattr__(dist_grad._local_tensor, 'data', grad_data) + + def _ensure_buffers_on_gpu(self) -> bool: + """Auto-reload any buffer on CPU back to GPU. + + Returns True if any buffer was moved (views were rebuilt). + """ + moved = False + for buf in (self.model_weight_buffer, self.main_weight_buffer, + self.main_grad_buffer, self.transpose_weight_buffer): + if buf is not None and buf._ensure_data_on_gpu(): + moved = True + if moved: + self._rebuild_dist_views() + return moved + def zero_grad(self, set_to_none: bool = True): """Zero the main gradient buffer and mark grads as zeroed.""" - self.release_grad_buffer() - if self.main_grad_buffer is not None: - self.main_grad_buffer.data.zero_() if set_to_none: for dist_param in self.dist_params: if dist_param.grad is not None: - del dist_param.grad + dist_param.grad = None if hasattr(dist_param, "decoupled_grad"): dist_param.decoupled_grad = None - self.is_zero_grad = True + self._maybe_free_grad_data() + else: + if self.main_grad_buffer is not None and self.main_grad_buffer.data is not None: + self.main_grad_buffer.data.zero_() + self._grad_buffer_is_fresh = True diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index 89b89f49a77..42e09d37bf2 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -3962,6 +3962,16 @@ def _add_distributed_args(parser): "reusing previously allocated buffers, rather than creating new buffers for each FSDP communication. " "This is required for user buffer registration and is enabled by default when using NCCL user buffers.", ) + group.add_argument( + '--fsdp-trace-pool', + action='store_true', + dest='fsdp_trace_pool', + help="Enable TracePoolAllocator for Megatron FSDP v2 communication buffers. " + "The TracePoolAllocator profiles one micro-batch to record allocation patterns, " + "then builds a static key-to-address plan that is reused across subsequent micro-batches. " + "This provides stable buffer addresses required for CUDA graph capture. " + "This flag only takes effect in the Megatron FSDP v2 path.", + ) group.add_argument( '--suggested-communication-unit-size', type=int, diff --git a/tests/unit_tests/distributed/megatron_fsdp/v2/test_dist_muon.py b/tests/unit_tests/distributed/megatron_fsdp/v2/test_dist_muon.py index 40aae5e8246..9378ee4f062 100644 --- a/tests/unit_tests/distributed/megatron_fsdp/v2/test_dist_muon.py +++ b/tests/unit_tests/distributed/megatron_fsdp/v2/test_dist_muon.py @@ -272,6 +272,7 @@ def test_dist_muon_matches_reference(strategy, nesterov, mesh_case, shapes): mesh=mesh, sharding_strategy=strategy, ) + pg._init_dist_grads() # Pass only the 2D-matrix dist_params + their grad DTensors (in production the # FullyShardV2MuonOptimizer wrapper does this filtering). The 1D bias is excluded, # so the optimizer never touches it. grad DTensors may be None on ranks whose shard diff --git a/tests/unit_tests/distributed/megatron_fsdp/v2/test_fully_shard.py b/tests/unit_tests/distributed/megatron_fsdp/v2/test_fully_shard.py index 82984f3563b..0958d60bea6 100644 --- a/tests/unit_tests/distributed/megatron_fsdp/v2/test_fully_shard.py +++ b/tests/unit_tests/distributed/megatron_fsdp/v2/test_fully_shard.py @@ -882,7 +882,6 @@ def test_no_params_module_ok(self): class TestCheckpoint: - @pytest.mark.skip(reason="Hangs. Debug in progress.") def test_get_state_dict_returns_dicts(self): """get_state_dict should return model and optimizer state dicts.""" torch.manual_seed(42) @@ -902,7 +901,6 @@ def test_get_state_dict_returns_dicts(self): assert isinstance(opt_sd, dict) assert len(model_sd) > 0, "Model state dict should not be empty" - @pytest.mark.skip(reason="Hangs. Debug in progress.") def test_get_state_dict_nested_fsdp(self): """get_state_dict should work with nested FSDP modules.""" torch.manual_seed(42) @@ -948,7 +946,6 @@ def test_get_state_dict_strict_all_dtensor(self): with pytest.raises(AssertionError, match="Expected all parameters to be DTensors"): get_state_dict(model, optimizer) - @pytest.mark.skip(reason="Hangs. Debug in progress.") def test_get_state_dict_llm_scenario(self): """Full LLM forward-backward-checkpoint cycle should work.""" torch.manual_seed(42) @@ -969,7 +966,6 @@ def test_get_state_dict_llm_scenario(self): assert len(model_sd) > 0 assert len(opt_sd) > 0 - @pytest.mark.skip(reason="Hangs. Debug in progress.") def test_get_state_dict_with_frozen_params(self): """get_state_dict should work with mixed frozen/trainable params.""" torch.manual_seed(42) diff --git a/tests/unit_tests/distributed/megatron_fsdp/v2/test_mcore_checkpoint.py b/tests/unit_tests/distributed/megatron_fsdp/v2/test_mcore_checkpoint.py index 812886a1441..a8c6ab2522c 100644 --- a/tests/unit_tests/distributed/megatron_fsdp/v2/test_mcore_checkpoint.py +++ b/tests/unit_tests/distributed/megatron_fsdp/v2/test_mcore_checkpoint.py @@ -406,8 +406,8 @@ def teardown_class(cls): # ---- MFSDP v2 → MFSDP v2 (round-trip) ---- pytest.param( "v2", - dict(data_parallel_sharding_strategy="optim_grads_params"), - dict(data_parallel_sharding_strategy="optim_grads_params"), + dict(data_parallel_sharding_strategy="optim_grads_params", fsdp_trace_pool=True), + dict(data_parallel_sharding_strategy="optim_grads_params", fsdp_trace_pool=True), id="v2_rt_optim_grads_params", ), pytest.param( diff --git a/tests/unit_tests/distributed/megatron_fsdp/v2/test_mcore_nd_parallel.py b/tests/unit_tests/distributed/megatron_fsdp/v2/test_mcore_nd_parallel.py index 7ee9f648b5d..054b01e7903 100644 --- a/tests/unit_tests/distributed/megatron_fsdp/v2/test_mcore_nd_parallel.py +++ b/tests/unit_tests/distributed/megatron_fsdp/v2/test_mcore_nd_parallel.py @@ -257,6 +257,7 @@ def _training_loop(seed=42, **kwargs): overlap_grad_reduce=True, use_megatron_fsdp_v2=True, gradient_accumulation_fusion=True, + fsdp_trace_pool=True, ), id="optim_grads_params_double_buffer", ), diff --git a/tests/unit_tests/distributed/megatron_fsdp/v2/test_param_group.py b/tests/unit_tests/distributed/megatron_fsdp/v2/test_param_group.py index 7c3957fc150..6c400931b9d 100644 --- a/tests/unit_tests/distributed/megatron_fsdp/v2/test_param_group.py +++ b/tests/unit_tests/distributed/megatron_fsdp/v2/test_param_group.py @@ -198,8 +198,7 @@ def test_init_buffers(strategy): if pg.requires_grad: assert pg.main_grad_buffer is not None assert pg.main_grad_buffer.is_distributed == g_dist - # Grad buffer must be zero-initialized - assert torch.all(pg.main_grad_buffer.data == 0) + assert pg.main_grad_buffer.data is None # lazy init torch.distributed.barrier() @@ -264,6 +263,7 @@ def test_reduce_grad(strategy): _, _, _, g_dist = _flags(strategy) for pg in groups: + pg._init_dist_grads() # lazily allocate grad buffer and dist_grads list gbuf = pg.main_grad_buffer if gbuf is None: # uint8 group has requires_grad=False, so no grad buffer From 0436cbd032ea550828311b2f77af9cec88f33e1a Mon Sep 17 00:00:00 2001 From: Jianbin Chang Date: Tue, 23 Jun 2026 20:45:00 +0800 Subject: [PATCH 25/28] Integrate M-FSDP v2 with 1F1B pipeline overlap; add lazy grad buffers, 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. --- .../distributed/fsdp/mcore_fsdp_adapter.py | 31 +- .../fsdp/src/megatron_fsdp/megatron_fsdp.py | 51 +- .../v2/design/1f1b_ep_overlap_fsdp_design.md | 544 ++++++++++++++ .../v2/{ => design}/allocator_design.md | 0 .../v2/{ => design}/cuda_graph_design.md | 4 +- .../megatron_fsdp/v2/{ => design}/design.md | 2 +- .../src/megatron_fsdp/v2/design/hooks_api.md | 136 ++++ .../{ => design}/lazy_grad_buffer_design.md | 0 .../mcore_fsdp_checkpoint_design.md | 24 +- .../v2/{ => design}/nvfp4_design.md | 4 +- .../v2/{ => design}/tp_support_design.md | 0 .../fsdp/src/megatron_fsdp/v2/dp_buffer.py | 9 +- .../fsdp/src/megatron_fsdp/v2/fsdp_module.py | 15 + .../fsdp/src/megatron_fsdp/v2/fully_shard.py | 40 +- .../fsdp/src/megatron_fsdp/v2/hooks.py | 676 +++++++++++------- .../fsdp/src/megatron_fsdp/v2/param_group.py | 1 - megatron/core/optimizer/clip_grads.py | 10 +- .../core/pipeline_parallel/combined_1f1b.py | 207 +++++- megatron/training/utils.py | 5 +- .../distributed/megatron_fsdp/utils.py | 12 +- .../v2/test_mcore_nd_parallel.py | 38 +- 21 files changed, 1473 insertions(+), 336 deletions(-) create mode 100644 megatron/core/distributed/fsdp/src/megatron_fsdp/v2/design/1f1b_ep_overlap_fsdp_design.md rename megatron/core/distributed/fsdp/src/megatron_fsdp/v2/{ => design}/allocator_design.md (100%) rename megatron/core/distributed/fsdp/src/megatron_fsdp/v2/{ => design}/cuda_graph_design.md (99%) rename megatron/core/distributed/fsdp/src/megatron_fsdp/v2/{ => design}/design.md (99%) create mode 100644 megatron/core/distributed/fsdp/src/megatron_fsdp/v2/design/hooks_api.md rename megatron/core/distributed/fsdp/src/megatron_fsdp/v2/{ => design}/lazy_grad_buffer_design.md (100%) rename megatron/core/distributed/fsdp/src/megatron_fsdp/v2/{ => design}/mcore_fsdp_checkpoint_design.md (97%) rename megatron/core/distributed/fsdp/src/megatron_fsdp/v2/{ => design}/nvfp4_design.md (99%) rename megatron/core/distributed/fsdp/src/megatron_fsdp/v2/{ => design}/tp_support_design.md (100%) diff --git a/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py b/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py index 51e852dd15d..6f389b11058 100644 --- a/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py +++ b/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py @@ -179,6 +179,27 @@ def __init__( self._annotate_tensor_parallelism(module) + if config.overlap_moe_expert_parallel_comm: + assert not ddp_config.fsdp_double_buffer, ( + "1F1B overlap with FSDP does not support double buffer. " + "Please set fsdp_double_buffer=False in the ddp config." + ) + assert config.cuda_graph_impl in ("none", "full_iteration"), ( + "1F1B overlap with FSDP does not support per-layer CUDA graphs " + f"(cuda_graph_impl={config.cuda_graph_impl!r}). " + "Use cuda_graph_impl='full_iteration' or disable CUDA graphs " + "(cuda_graph_impl='none')." + ) + + if ( + config.overlap_moe_expert_parallel_comm + and ddp_config.data_parallel_sharding_strategy == "optim_grads_params" + ): + assert self.fsdp_unit_modules == [TransformerLayer], ( + "EP overlap with FSDP currently requires fsdp_unit_modules " + f"to be [TransformerLayer], got {self.fsdp_unit_modules}." + ) + super().__init__( config=config, module=MegatronFSDP( @@ -192,7 +213,12 @@ def __init__( calculate_per_token_loss=config.calculate_per_token_loss, init_model_with_meta_device=config.init_model_with_meta_device, enable_fine_grained_param_gather_hook=( - config.fp8_recipe == "mxfp8" and ddp_config.fp8_param_gather + (config.fp8_recipe == "mxfp8" and ddp_config.fp8_param_gather) + or config.overlap_moe_expert_parallel_comm + ), + enable_fine_grained_param_gather_backward_hook=( + config.overlap_moe_expert_parallel_comm + and ddp_config.data_parallel_sharding_strategy == "optim_grads_params" ), ), ) @@ -273,6 +299,9 @@ def _init_with_fully_shard( "enable_async_reduce_grad": ddp_config.overlap_grad_reduce, "enable_trace_pool": ddp_config.fsdp_double_buffer or ddp_config.fsdp_trace_pool, "sharding_strategy": ddp_config.data_parallel_sharding_strategy, + "fine_grained_hooks": config.overlap_moe_expert_parallel_comm, + "skip_backward_callback": config.delay_wgrad_compute, + "skip_final_backward_callback": config.overlap_moe_expert_parallel_comm, } if config.calculate_per_token_loss: gradient_scaling_factor = None diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py index b589a5eece3..3b46521e45b 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py @@ -17,6 +17,7 @@ import logging from contextlib import contextmanager from enum import Enum, auto +from functools import partial from typing import Any, Dict, List, Optional, Tuple import torch @@ -46,11 +47,12 @@ try: # Default to Megatron-LM FW. - logger.info("Detected Megatron Core, using Megatron-FSDP with Megatron.") from megatron.core.distributed.distributed_data_parallel_config import ( DistributedDataParallelConfig, ) from megatron.core.utils import is_submodule + + logger.info("Detected Megatron Core, using Megatron-FSDP with Megatron.") except ImportError: # Megatron-LM is not installed, use Megatron-FSDP as a standalone module. logger.info("Megatron Core is not installed, Megatron-FSDP will run without Megatron Core.") @@ -214,6 +216,7 @@ def __init__( fsdp_db_use_persist_buf_on_alloc_fail: bool = False, disable_symmetric_registration: bool = False, enable_fine_grained_param_gather_hook: bool = False, + enable_fine_grained_param_gather_backward_hook: bool = False, report_nan_in_param_grad: bool = False, ): super().__init__() @@ -266,6 +269,9 @@ def __init__( self.calculate_per_token_loss = calculate_per_token_loss self.init_model_with_meta_device = init_model_with_meta_device self.enable_fine_grained_param_gather_hook = enable_fine_grained_param_gather_hook + self.enable_fine_grained_param_gather_backward_hook = ( + enable_fine_grained_param_gather_backward_hook + ) self.report_nan_in_param_grad = report_nan_in_param_grad # FSDPDistributedIndex stores the process groups and meshes used by Megatron-FSDP. @@ -693,18 +699,12 @@ def _process_post_backward_gradients(param_list): # Filter out shared parameters whose gradients are handled by the root hook. param_list = [p for p in param_list if not getattr(p, "_is_shared", False)] - # Filter out parameters whose gradient processing is deferred to a delayed - # wgrad accumulation hook (post_wgrad_grad_acc_hook). If skip_backward_post_hook - # is set but the delayed hook was never installed, process the parameter - # immediately as a safety fallback to avoid silently dropping gradients. - param_list = [ - p - for p in param_list - if not ( - getattr(p, 'skip_backward_post_hook', False) - and hasattr(p, 'post_wgrad_grad_acc_hook') - ) - ] + # Make sure for delayed wgrad params, the grad_acc_hooks are registered. + for p in param_list: + if getattr(p, 'skip_backward_post_hook', False): + assert hasattr( + p, 'post_wgrad_grad_acc_hook' + ), "Missing grad accumulation hook for delayed_wgrad_compute param." if not param_list: return @@ -884,7 +884,7 @@ def _pre_backward_param_unshard(module: nn.Module, *unused): self._root_pre_backward_hook_issued = False - def _root_pre_backward(module: nn.Module, *unused): + def _root_pre_backward(module: nn.Module, *unused, skip_backward_hook: bool = False): """Marks the module's training state as PRE_BACKWARD before the backprop, this function is registered on the root module. @@ -920,6 +920,8 @@ def _root_pre_backward(module: nn.Module, *unused): param.grad_added_to_main_grad = False # Queue the root post-backward hook to reduce leftover gradients after # the backward pass. + if skip_backward_hook: + return torch.autograd.Variable._execution_engine.queue_callback(_root_post_backward) @torch.compiler.disable @@ -1007,12 +1009,24 @@ def _register_pre_backward_param_unshard_hook(module): create_custom_backward_hook(module, _pre_backward_param_unshard) ) + # These hooks need to be exposed for manual management by 1F1B Overlapping + # and triggered by 1F1B Overlapped execution pipeline, except for + # `param_unshard` hook that needs to be installed at param level, + # such that non-overlapped params like embedding layer are also correctly + # unsharded. + self.post_forward_release_module = partial(_post_forward, input=None, output=None) + self.post_backward_release_module = _post_backward_release_module + self.pre_backward = partial(_root_pre_backward, module=None, skip_backward_hook=True) + self.post_backward = _root_post_backward + fsdp_modules = [] for name, module in root_module.named_modules(): # Set post backward hook for TE grouped gemm if enabled comm overlap setup_delayed_wgrad_acc_hook(module, _process_post_backward_gradients) if self.enable_fine_grained_param_gather_hook: _register_pre_forward_param_unshard_hook(module) + if self.enable_fine_grained_param_gather_backward_hook: + _register_pre_backward_param_unshard_hook(module) # Skip if the module is already registered in fsdp_modules. if any(is_submodule(module, fsdp_module) for fsdp_module in fsdp_modules): @@ -1067,7 +1081,11 @@ def _register_pre_backward_param_unshard_hook(module): continue self.grad_acc_hooks[f"grad_acc and reduce for {self.param_to_name[param]}"] = ( param.register_post_accumulate_grad_hook( - lambda p: _process_post_backward_gradients([p]) + lambda p: ( + None + if getattr(p, 'skip_backward_post_hook', False) + else _process_post_backward_gradients([p]) + ) ) ) @@ -1224,6 +1242,9 @@ def start_param_sync(self, *unused, force_sync: bool = False, force_dispatch: bo """ self._replace_param_with_raw_if_needed() + if self.data_parallel_sharding_strategy == "no_shard": + return + if not force_sync and self.ddp_config.overlap_param_gather: # All-gather the first bucket before the forward pass. if self.ddp_config.fsdp_all_gather_in_start_param_sync: diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/design/1f1b_ep_overlap_fsdp_design.md b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/design/1f1b_ep_overlap_fsdp_design.md new file mode 100644 index 00000000000..8f79ea469ac --- /dev/null +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/design/1f1b_ep_overlap_fsdp_design.md @@ -0,0 +1,544 @@ +# 1F1B (EP) Overlap — FSDP Integration Design + +This document describes the FSDP-side contract required by the 1F1B EP-overlap +schedule (`combined_1f1b`), and how Megatron FSDP v2 fulfils it. +v1 behaviour is documented in [§7 — v1 reference](#7-v1-reference-megatron_fsdp). + +--- + +## 0. Key concepts: `--overlap-moe-expert-parallel-comm` and `--delay-wgrad-compute` + +### `--overlap-moe-expert-parallel-comm` (EP overlap) + +Enables the **combined 1F1B schedule** (`combined_1f1b.py`). Instead of running +each `TransformerLayer` forward and backward sequentially, the schedule +interleaves the **sub-module operations** (attn, mlp, moe_dispatch, moe_combine) +of two different micro-batches: + +``` +comm_stream: combine_bwd │ dispatch_fwd → dispatch_bwd │ combine_fwd +comp_stream: attn_fwd │ mlp_bwd → mlp_bwd_dw → mlp_fwd │ attn_bwd → attn_bwd_dw + ── micro-batch N ──► ◄── micro-batch N-1 ── +``` + +- Forward micro-batch `N` (attn_fwd, mlp_fwd) runs **concurrently** with backward + micro-batch `N-1` (mlp_bwd, attn_bwd). +- MoE all-to-all communication (dispatch/combine) overlaps with computation on + the other stream. +- Global batch size and micro-batch size semantics are unchanged. You have + `mbs=4` — each forward or backward still operates on `mbs=4` tokens; the + "parallelism" is between **different** micro-batches. + +### `--delay-wgrad-compute` (delayed weight gradient) + +Instructs Transformer Engine layers to **postpone weight gradient computation** +until an explicit `backward_dw()` call. During `backward()`, only activation +gradients flow; weight gradients are computed later in `backward_dw()`. + +**Requires** `--overlap-moe-expert-parallel-comm` (asserted in +`transformer_config.py:2802-2804`). This lets the schedule overlap the delayed +wgrad kernels with other work (e.g., P2P communication): + +``` +comp_stream: ... mlp_bwd (act grads only) │ mlp_bwd_dw (wgrad) │ ... + ─── activation-grad-only ──► ─── weight grad ──► +``` + +Without `delay_wgrad_compute`, wgrads are computed inline during `backward()`. +With it, they move to `backward_dw()`, creating additional overlapping +opportunities but also requiring FSDP's gradient reduction to wait until after +`backward_dw()` — the central challenge for FSDP integration. + +### How they relate + +| Flag | Requires | Relationship | +|---|---|---| +| `overlap_moe_expert_parallel_comm` | — | Enables the combined 1F1B schedule. Without it, normal 1F1B or naive pipeline runs. | +| `delay_wgrad_compute` | `overlap_moe_expert_parallel_comm` | Moves wgrad computation from `backward()` to `backward_dw()`. FSDP must defer `reduce_grad()` until after `backward_dw()`. | + +--- + +## 1. Why the overlap schedule needs special FSDP handling + +**Normal FSDP flow** (hooks fire on `TransformerLayer`): + +``` +TransformerLayer.forward() + → pre-forward hook: unshard params + → actual compute + → post-forward hook: reshard params +backward() + → pre-backward hook: unshard params + → compute grads + → post-backward hook: reshard params + reduce-scatter grads +``` + +**EP overlap flow** (calls sub-modules directly, bypassing `TransformerLayer.forward()`): + +``` +combined_forward_backward_step() + → f_layer.attn.forward() ← no TransformerLayer hook fires + → b_layer.mlp.backward() ← no TransformerLayer hook fires + → f_layer.moe_dispatch.forward() + → ... +``` + +Because the schedule invokes sub-modules (`attn`, `moe_dispatch`, `mlp`, +`moe_combine`) directly, the FSDP hooks registered on `TransformerLayer` are +**never triggered**. FSDP must therefore expose a set of manual management +APIs that the schedule calls explicitly at the right moments. + +--- + +## 2. FSDP API contract + +### 2.1 Required attributes on the FSDP wrapper + +| Hook | Type | Required for | Semantics | v2 impl | +|---|---|---|---|---| +| Root backward-phase setup | callable | All sharding strategies | Set `backward_phase = True`, unshard root params for backward. Called once before the overlapped forward+backward run. | `mfsdp_pre_backward_setup` | +| Root backward finalization | callable | All sharding strategies | Handle any modules whose per-module post-backward was skipped, drain async reduce-grad events, reset state. Called once after the overlapped run. | `mfsdp_post_backward_final_callback` | +| Per-layer post-forward | callable | `optim_grads_params` only | Reshard (release all-gathered parameters) for one layer after its forward ops complete. Called on the last forward node. | `mfsdp_post_forward_hook` | +| Per-layer post-backward | callable | `optim_grads_params` only | **Reshard + reduce gradients** for one layer after its backward ops complete. Called on the last backward node. Copies `.grad` → main grad buffer → reduce-scatter → installs `dist_param.grad` as DTensor. | `mfsdp_post_backward_hook` | +| Gradient sync suppression | context manager | All sharding strategies | Suppress gradient reduce-scatter for inner micro-batches. | `no_sync` (nullcontext in v2) | +| Sharding strategy access | object | All | Provides `data_parallel_sharding_strategy` for the schedule to decide per-layer hooks. | `ddp_config` / `_fsdp_param_groups` | + +> **v1 vs v2 note**: In v1, the per-layer post-backward hook only releases +> parameters — gradient reduction happens separately via per-param +> `post_accumulate_grad_hook`. In v2, `mfsdp_post_backward_hook` performs +> both `reshard()` **and** `reduce_grad()` in a single call. Both v1 and v2 +> have a per-layer post-forward hook that only does reshard (no grad work). +> +> `_replace_param_with_raw_if_needed()` is **not required** in v2. See §2.2. + +### 2.2 Required hooks for fine-grained sub-module management + +When `overlap_moe_expert_parallel_comm=True`, two additional hook modes must be +enabled: + +| Hook mode | Required for | Effect | +|---|---|---| +| Fine-grained pre-forward unshard | All strategies | Register `mfsdp_forward_pre_hook` on **every sub-module** (not just FSDP units), because the schedule calls sub-modules directly. | +| Fine-grained pre-backward unshard | All strategies | Register `mfsdp_pre_backward_setup` via `register_multi_grad_hook` on each sub-module's output tensor, ensuring params are unsharded before backward compute. | + +> **Why `_replace_param_with_raw_if_needed()` is not required in v2**: v2 keeps +> DTensor params directly on the module and manages them via +> `unshard()`/`reshard()`. Fine-grained pre-forward hooks ensure parameters are +> all-gathered before each sub-module compute. v1 required the swap because its +> optimizer-facing params were separate from module params. + +### 2.3 Per-layer post-forward and post-backward hooks (`optim_grads_params` only) + +The schedule plan calls `set_fsdp_reshard_hooks(post_forward_hook, post_backward_hook)` +on each `TransformerLayerSchedulePlan` to wire: + +- **Post-forward**: attached to the **last forward node** + (`moe_combine` for MoE, `mlp` otherwise). Calls `mfsdp_post_forward_hook(layer)`, + which reshard parameters (release all-gathered buffer, install DTensor dist_params). + +- **Post-backward**: attached to the **last backward node** (`attn`). + Calls `mfsdp_post_backward_hook(layer)`, which does **both**: + 1. Reshard parameters (release all-gathered buffer, install DTensor dist_params). + 2. Reduce gradients: copy `.grad` → main grad buffer → reduce-scatter → + install `dist_param.grad = dist_grad` (DTensor). + +These are needed because the overlap schedule bypasses `TransformerLayer.forward()`, +so the normal forward/backward hooks never fire. + +--- + +## 3. Runtime call sequence (v2) + +### 3.1 Overall flow + +The schedule in `combined_1f1b_schedule_for_no_pipelining()` orchestrates four +phases per training step (file: `megatron/core/pipeline_parallel/combined_1f1b.py`): + +``` +combined_1f1b_schedule_for_no_pipelining(): +│ +├─ 1. combined_forward_backward_step( # first microbatch, fwd only +│ f_model=model, b_model=None, ... +│ ) +│ No FSDP root-level calls for this phase. Fine-grained pre-forward +│ hooks (registered on every sub-module) handle per-sub-module unshard. +│ +├─ 2. with no_sync_func(): # inner micro-batches +│ │ +│ └─ combined_forward_backward_step( +│ f_model=model, b_model=model, ... +│ ) +│ │ +│ ├─ 2a. pre_backward_fn(root_module) # mfsdp_pre_backward_setup +│ │ → ctx.backward_phase = True +│ │ → root_module.unshard(bwd_pass=True) +│ │ → skip_final_callback=True (no auto-enqueue) +│ │ +│ ├─ 2b. for each layer: layer_plan.set_fsdp_reshard_hooks( +│ │ mfsdp_post_forward_hook, mfsdp_post_backward_hook +│ │ ) +│ │ +│ ├─ 2c. TransformerModelChunkSchedulePlan.run(f, b, ...) +│ │ │ +│ │ │ During the run, per-layer hooks fire: +│ │ │ - After last fwd node: mfsdp_post_forward_hook(layer) +│ │ │ → layer.reshard() +│ │ │ - After last bwd node: mfsdp_post_backward_hook(layer) +│ │ │ → layer.reshard() + layer.reduce_grad() +│ │ │ +│ │ │ Also, the normal autograd post-backward hook +│ │ │ (RegisterFSDPBackwardFunction) is DISABLED for all FSDP +│ │ │ modules when delay_wgrad_compute=True. See §3.4. +│ │ +│ └─ 2d. mfsdp_post_backward_final_callback(root_module) +│ → Handles ALL fsdp modules whose per-module post-backward +│ was skipped (e.g., nested TEGroupedMLP expert modules). +│ Runs after all backward_dw() calls complete. +│ +└─ 3. combined_forward_backward_step( # last batch, bwd only + f_model=None, b_model=model, ... + ) + │ + ├─ 3a. pre_backward_fn(root_module) + ├─ 3b. TransformerModelChunkSchedulePlan.run(None, b, ...) + └─ 3c. mfsdp_post_backward_final_callback(root_module) +``` + +### 3.2 Single-layer execution order + +Inside `TransformerLayerSchedulePlan.run()` (file: +`megatron/core/models/common/model_chunk_schedule_plan.py`), each overlapped +layer pair executes: + +``` +line 263: b_layer.moe_combine.backward(grad) ← MoE combine backward +line 267: f_layer.attn.forward(f_input) ← attention forward +line 270: b_layer.mlp.backward(grad) ← autograd backward: + 1. fine-grained pre-bwd hook → unshard MLP + 2. TE mlp backward (act grads only if delay_wgrad) + 3. autograd post-bwd hook → SKIPPED (delay_wgrad) +line 274: f_layer.moe_dispatch.forward(f_input) ← MoE dispatch forward +line 277: b_layer.mlp.backward_dw() ← delayed TE mlp wgrad +line 278: b_layer.moe_dispatch.backward(grad) ← MoE dispatch backward +line 285: f_layer.mlp.forward(f_input) ← MLP forward +line 289: f_layer.moe_combine.forward(f_input) ← MoE combine forward +line 292: b_layer.attn.backward(grad) ← autograd backward: + 1. fine-grained pre-bwd hook → unshard attn + 2. TE attn backward (act grads only if delay_wgrad) + 3. autograd post-bwd hook → SKIPPED (delay_wgrad) +line 301: b_layer.attn.backward_dw() ← delayed TE attn wgrad + → fires mfsdp_post_backward_hook(layer) + → layer.reshard() + layer.reduce_grad() +``` + +### 3.3 Fine-grained pre-backward hook registration + +In v2, fine-grained hooks are registered via the `fine_grained_hooks` parameter +of `fully_shard()` (wired to `config.overlap_moe_expert_parallel_comm` in +`mcore_fsdp_adapter.py:276`). + +**Pre-forward** (`_register_forward_pre_hook(fine_grained=True)` in `fully_shard.py:124`): +- Registers `mfsdp_forward_pre_hook` on **every sub-module** of each FSDP unit. +- When the schedule calls `f_layer.attn.forward()`, the hook on the `attn` + sub-module fires → resolves parent FSDPModule → `unshard()`. + +**Pre-backward** (`_register_backward_pre_hook(fine_grained=True)` in `fully_shard.py:131`): +- Registers `mfsdp_pre_backward_setup` via `register_multi_grad_hook` on + every sub-module's output tensor. +- When backward reaches a sub-module, the hook fires → resolves parent + FSDPModule → `unshard(bwd_pass=True)` → resets `post_backward_issued`. + +### 3.4 Critical design constraint: `_register_backward_hook` and `delay_wgrad_compute` + +The normal autograd post-backward hook (`_register_backward_hook` at `fully_shard.py:141`) +inserts `RegisterFSDPBackwardFunction` into the autograd graph. During +backward, this function fires and calls `mfsdp_post_backward_hook(module)`, +which does `reshard()` + `reduce_grad()`. + +When `delay_wgrad_compute=True`, this hook fires **before** `backward_dw()` +completes the weight gradient computation. This causes `reduce_grad()` to +reduce only activation gradients, but not weight gradients. The subsequent +`backward_dw()` then writes to the already-resharded DTensor params, +corrupting `.grad`. + +**Fix**: When `skip_backward_callback=True` (wired to `config.delay_wgrad_compute` +in `mcore_fsdp_adapter.py:303`), the autograd `RegisterFSDPBackwardFunction` is +**skipped** for this module. Per-layer `reshard()` + `reduce_grad()` still fires +manually for the TransformerLayer via `set_fsdp_reshard_hooks` → +`mfsdp_post_backward_hook`. Remaining nested modules (e.g., TEGroupedMLP, root) +are handled by `mfsdp_post_backward_final_callback` (called at +`combined_1f1b.py:629`), which runs after all `backward_dw()` calls complete. + +For EP overlap **without** `delay_wgrad_compute`, the autograd hook fires at +the right time (all grads are ready inline during backward), so it is +**not** skipped. + +### 3.5 Hook fire count — what to expect + +`mfsdp_forward_pre_hook` fires on **every** `Module.__call__()`, not just on +the FSDP unit modules. Because the EP overlap schedule invokes sub-modules +directly, and fine-grained hooks are registered on all sub-modules, the **root +pre-forward hook fires multiple times per forward micro-batch** — once for each +root-child sub-module entered via `__call__`. + +In a typical GPT model, these root-child calls happen in: + +- `PreProcessNode.forward()` → `gpt_model._preprocess()`: + - `self.embedding(...)` — calls `embedding.__call__()` + - `self.rotary_pos_emb(...)` — calls `rotary_pos_emb.__call__()` +- `PostProcessNode.forward()` → `gpt_model._postprocess()`: + - `self.decoder.final_layernorm(...)` — calls `final_layernorm.__call__()` + +That is **3** root forward pre-hook fires per forward micro-batch. + +Similarly, the fine-grained **pre-backward** hook fires on each sub-module +whose backward is entered. For a root-child sub-module (e.g., `embedding`, +`final_layernorm`), this triggers `mfsdp_pre_backward_setup` which deduplicates +via `_fsdp_pre_backward_done` — so the root backward setup fires **once per +backward micro-batch**, on the first root-child sub-module whose backward runs. + +**Example**: World size `W=4`, `EP=2`, `mbs=4`, `global_batch_size=32`: + +``` +DP = W / EP = 2 (dense DP = 4, but micro-batching uses the inter-EP DP) +num_microbatches = 32 / (4 × 2) = 4 +``` + +The 1F1B schedule produces 5 phases — 4 forward + 4 backward. Expected hook +fires (root-level only): + +``` +Phase 0 (warmup): FWD[0] → F F F +Phase 1 (overlap): FWD[1] ─ BWD[0] → F F F B +Phase 2 (overlap): FWD[2] ─ BWD[1] → F F F B +Phase 3 (overlap): FWD[3] ─ BWD[2] → F F F B +Phase 4 (cleanup): BWD[3] → B + +Total per step: 12 F, 4 B +``` + +A typical log snippet would show: + +``` +=== Starting forward pass === (×3, phase 0) +=== Starting forward pass === (×3, phase 1) +=== Starting backward pass === +=== Starting forward pass === (×3, phase 2) +=== Starting forward pass === (×3, phase 3) +=== Starting backward pass === +=== Starting backward pass === (×2, phases 1-2-3-4 depending on dedup timing) +... +``` + +The **6 forward per 1 backward** pattern in the log happens naturally because +a typical print grouping captures ~2 forward phases worth of hooks before +each backward. This is expected and does **not** mean anything is wrong. The +root pre-forward hook is designed to be called repeatedly — it is idempotent +(both the root-level bookkeeping and the per-module `unshard()` are safe to +repeat). + +--- + +## 4. Hook behavior — v2 implementation + +The EP overlap schedule wires four hook functions defined in +`megatron/core/distributed/fsdp/src/megatron_fsdp/v2/hooks.py`. +Their schedule-side wiring is in `combined_1f1b.py` (see §3.1). + +### 4.1 Root backward-phase setup — `mfsdp_pre_backward_setup` + +- **File**: `megatron/core/distributed/fsdp/src/megatron_fsdp/v2/hooks.py:194` +- **When**: called once in `combined_forward_backward_step()` before the + schedule plan `.run()` (line 468 in `combined_1f1b.py`), with + `skip_final_callback=True`. +- **What it does**: + 1. Resolves target FSDPModule via `_find_fsdp_target(hook_module)`. + 2. Deduplicates via `_fsdp_pre_backward_done` flag. + 3. If root: sets `ctx.backward_phase = True`, advances backward module tracker. + 4. Does **not** auto-enqueue `mfsdp_post_backward_final_callback` + (`skip_final_callback=True`) — the schedule calls it manually. + 5. Calls `module.unshard(bwd_pass=True)` — all-gathers params for backward compute. + 6. Sets `module.post_backward_issued = False` — resets per-module bookkeeping. + 7. Resets TE gradient-accumulation fusion flags. + +### 4.2 Root backward finalization — `mfsdp_post_backward_final_callback` + +- **File**: `megatron/core/distributed/fsdp/src/megatron_fsdp/v2/hooks.py:252` +- **When**: called once in `combined_forward_backward_step()` after the + schedule plan `.run()` (line 629 in `combined_1f1b.py`). +- **What it does**: + 1. Iterates `reversed(ctx.forward_order)` — all FSDP modules in the tree. + 2. For each module with `post_backward_issued = False`: + - `module.reshard()` — release all-gathered buffers, install DTensor params. + - `module.reduce_grad()` — copy grads → main grad buffer → reduce-scatter → + set `dist_param.grad = dist_grad` (DTensor). + 3. Drains pending async reduce-grad events (`event.wait()` + release buffers). + 4. Resets root/context state (`backward_phase = False`, clears + `backward_done_modules`, `_fsdp_pre_backward_done` flags). + 5. Transitions bucket allocator from trace → optimized plan (first micro-batch). + +### 4.3 Per-layer post-forward (reshard only) — `mfsdp_post_forward_hook` + +- **File**: `megatron/core/distributed/fsdp/src/megatron_fsdp/v2/hooks.py:133` +- **When**: wired via `set_fsdp_reshard_hooks()`, fires after the last + forward node of each layer. +- **What it does**: + 1. Calls `module.reshard()` — release all-gathered params after forward. + 2. Has a backward-phase guard: if `ctx.backward_phase` is active and this + module is the current backward module, skips reshard (activation + recomputation case — params are still needed). + +### 4.4 Per-layer post-backward (reshard + reduce_grad) — `mfsdp_post_backward_hook` + +- **File**: `megatron/core/distributed/fsdp/src/megatron_fsdp/v2/hooks.py:225` +- **When**: wired via `set_fsdp_reshard_hooks()`, fires after the last + backward node (`attn`) of each layer. Also potentially fired by + `RegisterFSDPBackwardFunction` during autograd backward (but see §3.4). +- **What it does**: + 1. Adds module to `backward_done_modules`. + 2. Advances the backward module tracker. + 3. Calls `module.reshard()` — release all-gathered params. + 4. If sharding strategy is `optim_grads` or `optim_grads_params`: + calls `module.reduce_grad()` — copy grads → reduce-scatter → install DTensor grads. + 5. Sets `module.post_backward_issued = True`. + +### 4.5 Gradient sync suppression — `no_sync()` + +- **v2 status**: Returns `nullcontext` (no-op). + This means gradient reduce-scatter fires on every micro-batch, including + inner ones. The standard 1F1B pattern (suppress sync for inner + micro-batches, sync on the last) is not yet implemented for v2. + +--- + +## 5. Constraints enforced at init time + +When `overlap_moe_expert_parallel_comm=True`, the following constraints apply: + +| Constraint | Enforced? | Rationale | +|---|---|---| +| `overlap_moe_expert_parallel_comm` required by `delay_wgrad_compute` | ✅ `transformer_config.py:2802-2804` | Delayed wgrad only works in the 1F1B overlap schedule. | +| `delay_wgrad_compute` mutually exclusive with `overlap_dispatch_backward_with_experts_wgrad` | ✅ `transformer_config.py:2819-2822` | Two different wgrad-deferral strategies; choose one. | +| CUDA graph scope on `moe` and `mlp` blocked | ✅ `transformer_config.py:2790-2798` | Partial CUDA graph scopes conflict with fine-grained schedule. | +| TE >= 2.3.0 for `delay_wgrad_compute` | ✅ `megatron/training/arguments.py:1937` | `delay_wgrad_compute` kernel support requires TE >= 2.3.0. | +| Only `GPTModel` supported | ✅ `combined_1f1b.py:505` | Schedule plan build checks `isinstance(unwrapped_model, GPTModel)`. | +| Interleaved PP + FSDP blocked | ✅ `combined_1f1b.py:317-321` | Multi-chunk models not yet supported with EP overlap. | +| `fsdp_double_buffer = False` | ❌ not enforced in v2 | Double buffering is incompatible with per-sub-module parameter management. | +| `fsdp_unit_modules` compatibility | ❌ not enforced in v2 | The schedule expects each layer to be an FSDPModule (wrapped by `fully_shard()`), but this is not validated at init time. | + +--- + +## 6. Key code locations + +| Component | File | +|---|---| +| **v2 implementation** | | +| FSDP API definitions | `megatron/core/distributed/fsdp/src/megatron_fsdp/v2/hooks.py` | +| FSDP module (`FSDPModule`) | `megatron/core/distributed/fsdp/src/megatron_fsdp/v2/fsdp_module.py` | +| `fully_shard()` entry point | `megatron/core/distributed/fsdp/src/megatron_fsdp/v2/fully_shard.py` | +| Parameter groups & buffers | `megatron/core/distributed/fsdp/src/megatron_fsdp/v2/param_group.py` | +| FSDP adapter (config wiring) | `megatron/core/distributed/fsdp/mcore_fsdp_adapter.py` | +| **Shared (v1 + v2)** | | +| Schedule orchestration | `megatron/core/pipeline_parallel/combined_1f1b.py` | +| Schedule plan (layer-level) | `megatron/core/models/common/model_chunk_schedule_plan.py` | +| Schedule plan (chunk-level) | `megatron/core/models/common/model_chunk_schedule_plan.py` | +| Schedule dispatch | `megatron/core/pipeline_parallel/schedules.py` | +| Fine-grained callables | `megatron/core/models/gpt/fine_grained_callables.py` | +| Config definitions | `megatron/core/model_parallel_config.py`, `megatron/core/transformer/transformer_config.py` | +| **Tests** | | +| M-FSDP v2 EP overlap e2e | `tests/unit_tests/distributed/megatron_fsdp/v2/test_mcore_nd_parallel.py` | +| delay_wgrad_compute unit test | `tests/unit_tests/a2a_overlap/test_delay_wgrad_compute.py` | +| FSDP 1F1B overlap test | `tests/unit_tests/a2a_overlap/test_fsdp_1f1b_overlap.py` | + +--- + +## 7. v1 reference (`megatron_fsdp`) + +> This section describes the Megatron FSDP v1 (`megatron_fsdp`) implementation +> for reference. v1 uses different internal APIs than v2. + +### 7.1 Hook behavior — v1 + +> All v1 hooks live in `megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py`. + +#### 7.1.1 Root backward-phase setup — `pre_backward()` + +- **V1 impl**: `_root_pre_backward(module=None, skip_backward_hook=True)` +- **When**: called once in `combined_forward_backward_step()` before the + schedule plan `.run()`. +- **What it does**: + 1. Sets `_root_pre_backward_hook_issued = True` (idempotency guard). + 2. For `optim_grads_params`: sets all sub-module `_training_state` to + `PRE_BACKWARD`, marks all AG buckets as releasable. + 3. Tracks params that require gradient handling via + `_params_require_handle_grad`. + 4. Does NOT auto-enqueue `_root_post_backward` — the schedule calls it + manually via `post_backward()`. + +#### 7.1.2 Root backward finalization — `post_backward()` + +- **V1 impl**: `_root_post_backward()` +- **What it does**: + 1. Processes any remaining unhandled gradients. + 2. Launches async reduce-scatter for gradient-sharding strategies. + 3. Resets root state: `_root_pre_backward_hook_issued = False`, + increments `microbatch_count`. + +#### 7.1.3 Per-layer post-forward (reshard only) — `post_forward_release_module()` + +- **V1 impl**: `_post_forward(module, input=None, output=None)` +- **When**: wired via `set_fsdp_reshard_hooks()`, fires after the last + forward node of each layer. +- **What it does**: + 1. If `_training_state == PRE_BACKWARD`: lazy release (activation + recomputation case). + 2. Otherwise: release params via `release_module_parameters(module, bwd=False)`, + transition to `IDLE`. + +#### 7.1.4 Per-layer post-backward (release only, no grad reduction) — `post_backward_release_module()` + +- **V1 impl**: `_post_backward_release_module(module)` +- **When**: wired via `set_fsdp_reshard_hooks()`, fires after the last + backward node (`attn`) of each layer. +- **What it does**: + - Releases params for both backward and forward passes: + `release_module_parameters(module, bwd=True)` and + `release_module_parameters(module, bwd=False)`. + - Transitions all sub-modules to `IDLE` state. + +- **Key difference from v2**: v1's per-layer post-backward hook does **not** + perform gradient reduction. Gradient processing (reduce-scatter) is handled + separately by the per-param `post_accumulate_grad_hook` that fires + independently. In contrast, v2's `mfsdp_post_backward_hook` combines + **both** reshard and gradient reduction in a single call. + +- With `delay_wgrad_compute`: gradient processing for params with + `skip_backward_post_hook=True` is deferred to `backward_dw()` via + `setup_delayed_wgrad_acc_hook` (`megatron_fsdp.py:76-101`). + +#### 7.1.5 Gradient sync suppression — `no_sync()` + +- **V1 impl**: `MegatronFSDP.no_sync()` (context manager, sets + `is_last_microbatch = False` on enter, `True` on exit). +- **When**: wraps the inner micro-batches (all except first forward-only + and last backward-only). +- **Effect**: prevents gradient reduce-scatter for non-final micro-batches. + +#### 7.1.6 `_replace_param_with_raw_if_needed()` + +- **V1 impl**: `MegatronFSDP._replace_param_with_raw_if_needed()` +- **When**: called once at the start of the schedule (`combined_1f1b.py:179`), + before any layer access. +- **Effect**: swaps the distributed (optimizer-managed) `DTensor` parameters + back to raw `nn.Parameter` tensors so the schedule can call sub-modules + directly. In v1, optimizer-facing params are separate from module params, + so this swap is mandatory. v2 does not require this because DTensor + params live on the module directly and are managed via `unshard()`/`reshard()`. + +### 7.2 Key code locations (v1) + +| Component | File | +|---|---| +| FSDP API definitions | `megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py` | +| Delayed wgrad hook setup | `megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py:76-101` | +| Param attribute preservation | `megatron/core/distributed/fsdp/src/megatron_fsdp/param_and_grad_buffer.py:2743-2748` | diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/allocator_design.md b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/design/allocator_design.md similarity index 100% rename from megatron/core/distributed/fsdp/src/megatron_fsdp/v2/allocator_design.md rename to megatron/core/distributed/fsdp/src/megatron_fsdp/v2/design/allocator_design.md diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/cuda_graph_design.md b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/design/cuda_graph_design.md similarity index 99% rename from megatron/core/distributed/fsdp/src/megatron_fsdp/v2/cuda_graph_design.md rename to megatron/core/distributed/fsdp/src/megatron_fsdp/v2/design/cuda_graph_design.md index 637655353fb..99b9f7b5f2a 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/cuda_graph_design.md +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/design/cuda_graph_design.md @@ -7,11 +7,11 @@ mcore's CUDA graph system (`cuda_graph_impl="local"`, `cuda_graphs.py`) was designed for DDP's memory model — each layer receives freshly-allocated tensor -inputs/outputs. FSDP v2 shares the same pool-backed buffers across layers, +inputs/outputs. Megatron FSDP v2 shares the same pool-backed buffers across layers, and the FSDP hooks (unshard/reshard) are not captured in the graph in the same way. -This doc describes a CUDA graph system built INTO FSDP v2, using +This doc describes a CUDA graph system built INTO Megatron FSDP v2, using `TracePoolAllocator` as the stable-memory foundation. The user enables it with a single flag — everything else is automatic. diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/design.md b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/design/design.md similarity index 99% rename from megatron/core/distributed/fsdp/src/megatron_fsdp/v2/design.md rename to megatron/core/distributed/fsdp/src/megatron_fsdp/v2/design/design.md index 55148a74c69..fee7f75b521 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/design.md +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/design/design.md @@ -1,4 +1,4 @@ -# Design: Megatron-FSDP2 Implementation +# Design: Megatron FSDP v2 Implementation --- diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/design/hooks_api.md b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/design/hooks_api.md new file mode 100644 index 00000000000..f36c01b4f90 --- /dev/null +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/design/hooks_api.md @@ -0,0 +1,136 @@ +# Megatron FSDP v2 Hooks API + +The public API surface of `hooks.py` — functions callable by external code +(e.g. the 1F1B EP overlap schedule) for manual FSDP lifecycle management. + +All functions live in `megatron.core.distributed.fsdp.src.megatron_fsdp.v2.hooks`. + +--- + +## Target resolution + +### `_find_fsdp_target(hook_module) → Optional[FSDPModule]` + +Resolve the nearest parent FSDPModule for any module. + +| Condition | Returns | +|---|---| +| `isinstance(hook_module, FSDPModule)` | `hook_module` itself | +| `hasattr(hook_module, '_fsdp_parent_module')` | `hook_module._fsdp_parent_module` | +| Otherwise | `None` | + +`_fsdp_parent_module` is set on every non-FSDPModule sub-module during +`FSDPModule._init_fsdp_state()` (bottom-up pass, nearest ancestor wins). + +--- + +## Forward hooks + +### `mfsdp_forward_pre_hook(hook_module, args, kwargs)` + +Pre-forward: unshard parameters for the target FSDPModule. + +- Resolves target via `_find_fsdp_target`. No-op if `None`. +- Root phase: sets `ctx.forward_phase = True`, creates CG stream lazily. +- Unshards parameters (forward + optional backward pass). +- Frees stale grad data. +- **CUDA graph capture**: only when `isinstance(hook_module, FSDPModule)` (skipped for sub-modules). + +Accepts any module. Sub-modules delegate to parent via `_fsdp_parent_module`. + +**Repeatability**: This function may be called multiple times for the same target +when fine-grained hooks are active (e.g. sub-module forward triggers the hook, +followed by the enclosing FSDPModule's forward also triggering it). The +implementation MUST be safe to invoke repeatedly without observable overhead — +duplicate ``unshard()`` calls and idempotent bookkeeping must be effectively +free. + +### `mfsdp_post_forward_hook(module, *unused)` + +Post-forward: reshard parameters. + +- **Validates**: `TypeError` if `not isinstance(module, FSDPModule)`. +- Skips if `ctx.backward_phase` and this module is the current backward module + (activation recomputation: don't reshard before backward needs params). +- `module.reshard()`. + +--- + +## Backward hooks + +### `mfsdp_pre_backward_setup(hook_module, grads)` + +Pre-backward: root phase transition, unshard, TE gradient-fusion flags. + +- Resolves target via `_find_fsdp_target`. No-op if `None`. +- **Deduplication**: skips if `target._fsdp_pre_backward_done` is `True`. +- Root setup (first call): + - `ctx.backward_done_modules.clear()` + - `ctx.forward_phase = False`, `ctx.backward_phase = True` + - `ctx._advance_backward_module()` + - Enqueues `_register_post_backward_final_callback` (once). +- `target.unshard(bwd_pass=True)`. +- `target.post_backward_issued = False`. +- TE gradient-fusion: resets `grad_added_to_main_grad`, sets `overwrite_main_grad` for sharding strategies. +- `target._fsdp_pre_backward_done = True`. + +Accepts any module. Compatible with `register_multi_grad_hook` callback signature +`(module, grads)`. + +### `mfsdp_post_backward_hook(module)` + +Post-backward: mark module done, reshard, reduce gradients. + +- **Validates**: `TypeError` if `not isinstance(module, FSDPModule)`. +- `ctx.backward_done_modules.add(id(module))`. +- `ctx._advance_backward_module()`. +- `module.reshard()`. +- If any param group uses `optim_grads` or `optim_grads_params`: + `module.reduce_grad(async_op=ctx.enable_async_reduce_grad)`. +- `module.post_backward_issued = True`. + +### `mfsdp_post_backward_final_callback(root_module)` + +Finalise the backward pass for one micro-batch. + +- **Validates**: `TypeError` if not FSDPModule; `RuntimeError` if not root. +- Handles modules whose per-module post-backward was skipped (e.g. activation + recomputation): `reshard()` + `reduce_grad()`. +- Drains pending async reduce-grad events. +- Resets root context: `_post_backward_callback_queued`, `backward_phase`, + `backward_module`, `backward_done_modules`. +- Clears `_fsdp_pre_backward_done` on all FSDP modules in `forward_order`. +- First micro-batch only: transitions `TracePoolAllocator` from trace to + optimized plan. + +--- + +## Flag lifecycle + +``` +_fsdp_pre_backward_done: + False (init) ─► True (mfsdp_pre_backward_setup) ─► False (mfsdp_post_backward_final_callback) + +post_backward_issued: + False (init) ─► False (mfsdp_pre_backward_setup) ─► True (mfsdp_post_backward_hook) +``` + +All flags are initialised to `False` in `FSDPModule._init_fsdp_state()`. + +--- + +## Integration with the auto-hook path + +The registration functions (`_register_*`) use the public functions internally: + +| Registration function | Uses | +|---|---| +| `_register_forward_pre_hook` | `mfsdp_forward_pre_hook` | +| `_register_forward_hook` | `mfsdp_post_forward_hook` | +| `_register_backward_pre_hook` | `mfsdp_pre_backward_setup` (via `_create_custom_backward_hook`) | +| `_register_backward_pre_hook_fine_grained` | `mfsdp_pre_backward_setup` (via `_create_custom_backward_hook`) | +| `_register_backward_hook` | `mfsdp_post_backward_hook` (via `RegisterFSDPBackwardFunction`) | +| `_register_post_backward_final_callback` | `mfsdp_post_backward_final_callback` (via `Variable._execution_engine.queue_callback`) | + +This means the auto-hook path and the manual (1F1B overlap schedule) path share +identical logic — they just differ in *when* the functions are called. diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/lazy_grad_buffer_design.md b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/design/lazy_grad_buffer_design.md similarity index 100% rename from megatron/core/distributed/fsdp/src/megatron_fsdp/v2/lazy_grad_buffer_design.md rename to megatron/core/distributed/fsdp/src/megatron_fsdp/v2/design/lazy_grad_buffer_design.md diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/mcore_fsdp_checkpoint_design.md b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/design/mcore_fsdp_checkpoint_design.md similarity index 97% rename from megatron/core/distributed/fsdp/src/megatron_fsdp/v2/mcore_fsdp_checkpoint_design.md rename to megatron/core/distributed/fsdp/src/megatron_fsdp/v2/design/mcore_fsdp_checkpoint_design.md index 7b1051bc880..0011c0414c4 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/mcore_fsdp_checkpoint_design.md +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/design/mcore_fsdp_checkpoint_design.md @@ -24,10 +24,10 @@ the checkpoint save/load architecture and how it integrates with MCore's | Source Format | Target Format | Model | Optimizer | Notes | |---------------|--------------|-------|-----------|-------| -| MFSDP v2 | MFSDP v2 | ✓ | ✓ | Full round-trip and cross-setting (`optim_grads` ↔ `optim_grads_params`) | -| MFSDP v1 baseline | MFSDP v2 | ✗ | ✗ | Currently skipped in tests (``pytest.skip("v1 checkpoint format not available")``). Both use ``fsdp_dtensor`` format but key names may differ. Planned for future support. | -| ND-parallel (`torch_dist`) | MFSDP v2 | ✓ | ✓ (``fully_reshardable`` only) | Online conversion via `_load_torch_dist_into_megatron_fsdp_v2` in `checkpointing.py`. Expert weights are split from flattened multi-expert tensors. Optimizer states (`exp_avg`, `exp_avg_sq`) are loaded into V2's name-based format. Hi-precision FP32 optimizer param copies are used for model weights when available. ``dp_reshardable`` format is not supported (bucket-based layout incompatible with name-based V2 format). | -| ND-parallel (`torch`) | MFSDP v2 | ✗ | ✗ | Not supported (different serialization) | +| Megatron FSDP v2 | Megatron FSDP v2 | ✓ | ✓ | Full round-trip and cross-setting (`optim_grads` ↔ `optim_grads_params`) | +| MFSDP v1 baseline | Megatron FSDP v2 | ✗ | ✗ | Currently skipped in tests (``pytest.skip("v1 checkpoint format not available")``). Both use ``fsdp_dtensor`` format but key names may differ. Planned for future support. | +| ND-parallel (`torch_dist`) | Megatron FSDP v2 | ✓ | ✓ (``fully_reshardable`` only) | Online conversion via `_load_torch_dist_into_megatron_fsdp_v2` in `checkpointing.py`. Expert weights are split from flattened multi-expert tensors. Optimizer states (`exp_avg`, `exp_avg_sq`) are loaded into V2's name-based format. Hi-precision FP32 optimizer param copies are used for model weights when available. ``dp_reshardable`` format is not supported (bucket-based layout incompatible with name-based V2 format). | +| ND-parallel (`torch`) | Megatron FSDP v2 | ✗ | ✗ | Not supported (different serialization) | --- @@ -70,7 +70,7 @@ format for both paths. It uses DCP directly, storing each parameter as a `DTenso | `_maybe_wrap_as_uneven_dtensor(tensor, dist_param)` | Wrap a single plain tensor as an uneven DTensor if it matches the parameter's local shard shape; otherwise return unchanged. Shared by ``_build_dtensor_optim_sd`` and ``_preprocess_and_verify_v2_state_dict``. | | `load_torch_dist_into_fsdp_v2` | Entry point for online checkpoint conversion from legacy ``torch_dist`` format to ``fsdp_dtensor``. Five-phase pipeline: preprocess/verify, build name mapping, DCP load, expert params, verify. | | `add_module_prefix(state_dict)` | Add ``module.`` prefix to all state dict keys. Megatron FSDP v2 lacks ``MegatronFSDP`` wrapper so ``model.state_dict()`` keys have no prefix; this aligns with Megatron's checkpoint format. | -| `strip_module_prefix(state_dict)` | Remove ``module.`` prefix from state dict keys. Inverse of ``add_module_prefix``, used when loading checkpoint back into FSDP v2 model. | +| `strip_module_prefix(state_dict)` | Remove ``module.`` prefix from state dict keys. Inverse of ``add_module_prefix``, used when loading checkpoint back into Megatron FSDP v2 model. | | `get_model_state_dict(model)` | Get model state dict with ``module.`` prefix. Auto-detects whether prefix is already present; adds it if missing. | | `get_optimizer_state_dict(optimizer, is_loading)` | Get optimizer state dict via Path A (``sharded_param_state_fsdp_dtensor``). Delegates to ``optimizer.sharded_state_dict()`` with ``fsdp_dtensor`` sharding type. Returns ``{"state": ..., "param_to_group_meta": ...}``. | | `handle_fp8_extra_state_case(model_sd)` | Remove ``._extra_state`` keys from model state dict (FP8 artifact cleanup). | @@ -92,7 +92,7 @@ format for both paths. It uses DCP directly, storing each parameter as a `DTenso ### 3.2 Current Save Flow -> **Note:** This describes the current `generate_state_dict`-based flow. For FSDP v2, +> **Note:** This describes the current `generate_state_dict`-based flow. For Megatron FSDP v2, > ``_build_megatron_fsdp_v2_state_dict`` (Section 12) replaces this in the v2 code path. ``` @@ -120,7 +120,7 @@ save_checkpoint() ### 3.3 Current Load Flow -> **Note:** This describes the current `generate_state_dict`-based flow. For FSDP v2, +> **Note:** This describes the current `generate_state_dict`-based flow. For Megatron FSDP v2, > ``_build_megatron_fsdp_v2_state_dict`` (Section 12) replaces this in the v2 code path. ``` @@ -677,7 +677,7 @@ so the model state dict is complete. Torch_dist stores multi-layer and multi-expert tensors as a single fused tensor (e.g., ``decoder.layers.self_attention.linear_qkv.weight`` of shape ``(num_layers, H, 3*W)``). -FSDP v2 stores each layer as an individual DTensor +Megatron FSDP v2 stores each layer as an individual DTensor (``decoder.layers.0.self_attention.linear_qkv.weight`` of shape ``(H, 3*W)``). DCP cannot split one tensor into many, so Phase 4 handles this in `load_torch_dist_into_fsdp_v2`. @@ -817,7 +817,7 @@ provenance of chunk metadata when debugging size mismatches. | Source Tag | Set By | When | |---|---|---| -| `"init"` | `update_uneven_dtensor_chunk_metadata` (default) | FSDP v2 construction (`param_group.py`) | +| `"init"` | `update_uneven_dtensor_chunk_metadata` (default) | Megatron FSDP v2 construction (`param_group.py`) | | `"preprocess"` | `preprocess_state_dict_for_uneven_dtensor` | Save Phase 1 (re-computes via all_gather) | | `"propagate:"` | `copy_chunk_metadata` | Save Phase 2 (copied from model param; `` is the original tag) | | `"split"` | `split_dtensor` | Save Phase 3 (locally derived from parent's metadata) | @@ -865,7 +865,7 @@ and their consumers. | Model state dict keys | Keys have `module.` prefix from `MegatronFSDP` wrapper | Keys lack `module.` prefix; `add_module_prefix()` used for checkpoint alignment | | Model `state_dict_for_save_checkpoint` | `model.state_dict()` with `state_dict_pre_hook` | `model.state_dict()` (already DTensors) | | Optimizer buffer management | Megatron FSDP managed | Standard Torch optimizer managed | -| Gradient buffer | Megatron FSDP `param_and_grad_buffer` | None (FSDP v2 handles internally) | +| Gradient buffer | Megatron FSDP `param_and_grad_buffer` | None (Megatron FSDP v2 handles internally) | | Load model state dict | `module.load_state_dict(custom, strict)` with `_load_state_dict_post_hook` | `super().load_state_dict(state_dict, strict)` — after `strip_module_prefix()` | | Zero gradient | `model_chunk.zero_grad_buffer()` | `model_chunk._zero_grad_buffer()` | @@ -883,7 +883,7 @@ so the raw model's ``state_dict()`` produces keys **without** the prefix (e.g., - ``add_module_prefix(state_dict)`` — adds ``module.`` prefix to all keys before saving, aligning with Megatron's checkpoint format. - ``strip_module_prefix(state_dict)`` — removes ``module.`` prefix after loading, - aligning with the FSDP v2 model's expected key format. + aligning with the Megatron FSDP v2 model's expected key format. ``get_model_state_dict(model)`` auto-detects whether the prefix is present and adds it if missing. ``load_checkpoint`` strips the prefix back before calling @@ -899,7 +899,7 @@ ensures model and optimizer state dict keys are consistent in the checkpoint. ## 12. Implementation Checklist -### Phase 1: `distrib_optimizer.py` — FSDP v2 Guard Propagation +### Phase 1: `distrib_optimizer.py` — Megatron FSDP v2 Guard Propagation - [x] `__init__`: Guard early return with `use_megatron_fsdp or use_megatron_fsdp_v2` - [x] `state_dict()`: Return `self.optimizer.state_dict()` for FSDP paths diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/nvfp4_design.md b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/design/nvfp4_design.md similarity index 99% rename from megatron/core/distributed/fsdp/src/megatron_fsdp/v2/nvfp4_design.md rename to megatron/core/distributed/fsdp/src/megatron_fsdp/v2/design/nvfp4_design.md index 0610731f271..f364ce2508e 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/nvfp4_design.md +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/design/nvfp4_design.md @@ -20,7 +20,7 @@ FSDP path, following the same patterns used for FP8 support. - After optimizer step: `quantize_nvfp4_param_shard()` calls TE's `quantize_master_weights` to cast FP32 → NVFP4 in-place in the packed buffer. -### 2.2 What already works (FSDP v2 FP8 path) +### 2.2 What already works (Megatron FSDP v2 FP8 path) - `MixedPrecisionPolicy` (in `v2/mixed_precision.py`) handles FP8 / MXFP8 detection, buffer dtype, raw-data extraction, post-unshard @@ -259,7 +259,7 @@ No changes needed. Existing functions (`is_nvfp4tensor`, `quantize_nvfp4_param_s → TE creates NVFP4Tensor params (packed uint8, shape[-1] // 2) │ ▼ -[FSDP v2 ParameterGroup._init_buffers()] +[Megatron FSDP v2 ParameterGroup._init_buffers()] model_weight_buffer: DataParallelBuffer(uint8, packed shapes, distributed) main_weight_buffer: DataParallelBuffer(fp32, full shapes, distributed) main_grad_buffer: DataParallelBuffer(fp32, full shapes, distributed) diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/tp_support_design.md b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/design/tp_support_design.md similarity index 100% rename from megatron/core/distributed/fsdp/src/megatron_fsdp/v2/tp_support_design.md rename to megatron/core/distributed/fsdp/src/megatron_fsdp/v2/design/tp_support_design.md diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/dp_buffer.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/dp_buffer.py index 6d2d2bd226e..d49085bd7f7 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/dp_buffer.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/dp_buffer.py @@ -674,7 +674,6 @@ def fetch_buffer(self, *, as_shard: bool = False) -> torch.Tensor: @torch.no_grad() def reduce_grad( self, - grad_comm_dtype: Optional[torch.dtype] = None, overwrite_grad: bool = False, ): """Reduce gradients into the optimizer-facing local shard. @@ -690,7 +689,7 @@ def reduce_grad( if self.sharding_strategy in ("no_shard", "optim"): overwrite_grad = True - grad_comm_dtype = grad_comm_dtype or self.dtype + grad_comm_dtype = self.mp_policy.grad_comm_dtype or self.dtype if self.gradient_scaling_factor in (None, 1.0): op = torch.distributed.ReduceOp.SUM @@ -731,12 +730,10 @@ def reduce_grad( # accumulation buffer. The optimizer consumes only this rank's # virtual shard, so the one delayed RS writes directly into that # slice instead of accumulating into a separate shard buffer. - input_buffer = self.data + input_buffer = self.fetch_buffer() output_offset = sm.local_data_index - comm_input = ( - input_buffer if grad_comm_dtype == self.dtype else input_buffer.to(grad_comm_dtype) - ) + comm_input = input_buffer.to(grad_comm_dtype) if prescale: comm_input.mul_(self.gradient_scaling_factor) reduced_grad_shard = comm_input[output_offset : output_offset + sm.size] diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/fsdp_module.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/fsdp_module.py index 704e95a62c6..5beb2418367 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/fsdp_module.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/fsdp_module.py @@ -15,6 +15,7 @@ """FSDPModule implementation for Megatron-FSDP2.""" import logging +import weakref from contextlib import contextmanager, nullcontext from dataclasses import dataclass, field from typing import Any, Dict, List, Optional, Tuple @@ -610,8 +611,22 @@ def _init_fsdp_state( setattr(module, "_fsdp_module_idx", module_idx) setattr(module, "_fsdp_module_name", name) + module._fsdp_pre_backward_done = False + module.post_backward_issued = False module_idx += 1 + # Annotate every non-FSDPModule sub-module with its nearest parent + # FSDPModule. Process bottom-up (reverse forward_order) so that + # child FSDPModules claim their sub-modules before the root reaches + # them. + for module in reversed(forward_order): + for submodule in module.modules(): + if isinstance(submodule, FSDPModule): + continue + if hasattr(submodule, '_fsdp_parent_module'): + continue + submodule._fsdp_parent_module = weakref.ref(module) + if enable_cuda_graph: if len(forward_order) > 1: child_names = [ diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/fully_shard.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/fully_shard.py index 0f63e9b6fe8..4fb05118845 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/fully_shard.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/fully_shard.py @@ -59,6 +59,9 @@ def fully_shard( enable_trace_pool: bool = False, sharding_strategy: str = "optim_grads_params", enable_cuda_graph: bool = False, + fine_grained_hooks: bool = False, + skip_backward_callback: bool = False, # Skip autograd RegisterFSDPBackwardFunction. + skip_final_backward_callback: bool = False, ) -> nn.Module: """ Wrap a module with FSDP sharding semantics. @@ -69,6 +72,21 @@ def fully_shard( 3. Creates ParameterGroup for each group with dedicated buffers 4. Registers forward/backward hooks for unshard/reshard/reduce 5. Replaces module parameters with DTensor representations + + Args: + fine_grained_hooks: If ``True``, register pre-forward/backward hooks + on every sub-module (for EP-overlap / 1F1B schedules). + skip_backward_callback: If ``True``, skip the autograd post-backward + hook (``_register_backward_hook``). Per-layer reshard + reduce_grad + still fires via ``set_fsdp_reshard_hooks`` / ``mfsdp_post_backward_hook`` + for the TransformerLayer; remaining modules (e.g., TEGroupedMLP) are + handled by ``mfsdp_post_backward_final_callback`` after all + ``backward_dw()`` calls complete. Set when ``delay_wgrad_compute=True`` + because weight gradients are not ready during autograd backward. + skip_final_backward_callback: If ``True``, do not auto-enqueue + ``_register_post_backward_final_callback`` during the backward + pre-hook. The caller must invoke the final callback manually + (used by the 1F1B EP overlap schedule). """ if isinstance(module, FSDPModule): raise ValueError( @@ -115,12 +133,26 @@ def fully_shard( module._init_param_main_grad_func() _register_forward_pre_hook( - fsdp_module=module, - fine_grained=mp_policy.fine_grained_forward_hooks_required(module._fsdp_param_groups), + module, + fine_grained=( + fine_grained_hooks + or mp_policy.fine_grained_forward_hooks_required(module._fsdp_param_groups) + ), ) _register_forward_hook(module) - _register_backward_pre_hook(module) - _register_backward_hook(module) + _register_backward_pre_hook( + module, + fine_grained=fine_grained_hooks, + skip_final_callback=skip_final_backward_callback, + ) + # When delay_wgrad_compute is enabled, skip the autograd post-backward + # hook. Per-layer reshard+reduce_grad still fires via set_fsdp_reshard_hooks + # for TransformerLayer; remaining modules (TEGroupedMLP, root) are handled + # by mfsdp_post_backward_final_callback after all backward_dw() calls. + # + # When disabled, the autograd hook does both inline. + if not skip_backward_callback: + _register_backward_hook(module) module.reshard() diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/hooks.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/hooks.py index 738f6820526..6e31bbab9b1 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/hooks.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/hooks.py @@ -16,7 +16,7 @@ import functools import logging -from typing import Any, Callable, Dict, List, Tuple +from typing import Any, Callable, Dict, List, Optional, Tuple import torch import torch.nn as nn @@ -24,243 +24,464 @@ from torch.utils._pytree import tree_flatten, tree_map, tree_unflatten from .allocator import TracePoolAllocator +from .cuda_graph_runner import FSDPCudaGraphRunner from .fsdp_module import FSDPModule, _FSDPState from .utils import RegisterFSDPBackwardFunction -from .cuda_graph_runner import FSDPCudaGraphRunner logger = logging.getLogger(__name__) -def _register_forward_pre_hook(fsdp_module: FSDPModule, fine_grained: bool = False): - """Register a pre-forward hook on *hook_module*. +# --------------------------------------------------------------------------- +# Public helpers +# --------------------------------------------------------------------------- + + +def _find_fsdp_target(hook_module: nn.Module) -> Optional[FSDPModule]: + """Return the nearest parent FSDPModule for *hook_module*. + + Used by fine-grained hooks registered on sub-modules to resolve the + FSDPModule that owns the sub-module. The reference is stored as + ``_fsdp_parent_module`` during FSDP init (a ``weakref.ref`` to avoid + reference cycles). - Called before every ``forward()`` of the FSDP module — both in the - forward and backward (activation recomputation) passes. Handles - root-level phase bookkeeping, parameter unshard, and CUDA graph - capture (once per compatible module). + Returns: + The owning FSDPModule, or ``None`` if the module has no FSDP parent. """ + if isinstance(hook_module, FSDPModule): + return hook_module + parent_ref = getattr(hook_module, '_fsdp_parent_module', None) + if parent_ref is not None: + return parent_ref() + return None + + +def mfsdp_forward_pre_hook(hook_module: nn.Module, args: Any, kwargs: Any): + """Pre-forward hook for FSDP modules and fine-grained sub-modules. + + Resolves the target FSDPModule via :func:`_find_fsdp_target`, performs + parameter unshard, root-phase bookkeeping, and (for direct FSDPModule + calls only) CUDA graph capture. + + **Repeatability**: This function MUST be safe to call multiple times per + module without observable overhead. Fine-grained hook registration + (``_register_forward_pre_hook(fine_grained=True)``) installs the hook on + every sub-module of an FSDPModule. When a sub-module's ``forward()`` is + called, PyTorch triggers the pre-forward hook, which calls this function. + If the enclosing FSDPModule is also directly invoked (and its own pre-forward + hook fires), this function will be invoked again for the same target. + The implementation must handle this gracefully — duplicating a no-op + ``unshard()`` call or re-applying idempotent bookkeeping must not introduce + measurable latency. + """ + target = _find_fsdp_target(hook_module) + if target is None: + return - def forward_pre_hook(hook_module, args, kwargs): - ctx = fsdp_module._fsdp_root_context - assert not ctx.cuda_graph_active, ( - "hooks must not fire during CUDA graph capture" - ) + ctx = target._fsdp_root_context + assert not ctx.cuda_graph_active, ( + "hooks must not fire during CUDA graph capture" + ) - # ---- root: forward-phase setup (once per micro-batch) -------------- - if hook_module is fsdp_module and fsdp_module._fsdp_state._is_root: - if ctx.enable_cuda_graph and ctx.cuda_graph_stream is None: - ctx.cuda_graph_stream = torch.cuda.Stream() - torch.cuda.set_stream(ctx.cuda_graph_stream) - ctx.cuda_graph_pool = torch.cuda.graph_pool_handle() - ctx.forward_phase = True - ctx.backward_phase = False - - # ---- unshard parameters for this module --------------------------- - if ctx.backward_phase: - fsdp_module.unshard(async_op=ctx.enable_unshard_prefetch, bwd_pass=True) - fsdp_module.unshard(async_op=ctx.enable_unshard_prefetch, bwd_pass=False) - - for param_group in fsdp_module._fsdp_param_groups: - param_group._maybe_free_grad_data() - - # ---- CUDA graph capture (once per compatible module) -------------- - if ( - hook_module is fsdp_module - and fsdp_module._fsdp_state.enable_cuda_graph - and (not hasattr(fsdp_module, "_fsdp_cg_runner")) - and not ctx.backward_phase - and fsdp_module.cuda_graph_compatible - ): - if torch.distributed.get_rank() == 0: - logger.debug( - "Capturing CUDA graph for module %s (id=%s)", - fsdp_module._fsdp_module_name, - id(fsdp_module), - ) - cg_runner = FSDPCudaGraphRunner( - fsdp_module, graph_pool=ctx.cuda_graph_pool + # ---- root: forward-phase setup (once per micro-batch) ------------------ + if target._fsdp_state._is_root: + if ctx.enable_cuda_graph and ctx.cuda_graph_stream is None: + ctx.cuda_graph_stream = torch.cuda.Stream() + torch.cuda.set_stream(ctx.cuda_graph_stream) + ctx.cuda_graph_pool = torch.cuda.graph_pool_handle() + ctx.forward_phase = True + ctx.backward_phase = False + + # ---- unshard parameters for this module ------------------------------- + if ctx.backward_phase: + target.unshard(async_op=ctx.enable_unshard_prefetch, bwd_pass=True) + target.unshard(async_op=ctx.enable_unshard_prefetch, bwd_pass=False) + + # ---- free stale grad data (safe to repeat, idempotent) ---------------- + for param_group in target._fsdp_param_groups: + param_group._maybe_free_grad_data() + + # ---- CUDA graph capture (FSDPModule targets only) --------------------- + # Fine-grained hooks fire on sub-modules whose forward args differ from + # the FSDPModule's; CG capture is not meaningful there. + if ( + isinstance(hook_module, FSDPModule) + and target._fsdp_state.enable_cuda_graph + and (not hasattr(target, "_fsdp_cg_runner")) + and not ctx.backward_phase + and target.cuda_graph_compatible + ): + if torch.distributed.get_rank() == 0: + logger.debug( + "Capturing CUDA graph for module %s (id=%s)", + target._fsdp_module_name, + id(target), ) - cg_runner.capture_forward(*args, **kwargs) - cg_runner.install() - fsdp_module._fsdp_cg_runner = cg_runner - if torch.distributed.get_rank() == 0: - logger.debug( - "Captured CUDA graph for module %s (id=%s)", - fsdp_module._fsdp_module_name, - id(fsdp_module), - ) + cg_runner = FSDPCudaGraphRunner( + target, graph_pool=ctx.cuda_graph_pool + ) + cg_runner.capture_forward(*args, **kwargs) + cg_runner.install() + target._fsdp_cg_runner = cg_runner + if torch.distributed.get_rank() == 0: + logger.debug( + "Captured CUDA graph for module %s (id=%s)", + target._fsdp_module_name, + id(target), + ) + + +def mfsdp_post_forward_hook(module: nn.Module, *unused): + """Post-forward hook: reshard parameters. + + Only supports direct FSDPModule calls. Raises ``TypeError`` when + called with a non-FSDPModule (fine-grained path is not yet handled). + """ + if not isinstance(module, FSDPModule): + raise TypeError( + "mfsdp_post_forward_hook only supports FSDPModule, " + f"got {type(module).__name__}" + ) + ctx = module._fsdp_root_context + assert not ctx.cuda_graph_active, ( + "hooks must not fire during CUDA graph capture" + ) + if ctx.backward_phase and id(module) == ctx.backward_module: + return + module.reshard() + +# --------------------------------------------------------------------------- +# Hook registration +# --------------------------------------------------------------------------- + + +def _register_forward_pre_hook( + module: FSDPModule, fine_grained: bool = False +) -> None: + """Register a pre-forward hook on the FSDP module or its sub-modules. + + Args: + fsdp_module: The FSDP module to instrument. + fine_grained: If ``True``, register on every sub-module of + *fsdp_module* (for EP-overlap / 1F1B schedules). + ``_fsdp_parent_module`` must already be set on sub-modules + (done by :meth:`FSDPModule._init_fsdp_state`). + """ if fine_grained: - for submodule in fsdp_module.modules(): + for submodule in module.modules(): + fsdp_module = _find_fsdp_target(submodule) + if fsdp_module is None or fsdp_module is not module: + continue submodule.register_forward_pre_hook( - forward_pre_hook, prepend=True, with_kwargs=True, + mfsdp_forward_pre_hook, prepend=True, with_kwargs=True, ) else: - fsdp_module.register_forward_pre_hook( - forward_pre_hook, prepend=True, with_kwargs=True + module.register_forward_pre_hook( + mfsdp_forward_pre_hook, prepend=True, with_kwargs=True ) def _register_forward_hook(module: FSDPModule): """Register post-forward hook to reshard parameters.""" + module._mfsdp_forward_hook = module.register_forward_hook(mfsdp_post_forward_hook) - def reshard_param_groups(module, *unused): - ctx = module._fsdp_root_context - assert not ctx.cuda_graph_active, ( - "hooks must not fire during CUDA graph capture" - ) - if ctx.backward_phase and id(module) == ctx.backward_module: - return - module.reshard() - module._mfsdp_forward_hook = module.register_forward_hook(reshard_param_groups) +# --------------------------------------------------------------------------- +# Internal: backward hook helpers +# --------------------------------------------------------------------------- -def _register_backward_pre_hook(module: FSDPModule): - """ - Register backward pre-hook using multi-grad hooks on output tensors. +def mfsdp_pre_backward_setup( + hook_module: nn.Module, grads: Any = None, skip_final_callback: bool = False +): + """Pre-backward hook for FSDP modules and fine-grained sub-modules. - Attaches a ``register_multi_grad_hook`` to every tensor output of - ``module.forward()``. When autograd reaches this module during the - backward pass, the hook fires *before* the module's own backward, - giving FSDP a chance to unshard parameters for gradient computation. + Resolves the target FSDPModule via :func:`_find_fsdp_target`, performs + backward-phase root setup, parameter unshard, and TE gradient-fusion + bookkeeping. The ``_fsdp_pre_backward_done`` flag prevents redundant + calls when multiple sub-modules share the same parent. + + Compatible with ``register_multi_grad_hook`` callback signature + (module, grads). + + Args: + hook_module: Module whose backward pass is about to start. + grads: Gradients from ``register_multi_grad_hook`` (unused). + skip_final_callback: If ``True``, do **not** auto-enqueue + ``mfsdp_post_backward_final_callback``. The caller is + responsible for calling it manually (used by the 1F1B EP + overlap schedule). """ + target = _find_fsdp_target(hook_module) + if target is None: + return + if target._fsdp_pre_backward_done: + return - def create_custom_backward_hook( - module: FSDPModule, custom_backward_handler: Callable - ): - """Wrap *module* so that ``custom_backward_handler`` fires as a - pre-backward hook via ``register_multi_grad_hook``.""" + _pre_backward_setup(target, skip_final_callback=skip_final_callback) + target._fsdp_pre_backward_done = True - @torch.compiler.disable - def forward_hook(_module, inputs, output): - assert not _module._fsdp_root_context.cuda_graph_active, ( - "hooks must not fire during CUDA graph capture" - ) - # ``view_as`` ensures the autograd graph sees a distinct tensor - # object even when the module returns a view of an input. - output = tree_map( - lambda t: t.view_as(t) if torch.is_tensor(t) else t, output - ) - output_list = [] - if isinstance(output, torch.Tensor): - output_list = [output] - elif isinstance(output, (tuple, list)): - output_list = [t for t in output if isinstance(t, torch.Tensor)] +def mfsdp_post_backward_hook(module: nn.Module): + """Post-backward hook: reshard parameters and reduce gradients. - torch.autograd.graph.register_multi_grad_hook( - output_list, - lambda grads: custom_backward_handler(_module, grads), - mode="any", - ) - return output - - return module.register_forward_hook(forward_hook) - - def pre_backward_hook(module: FSDPModule, grads): - """Pre-backward callback for a single FSDP module. - - Invoked by ``register_multi_grad_hook`` when autograd reaches - this module during the backward pass — *before* the module's - own ``backward()`` runs. Executed in reverse forward order. - - Execution flow - -------------- - 1. **Assert** not inside CUDA graph capture. - 2. **Root setup** (first module in backward): - clear backward tracking, switch phase flags, advance the - backward module cursor, and enqueue the post-backward final - callback (once). - 3. **Unshard** parameters for this module (bwd pass). - 4. **Reset per-module state**: mark ``post_backward_issued`` - as not yet handled, and reset Transformer Engine gradient - accumulation flags. - 5. **TE wgrad fusion** (``optim_grads`` / ``optim_grads_params``): - set ``overwrite_main_grad`` so TE writes weight gradients - directly into ``param.main_grad`` (overwrite, not accumulate) - to avoid double-counting across micro-batches. - 6. **Unshard main_grad_buffer** if the parameter group owns one. - """ - ctx = module._fsdp_root_context - assert not ctx.cuda_graph_active, ( - "hooks must not fire during CUDA graph capture" + Only supports direct FSDPModule calls. Raises ``TypeError`` when + called with a non-FSDPModule (fine-grained path is not yet handled). + """ + if not isinstance(module, FSDPModule): + raise TypeError( + "mfsdp_post_backward_hook only supports FSDPModule, " + f"got {type(module).__name__}" ) - - # ---- root: backward-phase setup ----------------------------------- - if module._fsdp_state._is_root: - ctx.backward_done_modules.clear() - ctx.forward_phase = False - ctx.backward_phase = True - ctx._advance_backward_module() - if not module._fsdp_state._post_backward_callback_queued: - _register_post_backward_final_callback(module._fsdp_state, module) - - # ---- unshard params for backward compute -------------------------- - module.unshard(async_op=ctx.enable_unshard_prefetch, bwd_pass=True) - - # ---- reset per-module bookkeeping --------------------------------- - # ``post_backward_issued`` guards against the post-backward path - # (reshard + reduce_grad) being skipped. - setattr(module, "post_backward_issued", False) - - # ---- Transformer Engine gradient-accumulation fusion --------------- - # - # In the eager (trace) backward, TE sets - # ``grad_added_to_main_grad = True`` on each param whose weight - # gradient it writes directly to ``main_grad``. Under CUDA graph - # replay the GPU kernel still runs, but the Python-side - # ``setattr`` is not part of the graph. We recorded the per-param - # flag after the trace backward (see ``_post_backward_final_callback``) - # and restore it here so that ``reduce_grad`` knows TE already - # populated ``main_grad``. - for param_group in module._fsdp_param_groups: - for param in param_group.params: - setattr(param, "grad_added_to_main_grad", False) - if param_group.sharding_strategy in ( - "optim_grads_params", - "optim_grads", - ): - # TE's backward kernel writes weight gradients directly - # into param.main_grad. By default TE *accumulates* - # (adds) into main_grad, which silently doubles - # gradients when the buffer isn't zeroed between - # micro-batches. ``overwrite_main_grad`` tells TE to - # overwrite instead. - setattr(param, "overwrite_main_grad", True) - if param_group.main_grad_buffer is not None: - param_group._init_dist_grads() - param_group.main_grad_buffer.fetch_buffer() - - module._mfsdp_backward_pre_hook = create_custom_backward_hook( - module, custom_backward_handler=pre_backward_hook + ctx = module._fsdp_root_context + assert not ctx.cuda_graph_active, ( + "hooks must not fire during CUDA graph capture" ) + for submodule in module._get_fsdp_modules(recursive=True): + if submodule.post_backward_issued: + continue + ctx.backward_done_modules.add(id(submodule)) + submodule.reshard() + if any( + param_group.sharding_strategy in ("optim_grads", "optim_grads_params") + for param_group in submodule._fsdp_param_groups + ): + submodule.reduce_grad(async_op=ctx.enable_async_reduce_grad) + submodule.post_backward_issued = True + ctx._advance_backward_module() -def _register_backward_hook(module: FSDPModule): - """ - Register backward hook using autograd Function. - This inserts a RegisterFSDPBackwardFunction in the backward pass - that triggers reshard and reduce_grad after gradients are computed. - """ +def mfsdp_post_backward_final_callback(root_module: nn.Module): + """Finalise the backward pass: drain skipped modules, reset state, + clear fine-grained flags, and (on the first micro-batch) transition + the bucket allocator from trace to optimized plan. - def post_backward(module: FSDPModule): - """Hook called after backward pass for this module.""" - ctx = module._fsdp_root_context - assert not ctx.cuda_graph_active, ( - "hooks must not fire during CUDA graph capture" + Only supports the root FSDP module. Raises ``TypeError`` if + *root_module* is not an FSDPModule, or ``RuntimeError`` if it is + not marked as root. + """ + if not isinstance(root_module, FSDPModule): + raise TypeError( + "mfsdp_post_backward_final_callback only supports FSDPModule, " + f"got {type(root_module).__name__}" ) - ctx.backward_done_modules.add(id(module)) - ctx._advance_backward_module() + if not root_module._fsdp_state._is_root: + raise RuntimeError( + "mfsdp_post_backward_final_callback requires root FSDP module" + ) + + ctx = root_module._fsdp_root_context + assert not ctx.cuda_graph_active, ( + "hooks must not fire during CUDA graph capture" + ) + + # ---- handle modules whose per-module post-backward was skipped ---- + for module in reversed(ctx.forward_order): + if module.post_backward_issued: + continue module.reshard() if any( param_group.sharding_strategy in ("optim_grads", "optim_grads_params") for param_group in module._fsdp_param_groups ): module.reduce_grad(async_op=ctx.enable_async_reduce_grad) - module.post_backward_issued = True + # ---- drain pending async reduce-grad events ----------------------- + stream = ctx.rs_stream + for buckets in ctx.reduce_grad_buckets.values(): + while len(buckets) > 0: + event, param_group = buckets.pop() + event.wait() + param_group.release_grad_buffer() + torch.cuda.current_stream().wait_stream(stream) + + # ---- reset root / context state for the next micro-batch ---------- + root_module._fsdp_state._post_backward_callback_queued = False + ctx.backward_phase = False + ctx.backward_module = None + ctx.backward_done_modules.clear() + + # ---- clear fine-grained pre-backward flags ------------------------- + for module in ctx.forward_order: + module._fsdp_pre_backward_done = False + + # ---- trace → optimized transition (first micro-batch only) -------- + if isinstance(ctx.bucket_allocator, TracePoolAllocator): + bucket_alloc = ctx.bucket_allocator + if bucket_alloc.phase == "trace": + if torch.distributed.get_rank() == 0: + logger.debug(bucket_alloc.dump_trace()) + for m in ctx.forward_order: + logger.debug( + f"module_id={id(m)}, module_name={m._fsdp_module_name}" + ) + bucket_alloc.plan() + elif bucket_alloc.phase != "optimized": + raise ValueError( + f"Unexpected bucket allocator phase: {bucket_alloc.phase}" + ) + + +# --------------------------------------------------------------------------- +# Internal: backward hook helpers +# --------------------------------------------------------------------------- + + +def _create_custom_backward_hook( + module: nn.Module, + custom_backward_handler: Callable, + ctx_module: Optional[nn.Module] = None, +): + """Wrap *module* so that ``custom_backward_handler`` fires as a + pre-backward hook via ``register_multi_grad_hook``. + + Args: + module: Module whose output tensors are instrumented. + custom_backward_handler: Callback invoked when backward reaches + this module. + ctx_module: Module whose ``_fsdp_root_context`` is checked for + CUDA-graph safety. Defaults to *module*. + """ + _ctx_source = ctx_module if ctx_module is not None else module + + @torch.compiler.disable + def forward_hook(_module, inputs, output): + if hasattr(_ctx_source, '_fsdp_root_context'): + assert not _ctx_source._fsdp_root_context.cuda_graph_active, ( + "hooks must not fire during CUDA graph capture" + ) + output = tree_map( + lambda t: t.view_as(t) if torch.is_tensor(t) else t, output + ) + + output_list = [] + if isinstance(output, torch.Tensor): + output_list = [output] + elif isinstance(output, (tuple, list)): + output_list = [t for t in output if isinstance(t, torch.Tensor)] + + torch.autograd.graph.register_multi_grad_hook( + output_list, + lambda grads: custom_backward_handler(_module, grads), + mode="any", + ) + return output + + return module.register_forward_hook(forward_hook) + + +def _pre_backward_setup( + module: FSDPModule, skip_final_callback: bool = False +): + """Shared pre-backward logic: root setup, unshard, TE flags. + + Used by both the normal and fine-grained backward pre-hook paths. + + Args: + module: The FSDPModule whose backward is starting. + skip_final_callback: If ``True``, do not enqueue the post-backward + final callback. The caller must call + ``mfsdp_post_backward_final_callback`` manually later. + + .. note:: + + When CUDA graph is enabled, TE wgrad fusion writes directly into + ``param.main_grad`` during the trace (eager) backward. Under graph + replay only the GPU kernel runs — the Python-side flags that mark + ``grad_added_to_main_grad`` are not part of the graph. We eagerly + allocate the main gradient buffer and its full unsharded fetch-buffer + here so that memory addresses are fixed across graph replay iterations. + """ + ctx = module._fsdp_root_context + assert not ctx.cuda_graph_active, ( + "hooks must not fire during CUDA graph capture" + ) + + # ---- root: backward-phase setup ----------------------------------- + if module._fsdp_state._is_root: + ctx.backward_done_modules.clear() + ctx.forward_phase = False + ctx.backward_phase = True + ctx._advance_backward_module() + if not skip_final_callback and not module._fsdp_state._post_backward_callback_queued: + _register_post_backward_final_callback(module._fsdp_state, module) + + # ---- unshard params for backward compute -------------------------- + module.unshard(async_op=ctx.enable_unshard_prefetch, bwd_pass=True) + + # ---- reset per-module bookkeeping --------------------------------- + module.post_backward_issued = False + + # ---- Transformer Engine gradient-accumulation fusion --------------- + for param_group in module._fsdp_param_groups: + for param in param_group.params: + param.grad_added_to_main_grad = False + if param_group.sharding_strategy in ( + "optim_grads_params", + "optim_grads", + ): + param.overwrite_main_grad = True + # CUDA graph + TE wgrad fusion: during graph capture the eager backward + # runs once and TE sets grad_added_to_main_grad=True on each param it + # writes to. Under replay only the GPU kernel runs — the Python-side + # setattr is not part of the graph. We must allocate the main gradient + # buffer and its full unsharded fetch-buffer BEFORE capture so that + # memory addresses are fixed across replay iterations. Without this, + # TE would write to stale or uninitialised buffer addresses on replay. + if module._fsdp_state.enable_cuda_graph and param_group.main_grad_buffer is not None: + param_group._init_dist_grads() + param_group.main_grad_buffer.fetch_buffer() + + return ctx + + +# --------------------------------------------------------------------------- +# Backward hook registration +# --------------------------------------------------------------------------- + + +def _register_backward_pre_hook( + module: FSDPModule, + fine_grained: bool = False, + skip_final_callback: bool = False, +): + """Register backward pre-hook using multi-grad hooks on output tensors. + + Attaches a ``register_multi_grad_hook`` to every tensor output of + ``module.forward()``. When autograd reaches this module during the + backward pass, the hook fires *before* the module's own backward, + giving FSDP a chance to unshard parameters for gradient computation. + """ + if fine_grained: + for submodule in module.modules(): + fsdp_module = _find_fsdp_target(submodule) + if fsdp_module is None or fsdp_module is not module: + continue + submodule._mfsdp_backward_pre_hook = _create_custom_backward_hook( + submodule, + custom_backward_handler=lambda m, g: mfsdp_pre_backward_setup( + m, g, skip_final_callback=skip_final_callback + ), + ctx_module=module, + ) + return + + module._mfsdp_backward_pre_hook = _create_custom_backward_hook( + module, custom_backward_handler=lambda m, g: mfsdp_pre_backward_setup( + m, g, skip_final_callback=skip_final_callback + ), + ) + + +def _register_backward_hook(module: FSDPModule): + """ + Register backward hook using autograd Function. + + This inserts a RegisterFSDPBackwardFunction in the backward pass + that triggers ``mfsdp_post_backward_hook`` after gradients are + computed — resharding parameters and reducing gradients. + """ @torch.compiler.disable def _register_post_backward_hook( post_backward_hook: Callable, @@ -314,10 +535,16 @@ def _register_post_backward_hook( return args, kwargs module._mfsdp_backward_hook = module.register_forward_pre_hook( - functools.partial(_register_post_backward_hook, post_backward), with_kwargs=True + functools.partial(_register_post_backward_hook, mfsdp_post_backward_hook), + with_kwargs=True, ) +# --------------------------------------------------------------------------- +# Post-backward final callback +# --------------------------------------------------------------------------- + + def _register_post_backward_final_callback( state: _FSDPState, module: nn.Module ) -> None: @@ -326,70 +553,13 @@ def _register_post_backward_final_callback( backward pass has completed. Registered once by the root FSDP module (avoids duplicates). - The callback: - - Reshards and reduces gradients for any module whose per-module - post-backward hook was silently skipped (e.g. activation - recomputation). - - Waits for all async reduce-grad operations to finish. - - Resets root / context state for the next micro-batch. - - On the first call (trace → optimized transition): builds the - pool plan and records TE wgrad-fusion flags for CUDA graph restore. + Delegates to :func:`mfsdp_post_backward_final_callback`. """ assert state._is_root, "Only root FSDP should register post-backward callback" if state._post_backward_callback_queued: return - def _post_backward_final_callback(root_state: _FSDPState, root_module: nn.Module): - """Engine callback — the last thing autograd runs after the backward - pass of every micro-batch.""" - ctx = root_module._fsdp_root_context - assert not ctx.cuda_graph_active, ( - "hooks must not fire during CUDA graph capture" - ) - - # ---- handle modules whose per-module post-backward was skipped ---- - for module in reversed(ctx.forward_order): - if getattr(module, "post_backward_issued", False): - continue - module.reshard() - if any( - param_group.sharding_strategy in ("optim_grads", "optim_grads_params") - for param_group in module._fsdp_param_groups - ): - module.reduce_grad(async_op=ctx.enable_async_reduce_grad) - - # ---- drain pending async reduce-grad events ----------------------- - stream = ctx.rs_stream - for buckets in ctx.reduce_grad_buckets.values(): - while len(buckets) > 0: - event, param_group = buckets.pop() - event.wait() - param_group.release_grad_buffer() - torch.cuda.current_stream().wait_stream(stream) - - # ---- reset root / context state for the next micro-batch ---------- - root_state._post_backward_callback_queued = False - ctx.backward_phase = False - ctx.backward_module = None - ctx.backward_done_modules.clear() - - # ---- trace → optimized transition (first micro-batch only) -------- - if isinstance(ctx.bucket_allocator, TracePoolAllocator): - bucket_alloc = ctx.bucket_allocator - if bucket_alloc.phase == "trace": - if torch.distributed.get_rank() == 0: - logger.debug(bucket_alloc.dump_trace()) - for m in ctx.forward_order: - logger.debug( - f"module_id={id(m)}, module_name={m._fsdp_module_name}" - ) - bucket_alloc.plan() - elif bucket_alloc.phase != "optimized": - raise ValueError( - f"Unexpected bucket allocator phase: {bucket_alloc.phase}" - ) - state._post_backward_callback_queued = True Variable._execution_engine.queue_callback( - functools.partial(_post_backward_final_callback, state, module) + functools.partial(mfsdp_post_backward_final_callback, module) ) diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/param_group.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/param_group.py index 66e37e74dbb..7dc9ba26777 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/param_group.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/param_group.py @@ -268,7 +268,6 @@ def reduce_grad(self): # so the first reduce_grad after either event overwrites instead of # accumulating — no stale data from uninitialised or zeroed buffers. self.main_grad_buffer.reduce_grad( - grad_comm_dtype=self.mp_policy.grad_comm_dtype, overwrite_grad=self._grad_buffer_is_fresh, ) self._grad_buffer_is_fresh = False diff --git a/megatron/core/optimizer/clip_grads.py b/megatron/core/optimizer/clip_grads.py index 01095a50154..372d33b926e 100644 --- a/megatron/core/optimizer/clip_grads.py +++ b/megatron/core/optimizer/clip_grads.py @@ -235,11 +235,13 @@ def count_zeros_fp32( use_megatron_fsdp = False for param in parameters: if getattr(param, "__fsdp_param__", False) and param.grad is not None: - # If the parameter is managed by Megatron FSDP, we need to handle it differently. use_megatron_fsdp = True - grad = param.grad._local_tensor - num_zeros = grad.numel() - torch.count_nonzero(grad) - total_num_zeros += num_zeros + is_not_shared = param_is_not_shared(param) + is_not_tp_duplicate = param_is_not_tensor_parallel_duplicate(param, tp_group=tp_group) + if is_not_shared and is_not_tp_duplicate: + grad = param.grad._local_tensor + num_zeros = grad.numel() - torch.count_nonzero(grad) + total_num_zeros += num_zeros continue grad_attr = "decoupled_grad" if use_decoupled_grad else "grad" diff --git a/megatron/core/pipeline_parallel/combined_1f1b.py b/megatron/core/pipeline_parallel/combined_1f1b.py index b1ebbb876ff..4a05131dc30 100644 --- a/megatron/core/pipeline_parallel/combined_1f1b.py +++ b/megatron/core/pipeline_parallel/combined_1f1b.py @@ -2,9 +2,10 @@ import contextlib from contextlib import nullcontext -from typing import List, Union +from typing import List, Optional, Union import torch +import torch.nn as nn from megatron.core.distributed.fsdp.src.megatron_fsdp.utils import find_megatron_fsdp from megatron.core.enums import Fp8Recipe @@ -21,6 +22,128 @@ Shape = Union[List[int], torch.Size] +# --------------------------------------------------------------------------- +# FSDP version-agnostic helpers +# --------------------------------------------------------------------------- + + +def _find_mfsdp_root_module(model: nn.Module) -> Optional[nn.Module]: + """Return the root Megatron FSDP module, or ``None`` if not using M-FSDP. + + Supports both v1 (``MegatronFSDP``) and v2 (``FSDPModule``) via duck + typing. For v2 the model may be wrapped in ``FullyShardedDataParallel``, + so we walk the sub-modules. + """ + # v1: the model object itself may be the MegatronFSDP wrapper + v1 = find_megatron_fsdp(model) + if v1 is not None: + return v1 + + # v2: walk sub-modules to find the root FSDPModule + v2_modules: List[nn.Module] = [] + for child in model.modules(): + if hasattr(child, '_fsdp_state'): + v2_modules.append(child) + if getattr(child._fsdp_state, '_is_root', False): + return child + + if v2_modules: + raise RuntimeError( + "Found M-FSDP v2 modules but none marked as root. " + "Ensure the root FSDP module has _fsdp_state._is_root = True. " + "This is normally set by fully_shard() on the outermost module. " + f"v2 modules found: {[m.__class__.__name__ for m in v2_modules[:5]]}" + ) + + return None + + +def _get_mfsdp_post_backward_final_callback(root_module: nn.Module): + """Return a ``post_backward_final_callback`` callable for *root_module*. + + v1: ``root_module.post_backward`` (MegatronFSDP._root_post_backward). + v2: ``mfsdp_post_backward_final_callback`` from the v2 hooks module. + """ + # v2 + if hasattr(root_module, '_fsdp_root_context'): + from megatron.core.distributed.fsdp.src.megatron_fsdp.v2.hooks import ( + mfsdp_post_backward_final_callback, + ) + return mfsdp_post_backward_final_callback + + # v1 + return root_module.post_backward + + +def _get_mfsdp_pre_backward_setup(root_module: nn.Module): + """Return a ``pre_backward_setup(hook_module, grads, *, skip_final_callback)`` + callable for *root_module*. + + v1: calls ``root_module.pre_backward()`` (which has *skip_backward_hook* + semantics baked in via the stored partial). + v2: ``mfsdp_pre_backward_setup`` from the v2 hooks module. + """ + # v2 + if hasattr(root_module, '_fsdp_root_context'): + from megatron.core.distributed.fsdp.src.megatron_fsdp.v2.hooks import ( + mfsdp_pre_backward_setup, + ) + + def _v2_pre_backward(hook_module, grads=None, *, skip_final_callback=True): + mfsdp_pre_backward_setup( + hook_module, grads, skip_final_callback=skip_final_callback + ) + + return _v2_pre_backward + + # v1: MegatronFSDP.pre_backward is a partial of _root_pre_backward with + # skip_backward_hook=True — exactly what the overlap schedule needs. + return root_module.pre_backward + + +def _get_mfsdp_reshard_hooks(root_module: nn.Module): + """Return ``(post_forward_hook, post_backward_hook)`` for per-layer + parameter release in the overlap schedule. + + v1: ``(post_forward_release_module, post_backward_release_module)``. + v2: ``(mfsdp_post_forward_hook, mfsdp_post_backward_hook)`` — these + assert ``isinstance(FSDPModule)``, which is satisfied because the + schedule plan layer is an FSDPModule. + """ + # v2 + if hasattr(root_module, '_fsdp_root_context'): + from megatron.core.distributed.fsdp.src.megatron_fsdp.v2.hooks import ( + mfsdp_post_backward_hook, + mfsdp_post_forward_hook, + ) + return mfsdp_post_forward_hook, mfsdp_post_backward_hook + + # v1 + return root_module.post_forward_release_module, root_module.post_backward_release_module + + +def _get_mfsdp_sharding_strategy(root_module: nn.Module) -> Optional[str]: + """Return the data-parallel sharding strategy string, or ``None``.""" + # v2 + if hasattr(root_module, '_fsdp_root_context'): + for child in root_module.modules(): + if hasattr(child, '_fsdp_param_groups'): + for pg in child._fsdp_param_groups: + return pg.sharding_strategy + return "no_shard" + + # v1 + if hasattr(root_module, 'ddp_config'): + return root_module.ddp_config.data_parallel_sharding_strategy + + return None + + +# --------------------------------------------------------------------------- +# Schedules +# --------------------------------------------------------------------------- + + def combined_1f1b_schedule_for_no_pipelining( forward_step_func, data_iterator, @@ -54,13 +177,17 @@ def combined_1f1b_schedule_for_no_pipelining( """ set_streams(high_priority=config.high_priority_a2a_comm_stream) - fsdp_wrapper = find_megatron_fsdp(model) - if fsdp_wrapper is not None: - # The overlap schedule bypasses MegatronFSDP.forward(), which normally - # swaps distributed (optimizer-managed) parameters back to raw parameters. - # We must do this explicitly before the schedule accesses layers directly. - fsdp_wrapper._replace_param_with_raw_if_needed() + # Resolve FSDP root module (v1 or v2) + root_module = _find_mfsdp_root_module(model) + is_v1 = root_module is not None and hasattr(root_module, 'ddp_config') + is_v2 = root_module is not None and hasattr(root_module, '_fsdp_root_context') + is_mfsdp = is_v1 or is_v2 + + if is_v1: + # v1: swap distributed (optimizer-managed) params → raw params before + # the schedule accesses layers directly. + root_module._replace_param_with_raw_if_needed() # The forward step for the first microbatch is executed alone, no a2a overlapping output_tensor, num_tokens, _ = combined_forward_backward_step( @@ -79,6 +206,9 @@ def combined_1f1b_schedule_for_no_pipelining( checkpoint_activations_microbatch=None, is_first_microbatch=check_first_val_step(True), current_microbatch=0, + root_module=root_module, + is_mfsdp=is_mfsdp, + is_v1=is_v1, ) # The forward step is executed in parallel with the backward step of another microbatch # EP A2A in forward step is hidden by the attention/mlp computation in the backward step @@ -102,7 +232,9 @@ def combined_1f1b_schedule_for_no_pipelining( checkpoint_activations_microbatch=None, is_first_microbatch=check_first_val_step((i + 1) == 0), current_microbatch=(i + 1), - fsdp_wrapper=fsdp_wrapper, + root_module=root_module, + is_mfsdp=is_mfsdp, + is_v1=is_v1, ) total_num_tokens += num_tokens # The backward step for the last microbatch is executed alone, no a2a overlapping @@ -119,7 +251,9 @@ def combined_1f1b_schedule_for_no_pipelining( output_tensor, # b_output_tensor output_tensor_grad, # b_output_tensor_grad config, - fsdp_wrapper=fsdp_wrapper, + root_module=root_module, + is_mfsdp=is_mfsdp, + is_v1=is_v1, ) return forward_data_store, total_num_tokens @@ -191,12 +325,10 @@ def combined_1f1b_schedule_for_interleaved_pipelining(): set_streams(high_priority=config.high_priority_a2a_comm_stream) - # Interleaved pipeline with FSDP(optim_grads_params) is not yet supported: - # _replace_param_with_raw_if_needed() and root pre/post_backward() are not - # handled for multi-chunk models in this path. + # Interleaved pipeline with FSDP(optim_grads_params) is not yet supported. if isinstance(model, (list, tuple)): for m in model: - assert find_megatron_fsdp(m) is None, ( + assert find_megatron_fsdp(m) is None and _find_mfsdp_root_module(m) is None, ( "EP overlap 1F1B with FSDP is not supported for interleaved " "pipeline parallelism (virtual_pipeline_model_parallel_size > 1). " "Use pipeline_model_parallel_size=1 or disable FSDP." @@ -289,7 +421,9 @@ def combined_forward_backward_step( is_first_microbatch=False, current_microbatch=None, encoder_decoder_xattn=False, - fsdp_wrapper=None, + root_module=None, + is_mfsdp=False, + is_v1=False, ): """Merged forward and backward step for combined 1f1b scheduler. @@ -304,6 +438,12 @@ def combined_forward_backward_step( post_forward (callable): The function to call after the forward_step. post_backward (callable): The function to call after the backward_step. + root_module: Root Megatron FSDP module (v1 ``MegatronFSDP`` or v2 root + ``FSDPModule``), or ``None`` if FSDP is not in use. + is_mfsdp: ``True`` if FSDP is active (v1 or v2). + is_v1: ``True`` if using v1 (``MegatronFSDP`` with ``ddp_config``); + ``False`` for v2. + Returns: forward_output_tensor (Tensor or list[Tensor]): The output object(s) from the forward step. forward_num_tokens (Tensor): The number of tokens. @@ -335,8 +475,12 @@ def forward_backward_step(): from .schedules import set_current_microbatch - if fsdp_wrapper is not None and b_model is not None: - fsdp_wrapper.pre_backward() + if is_mfsdp and b_model is not None: + if is_v1: + root_module.pre_backward() + else: + pre_backward_fn = _get_mfsdp_pre_backward_setup(root_module) + pre_backward_fn(root_module, skip_final_callback=True) if f_model is not None and config.timers is not None: config.timers('forward-compute', log_level=2).start() @@ -387,20 +531,15 @@ def forward_backward_step(): # schedule bypasses normal FSDP forward/backward hooks, so we release # each layer's all-gathered parameters explicitly after its compute. # Only needed for optim_grads_params strategy (where params are sharded). - forward_fsdp_wrapper = find_megatron_fsdp(f_model) - if ( - forward_fsdp_wrapper is not None - and forward_fsdp_wrapper.ddp_config.data_parallel_sharding_strategy - == "optim_grads_params" - ): - for i in range(f_schedule_plan.num_layers()): - layer_plan = f_schedule_plan.get_layer(i) - # Wire explicit FSDP reshard hooks for the EP-overlap schedule, - # which bypasses the normal TransformerLayer-level FSDP hooks. - layer_plan.set_fsdp_reshard_hooks( - forward_fsdp_wrapper.post_forward_release_module, - forward_fsdp_wrapper.post_backward_release_module, - ) + if is_mfsdp: + f_root = _find_mfsdp_root_module(f_model) + if f_root is not None: + sharding_strategy = _get_mfsdp_sharding_strategy(f_root) + if sharding_strategy == "optim_grads_params": + post_fwd, post_bwd = _get_mfsdp_reshard_hooks(f_root) + for i in range(f_schedule_plan.num_layers()): + layer_plan = f_schedule_plan.get_layer(i) + layer_plan.set_fsdp_reshard_hooks(post_fwd, post_bwd) # backward preprocess, the same as the backward_step() unwrap_input_tensor_grad = False @@ -497,7 +636,11 @@ def forward_backward_step(): if unwrap_input_tensor_grad: input_tensor_grad = input_tensor_grad[0] - if fsdp_wrapper is not None and b_model is not None: - fsdp_wrapper.post_backward() + if is_mfsdp and b_model is not None: + if is_v1: + root_module.post_backward() + else: + post_backward_fn = _get_mfsdp_post_backward_final_callback(root_module) + post_backward_fn(root_module) return output_tensor, num_tokens, input_tensor_grad diff --git a/megatron/training/utils.py b/megatron/training/utils.py index 0e05c11e52a..ebfee1b0ab9 100644 --- a/megatron/training/utils.py +++ b/megatron/training/utils.py @@ -199,7 +199,10 @@ def calc_dtensor_params_l2_norm(params): """Calculate l2 norm of DTensor parameters.""" params_data = defaultdict(list) for param in params: - params_data[param._spec].append(param._local_tensor) + local_tensor = param._local_tensor + if local_tensor.dtype != torch.float32 and local_tensor.numel() > 0: + local_tensor = local_tensor.float() + params_data[param._spec].append(local_tensor) total_norm_2 = torch.zeros((1,), dtype=torch.float32, device='cuda') dummy_overflow_buf = torch.zeros((1,), dtype=torch.int, device='cuda') diff --git a/tests/unit_tests/distributed/megatron_fsdp/utils.py b/tests/unit_tests/distributed/megatron_fsdp/utils.py index 18a2da63786..975d573f120 100644 --- a/tests/unit_tests/distributed/megatron_fsdp/utils.py +++ b/tests/unit_tests/distributed/megatron_fsdp/utils.py @@ -163,7 +163,7 @@ def __getitem__(self, idx): } -def _forward_step_func(data_iterator, model, device="cuda"): +def _forward_step_func(data_iterator, model, device="cuda", return_schedule_plan=False): def loss_func(loss_mask: torch.Tensor, output_tensor: torch.Tensor): @@ -191,6 +191,16 @@ def loss_func(loss_mask: torch.Tensor, output_tensor: torch.Tensor): ) position_ids = data["position_ids"].to(device, non_blocking=True) + if return_schedule_plan: + schedule_plan = model.build_schedule_plan( + input_ids=tokens, + position_ids=position_ids, + attention_mask=attention_mask, + labels=labels, + loss_mask=loss_mask, + ) + return schedule_plan, partial(loss_func, loss_mask) + output_tensor = model(tokens, position_ids, attention_mask, labels=labels) return output_tensor, partial(loss_func, loss_mask) diff --git a/tests/unit_tests/distributed/megatron_fsdp/v2/test_mcore_nd_parallel.py b/tests/unit_tests/distributed/megatron_fsdp/v2/test_mcore_nd_parallel.py index 054b01e7903..78c0cb7036e 100644 --- a/tests/unit_tests/distributed/megatron_fsdp/v2/test_mcore_nd_parallel.py +++ b/tests/unit_tests/distributed/megatron_fsdp/v2/test_mcore_nd_parallel.py @@ -1,5 +1,6 @@ # Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. import copy +import time import pytest import torch @@ -185,6 +186,7 @@ def _training_loop(seed=42, **kwargs): padded_vocab_size=VOCAB_SIZE, seq_length=MAX_SEQ_LEN, sequence_parallel=TP > 1, + expert_model_parallel_size=EP, tensor_model_parallel_size=TP, pipeline_model_parallel_size=PP, num_layers_per_virtual_pipeline_stage=VPP, @@ -211,7 +213,8 @@ def _training_loop(seed=42, **kwargs): param_snapshots = [] # Training loop - for _ in range(NUM_TRAINING_STEPS): + for step in range(NUM_TRAINING_STEPS): + t0 = time.time() optim.zero_grad() output = pretrain_forward_backward( model=model_chunks, @@ -226,6 +229,18 @@ def _training_loop(seed=42, **kwargs): # Collect loss outputs.append(output[-1]) + if torch.distributed.get_rank() == 0: + elapsed = time.time() - t0 + mem_alloc = torch.cuda.memory_allocated() / 1024**3 + mem_reserved = torch.cuda.max_memory_reserved() / 1024**3 + print( + f"[Step {step + 1}/{NUM_TRAINING_STEPS}] " + f"loss={output[-1]['lm loss'].item():.6f} " + f"time={elapsed:.2f}s " + f"mem_alloc={mem_alloc:.2f}GiB " + f"mem_reserved_max={mem_reserved:.2f}GiB" + ) + torch.cuda.reset_peak_memory_stats() if capture_param_snapshots: param_snapshots.append( TestMegatronFSDPE2E._capture_named_params(model_chunks) @@ -290,6 +305,21 @@ def _training_loop(seed=42, **kwargs): ), id="optim_grads_params_nvfp4_param_gather", ), + pytest.param( + dict( + bf16=True, + data_parallel_sharding_strategy="optim_grads_params", + fp8="e4m3", + fp8_param_gather=True, + fp8_recipe="mxfp8", + moe_grouped_gemm=True, + use_megatron_fsdp_v2=True, + moe_token_dispatcher_type="alltoall", + overlap_moe_expert_parallel_comm=True, + delay_wgrad_compute=True, + ), + id="ep_overlap-optim_grads_params", + ), ], ) def test_compatible_with_nd_parallel(self, ref_cache, nd_topology, spec_configs): @@ -314,6 +344,12 @@ def test_compatible_with_nd_parallel(self, ref_cache, nd_topology, spec_configs) except ImportError: pytest.skip("NVFP4 support check requires Transformer Engine >= 2.7.0.dev0") + if spec_configs.get("overlap_moe_expert_parallel_comm"): + from megatron.core.utils import is_te_min_version + + if not is_te_min_version("2.3.0"): + pytest.skip("EP overlap requires Transformer Engine >= 2.3.0") + reference_kind = "distopt" ref_cache_key = ( reference_kind, From 35be05931379db7672679d4393d2757e855758f8 Mon Sep 17 00:00:00 2001 From: Jianbin Chang Date: Tue, 30 Jun 2026 18:29:38 +0800 Subject: [PATCH 26/28] CUDA graph capture via vendored te-graph-runtime with `capture_time_hooks` for MFSDP v2 (#21) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 … --- examples/qwenimage_mfsdp/README.md | 168 ++ examples/qwenimage_mfsdp/test_qwenimage.py | 626 ++++++ .../fsdp/src/megatron_fsdp/utils.py | 2 + .../fsdp/src/megatron_fsdp/v2/README.md | 62 +- .../v2/cuda_graph_memory_analysis.md | 577 +++++ .../src/megatron_fsdp/v2/cuda_graph_runner.py | 608 ++--- .../v2/design/cuda_graph_design.md | 460 ++-- .../fsdp/src/megatron_fsdp/v2/dp_buffer.py | 78 +- .../fsdp/src/megatron_fsdp/v2/fsdp_module.py | 96 +- .../fsdp/src/megatron_fsdp/v2/fully_shard.py | 5 +- .../fsdp/src/megatron_fsdp/v2/hooks.py | 71 +- .../src/megatron_fsdp/v2/mixed_precision.py | 4 +- .../fsdp/src/megatron_fsdp/v2/param_group.py | 28 +- .../v2/te_graph_runtime/README.md | 57 + .../v2/te_graph_runtime/__init__.py | 20 + .../v2/te_graph_runtime/graph.py | 1962 +++++++++++++++++ .../v2/test_cuda_graph_runner.py | 94 + .../megatron_fsdp/v2/test_dist_muon.py | 2 +- 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 create mode 100644 tests/unit_tests/distributed/megatron_fsdp/v2/test_cuda_graph_runner.py diff --git a/examples/qwenimage_mfsdp/README.md b/examples/qwenimage_mfsdp/README.md new file mode 100644 index 00000000000..d16b8c9e722 --- /dev/null +++ b/examples/qwenimage_mfsdp/README.md @@ -0,0 +1,168 @@ +# FSDP1 vs Megatron-FSDP — QwenImage Benchmark + +Minimal repro comparing PyTorch FSDP1 against Megatron-FSDP on +`QwenImageTransformer2DModel`. All code is self-contained in `test_qwenimage.py`; +only stock packages are required below. + +## Environment Setup (run once) + +```bash +pip install "diffusers>=0.37.0" # QwenImage model +pip install megatron-fsdp # if not installed in repo +pip install huggingface_hub # model download + +# Flash attention (pick one tier): +# Tier 1: FA3 — best perf, install from PyPI +pip install flash_attn_interface +# Tier 2: FA2 — if FA3 unavailable +pip install flash-attn --no-build-isolation +# Tier 3: no flash-attn at all → use --attention native below +``` + +## Download Model (run once) + +```bash +hf download Qwen/Qwen-Image \ + --include "transformer/*" \ + --local-dir /tmp/qwen-image +``` + +## Run + +### Single node, 4 GPU (full shard) + +```bash +# Tier 1: FA3 (_flash_3) +torchrun --nproc_per_node=4 test_qwenimage.py \ + --backend mfsdp --sharding full \ + --pretrained_model_name_or_path /tmp/qwen-image \ + --num_gpus_per_node 4 --batch_size 4 --height 512 --width 512 \ + --attention _flash_3 --compile --bench_steps 20 --warmup_steps 3 + +# Tier 2: FA2 (flash) — if FA3 unavailable +nsys profile \ +torchrun --nproc_per_node=4 test_qwenimage.py \ + --backend mfsdp --sharding full \ + --pretrained_model_name_or_path /tmp/qwen-image \ + --num_gpus_per_node 4 --batch_size 4 --height 512 --width 512 \ + --attention flash --compile --bench_steps 3 --warmup_steps 1 + +# Tier 3: native attention — no flash-attn needed +torchrun --nproc_per_node=4 test_qwenimage.py \ + --backend mfsdp --sharding full \ + --pretrained_model_name_or_path /tmp/qwen-image \ + --num_gpus_per_node 4 --batch_size 4 --height 512 --width 512 \ + --attention native --compile --bench_steps 20 --warmup_steps 3 + +# PyTorch FSDP1 for comparison (swap --backend) +nsys profile \ +torchrun --nproc_per_node=4 test_qwenimage.py \ + --backend fsdp1 --sharding full \ + --pretrained_model_name_or_path /tmp/qwen-image \ + --num_gpus_per_node 4 --batch_size 4 --height 512 --width 512 \ + --attention flash --compile --bench_steps 20 --warmup_steps 3 +``` + +### Single node, 8 GPU (hybrid shard) + +```bash +torchrun --nproc_per_node=8 test_qwenimage.py \ + --backend mfsdp --sharding hybrid \ + --pretrained_model_name_or_path /tmp/qwen-image \ + --batch_size 4 --height 512 --width 512 \ + --attention _flash_3 --compile --bench_steps 20 --warmup_steps 3 +``` + +### With numerical verification (adds sync overhead) + +```bash +torchrun --nproc_per_node=4 test_qwenimage.py \ + --backend mfsdp --sharding full \ + --pretrained_model_name_or_path /tmp/qwen-image \ + --num_gpus_per_node 4 --batch_size 4 --height 512 --width 512 \ + --attention flash --compile --verify +``` + +### Multi-node (2+ nodes) + +```bash +torchrun --nnodes=$NNODES --node_rank=$NODE_RANK \ + --nproc_per_node=8 --master_addr=$MASTER_ADDR --master_port=$MASTER_PORT \ + test_qwenimage.py \ + --backend mfsdp --sharding hybrid \ + --pretrained_model_name_or_path /tmp/qwen-image \ + --batch_size 4 --height 512 --width 512 \ + --attention _flash_3 --compile --bench_steps 20 --warmup_steps 3 +``` + +## Benchmarks + +`QwenImageTransformer2DModel`, `bs=4`, `512×512`, `bf16`, `torch.compile`, FA2. +`[mfsdpv2+cg]` uses `--cuda-graph --trace-pool`. + +| Backend | 8×H100 | 4×GB200 | +|---------|--------|---------| +| **fsdp1** | 729 ms / 60.2 GB | 679 ms / 75.4 GB | +| **mfsdpv2** | 769 ms / 59.3 GB | 647 ms / 74.7 GB | +| **mfsdpv2+cg** | **674 ms** / 68.3 GB | **364 ms** / 88.7 GB | + +CG delivers **11% faster** on H100 and **44% faster** on GB200 at the cost +of higher peak memory (pool-backed graph buffers). + +## torch FSDP1 (reference API) + +```python +mesh = init_device_mesh("cuda", (num_nodes, gpus_per_node), + mesh_dim_names=("replicate", "shard")) +mp = MixedPrecision(param_dtype=bf16, reduce_dtype=bf16, buffer_dtype=bf16) +FSDP(model, device_mesh=mesh, sharding_strategy=ShardingStrategy.HYBRID_SHARD, + mixed_precision=mp, + auto_wrap_policy=ModuleWrapPolicy({QwenImageTransformerBlock}), + backward_prefetch=BackwardPrefetch.BACKWARD_PRE, forward_prefetch=True, + use_orig_params=True, limit_all_gathers=True) +``` + +## Megatron-FSDP (reference API) + +```python +mesh = init_device_mesh("cuda", (num_nodes, gpus_per_node), + mesh_dim_names=("dp_outer", "dp_shard")) +mesh[("dp_outer", "dp_shard")]._flatten("hsdp") +mp = MixedPrecisionPolicy(main_params_dtype=bf16, main_grads_dtype=bf16, + grad_comm_dtype=bf16) +fully_shard_model( + module=model, fsdp_unit_modules=[QwenImageTransformerBlock], + device_mesh=mesh, dp_shard_dim="dp_shard", + dp_outer_dim="dp_outer", hybrid_fsdp_group=mesh["hsdp"].get_group(), + zero_dp_strategy=3, # ZeRO-3 within node + outer_dp_sharding_strategy=0, # REPLICATE across nodes (match FSDP1 HYBRID_SHARD) + sync_model_each_microbatch=True, + overlap_grad_reduce=True, overlap_param_gather=True, + mixed_precision_policy=mp, +) + +# Optimizer +fully_shard_optimizer(optim) +optim.step(sync_grad_before_optimizer_step=True, install_optimized_model_weights=True) +optim.zero_grad(set_to_none=True, zero_grad_buffer=True) +``` + +## Notes + +- Per-block `torch.compile` (identical for both backends): + ```python + for blk in model.module.transformer_blocks: + blk.compile() + ``` +- `--verify` gradient probe — mfsdp does not expose gradients via `param.grad`: + ```python + def _local_grad(p): + g = p.grad + if g is None and hasattr(p, "get_main_grad"): g = p.get_main_grad() + if g is None: g = getattr(p, "main_grad", None) + if g is None: g = getattr(p, "decoupled_grad", None) + if hasattr(g, "to_local"): g = g.to_local() + return g + ``` +- FA3 autograd shim is applied inline by the script — no manual patching required. +- `hybrid` sharding on a single node falls back to `full` shard automatically. diff --git a/examples/qwenimage_mfsdp/test_qwenimage.py b/examples/qwenimage_mfsdp/test_qwenimage.py new file mode 100644 index 00000000000..ee1c46fd293 --- /dev/null +++ b/examples/qwenimage_mfsdp/test_qwenimage.py @@ -0,0 +1,626 @@ +""" +Single-file end-to-end training test for QwenImageTransformer2DModel +with Megatron FSDP v1/v2 and PyTorch FSDP1. + +Self-contained — requires only stock packages (`torch`, `diffusers`, +`megatron-fsdp`). The transformer body is the upstream diffusers model. + +Features: + - FSDP1 / Megatron FSDP v1 / v2 backends + - Per-block torch.compile + - CUDA graph capture (v2 with --cuda-graph) + - TracePoolAllocator (v2 with --trace-pool) + - Gradient checkpointing (--gradient_checkpointing) + - FA2 / FA3 / native attention backends + - Memory history OOM dump (--record-memory-history) + - Per-step GPU memory tracking + - NVML real GPU memory logging + - Numerical verification probe (--verify) + - Nsight profiling support (--cuda_profiler_capture) + +Usage (per node): + torchrun --nnodes=$NNODES --node_rank=$NODE_RANK \ + --nproc_per_node=8 --master_addr=$MASTER_ADDR --master_port=$MASTER_PORT \ + test_qwenimage.py \ + --backend {fsdp1|mfsdp|mfsdpv2} \ + --sharding {full|hybrid} \ + --batch_size 4 --height 512 --width 512 \ + --bench_steps 20 --warmup_steps 3 \ + --compile --cuda-graph --trace-pool + +Reports per-step latency, peak GPU memory, and per-step memory summary. +`--verify` probes global loss + grad-norm for numerical agreement across +backends (adds sync overhead — only relative deltas are meaningful). + +Note: mfsdpv2 uses a 1D device mesh (no HSDP). `--sharding hybrid` with +mfsdpv2 falls back to full sharding on the 1D mesh. +""" +import argparse +import atexit +from contextlib import contextmanager +import os +from pathlib import Path +import sys +import time +import traceback + +import torch +import torch.distributed as dist +from torch.distributed.device_mesh import init_device_mesh + +from diffusers.models.transformers.transformer_qwenimage import ( + QwenImageTransformer2DModel, + QwenImageTransformerBlock, +) +from diffusers.models.attention_dispatch import attention_backend + +import logging +logging.basicConfig( + level=logging.DEBUG, + format="%(asctime)s %(name)s %(levelname)s %(message)s", +) + + + +# ---------- FA3 training shim ------------------------------------------------ +# diffusers 0.37.x registers `_flash_3` via a custom-op that only has +# forward + fake — there is no autograd registration (see the TODO in +# diffusers/models/attention_dispatch.py pointing at the unmerged +# Dao-AILab/flash-attention#1590), so training trips +# "Trying to backward through _diffusers_flash_attn_3._flash_attn_forward.default +# but no autograd formula was registered." +# `flash_attn_interface.flash_attn_func` from the FA3 wheel itself IS +# autograd-aware. Patching the registry entry to call it directly is the +# smallest change that makes `attention_backend("_flash_3")` usable in training, +# matching the FA3 path our internal fork (model/qwen_image/modules/attention_utils.py) +# already relies on for H800 training. +def _enable_fa3_training_patch(): + import flash_attn_interface + from diffusers.models.attention_dispatch import ( + _AttentionBackendRegistry, AttentionBackendName, + ) + + def _flash_attention_3_training( + query, key, value, attn_mask=None, scale=None, + is_causal=False, return_lse=False, _parallel_config=None, + ): + if attn_mask is not None: + raise ValueError("`attn_mask` is not supported for flash-attn 3.") + out = flash_attn_interface.flash_attn_func( + q=query, k=key, v=value, + softmax_scale=scale, causal=is_causal, + return_attn_probs=return_lse, + ) + if return_lse: + o, lse, *_ = out + return o, lse.permute(0, 2, 1) + return out[0] if isinstance(out, tuple) else out + + _AttentionBackendRegistry._backends[AttentionBackendName._FLASH_3] = ( + _flash_attention_3_training + ) + _AttentionBackendRegistry._supported_arg_names[AttentionBackendName._FLASH_3] = ( + set(_flash_attention_3_training.__code__.co_varnames[ + :_flash_attention_3_training.__code__.co_argcount]) + ) + + +# ----------------------------- nvtx ----------------------------- +def nvtx_range_start(name): + if hasattr(torch.cuda.nvtx, "range_start"): + return torch.cuda.nvtx.range_start(name) + torch.cuda.nvtx.range_push(name) + return None + + +def nvtx_range_end(handle): + if handle is None: + torch.cuda.nvtx.range_pop() + else: + torch.cuda.nvtx.range_end(handle) + + +@contextmanager +def nvtx_range(name): + handle = nvtx_range_start(name) + try: + yield + finally: + nvtx_range_end(handle) + + +_fwd_nvtx_handles = [] +_bwd_nvtx_handles = [] + + +class _NvtxFwdStartBwdEnd(torch.autograd.Function): + @staticmethod + def forward(ctx, x): + _fwd_nvtx_handles.append(nvtx_range_start("fwd")) + return x + + @staticmethod + def backward(ctx, grad): + if _bwd_nvtx_handles: + nvtx_range_end(_bwd_nvtx_handles.pop()) + return grad + + +class _NvtxFwdEndBwdStart(torch.autograd.Function): + @staticmethod + def forward(ctx, x): + if _fwd_nvtx_handles: + nvtx_range_end(_fwd_nvtx_handles.pop()) + return x + + @staticmethod + def backward(ctx, grad): + _bwd_nvtx_handles.append(nvtx_range_start("bwd")) + return grad + + +def _looks_like_vit_module(name, module): + lowered_name = name.lower() + lowered_class = module.__class__.__name__.lower() + compact_name = lowered_name.replace(".", "_") + compact_class = lowered_class.replace("_", "") + tokens = set(compact_name.split("_")) + return ( + "vit" in tokens + or "visiontransformer" in compact_class + or "visiontransformer" in compact_name.replace("_", "") + or "vision_model" in lowered_name + or "vision_tower" in lowered_name + or "image_encoder" in lowered_name + or lowered_name.endswith("visual") + ) + + +def install_vit_nvtx_tag(model): + for name, module in model.named_modules(): + if not name or not _looks_like_vit_module(name, module): + continue + original_forward = module.forward + + def tagged_forward(*args, _original_forward=original_forward, **kwargs): + with nvtx_range("vit_fwd"): + return _original_forward(*args, **kwargs) + + module.forward = tagged_forward + return name + return None + + +# ------------------------ memory history / OOM dump ------------------------ +class MemoryHistoryManager: + """Records CUDA memory history and dumps snapshots on OOM and/or exit.""" + + def __init__(self, out_dir: str, oom_only: bool = False): + self._out_dir = out_dir + self._oom_only = oom_only + self._dumped = False + + def start(self): + Path(self._out_dir).mkdir(parents=True, exist_ok=True) + torch.cuda.memory._record_memory_history( + max_entries=200000, stacks="all", + ) + rank = torch.distributed.get_rank() if torch.distributed.is_initialized() else 0 + if rank == 0: + print(f"[rank0] Memory history recording enabled, " + f"dump dir={self._out_dir} oom_only={self._oom_only}") + + def dump(self, tag: str = ""): + if self._dumped: + return + self._dumped = True + rank = torch.distributed.get_rank() if torch.distributed.is_initialized() else 0 + suffix = f"_{tag}" if tag else "" + path = os.path.join(self._out_dir, f"memory_snapshot_rank{rank}{suffix}.pickle") + try: + torch.cuda.memory._dump_snapshot(path) + if rank == 0: + print(f"[rank0] Memory snapshot dumped: {path}") + except Exception as e: + if rank == 0: + print(f"[rank0] Memory snapshot dump failed: {e}") + + def stop(self): + torch.cuda.memory._record_memory_history(enabled=None) + + def dump_on_normal_exit(self): + if self._oom_only or self._dumped: + return + self.dump() + + def dump_on_oom(self): + self.dump(tag="OOM") + + +# ------------------------------------------------------------------ +def parse(): + p = argparse.ArgumentParser() + p.add_argument("--backend", required=True, choices=["fsdp1", "mfsdp", "mfsdpv2"]) + p.add_argument("--sharding", default="hybrid", choices=["full", "hybrid"]) + p.add_argument("--pretrained_model_name_or_path", required=True, + help="HF repo or local dir containing the QwenImage transformer.") + p.add_argument("--subfolder", default="transformer") + p.add_argument("--attention", default="_flash_3", + help="diffusers attention backend name (e.g. _flash_3, flash, native).") + p.add_argument("--mixed_precision", default="bf16", choices=["bf16", "fp32"]) + p.add_argument("--batch_size", type=int, default=4) + p.add_argument("--height", type=int, default=512) + p.add_argument("--width", type=int, default=512) + p.add_argument("--instruction_seq_len", type=int, default=64) + p.add_argument("--vl_patch_size", type=int, default=14) + p.add_argument("--vl_merge_size", type=int, default=2) + p.add_argument("--bench_steps", type=int, default=20) + p.add_argument("--warmup_steps", type=int, default=3) + p.add_argument("--num_gpus_per_node", type=int, default=8) + p.add_argument("--lr", type=float, default=1e-5) + p.add_argument("--seed", type=int, default=1234) + p.add_argument("--compile", action="store_true", + help="Per-block torch.compile on transformer_blocks.") + p.add_argument("--trace-pool", action="store_true", + help="Use TracePoolAllocator for stable buffer addresses (mfsdpv2 only).") + p.add_argument("--cuda-graph", action="store_true", + help="Enable CUDA graph capture on leaf FSDP modules (mfsdpv2 only).") + p.add_argument("--gradient_checkpointing", action="store_true") + p.add_argument("--verify", action="store_true", + help="Cross-rank gloss + global grad-norm probe; timing INVALID for this run.") + p.add_argument("--cuda_profiler_capture", action="store_true", + help="Bracket bench steps with cudaProfilerStart/Stop for Nsight capture ranges.") + p.add_argument("--record-memory-history", type=str, default=None, metavar="DIR", + help="Enable CUDA memory recording (max_entries=200000). " + "Dumps a snapshot to DIR on normal exit AND on OOM. " + "Files: memory_snapshot_rank{N}.pickle. " + "Loadable at https://pytorch.org/memory_viz.") + p.add_argument("--record-memory-history-oom-only", action="store_true", + help="When set with --record-memory-history, only dump the " + "snapshot on OOM (skip normal exit dump).") + return p.parse_args() + + +# ------------------------ FSDP wrap helpers ------------------------ +def wrap_fsdp1(model, world_size, num_gpus_per_node, dtype, sharding, local_rank): + from torch.distributed.fsdp import ( + FullyShardedDataParallel as FSDP, + ShardingStrategy, MixedPrecision, BackwardPrefetch, + ) + from torch.distributed.fsdp.wrap import ModuleWrapPolicy + + use_hsdp = sharding == "hybrid" and (world_size // num_gpus_per_node) > 1 + if use_hsdp: + num_nodes = world_size // num_gpus_per_node + mesh = init_device_mesh("cuda", (num_nodes, num_gpus_per_node), + mesh_dim_names=("replicate", "shard")) + strat = ShardingStrategy.HYBRID_SHARD + else: + mesh = init_device_mesh("cuda", (world_size,), mesh_dim_names=("shard",)) + strat = ShardingStrategy.FULL_SHARD + + mp = MixedPrecision(param_dtype=dtype, reduce_dtype=dtype, buffer_dtype=dtype) + return FSDP( + model, device_mesh=mesh, sharding_strategy=strat, mixed_precision=mp, + auto_wrap_policy=ModuleWrapPolicy({QwenImageTransformerBlock}), + backward_prefetch=BackwardPrefetch.BACKWARD_PRE, forward_prefetch=True, + use_orig_params=True, device_id=local_rank, limit_all_gathers=True, + ) + + +def wrap_mfsdp(model, world_size, num_gpus_per_node, dtype, sharding, + sync_each_microbatch): + from megatron_fsdp import fully_shard_model, MixedPrecisionPolicy + + use_hsdp = sharding == "hybrid" and (world_size // num_gpus_per_node) > 1 + if use_hsdp: + num_nodes = world_size // num_gpus_per_node + mesh = init_device_mesh("cuda", (num_nodes, num_gpus_per_node), + mesh_dim_names=("dp_outer", "dp_shard")) + sub = mesh[("dp_outer", "dp_shard")] + (sub._flatten if hasattr(sub, "_flatten") else sub.flatten)("hsdp") + hsdp_kwargs = dict(dp_outer_dim="dp_outer", + hybrid_fsdp_group=mesh["hsdp"].get_group(), + outer_dp_sharding_strategy=0) # REPLICATE outer + else: + mesh = init_device_mesh("cuda", (world_size,), mesh_dim_names=("dp_shard",)) + hsdp_kwargs = {} + + mp = MixedPrecisionPolicy(main_params_dtype=dtype, main_grads_dtype=dtype, + grad_comm_dtype=dtype) + return fully_shard_model( + module=model, + fsdp_unit_modules=[QwenImageTransformerBlock], + device_mesh=mesh, dp_shard_dim="dp_shard", + zero_dp_strategy=3, + sync_model_each_microbatch=sync_each_microbatch, + overlap_grad_reduce=True, overlap_param_gather=True, + mixed_precision_policy=mp, **hsdp_kwargs, + ) + + +def wrap_mfsdpv2(model, world_size, num_gpus_per_node, dtype, sharding, + enable_trace_pool=False, enable_cuda_graph=False): + from megatron.core.distributed.fsdp.src.megatron_fsdp.v2 import ( + fully_shard, MixedPrecisionPolicy, + ) + + use_hsdp = sharding == "hybrid" and (world_size // num_gpus_per_node) > 1 + if use_hsdp: + mesh = init_device_mesh("cuda", (world_size,), mesh_dim_names=("dp_shard",)) + shard_strategy = "optim_grads_params" + if dist.get_rank() == 0: + print("[mfsdpv2] HSDP not supported; falling back to full sharding on 1D mesh") + else: + mesh = init_device_mesh("cuda", (world_size,), mesh_dim_names=("dp_shard",)) + shard_strategy = "optim_grads_params" + + mp = MixedPrecisionPolicy(main_params_dtype=dtype, main_grads_dtype=dtype, + grad_comm_dtype=dtype) + + for blk in model.transformer_blocks: + fully_shard(blk, mesh=mesh, mp_policy=mp, + sharding_strategy=shard_strategy, + enable_unshard_prefetch=True, + enable_async_reduce_grad=True, + enable_trace_pool=enable_trace_pool, + enable_cuda_graph=enable_cuda_graph) + + fully_shard(model, mesh=mesh, mp_policy=mp, + sharding_strategy=shard_strategy, + enable_unshard_prefetch=True, + enable_async_reduce_grad=True, + enable_trace_pool=enable_trace_pool) + + return model + + +# ------------------------ verify probe ------------------------ +def _local_grad(p): + """Read this param's local-shard gradient across both backends. + + fsdp1 (use_orig_params): p.grad is the local flat shard. + mfsdp: grad lives in an internal param_and_grad_buffer; p.grad is None. + Try .get_main_grad() -> .main_grad -> .decoupled_grad. + """ + g = p.grad + if g is None: + getter = getattr(p, "get_main_grad", None) + if getter is not None: + try: g = getter() + except Exception: g = None + if g is None: g = getattr(p, "main_grad", None) + if g is None: g = getattr(p, "decoupled_grad", None) + if g is None: return None + if hasattr(g, "to_local"): g = g.to_local() + return g + + +@torch.no_grad() +def verify_stats(model, loss, device): + """Global gloss (mean over ranks) + global grad-norm (over all shards/ranks).""" + torch.cuda.synchronize() # mfsdp reduce-scatters on a side stream + sumsq = torch.zeros((), device=device, dtype=torch.float32) + n = 0 + for p in model.parameters(): + if not p.requires_grad: continue + g = _local_grad(p) + if g is None: continue + sumsq += g.detach().float().pow(2).sum() + n += 1 + dist.all_reduce(sumsq, op=dist.ReduceOp.SUM) + gloss = loss.detach().float().clone() + dist.all_reduce(gloss, op=dist.ReduceOp.SUM) + return (gloss / dist.get_world_size()).item(), sumsq.sqrt().item(), n + + +# ------------------------ synthetic batch ------------------------ +def qwen25vl_vision_tokens(h, w, patch=14, merge=2): + f = patch * merge + return max(1, round(h/f)) * max(1, round(w/f)) + + +def make_batch(args, device, dtype, gen): + B = args.batch_size + H, W = args.height // 16, args.width // 16 # 16x packing (VAE 8x + 2x2 patchify) + seq = H * W + txt_len = args.instruction_seq_len + qwen25vl_vision_tokens( + args.height, args.width, args.vl_patch_size, args.vl_merge_size) + hidden = torch.randn(B, seq, 64, device=device, dtype=dtype, generator=gen) + timestep = torch.rand(B, device=device, dtype=dtype, generator=gen) * 1000 + prompt = torch.randn(B, txt_len, 3584, device=device, dtype=dtype, generator=gen) + img_shapes = [[(1, H, W)]] * B + # No padding in synthetic batch → omit attention mask so FA3 path is reachable. + return hidden, timestep, prompt, img_shapes, [txt_len] * B + + +# ----------------------------- main ----------------------------- +def main(): + args = parse() + if args.attention == "_flash_3": + _enable_fa3_training_patch() + rank, local_rank, world = (int(os.environ[k]) for k in + ("RANK", "LOCAL_RANK", "WORLD_SIZE")) + dist.init_process_group("nccl") + torch.cuda.set_device(local_rank) + device = torch.device(f"cuda:{local_rank}") + dtype = torch.bfloat16 if args.mixed_precision == "bf16" else torch.float32 + + if rank == 0: + txt_len = args.instruction_seq_len + qwen25vl_vision_tokens( + args.height, args.width, args.vl_patch_size, args.vl_merge_size) + print(f"[{args.backend}] world={world} dtype={dtype} bs={args.batch_size} " + f"img={args.height}x{args.width} txt={txt_len} " + f"sharding={args.sharding} compile={args.compile} gc={args.gradient_checkpointing}") + + model = QwenImageTransformer2DModel.from_pretrained( + args.pretrained_model_name_or_path, subfolder=args.subfolder, torch_dtype=dtype, + ).to(dtype) + if args.gradient_checkpointing: + model.enable_gradient_checkpointing() + model.train() + + if args.backend == "fsdp1": + model = wrap_fsdp1(model, world, args.num_gpus_per_node, dtype, + args.sharding, local_rank) + elif args.backend == "mfsdpv2": + model = wrap_mfsdpv2(model, world, args.num_gpus_per_node, dtype, + args.sharding, enable_trace_pool=args.trace_pool, + enable_cuda_graph=args.cuda_graph) + else: + # verify needs grads finished before optimizer.step(); benchmark path keeps overlap + model = wrap_mfsdp(model, world, args.num_gpus_per_node, dtype, + args.sharding, sync_each_microbatch=True) + + if args.compile: + inner = model.module if hasattr(model, "module") else model + for blk in inner.transformer_blocks: + blk.compile() + + tagged_vit = install_vit_nvtx_tag(model) + if rank == 0: + if tagged_vit is not None: + print(f"[nvtx] vit_fwd tagged module: {tagged_vit}") + else: + print("[nvtx] no ViT-like module found; vit_fwd tag skipped") + + optim = torch.optim.AdamW([p for p in model.parameters() if p.requires_grad], + lr=args.lr, fused=True) + if args.backend == "mfsdp": + from megatron_fsdp import fully_shard_optimizer + fully_shard_optimizer(optim) + + torch.cuda.empty_cache() + torch.cuda.reset_peak_memory_stats() + + # ---- memory history recording ---- + mem_mgr = None + if args.record_memory_history: + mem_mgr = MemoryHistoryManager( + out_dir=args.record_memory_history, + oom_only=args.record_memory_history_oom_only, + ) + mem_mgr.start() + # Dump on normal exit via atexit (unless oom_only) + if not args.record_memory_history_oom_only: + atexit.register(mem_mgr.dump_on_normal_exit) + + _mem_log("after_model_init", rank=rank) + + gen = torch.Generator(device=device).manual_seed(args.seed + rank) + step_times = [] + cuda_profiler_active = False + oom_occurred = False + with attention_backend(args.attention): + try: + for step in range(args.warmup_steps + args.bench_steps): + if args.cuda_profiler_capture and step == args.warmup_steps: + torch.cuda.synchronize() + dist.barrier() + torch.cuda.profiler.start() + cuda_profiler_active = True + if rank == 0: + print("[nsys] cudaProfilerStart") + hs, ts, pe, img_shapes, txt_lens = make_batch(args, device, dtype, gen) + hs.requires_grad_(True) + torch.cuda.synchronize(); dist.barrier() + t0 = time.perf_counter() + + hs = _NvtxFwdStartBwdEnd.apply(hs) + t_fwd = time.perf_counter() + with torch.amp.autocast("cuda", dtype=dtype): + out = model(hidden_states=hs, timestep=ts / 1000, + encoder_hidden_states=pe, + img_shapes=img_shapes, txt_seq_lens=txt_lens, return_dict=False) + pred = (out[0] if isinstance(out, tuple) else out)[:, :hs.size(1)] + pred = _NvtxFwdEndBwdStart.apply(pred) + loss = pred.float().pow(2).mean() + loss.backward() + t_bwd = time.perf_counter() + + if step < args.warmup_steps: + _mem_log(f"step={step} fwd_bwd fwd_ms={((t_bwd - t_fwd) * 1000):.1f} " + f"bwd_ms={((time.perf_counter() - t_bwd) * 1000):.1f}", rank=rank) + + if args.verify: + gloss, gnorm, n = verify_stats(model, loss, device) + + with nvtx_range("optimizer"): + if args.backend == "mfsdp": + optim.step(sync_grad_before_optimizer_step=True, + install_optimized_model_weights=True) + optim.zero_grad(set_to_none=True, zero_grad_buffer=True) + elif args.backend == "mfsdpv2": + optim.step() + optim.zero_grad(set_to_none=True) + else: + optim.step(); optim.zero_grad(set_to_none=True) + + torch.cuda.synchronize() + dt = time.perf_counter() - t0 + if step >= args.warmup_steps: + step_times.append(dt) + if rank == 0: + tag = "warmup" if step < args.warmup_steps else "bench " + if args.verify: + print(f"[{args.backend}] {tag} step {step:3d} | VERIFY (timing invalid) | " + f"gloss={gloss:.6f} | gnorm={gnorm:.4f} | n_grad={n}") + else: + print(f"[{args.backend}] {tag} step {step:3d} | {dt*1000:8.2f} ms | " + f"loss={loss.item():.4f}") + + except (torch.cuda.OutOfMemoryError, RuntimeError) as exc: + oom_occurred = True + if rank == 0: + print(f"\n[rank0] OOM / CUDA error at step {step}: {exc}") + traceback.print_exc() + if mem_mgr is not None: + mem_mgr.dump_on_oom() + raise + + if cuda_profiler_active: + torch.cuda.synchronize() + dist.barrier() + torch.cuda.profiler.stop() + if rank == 0: + print("[nsys] cudaProfilerStop") + + peak = torch.tensor([torch.cuda.max_memory_allocated() / 1e9], device=device) + dist.all_reduce(peak, op=dist.ReduceOp.MAX) + if rank == 0: + avg = sum(step_times) / max(1, len(step_times)) * 1000 + print(f"\n[{args.backend}] avg step (n={len(step_times)}): {avg:.2f} ms | " + f"peak mem: {peak.item():.2f} GB") + + _mem_log("final", rank=rank) + + if mem_mgr is not None: + mem_mgr.dump_on_normal_exit() + mem_mgr.stop() + + dist.destroy_process_group() + + +def _fmt_bytes(n: int) -> str: + for power, suffix in [(4, "TB"), (3, "GB"), (2, "MB"), (1, "KB"), (0, "B")]: + unit = 1024 ** power + if n >= unit: + return f"{n / unit:.2f} {suffix}" + return f"{n} B" + + +def _mem_log(tag="", rank=None): + """Log CUDA memory stats for the current rank.""" + if rank is None: + rank = torch.distributed.get_rank() + alloc = torch.cuda.memory_allocated() + max_alloc = torch.cuda.max_memory_allocated() + reserved = torch.cuda.memory_reserved() + max_reserved = torch.cuda.max_memory_reserved() + prefix = f"[rank{rank}] {tag}" if tag else f"[rank{rank}]" + print(f"{prefix} alloc={_fmt_bytes(alloc)} max_alloc={_fmt_bytes(max_alloc)} " + f"reserved={_fmt_bytes(reserved)} max_reserved={_fmt_bytes(max_reserved)}") + + +if __name__ == "__main__": + main() diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/utils.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/utils.py index 50334718758..53b4aedec08 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/utils.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/utils.py @@ -846,5 +846,7 @@ def using_tensor_parallel(dist_index, is_expert_parallel: bool = False) -> bool: """ Check if tensor parallelism is being used based on the distributed index. """ + if dist_index.tp_dim is None: + return False tp_mesh = dist_index.get_submesh(dist_index.tp_dim, is_expert_parallel=is_expert_parallel) return tp_mesh.mesh.numel() > 1 diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/README.md b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/README.md index 0dffff686f6..44624770de4 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/README.md +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/README.md @@ -46,7 +46,7 @@ Controls parameter/gradient dtypes and communication precision. | Field | Default | Purpose | |-------|---------|---------| -| `main_params_dtype` | ``None`` | Dtype for optimizer main-weight buffer. ``None`` = no separate buffer, optimizer mutates model weights directly. Set to ``torch.float32`` for quantized models (FP8/NVFP4) so the optimizer works on high-precision copies. | +| `main_params_dtype` | ``None`` | Dtype for optimizer main-weight buffer. ``None`` = no separate buffer, optimizer mutates model weights directly. Set to ``torch.float32`` for quantized models (FP8/NVFP4) so the optimizer works on high-precision copies. When this equals the model-weight dtype and the sharding layout matches, the separate buffer is skipped automatically to avoid a redundant copy. | | `main_grads_dtype` | ``None`` | Dtype for optimizer main-grad buffer. When ``None`` and ``use_decoupled_grad=False``, aligns with ``main_params_dtype``. Otherwise falls back to ``param.dtype``. | | `grad_comm_dtype` | ``None`` | Dtype for gradient reduce-scatter communication. ``None`` = use ``main_grads_dtype``. | | `use_decoupled_grad` | ``False`` | When ``False``, ``main_grads_dtype`` is inferred from ``main_params_dtype`` so the optimizer operates in a consistent precision context. | @@ -102,7 +102,12 @@ Mixin class added to wrapped modules. Methods: > **Experimental** — CUDA graph support in Megatron FSDP v2 is an experimental > feature. The API and behaviour may change in future releases without notice. -FSDP v2 supports transparent CUDA graph capture for individual FSDP modules. +**Why MFSDP v2 can support CUDA graphs.** The [`TracePoolAllocator`](allocator.py) +pre-plans every parameter/gradient buffer slot at a fixed offset in a +persistent pool tensor during the first (trace) micro-batch. Once `plan()` +has committed those offsets, all buffer addresses become deterministic across +micro-batches — the stable memory foundation that CUDA graph capture requires. + Enable it per-module with ``enable_cuda_graph=True``: ```python @@ -111,23 +116,19 @@ for layer in model.layers: fully_shard(model) # root without CUDA graph ``` -Capture happens automatically on the first optimized forward pass -(after the trace pool has been planned). Subsequent forward passes replay -the captured graph, bypassing FSDP hooks entirely for that module. - **How it works:** -1. The forward pre-hook detects ``enable_cuda_graph=True`` and creates an - ``FSDPCudaGraphRunner`` for the module. -2. The runner pops FSDP hooks, warms up via ``make_graphed_callables``, - and captures the module's ``forward()`` (without FSDP collectives) into a - CUDA graph. -3. After capture, ``install()`` patches ``module.forward`` to a wrapper - that replays the graph on subsequent calls. -4. After capture, ``install()`` patches ``module.forward`` to a wrapper that - replays the graph on subsequent calls. During replay, the patched forward - calls the graphed callable directly, bypassing FSDP forward hooks entirely. - Backward hooks (unshard/reshard/reduce_grad) still fire normally. +1. The first optimized forward pass records sample arguments for each + eligible module. +2. After the first backward completes, a single batch call to + `te-graph-runtime`'s `make_graphed_callables` captures forward + backward + graphs for all modules in correct order (fwds in forward-module order, + bwds in reverse) using the shared trace pool. +3. FSDP unshard/reshard hooks run **outside** the CUDA graph capture via + `capture_time_hooks` — they are never graphed. During replay they + fire normally around the graphed forward/backward. +4. Replay runs entirely through the captured graphs — no Python hooks fire + inside the graphed region. **Limitation — nesting:** A parent FSDP module that contains other FSDP modules as children **cannot** use ``enable_cuda_graph=True``. Only leaf @@ -135,7 +136,7 @@ FSDP modules (those without FSDP children) are eligible. Attempting to enable CUDA graph on a module with FSDP children raises a ``RuntimeError``. ```python -# OK — layers are leaf FSDP modules (no FSDP children inside them) +# OK — layers are leaf FSDP modules for layer in model.layers: fully_shard(layer, enable_cuda_graph=True) @@ -143,6 +144,8 @@ for layer in model.layers: fully_shard(model, enable_cuda_graph=True) # raises RuntimeError ``` +See [`design/cuda_graph_design.md`](design/cuda_graph_design.md) for the full architecture. + ### DataParallelBuffer Flat buffer managing (a shard of) parameter/gradient data: @@ -203,14 +206,14 @@ overlap (prefetch/unshard pipelining) is not applicable. ### CUDA Graph -- **Experimental.** CUDA graph support is experimental. The API and behavior - may change in future releases without notice. Enable via - ``enable_cuda_graph=True`` on leaf FSDP modules only. -- **Nesting not supported.** A parent FSDP module that contains other FSDP - modules as children cannot use ``enable_cuda_graph=True``. Only leaf FSDP - modules (those without FSDP children) are eligible for CUDA graph capture. - Attempting to enable CUDA graph on a nested FSDP module raises a - ``RuntimeError``. +- **Experimental.** Enable via ``enable_cuda_graph=True`` on leaf FSDP modules. + Built on vendored [te-graph-runtime](https://github.com/buptzyb/te-graph-runtime) + with local modifications. See [`design/cuda_graph_design.md`](design/cuda_graph_design.md). +- **Requires `TracePoolAllocator`.** CUDA graph capture depends on the + deterministic buffer addresses provided by the trace pool; modules must be + wrapped with ``enable_trace_pool=True``. +- **Nesting not supported.** Only leaf FSDP modules (those without FSDP children) + are eligible for capture. ### `fully_shard()` API Parameters @@ -316,8 +319,11 @@ torchrun --nproc_per_node=2 examples/megatron_fsdp/fsdp_toy.py \ compute operations (all-gather, reduce-scatter) run on side streams (`ag_stream`, `rs_stream`). CUDA events inserted at the boundary between allocation and compute, and between compute and free, guarantee ordering - without ``record_stream``. Keeping allocations on the default stream avoids - non-deterministic caching-allocator behaviour and peak memory regressions. + without ``record_stream``. ``record_stream`` is intentionally avoided + because it forces the caching allocator to hold memory blocks until the + recorded stream finishes, preventing reuse across iterations and causing + significant peak memory regressions + ([discussion](https://dev-discuss.pytorch.org/t/1486)). ## Unit Tests diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/cuda_graph_memory_analysis.md b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/cuda_graph_memory_analysis.md new file mode 100644 index 00000000000..4ef069d3117 --- /dev/null +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/cuda_graph_memory_analysis.md @@ -0,0 +1,577 @@ +# CUDA Graph Memory Overhead — Analysis & Findings + +> **Status:** experimental findings from profiling the per-layer graph +> capture path. Numbers come from a QwenImage diffusers model wrapped +> with Megatron FSDP v2 + `enable_cuda_graph=True`. The qualitative +> conclusions hold for any per-layer graph capture that pops torch +> forward hooks before capture. + +## TL;DR + +CG capture costs **+1.222 GB / captured layer** vs no-CG, in steady +state. There are **two independent sources**, both rooted in the same +root cause: **CG capture pops all module hooks (including torch._dynamo +inductor hooks), so the captured forward sees *un-fused* python-level +ops**. + +| Source | Per-layer | % of total | +|---|---:|---:| +| Framed activations (inductor fusion skipped → one SavedVariable per python op instead of one per fused region) | +0.524 GB | 43 % | +| No-frame cuBLAS / cuDNN workspaces (no fusion → one workspace per matmul, never recycled across layers in the graph pool) | +0.698 GB | 57 % | +| **Total** | **+1.222 GB** | **100 %** | + +Neither source is "saved activations that backward would have freed" — +that was the wrong framing in earlier drafts. Both sources survive in +the graph pool forever because graph memory is pinned for replay, but +their *magnitude* would be much smaller if CG capture went through +inductor's fused kernels rather than the raw python ops. + +## How we measured (apples-to-apples) + +Two env-var gated debug modes in `FSDPCudaGraphRunner` +(see `cuda_graph_runner.py`): + +| `MFSDP_CG_MEM_DEBUG` | Behaviour | +|---|---| +| `cg` | Capture CG normally. For the first `MFSDP_CG_MEM_SNAP_LAYERS` layers, record stack-aware `torch.cuda.memory._record_memory_history` around the forward-graph capture and dump `cg_layer{N}_rank{R}.pickle` immediately after capture. | +| `nocg` | Skip CG capture entirely. `capture_forward` just stores debug state; `install()` **patches `module.forward`** with a wrapper (`_debug_recorded_forward`) that records + snapshots the **real** forward — the one that runs after `forward_pre_hook` returns. So `nocg_layer{N}_rank{R}.pickle` captures the actual training forward, with hooks (and inductor) intact. | +| _(unset)_ | Normal production. | + +Run separately: + +```bash +MFSDP_CG_MEM_DEBUG=cg MFSDP_CG_MEM_SNAP_LAYERS=3 MFSDP_CG_MEM_SNAPSHOT_DIR=/tmp/cg_mem python train.py +MFSDP_CG_MEM_DEBUG=nocg MFSDP_CG_MEM_SNAP_LAYERS=3 MFSDP_CG_MEM_SNAPSHOT_DIR=/tmp/cg_mem python train.py +``` + +Per-layer peak (`peak_alloc`, `peak_reserved`, `post`) is logged for +each captured layer in both modes. Snapshots are loadable at +https://pytorch.org/memory_viz. + +### Why the nocg path patches `module.forward` + +Earlier versions of the debug code ran an *extra* forward inside +`capture_forward` to take the snapshot. That created a measurement +artifact: the extra forward's autograd tape was destroyed when +`_debug_run_eager` returned, freeing its activations into the caching +allocator's free list. The real forward (the one that runs after the +pre-hook returns) then *reused* those freed addresses, making no-CG +appear to grow at ~0.5 GB/layer instead of its true ~1.2 GB/layer +forward-only rate. + +The current design patches `module.forward` itself so the snapshot is +taken during the actual training forward. This means: + +- The autograd tape accumulates naturally across layers (as in real + training). +- Inductor is active (its dynamo hooks survive because `capture_forward` + returns early in `nocg` mode, before `_pop_hooks_recursive` runs). +- The snapshot captures true per-layer forward cost. + +## The numbers + +### Per-layer live memory (total active bytes across all pools) + +| Layer | CG total | NOCG total | CG − NOCG | +|---:|---:|---:|---:| +| 0 | 52.055 GB | 51.450 GB | **+0.606 GB** | +| 1 | 53.771 GB | 51.943 GB | **+1.828 GB** | +| 2 | 55.487 GB | 52.437 GB | **+3.050 GB** | + +### Per-layer growth rate (averaged over L0→L1 and L1→L2) + +| Component | CG | NOCG | Savings | +|---|---:|---:|---:| +| Framed activations (graph pool in CG, caching pool in NOCG) | +1.176 GB | +0.652 GB | **0.524 GB** | +| No-frame workspaces (caching pool, both modes) | +0.540 GB | −0.158 GB | **0.698 GB** | +| **Total** | **+1.716 GB** | **+0.494 GB** | **1.222 GB** | + +Note: NOCG's no-frame caching pool actually **shrinks** by 158 MB/layer. +That is the inductor-fused forward releasing workspace back to the +caching allocator after each fused kernel completes; the only thing +that persists is autograd-saved activations. + +## Source A — framed activations (0.524 GB/layer) + +### Per-operator-family breakdown at layer 0 + +| Family | CG graph pool | NOCG caching pool | Diff | +|---|---:|---:|---:| +| Linear (`linear.py:134`) | 0.397 GB | 0.000 GB | +0.397 | +| LayerNorm (`functional.py:2935`, `normalization.py:554/555/560`) | 0.454 GB | 0.000 GB | +0.454 | +| GELU (`activations.py:85`) | 0.129 GB | 0.000 GB | +0.129 | +| QwenImage attention processor (`transformer_qwenimage.py:558/559/560/582/733/739`) | 0.163 GB | 0.000 GB | +0.163 | +| Attention (`flash_attn_interface.py:96`) | 0.034 GB | 0.034 GB | 0.000 | +| Inductor triton kernel (`ctjbu2*.py`) | 0.000 GB | 0.618 GB | −0.618 | +| **Total** | **1.176 GB** | **0.652 GB** | **+0.524 GB** | + +### Why the discrepancy + +CG's call stacks show raw python ops: +``` +linear.py:134 forward + ↑ module.py:1789 _call_impl + ↑ activations.py:88 forward (diffusers GEGLU) + ↑ attention.py:1741 forward + ↑ transformer_qwenimage.py:732 forward + ↑ cuda_graph_runner.py:516 capture_forward + ↑ hooks.py:84 forward_pre_hook +``` + +NOCG's call stacks show inductor triton kernels: +``` +ctjbu2xj2tikjgdogi372mldv5x4h2vgi5fvncjsn3yscjbwxopj.py:1631 call ← inductor-fused kernel + ↑ activations.py:88 forward + ↑ attention.py:1741 forward + ↑ transformer_qwenimage.py:732 forward + ↑ utils.py:3331 run ← inductor runtime + ↑ output_code.py:638 __call__ ← inductor output code + ↑ runtime_wrappers.py:930 inner_fn +``` + +100 % of NOCG framed bytes (45/45 blocks, 0.652 GB) come from inductor +kernels. 0 % of CG framed bytes touch inductor (0/68 blocks). This is +because `_pop_hooks_recursive` (called at `cuda_graph_runner.py:438`) +removes the dynamo hooks that drive inductor compilation. CG capture +runs the un-fused python-level forward; each python op saves its own +input tensor for backward, producing the per-call-site breakdown above. + +The inductor path, by contrast, fuses sequences of ops (Linear + +LayerNorm + GELU + ...) into a single triton kernel. Saved-tensor +points become fusion-region boundaries rather than per-op boundaries, +so far fewer SavedVariables are allocated. + +## Source B — no-frame workspaces (0.698 GB/layer) + +### Layer-over-layer delta in no-frame caching blocks + +| Mode | L0→L1 gone | L0→L1 new | Net | L1→L2 gone | L1→L2 new | Net | +|---|---|---|---:|---|---|---:| +| NOCG | 1 block / −0.158 GB | 0 / 0 GB | **−0.158 GB** | 1 / −0.158 GB | 0 / 0 GB | **−0.158 GB** | +| CG | 2 blocks / −0.032 GB | 52 / +0.706 GB | **+0.674 GB** | 2 / −0.032 GB | 52 / +0.707 GB | **+0.674 GB** | + +- **NOCG** actually **frees** one ~158 MB no-frame block per layer — + inductor's fused kernel allocates a single workspace and releases it + when the kernel completes. Inductor's per-call workspace is not + captured by autograd, so it returns to the caching allocator's free + list. +- **CG** adds **+52 no-frame blocks per layer (0.706 GB)** that never + get recycled. These are cuBLAS / cuDNN workspaces — every captured + python-level matmul and convolution grabs its own workspace. The + graph capture prevents workspace reuse across layers (the kernels + baked into the graph reference those workspace addresses for replay). + +The 52-blocks-per-layer pattern is highly regular (predictable block +sizes: 2 × 96 MB, 2 × 36 MB, 4 × 33 MB, 9 × 24 MB, 9 × 9 MB, ...), +suggesting one workspace per matmul/conv per layer. With inductor +fusion, far fewer separate matmuls happen, so far fewer workspaces are +needed. + +## Why CG capture bypasses inductor + +`hooks.py:84` triggers CG capture inside the FSDP `forward_pre_hook` +(which is `@torch.compiler.disable`d, so dynamo treats it as a graph +break). Inside the pre-hook, `capture_forward` (line 438) does two +things: + +1. **Pops all hooks recursively** via `_pop_hooks_recursive`. This + removes the FSDP `forward_pre_hook` / `forward_hook` that triggered + capture — necessary, otherwise they would re-fire during the + captured forward and race with the training loop. + +2. **Invokes the forward body directly** via `_call_module` → + `self._module.forward(**kwargs)`. This bypasses `nn.Module.__call__` + (and thus `_call_impl`), which is where PyTorch 2.x installs the + dynamo dispatch. + +The previous version of this document blamed step 1 — claiming that +`_pop_hooks_recursive` removes "the dynamo/inductor hooks." **This is +incorrect.** In PyTorch 2.4, `module.compile()` does NOT install any +hooks; it sets `module._compiled_call_impl = torch.compile(_call_impl)` +(a `__call__`-level dispatch). Hooks popping is therefore orthogonal +to inductor: removing `_forward_pre_hooks` / `_forward_hooks` etc. +has no effect on whether dynamo fires. + +The actual cause is step 2 — `_call_module` calls `.forward()` +directly, never going through `_wrapped_call_impl` / +`_compiled_call_impl`. Furthermore, even if the call site were +switched to `self._module(**kwargs)`, the dispatch happens INSIDE a +dynamo graph break (the FSDP pre-hook is `@torch.compiler.disable`), +so dynamo would not have an active trace to fall back into. The +captured forward therefore runs the raw python body, emitting eager +aten ops + cuBLAS workspaces (Source B) and per-op SavedVariables +(Source A). No inductor fusion, no triton kernels. + +## Practical levers + +### Lever 1 — Capture a torch.compile()'d forward body (biggest win) + +Compile `module.forward` directly inside `capture_forward` so the +warmup below populates dynamo's cache for the forward BODY (not +`_call_impl`), and the capture (inside `torch.cuda.graph(pool=...)`) +runs the cached inductor code. Triton kernel launches are +CUDA-graph-capturable, so the resulting CG graph contains the same +fused kernels as the no-CG path — eliminating essentially all of the +1.222 GB/layer overhead (both Source A and Source B). + +Open questions: +- Does dynamo's cached lookup fire inside `torch.cuda.graph()` context? + Initial source inspection suggests yes (no global "skip during + capture" guard in `torch/_dynamo/eval_frame.py`); needs on-GPU + confirmation. +- Does inductor's allocator play nice with `torch.cuda.graph(pool=...)`? + Inductor-allocated workspaces would land in the graph pool and be + reused on replay — desirable. + +Implementation sketch (CURRENT FIX, env-gated on +`MFSDP_CG_COMPILE_FWD=1`): + +```python +# Inside capture_forward, after _pop_hooks_recursive and gc.freeze, +# BEFORE the warmup loop: +if _CG_COMPILE_FWD and not hasattr(self._module.forward, "get_compiler_config"): + self._orig_fwd_body = self._module.forward + self._module.forward = torch.compile(self._orig_fwd_body) + self._captured_fwd_was_compiled = True + +# Warmup (3 iters): _call_module → self._module.forward (compiled). +# Iter 1 triggers dynamo compile; iters 2–3 use the cache. + +# Capture: _call_module inside torch.cuda.graph(pool=...). +# Compiled code runs from dynamo cache → triton kernels fire → captured. + +# Finally: restore self._module.forward = self._orig_fwd_body so +# install() sees the user-written body and substitutes _patched_fwd +# (the CG replay function) over it. +``` + +User-visible effect: +- `blk(*args)` → dynamo's compiled `_call_impl` (user's `blk.compile()`) + → forwards call to `_patched_fwd` → `_CudaGraphFunction.apply` + → `fwd_graph.replay()` fires the captured triton kernels. + +The user's outer `blk.compile()` is still useful: it fuses the +pre/post-hook graph breaks, the autograd-Function boundary, and any +ops outside the FSDP block. + +### Lever 2 — Activation recomputation on the captured block + +The 0.524 GB/layer of framed activations (Source A) is forward +intermediates saved for backward. If `module.forward` is wrapped in +`torch.utils.checkpoint.checkpoint`, the captured forward graph can +free its intermediates before end-of-capture. Per-layer pinned +footprint from Source A drops from 0.524 GB to ~50 MB (just block I/O). + +This does NOT help Source B (cuBLAS workspaces) — those are captured +regardless of whether the forward is checkpointed. And it adds +compute (forward runs twice). + +### Lever 3 — Replace diffusers RMSNorm + Linear + GELU with TE fused kernels + +TE fuses norm + linear + activation into one kernel that uses +kernel-private workspace. The Linear (0.397 GB) + LayerNorm (0.454 GB) ++ GELU (0.129 GB) buckets in Source A collapse into a handful of TE +fused ops; pinned bytes drop proportionally. Also reduces Source B +because there are fewer separate matmul calls. + +### Lever impact summary + +| Lever | Source A | Source B | Total | Notes | +|---|---:|---:|---:|---| +| Lever 1 (capture through inductor) | −0.524 GB | −0.698 GB | **−1.222 GB** | Recovers essentially all overhead. Implementation risk in graph-pool / inductor-allocator interaction. | +| Lever 2 (activation recompute) | −0.474 GB | 0 | **−0.474 GB** | Adds compute. Standard pattern in `te.fp8_checkpoint`. | +| Lever 3 (TE fused ops) | −0.5 to −0.6 GB | −0.3 to −0.4 GB | **−0.8 to −1.0 GB** | Direct replacement; works alongside Lever 1 or 2. | +| L1 + L2 combined | −0.524 GB* | −0.698 GB** | **−1.222 GB** | *recompute on fused kernels saves less per layer, but the per-fused-region count is what matters; **inductor workspace recycling still helps. | + +Lever 1 is the clearly preferred direction — it fixes the root cause +rather than treating the symptoms. + +## What is *not* the problem + +Earlier drafts of this analysis attributed CG overhead to "saved +activations that backward would normally free." That framing is +incomplete: + +- **NOCG also saves activations for backward** — its 0.652 GB/layer of + framed caching-pool bytes are SavedVariables too. Per-layer growth + in NOCG is real (0.494 GB/layer total); CG's overhead is on top of + that, not instead of it. +- After backward, NOCG's framed activations are freed by the caching + allocator and the addresses are reused next microbatch. CG's framed + activations are pinned for replay. But the *magnitude* of CG's + pinned bytes is inflated by the inductor-bypass: if both paths + captured the same fused-kernel forward, CG's per-layer pin would be + 0.652 GB, not 1.176 GB. +- The dominant saving (0.698 GB/layer, 57 %) is **not about + SavedVariables at all** — it is about cuBLAS/cuDNN workspaces that + accumulate because un-fused python-level matmuls each grab their own. + +## Methodology artefacts removed + +The current nocg measurement path (patching `module.forward` with +`_debug_recorded_forward`) addresses three artefacts that affected +earlier measurements: + +1. **Tape destruction before snapshot (early `_debug_run_eager`).** + The first version did `del out, flat_out` before taking the + snapshot, which collapsed the autograd tape and made NOCG look like + it held 0 activations. Fixed by keeping `out` alive in scope until + the function returns. +2. **Extra forward before real forward (later `_debug_run_eager`).** + The second version ran an extra forward inside `capture_forward`, + which polluted the caching allocator's free list and made NOCG look + like it grew only 0.5 GB/layer. Fixed by patching `module.forward` + so the snapshot is taken during the real training forward. +3. **`empty_cache()` in the debug path.** An earlier version called + `torch.cuda.empty_cache()` after each layer's snapshot, which + perturbed the caching allocator's free list between layers. + Removed. + +## Files & tooling reference + +| Artifact | Location / how to produce | +|---|---| +| Debug mode implementation | `cuda_graph_runner.py` — `_CG_MEM_MODE`, `_debug_recorded_forward`, the nocg branch in `capture_forward`, the nocg branch in `install`. | +| Snapshot dumps | `cg_layer{N}_rank{R}.pickle`, `nocg_layer{N}_rank{R}.pickle` | +| Snapshot visualization | https://pytorch.org/memory_viz (drag pickles in) | +| Env vars | `MFSDP_CG_MEM_DEBUG={cg,nocg}`, `MFSDP_CG_MEM_SNAP_LAYERS=N`, `MFSDP_CG_MEM_SNAPSHOT_DIR=` | +| Snapshot analysis scripts | ad-hoc Python (see git history of this doc for `rigorous_compare.py`, `confirm_theory.py`, `layer_delta.py`). | + +## Open questions + +- Can `_pop_hooks_recursive` be made to selectively pop only FSDP + hooks, leaving dynamo/inductor hooks intact? If so, Lever 1 becomes + a single-line fix. +- Does `torch.compile(module)` before CG capture produce a forward that + plays nicely with `torch.cuda.graph(pool=...)`? Specifically, does + inductor's allocator honor the pool argument? +- The 52-blocks-per-layer no-frame pattern in CG is highly regular — + cataloguing which python ops each block corresponds to would let us + estimate Lever 3's impact per-op rather than per-family. +- For non-attention modules (e.g. MLP-only blocks), is the + inductor-bypass effect smaller (fewer separate matmuls to fuse)? + +--- + +## Full-snapshot analysis (post-Lever-1) + +After enabling `MFSDP_CG_COMPILE_FWD=1` (Lever 1 — `torch.compile` the +forward body during capture), a full end-of-training-step snapshot was +captured for both NOCG and CG+compile. These snapshots include the +optimizer state, gradient buffers, and the captured graphs for all 60 +transformer layers, giving an apples-to-apples comparison of the +*total* memory cost. + +Files: +- `memory_snapshot_nocg.pickle` — no CG (caching allocator only) +- `memory_snapshot_cg_compile.pickle` — CG + `torch.compile` fusion + +### Top-line: CG+compile costs +50 GB vs NOCG (101 GB vs 51 GB) + +| Category (pool) | NOCG | CG+C | Δ GB | +|---|---:|---:|---:| +| inductor:attn/blk (graph) | 0.000 | 21.124 | **+21.124** | +| inductor:activations (graph) | 0.000 | 18.741 | **+18.741** | +| no_frames (caching) — model weights/optimizer | 20.441 | 20.441 | 0 | +| adam.py (caching) — optimizer state | 20.431 | 20.431 | 0 | +| param_group.py (caching) — FSDP grad buffers | 10.216 | 10.215 | 0 | +| cg_runner (caching) — static_inputs | 0.000 | 4.166 | **+4.166** | +| fsdp_allocator (caching) — gradient bucket coloring | 0.000 | 3.561 | **+3.561** | +| flash_attn (graph) | 0.000 | 2.115 | **+2.115** | +| inductor (caching) — warmup allocations | 0.000 | 0.272 | +0.272 | +| **TOTAL** | **51.115** | **101.093** | **+49.979** | + +Model + optimizer + grad buffers (the three "no Δ" rows) account for +51 GB in both runs — that's the irreducible baseline. CG+compile +adds **+50 GB on top**, all in 5 categories broken down by source +below. + +### Big bucket 1: inductor graph-pool capture — +39.865 GB (+0.664 GB/layer) + +By far the largest CG overhead. All inductor triton-kernel allocations +during the captured forward go to the graph pool `(0, 1)`, where they +get pinned for replay. The 60-layer model accumulates 2552 such +blocks, totalling ~40 GB. Inductor aggregations grouped by user-code +call site: + +| Role (user-code frame) | Blocks | GB total | MB/layer | +|---|---:|---:|---:| +| activations.py:88 (GELU/SiLU outputs saved for backward) | 600 | 18.741 | 312.4 | +| attention linear / QKV / proj (transformer_qwenimage.py:520/558-586) | 1020 | 16.716 | 278.6 | +| transformer block fwd (transformer_qwenimage.py:689/733) | 332 | 2.254 | 37.6 | +| attention module (attention.py:1741) | 120 | 2.082 | 34.7 | +| apply_rotary_emb + modulation | 480 | 0.071 | 1.2 | +| **Total inductor at (0, 1)** | **2552** | **39.865** | **664.4** | + +**Why these accumulate per layer.** In NOCG, inductor's triton kernels +allocate intermediate workspaces that get released to the caching +allocator's free list as soon as each kernel completes; the next layer +reuses the same addresses. In CG, the captured graph *references* +every intermediate for replay, so the graph pool can't recycle them +between layers — the per-layer activations/workspaces accumulate +linearly with layer count. + +**The 18.7 GB GELU/SiLU activation row** is the dominant contributor. +These are activation outputs saved on the autograd tape for backward +(inductor emits them as its own kernel outputs). Two inductor call +sites generate 60 blocks each of 100 MB/layer (= 6 GB per call site, +12 GB total for the pair). Activation checkpointing inside the +compiled forward (Lever 2 below) would free these at end-of-capture, +reducing the row to ~0. + +### Big bucket 2: cg_runner static_inputs — +4.166 GB (caching pool) + +Allocated at `cuda_graph_runner.py:419`: +```python +t.clone().detach().requires_grad_(t.requires_grad) for t in flat_live +``` + +These are the per-layer **static input buffers** — one set per captured +layer, ~70 MB per layer × 60 layers = 4.2 GB. They live in the caching +pool (allocated outside `torch.cuda.graph`), so the per-layer snapshots +correctly show no graph-pool growth from these. + +Each static input is a clone of the user's actual forward input, used +to feed the captured graph during replay (`_CudaGraphFunction.forward` +copies live values into them before `fwd_graph.replay()`). + +**Possible reduction:** if multiple layers' static inputs share the +same shape/dtype, they could share a single static buffer per +(shape, dtype) pair via the graph pool — but this would require +reordering captures and is invasive. + +### Big bucket 3: FSDP gradient bucket coloring — +3.561 GB (caching pool) + +Allocated at `allocator.py:438` (`_color_and_allocate_slots`). Seven +huge blocks (~508 MB each) allocated only in CG mode. + +These are the FSDP `TracePoolAllocator` reduce-scatter gradient +buckets, allocated during the post-backward final callback. In CG mode +the trace→optimized transition happens after the first backward, and +the bucket allocation happens then; in NOCG mode (or pre-CG) the +allocator either hasn't planned yet or uses a different allocation +strategy that doesn't materialise these blocks. + +**Worth investigating:** the same bucket allocation may be avoidable +under CG by re-using the existing grad buffer storage as the bucket, +rather than allocating dedicated reduce-scatter buckets. + +### Big bucket 4: flash_attn graph pool — +2.115 GB + +180 blocks at ~35 MB/layer (graph pool). Flash attention's forward +intermediates (V matrix accumulators, softmax workspace) are captured +into the graph pool like inductor's allocations. These are +forward-only and don't need to persist for backward — but once the +graph captures them, they're pinned for replay. + +### Small bucket 5: inductor warmup — +0.272 GB (caching pool) + +Inductor allocations during the warmup iterations (before graph +capture starts) that did NOT get released to the free list. These +show up at `ctjbu2*.py` (30 blocks, 0.144 GB) and a second inductor +code file `cwy3f5*.py` (15 blocks, 0.128 GB). Tiny compared to the +graph-pool bucket, but indicates inductor keeps some persistent state +across the warmup→capture transition. + +### Per-layer vs total: reconciliation + +Earlier per-layer snapshots (cg_layer{0,1,2}_rank0_compile.pickle) +showed **+1.286 GB/layer** growth, but the full snapshot shows +**+50 GB / 60 layers = +0.833 GB/layer**. The discrepancy comes from +capturing different points in time: + +| Captured at | Per-layer snapshots | Full snapshots | +|---|---|---| +| When | end of forward capture for layer N | end of full training step (after backward + bucket coloring) | +| Layer 0..N forward graphs in pool | yes (cumulative) | yes (cumulative) | +| Backward graphs | no | yes (lazy, first backward) | +| Bucket coloring | no | yes | +| Optimizer state | no | yes (same in both) | + +The +1.286 GB/layer figure represents **forward-capture** growth only; +the +0.833 GB/layer figure amortises the one-time backward + bucket +costs across the whole step. Both are correct for their respective +points in time. + +## Updated lever priorities + +Given the full-snapshot breakdown, the lever priority order shifts: + +### Lever 2 (was Lever 2): no_grad forward capture — biggest single win + +**Expected savings: ~18.7 GB (the inductor activations row).** + +The captured forward graph accumulates SavedVariables for backward — +the 18.7 GB inductor-activations row. These are dead weight at replay +time because `_capture_backward_and_run` already re-runs the forward +with grad enabled (via `replay_inputs = t.detach().clone().requires_grad_(...)`) +to build a fresh autograd tape for the bwd_graph. + +Wrapping the capture call with `torch.no_grad()` drops all +SavedVariables from the fwd graph. This is strictly better than +`torch.utils.checkpoint` (Lever 2's original form): +- same fwd-graph savings as checkpoint +- no extra recompute in bwd capture (checkpoint's unpack hook + would re-run the forward during the runner's bwd-capture re-run, + doubling the recompute) +- no bwd-graph growth +- bonus: capture-time behavior matches replay-time behavior. + `_CudaGraphFunction.forward` is a `torch.autograd.Function`, and + PyTorch runs autograd Function forward methods with grad + disabled — so `fwd_graph.replay()` already executes under no_grad + at runtime. The legacy grad-enabled capture was inconsistent. + +**Status: IMPLEMENTED (default ON).** Gated by +`MFSDP_CG_NO_GRAD_FWD` (default `1`; set to `0` for legacy +behavior). The wrap is applied only to the capture call, not to +the warmup (which still needs grad to settle FP8 scales via +`torch.autograd.grad`): + +```python +# Inside capture_forward, around the CUDA graph capture: +no_grad_ctx = torch.no_grad() if _CG_NO_GRAD_FWD else contextlib.nullcontext() +with torch.cuda.graph(self.fwd_graph, pool=..., stream=...): + with no_grad_ctx: + out = self._call_module(static_inputs, tensor_names, frozen_kwargs) +``` + +Backward graph capture (`_capture_backward_and_run`) is unchanged: +it still re-runs the forward with `requires_grad_(True)` and grad +enabled, building its own autograd tape that gets captured into +`bwd_graph`. The bwd graph size is unaffected. + +For the historical `torch.utils.checkpoint`-based approach (kept +for A/B testing), see `cuda_graph_checkpoint_design.md` and the +`MFSDP_CG_USE_CHECKPOINT` env var. + +### Lever 4 (new): Share static input buffers across layers + +**Expected savings: ~3-4 GB (the cg_runner row).** + +If the same shapes recur across layers (true for transformer blocks +with identical attention dims), allocate one set of static I/O buffers +per (shape, dtype) pair instead of per layer. Live inputs get copied +into the shared buffer before each replay. + +### Lever 5 (new): Avoid duplicate gradient bucketing in CG + +**Expected savings: ~3.5 GB (the fsdp_allocator row).** + +Investigate why the FSDP `TracePoolAllocator` materialises 7 large +reduce-scatter buckets only in CG mode. If the existing grad buffer +storage can be used directly for reduce-scatter (without a separate +bucket allocation), this disappears. + +### Lever 6 (new): flash_attn recompute + +**Expected savings: ~2 GB (the flash_attn row).** + +Flash attention's forward intermediates account for 2.1 GB in the +graph pool. A recompute variant (re-run flash-attn forward in +backward) would let the captured graph free them. + +### Lever 3 (carried forward): TE fused ops + +Still relevant for the inductor attention-linear bucket (16.7 GB), but +Lever 2 + Lever 4 + Lever 5 together would already close 25 GB of the +50 GB gap. diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/cuda_graph_runner.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/cuda_graph_runner.py index ebcbb297580..5f3a7ba5329 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/cuda_graph_runner.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/cuda_graph_runner.py @@ -12,22 +12,90 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""CUDA graph capture / replay for individual FSDP modules.""" +"""CUDA graph capture / replay for individual FSDP v2 modules. + +Built on ``te_graph_runtime.make_graphed_callables`` which supports +``capture_time_hooks`` — hooks that run outside CUDA graph capture (for +FSDP unshard / reshard) and are not replayed. ``sample_kwargs`` is used +so modules receive keyword arguments natively. + +A single ``CudaGraphRunner`` instance is stored on the root context and +orchestrates: + + 1. Recording sample args for each eligible FSDP module during the + first optimized forward pass. + 2. Calling ``make_graphed_callables`` with all modules and + ``capture_time_hooks`` that perform unshard / reshard. +""" # noqa: E501 import inspect -import gc -import warnings +import logging +import os from collections import OrderedDict from typing import Any, Dict, List, Optional, Tuple import torch +logger = logging.getLogger(__name__) + +# --------------------------------------------------------------------------- +# NVML memory helper (real GPU memory, not just torch allocator view) +# --------------------------------------------------------------------------- + + +def _nvml_device_memory(device: Optional[int] = None) -> Optional[Tuple[int, int]]: + """Return (used_MiB, total_MiB) from NVML, or None if unavailable.""" + try: + import pynvml + except ImportError: + return None + try: + pynvml.nvmlInit() + except pynvml.NVMLError: + return None + try: + if device is None: + device = torch.cuda.current_device() + handle = pynvml.nvmlDeviceGetHandleByIndex(device) + info = pynvml.nvmlDeviceGetMemoryInfo(handle) + return (info.used // (1024 * 1024), info.total // (1024 * 1024)) + except Exception: + return None + + +def _mem_snapshot() -> Dict[str, int]: + """Capture a snapshot of memory counters across torch and NVML.""" + snap = { + "torch_alloc": torch.cuda.memory_allocated() // 1_000_000, + "torch_reserved": torch.cuda.memory_reserved() // 1_000_000, + } + nvml = _nvml_device_memory() + if nvml is not None: + snap["nvml_used"] = nvml[0] + snap["nvml_total"] = nvml[1] + return snap + + +def _fmt_mem_snapshot(before: Dict[str, int], after: Dict[str, int], peak_alloc: int) -> str: + """Format memory diff as a human-readable string.""" + parts = [ + f"torch_alloc {before['torch_alloc']}→{after['torch_alloc']} MB " + f"(Δ{after['torch_alloc'] - before['torch_alloc']:+d})", + f"torch_reserved {before['torch_reserved']}→{after['torch_reserved']} MB " + f"(Δ{after['torch_reserved'] - before['torch_reserved']:+d})", + f"peak_alloc {peak_alloc // 1_000_000} MB", + ] + if "nvml_used" in before: + parts.append( + f"nvml_used {before['nvml_used']}→{after['nvml_used']} MB " + f"(Δ{after['nvml_used'] - before['nvml_used']:+d})" + ) + return " ".join(parts) -# ------------------------------------------------------------------ -# Helpers -# ------------------------------------------------------------------ +# --------------------------------------------------------------------------- +# Hook save / restore +# --------------------------------------------------------------------------- -# All known hook attributes across PyTorch versions (including 2.x additions) _HOOK_ATTRS = [ "_forward_pre_hooks", "_forward_hooks", @@ -41,333 +109,275 @@ ] -def _get_forward_param_names(module: torch.nn.Module) -> List[str]: - """Return the ordered parameter names of module.forward (excluding 'self').""" - sig = inspect.signature(module.forward) - return [ - name - for name, p in sig.parameters.items() - if name != "self" - and p.kind - in ( - inspect.Parameter.POSITIONAL_OR_KEYWORD, - inspect.Parameter.KEYWORD_ONLY, - inspect.Parameter.POSITIONAL_ONLY, - ) - ] +def _pop_all_hooks(module): + saved = [] + for sub in module.modules(): + snap = {} + for attr in _HOOK_ATTRS: + if hasattr(sub, attr): + snap[attr] = getattr(sub, attr) + setattr(sub, attr, OrderedDict()) + saved.append((sub, snap)) + return saved -class _ForwardShim(torch.nn.Module): - """Wraps module.forward so that non-tensor kwargs are frozen at - capture time and tensor inputs are passed positionally in signature - order. +def _restore_all_hooks(saved): + for sub, snap in saved: + for name, value in snap.items(): + if value is not None: + setattr(sub, name, value) - Handles None outputs by filtering them out before returning to - make_graphed_callables, and records their positions so they can be - restored during replay. - """ - def __init__( - self, - module: torch.nn.Module, - tensor_param_names: List[str], - frozen_kwargs: dict, - ): - super().__init__() - self.module = module - self.tensor_param_names = tensor_param_names - self.frozen_kwargs = frozen_kwargs - - # Populated on first forward: tracks which output positions are None - self._none_mask: Optional[List[bool]] = None - # Whether the original output was a tuple (vs single tensor) - self._output_is_tuple: bool = True - - def forward(self, *flat_tensor_args): - kwargs = dict(zip(self.tensor_param_names, flat_tensor_args)) - kwargs.update(self.frozen_kwargs) - outputs = self.module.forward(**kwargs) - - # Handle single tensor output - if not isinstance(outputs, tuple): - self._output_is_tuple = False - if outputs is None: - raise RuntimeError( - "Module returned None as its only output — cannot capture " - "with CUDA graphs. The module must return at least one tensor." +def _prepare_compiled_modules_for_capture(modules): + """Convert ``Module.compile()`` modules to compiled forward bodies. + + ``nn.Module.compile()`` compiles ``Module._call_impl``, which includes + module-hook dispatch. FSDP removes those hooks and replaces them with + ``capture_time_hooks`` while building its explicit CUDA graphs. Keeping + the compiled ``_call_impl`` can therefore trigger a guard failure and a + lazy recompile inside CUDA stream capture. + + Compile the forward body instead, with Inductor CUDA graphs disabled so + that the FSDP runner remains the sole CUDA-graph owner. The returned state + is only for rollback if explicit graph capture fails; after successful + installation, the stale compiled ``_call_impl`` must remain disabled. + """ + saved = [] + try: + for module in modules: + compiled_call_impl = getattr(module, "_compiled_call_impl", None) + if compiled_call_impl is None: + continue + + original_forward = module.forward + saved.append((module, original_forward, compiled_call_impl)) + module._compiled_call_impl = None + + # Avoid wrapping a forward body that the user already compiled + # directly. This branch mainly handles ``module.compile()``. + if not hasattr(original_forward, "_torchdynamo_orig_callable"): + module.forward = torch.compile( + original_forward, + dynamic=False, + options={"triton.cudagraphs": False}, ) - return outputs - - # Handle tuple outputs — filter out None values - self._output_is_tuple = True - if self._none_mask is None: - self._none_mask = [o is None for o in outputs] - - filtered = tuple(o for o in outputs if o is not None) - if len(filtered) == 0: - raise RuntimeError( - "Module returned a tuple of all None values — cannot capture " - "with CUDA graphs. At least one output must be a tensor." - ) - if len(filtered) == 1: - return filtered[0] - return filtered - - def restore_none_positions(self, graph_output) -> Any: - """Re-insert None values at their original positions after graph replay.""" - if not self._output_is_tuple or self._none_mask is None: - return graph_output - - # If all non-None outputs were collapsed to a single tensor - if not isinstance(graph_output, tuple): - graph_output = (graph_output,) - - full: List[Any] = [] - tensor_iter = iter(graph_output) - for is_none in self._none_mask: - if is_none: - full.append(None) - else: - full.append(next(tensor_iter)) - return tuple(full) - - -def _pop_hooks(module: torch.nn.Module) -> Dict[str, Any]: - """Remove all hooks from *module* (non-recursive) and return a snapshot.""" - saved: Dict[str, Any] = {} - for attr in _HOOK_ATTRS: - if hasattr(module, attr): - saved[attr] = getattr(module, attr) - setattr(module, attr, OrderedDict()) + except Exception: + _restore_compiled_modules_after_capture_failure(saved) + raise return saved -def _pop_hooks_recursive(module: torch.nn.Module) -> List[Tuple[torch.nn.Module, Dict[str, Any]]]: - """Remove all hooks from *module* and all its submodules recursively. +def _restore_compiled_modules_after_capture_failure(saved): + """Restore module-level compilation when explicit capture fails.""" + for module, original_forward, compiled_call_impl in saved: + module.forward = original_forward + module._compiled_call_impl = compiled_call_impl + + +class CudaGraphRunner: + """Orchestrates per-module sample-arg recording and batch graph capture. - Returns a list of (submodule, saved_hooks) tuples for restore. - Using direct module references instead of id() for safety. + Created once by the root forward pre-hook and stored on + ``ctx.cuda_graph_runner``. """ - saved: List[Tuple[torch.nn.Module, Dict[str, Any]]] = [] - for submodule in module.modules(): - saved.append((submodule, _pop_hooks(submodule))) - return saved + def __init__(self, graph_pool: Any, num_warmup_iters: int = 3): + self._graph_pool = graph_pool + self._num_warmup = num_warmup_iters + self._captured = False -def _restore_hooks(module: torch.nn.Module, saved: Dict[str, Any]) -> None: - """Put the hooks back exactly as they were.""" - for name, value in saved.items(): - if value is not None: - setattr(module, name, value) + # Per-module state recorded during the first optimized forward. + self._sample_args: Dict[int, Tuple] = {} + self._sample_kwargs: Dict[int, Dict[str, Any]] = {} + self._modules_ordered: List[torch.nn.Module] = [] + self._compiled_module_state = [] + # ---- called from hooks ------------------------------------------------ -def _restore_hooks_recursive( - module: torch.nn.Module, saved: List[Tuple[torch.nn.Module, Dict[str, Any]]] -) -> None: - """Restore hooks for all submodules saved by ``_pop_hooks_recursive``.""" - for submodule, sub_saved in saved: - _restore_hooks(submodule, sub_saved) + def record_module( + self, module: torch.nn.Module, args: Tuple, kwargs: Dict[str, Any] + ) -> None: + """Record sample args for *module* during the first optimized forward.""" + if self._captured: + return + mid = id(module) + if mid in self._sample_args: + return + # Normalize Module.compile() before capture setup. te-graph-runtime + # detects this compiled forward body and warms the capture-equivalent + # hook specialization before entering torch.cuda.graph. + self._compiled_module_state.extend( + _prepare_compiled_modules_for_capture([module]) + ) -# ------------------------------------------------------------------ -# Runner -# ------------------------------------------------------------------ + sig = inspect.signature(module.forward) + has_self = "self" in sig.parameters + bound = ( + sig.bind(module, *args, **kwargs) + if has_self + else sig.bind(*args, **kwargs) + ) + all_kwargs = { + n: bound.arguments[n] + for n in bound.arguments + if not (has_self and n == "self") + } + self._sample_args[mid] = tuple() # all via kwargs + self._sample_kwargs[mid] = all_kwargs + self._modules_ordered.append(module) + n_tensor = sum( + 1 for v in all_kwargs.values() if isinstance(v, torch.Tensor) + ) + if torch.distributed.is_initialized() and torch.distributed.get_rank() == 0: + logger.info( + "CudaGraphRunner: recorded module %s (id=%s), " + "%d kwargs (%d tensor)", + getattr(module, "_fsdp_module_name", module.__class__.__name__), + id(module), + len(all_kwargs), n_tensor, + ) -class FSDPCudaGraphRunner: - """Captures a forward+backward CUDA graph for one FSDP module. + def capture_and_install( + self, root_module: torch.nn.Module, + capture_stream: Optional[torch.cuda.Stream] = None, + ) -> None: + """Capture all graphs + install wrappers on recorded modules.""" + if self._captured or not self._modules_ordered: + return + self._captured = True - During capture hooks are temporarily removed so the graph records - only the user's ``forward()``, not FSDP all-gather / reduce-scatter - collectives. FSDP side streams are disabled for the capture region. + modules = self._modules_ordered + n = len(modules) - Parameters: - fsdp_module: The FSDP module to capture. - gc_freeze: If True (default), call ``gc.collect()`` and ``gc.freeze()`` - before capture to prevent Python GC from stalling replay. - graph_pool: Optional shared CUDA graph memory pool handle obtained - via ``torch.cuda.graph_pool_handle()``. When provided, the - ``CUDAGraph`` is created with this handle so multiple FSDP - modules share the same backing memory pool, reducing total - GPU memory consumption. + if torch.distributed.is_initialized() and torch.distributed.get_rank() == 0: + logger.info("CudaGraphRunner: capturing %d modules", n) - Usage:: + # Prefer installed te-graph-runtime (https://github.com/buptzyb/te-graph-runtime), + # fall back to vendored copy. + try: + from te_graph_runtime import make_graphed_callables + except ImportError: + from .te_graph_runtime import make_graphed_callables + + sample_args_list: List[Tuple] = [] + sample_kwargs_list: List[Dict[str, Any]] = [] + capture_hooks: List[Dict] = [] + + for m in modules: + mid = id(m) + # Clone tensor values so warmup gets fresh leaves without + # residual autograd state from the first forward+backward. + args = tuple( + v.detach().clone().requires_grad_(v.requires_grad) + if isinstance(v, torch.Tensor) else v + for v in self._sample_args[mid] + ) + kw = { + k: v.detach().clone().requires_grad_(v.requires_grad) + if isinstance(v, torch.Tensor) else v + for k, v in self._sample_kwargs[mid].items() + } + sample_args_list.append(args) + sample_kwargs_list.append(kw) + + capture_hooks.append({ + "forward_pre_hooks": {0: _make_fwd_pre_hook(m)}, + "forward_pre_hooks_with_kwargs": {0: True}, + "forward_hooks": {0: _make_fwd_post_hook(m)}, + "forward_hooks_with_kwargs": {0: True}, + "backward_pre_hooks": {0: _make_bwd_pre_hook(m)}, + "backward_hooks": {0: _make_bwd_post_hook(m)}, + }) + + self._sample_args.clear() + self._sample_kwargs.clear() + + compiled_module_state = self._compiled_module_state + if compiled_module_state and ( + not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0 + ): + logger.info( + "CudaGraphRunner: converted %d Module.compile() wrappers to " + "compiled forward bodies", + len(compiled_module_state), + ) - runner = FSDPCudaGraphRunner(my_fsdp_module) - runner.capture_forward(sample_input) - runner.install() # patches module.forward - output = my_fsdp_module(input_batch) # replays graph, no hooks - runner.uninstall() # restore original behaviour - """ + # Pop real FSDP hooks so make_graphed_callables passes its assertion. + # capture_time_hooks handle unshard/reshard during warmup + capture. + saved_hooks = _pop_all_hooks(root_module) - def __init__( - self, - fsdp_module: torch.nn.Module, - gc_freeze: bool = True, - graph_pool: Optional[Any] = None, - ): - warnings.warn( - "FSDPCudaGraphRunner is an experimental feature. The API and " - "behaviour may change in future releases without notice.", - FutureWarning, - stacklevel=2, - ) - self._module: torch.nn.Module = fsdp_module - self._gc_freeze: bool = gc_freeze - self._graph_pool: Optional[int] = graph_pool + try: + torch.cuda.reset_peak_memory_stats() + _mem_before = _mem_snapshot() + + graphed = make_graphed_callables( + tuple(modules), + tuple(sample_args_list), + num_warmup_iters=self._num_warmup, + sample_kwargs=tuple(sample_kwargs_list), + pool=self._graph_pool, + capture_time_hooks=capture_hooks, + capture_stream=capture_stream, + ) + except Exception: + _restore_compiled_modules_after_capture_failure(compiled_module_state) + raise + finally: + _restore_all_hooks(saved_hooks) - # Will hold the callable returned by make_graphed_callables - self._graphed: Optional[Any] = None + _mem_after = _mem_snapshot() + _peak_alloc = torch.cuda.max_memory_allocated() - self._orig_fwd: Optional[Any] = None - self._use_cuda_graph: bool = False - self._captured: bool = False + if torch.distributed.is_initialized() and torch.distributed.get_rank() == 0: + logger.info( + "CudaGraphRunner: %d modules captured %s", + n, + _fmt_mem_snapshot(_mem_before, _mem_after, _peak_alloc), + ) - # Saved during capture for install() replay flattening - self._tensor_param_names: List[str] = [] - self._frozen_kwargs: Dict[str, Any] = {} + if not isinstance(graphed, tuple): + graphed = (graphed,) - # The shim used during capture (needed for None restoration) - self._shim: Optional[_ForwardShim] = None + # make_graphed_callables already replaced module.forward with + # the graphed version that handles kwargs natively. + for module in modules: + module._fsdp_cg_installed = True + self._compiled_module_state = [] - # ------------------------------------------------------------------ - # 1. Capture - # ------------------------------------------------------------------ + if torch.distributed.is_initialized() and torch.distributed.get_rank() == 0: + logger.info("CudaGraphRunner: installed CUDA graphs on %d modules", n) - def capture_forward( - self, - *sample_args, - **sample_kwargs, - ) -> None: - assert self._module.cuda_graph_compatible, ( - "CUDA graph capture requires TracePoolAllocator in optimized phase" - ) - # Introspect the module's forward signature - param_names = _get_forward_param_names(self._module.__class__) - - # Separate tensor vs non-tensor inputs - bound = {} - for i, val in enumerate(sample_args): - if i < len(param_names): - bound[param_names[i]] = val - bound.update(sample_kwargs) - - tensor_names = [ - n for n in param_names if n in bound and isinstance(bound[n], torch.Tensor) - ] - frozen_kwargs = { - n: v for n, v in bound.items() if not isinstance(v, torch.Tensor) - } - flat_sample = tuple( - bound[n].clone().detach().requires_grad_(True) for n in tensor_names - ) +# --------------------------------------------------------------------------- +# capture_time_hooks (unshard / reshard outside graph, not replayed) +# --------------------------------------------------------------------------- - # Build shim (handles None filtering) - shim = _ForwardShim(self._module, tensor_names, frozen_kwargs) - for param in self._module.parameters(): - param.grad = None +def _make_fwd_pre_hook(module): + def hook(mod, args, kwargs): + module.unshard() + return hook - # For gradient accumulate fusion - self.unshard_main_grad_buffer() - if self._gc_freeze: - gc.collect() - gc.freeze() +def _make_fwd_post_hook(module): + def hook(mod, args, kwargs, output): + module.reshard() + return hook - # Disable side-stream collectives during capture so every CUDA - # operation lands on the default (capture) stream. - saved_hooks = _pop_hooks_recursive(self._module) - ctx = self._module._fsdp_root_context - ctx.cuda_graph_active = True - try: - torch.cuda.synchronize() - self._graphed = torch.cuda.make_graphed_callables( - shim, - sample_args=flat_sample, - num_warmup_iters=3, - allow_unused_input=True, - # pool=self._graph_pool, - ) - finally: - ctx.cuda_graph_active = False - _restore_hooks_recursive(self._module, saved_hooks) - self.reshard_main_grad_buffer() - self._shim = shim - self._tensor_param_names = tensor_names - self._frozen_kwargs = frozen_kwargs - self._captured = True +def _make_bwd_pre_hook(module): + def hook(mod, grad_output): + module.unshard(bwd_pass=True) + return hook - # ------------------------------------------------------------------ - # 2. Install / uninstall the patched forward - # ------------------------------------------------------------------ - def install(self) -> None: - if not self._captured: - raise RuntimeError("Call capture_forward() first") - if self._orig_fwd is not None: - return - self._orig_fwd = self._module.forward - graphed = self._graphed - shim = self._shim - param_names = _get_forward_param_names(self._module.__class__) - tensor_names = self._tensor_param_names - - def _patched_fwd(*args, **kwargs): - if self._use_cuda_graph: - bound = {} - for i, val in enumerate(args): - if i < len(param_names): - bound[param_names[i]] = val - bound.update(kwargs) - flat = tuple(bound[n] for n in tensor_names) - result = graphed(*flat) - # Restore None positions that were filtered during capture - return shim.restore_none_positions(result) - return self._orig_fwd(*args, **kwargs) - - self._module.forward = _patched_fwd - self._use_cuda_graph = True - - def uninstall(self) -> None: - """Restore the original ``forward``.""" - if self._orig_fwd is None: - return - self._module.forward = self._orig_fwd - self._orig_fwd = None - self._use_cuda_graph = False - - # ------------------------------------------------------------------ - # 3. Properties - # ------------------------------------------------------------------ - - @property - def captured(self) -> bool: - """True if ``capture_forward`` has been called successfully.""" - return self._captured - - @property - def using_cuda_graph(self) -> bool: - """True if the patch is currently active (install() called).""" - return self._use_cuda_graph - - def reset(self) -> None: - """Uninstall the patch and allow a fresh capture later.""" - self.uninstall() - self._graphed = None - self._shim = None - self._captured = False - self._graph = None - - def unshard_main_grad_buffer(self): - """Unshard the main grad buffer for all param groups.""" - for group in self._module._fsdp_param_groups: - if hasattr(group, "main_grad_buffer"): - group.main_grad_buffer.fetch_buffer() - - def reshard_main_grad_buffer(self): - """Reshard the main grad buffer for all param groups.""" - for group in self._module._fsdp_param_groups: - group.release_grad_buffer() +def _make_bwd_post_hook(module): + def hook(mod, grad_input, grad_output): + module.reshard() + # Clear grad to avoid memory leak in CUDA graph capture. + for param_group in module._fsdp_param_groups: + for param in param_group.params: + param.grad = None + return hook diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/design/cuda_graph_design.md b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/design/cuda_graph_design.md index 99b9f7b5f2a..09fd399a19f 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/design/cuda_graph_design.md +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/design/cuda_graph_design.md @@ -3,17 +3,24 @@ > **Experimental** — CUDA graph support in Megatron FSDP v2 is an experimental > feature. The API and behaviour may change in future releases without notice. +## Acknowledgements + +Built on [te-graph-runtime](https://github.com/buptzyb/te-graph-runtime) by +[@buptzyb](https://github.com/buptzyb) (Robin Zhang) — a standalone +TE-compatible `make_graphed_callables` with `capture_time_hooks` support. +Vendored under `te_graph_runtime/` with local modifications for FSDP v2. + ## 1. Motivation -mcore's CUDA graph system (`cuda_graph_impl="local"`, `cuda_graphs.py`) was -designed for DDP's memory model — each layer receives freshly-allocated tensor -inputs/outputs. Megatron FSDP v2 shares the same pool-backed buffers across layers, -and the FSDP hooks (unshard/reshard) are not captured in the graph in the -same way. +MCore's CUDA graph system was designed for DDP's memory model — each layer +receives freshly-allocated tensor inputs/outputs. Megatron FSDP v2 shares +pool-backed buffers across layers, and FSDP hooks (unshard/reshard) are not +part of the graphed region. -This doc describes a CUDA graph system built INTO Megatron FSDP v2, using -`TracePoolAllocator` as the stable-memory foundation. The user enables it -with a single flag — everything else is automatic. +This system uses `te-graph-runtime`'s `make_graphed_callables` with +`capture_time_hooks` to run FSDP unshard/reshard **outside** CUDA graph +capture — they are never captured or replayed. A single batch call captures +all eligible modules in forward execution order with a shared memory pool. ## 2. One knob @@ -21,21 +28,70 @@ with a single flag — everything else is automatic. fully_shard(module, enable_cuda_graph=True) ``` -No `--cuda-graph-warmup-steps`, no `--cuda-graph-scope`, no coordination with -the pipeline schedule. The system automatically captures and replays on the -first optimized microbatch. +## 3. Architecture + +### 3.0 Why te-graph-runtime instead of torch.cuda.make_graphed_callables? + +| Feature | `torch.cuda.make_graphed_callables` | `te-graph-runtime` | +|---------|-------------------------------------|-------------------| +| **Module hooks** | Asserts modules have NO hooks at capture time | `capture_time_hooks` — hooks that run outside graph capture, not registered on modules | +| **Keyword args** | Positional tensor args only | `sample_kwargs` passes keyword tensor args natively | +| **TE/FP8** | No TE-specific support | TE FP8/FP4 recipes, amax reduction, delayed wgrad, RNG tracker state | +| **Parameter grad lifetime** | Standard PyTorch lifetime | TE PR #2937 grad lifetime fix | + +FSDP v2 modules require unshard/reshard hooks around every forward/backward — +even during warmup and capture. PyTorch's version rejects modules with hooks. +`te-graph-runtime`'s `capture_time_hooks` runs these essential operations +outside the CUDA graph region without violating the no-hooks assertion. + +### 3.1 te-graph-runtime + +`make_graphed_callables` wraps each module with a `Graphed` autograd Function +that replays forward and backward CUDA graphs. It handles warmup, capture +order (fwds in forward-module order, bwds in reverse), shared memory pool, and +autograd wiring internally. + +**Warmup and capture share the same CUDA stream for activation reuse** — intermediate +tensors (cuDNN/cuBLAS workspaces, attention intermediates) allocated during +warmup stay at the same addresses in the same CUDA context, so the capture +phase reuses them instead of freeing and reallocating. This saves significant +GPU memory vs. a throwaway warmup stream that would fragment the caching +allocator with stale allocations. + +Key features we rely on: + +- **`capture_time_hooks`** — hook functions that run during warmup + capture + but are NOT captured into the CUDA graph. Used for FSDP unshard/reshard. +- **`sample_kwargs`** — keyword tensor arguments passed natively to modules, + avoiding positional shim complexity. +- **`pool`** — shared `graph_pool_handle` for all modules. +- **`capture_stream`** — optional parameter to serialize captures on the same + stream (important for shared-pool correctness). + +### 3.2 CudaGraphRunner + +A lightweight orchestrator stored on `ctx.cuda_graph_runner`: + +```python +class CudaGraphRunner: + def record_module(self, module, args, kwargs): + """Record sample args during the first optimized forward.""" + def capture_and_install(self, root_module, capture_stream=None): + """Pop hooks → make_graphed_callables → restore hooks.""" +``` + +### 3.3 capture_time_hooks -## 3. Why TracePoolAllocator is the enabler +Four minimal hooks attached via `capture_time_hooks` to each module: -CUDA graphs require **stable buffer addresses**. After `plan()` allocates -the pool tensor, every slot has a fixed offset. The returned views' addresses -are deterministic every micro-batch. +| Hook | Trigger | Action | +|------|---------|--------| +| `forward_pre_hooks_with_kwargs` | Before warmup/capture forward | `module.unshard()` | +| `forward_hooks_with_kwargs` | After warmup/capture forward | `module.reshard()` | +| `backward_pre_hooks` | Before warmup/capture backward | `module.unshard(bwd_pass=True)` | +| `backward_hooks` | After warmup/capture backward | `module.reshard()` + clear `param.grad` | -During graph replay, the graphed CUDA kernels operate on the exact addresses -that were recorded during capture. The allocator is **not called** for -graphed modules during replay — the graph uses the addresses directly. Pool -slots are pre-allocated by `plan()`, so no dynamic allocation occurs inside -the graph region. +These run outside `torch.cuda.graph()` — unshard/reshard are never captured. ## 4. Lifecycle @@ -43,308 +99,150 @@ the graph region. ╔═══════════════════════════════════════════════════════════════╗ ║ Stage 1 — Microbatch 0 (trace) ║ ║ ║ -║ forward: _trace_allocate / _trace_free → trace recorded ║ -║ backward: _trace_allocate / _trace_free → trace continues ║ -║ plan(): pool built, _seq_ops populated, phase="optimized" ║ +║ forward: _trace_allocate / _trace_free → trace recorded ║ +║ backward: _trace_allocate / _trace_free → trace continues ║ +║ plan(): pool built, phase="optimized" ║ ║ snapshot_slots(): freeze slot state for replay ║ ╚═══════════════════════════════════════════════════════════════╝ │ ╔═══════════════════════════════════════════════════════════════╗ -║ Stage 2 — Microbatch 1+ (capture + eager forward) ║ +║ Stage 2 — Microbatch 1 (record sample args + eager) ║ ║ ║ ║ root_forward_pre_hook: ║ -║ reset_cursor() ║ -║ restore_slots() (restore capture-time slot state) ║ +║ Creates shared graph_pool_handle + capture_stream ║ +║ reset_cursor(); restore_slots() ║ +║ ║ ║ forward (per graphed FSDPModule): ║ -║ forward_pre_hook detects enable_cuda_graph=True: ║ -║ 1. Creates FSDPCudaGraphRunner(module) ║ -║ 2. Runner sets cuda_graph_active=True, pops hooks ║ -║ 3. Calls torch.cuda.make_graphed_callables(shim, ...) ║ -║ → warms up, captures forward+backward graph ║ -║ 4. Runner restores hooks, clears cuda_graph_active ║ -║ 5. install() patches module.forward to graph replay ║ -║ Module then runs _orig_forward eagerly (this call) ║ -║ backward: ║ -║ post_backward hooks fire normally ║ -║ reshard / reduce_grad run normally ║ -║ post-bwd: ║ -║ enable_flexible_mode() ║ +║ forward_pre_hook: unshard params ║ +║ CudaGraphRunner.record_module() — records sample kwargs ║ +║ module.forward() runs eagerly ║ +║ post_forward_hook: reshard ║ +║ ║ +║ backward: eager ║ +║ pre_backward: unshard (bwd_pass) ║ +║ post_backward: reshard + reduce_grad ║ +║ ║ +║ post_backward_final_callback: ║ +║ trace → optimized transition (plan) ║ +║ _maybe_capture_cuda_graphs() ║ +║ → CudaGraphRunner.capture_and_install() ║ +║ 1. Clone sample kwargs (fresh leaves) ║ +║ 2. Pop all FSDP hooks from module tree ║ +║ 3. make_graphed_callables(tuple(modules), ...) ║ +║ → warmup + forward capture + backward capture ║ +║ → replaces module.forward with graphed version ║ +║ 4. Restore FSDP hooks ║ +║ 5. Mark modules as _fsdp_cg_installed ║ ╚═══════════════════════════════════════════════════════════════╝ │ ╔═══════════════════════════════════════════════════════════════╗ ║ Stage 3 — Microbatch 2+ (replay) ║ ║ ║ ║ root_forward_pre_hook: ║ -║ reset_cursor() ║ -║ restore_slots() ║ +║ reset_cursor(); restore_slots() ║ +║ ║ ║ forward (per graphed FSDPModule): ║ -║ forward_pre_hook: runner already exists → skip capture ║ -║ patched forward: graphed(*flat) → replays captured graph ║ -║ → forward+backward complete within graph callable ║ -║ → NO Python FSDP hooks fire during graph ║ -║ → allocator NOT called for graphed modules ║ -║ forward (per non-graphed FSDPModule): ║ -║ normal eager: hooks fire → alloc calls ║ +║ forward_pre_hook: unshard; _fsdp_cg_installed → skip record ║ +║ module.forward (graphed) → Graphed.apply → fwd_graph.replay ║ +║ post_forward_hook: reshard ║ +║ ║ ║ backward: ║ -║ post_backward hooks fire normally (after graph returns) ║ -║ reshard / reduce_grad run normally ║ -║ post-bwd: ║ -║ enable_flexible_mode() ║ +║ pre_backward: unshard (bwd_pass) ║ +║ Graphed.backward → bwd_graph.replay ║ +║ post_backward: reshard + reduce_grad ║ ╚═══════════════════════════════════════════════════════════════╝ ``` -## 5. FSDPCudaGraphRunner +## 5. Hook strategy -The runner encapsulates per-module CUDA graph capture and replay: +### 5.1 During capture (`make_graphed_callables`) -```python -class FSDPCudaGraphRunner: - def capture_forward(self, *sample_args, **sample_kwargs): - # 1. Build a _ForwardShim that freezes non-tensor kwargs - # 2. Disable side-stream collectives on the root context - # 3. Pop FSDP hooks from the module (save/restore) - # 4. Call torch.cuda.make_graphed_callables(shim, ...) - # → handles warmup (3 iters) + capture internally - # 5. Restore hooks and side-stream settings - - def install(self): - # Patches module.forward to a wrapper that flattens - # args/kwargs → positional tensors and calls graphed(*flat) -``` - -**Why `make_graphed_callables`?** PyTorch's built-in utility handles: -- Warmup iterations to settle cuDNN/cuBLAS auto-tuning and TE FP8 scales -- Memory pool allocation for captured tensors -- AccumulateGrad node creation on the capture stream (avoids - `cudaErrorStreamCaptureImplicit` by keeping `.grad` tensors alive from - warmup through capture) +Real FSDP hooks are **popped** from the entire module tree before calling +`make_graphed_callables` (it asserts modules have no hooks). `capture_time_hooks` +handle unshard/reshard during warmup, forward capture, and backward capture. -**Why pop hooks?** The graph records the module's `forward()` only — no -FSDP all-gather or reduce-scatter collectives. Hooks are restored -immediately after capture, before the eager forward call. +### 5.2 During replay -## 6. Per-FSDPModule selectivity +Real FSDP hooks are **restored** after capture. They fire normally around +`module.__call__`: -```python -# Only specific leaf layers graphed -for layer in model.layers: - fully_shard(layer, enable_cuda_graph=True) -fully_shard(model, enable_cuda_graph=False) -``` +- `forward_pre_hook` → unshard +- `module.forward` (graphed) → fwd replay +- `forward_hook` → reshard +- `backward_pre_hook` → unshard (bwd_pass) +- `Graphed.backward` → bwd replay +- `backward_hook` → reshard + reduce_grad -Each `FSDPModule` carries a flag in `_FSDPState`: +## 6. Capture stream & shared pool ```python -class _FSDPState: - enable_cuda_graph: bool = False +# hooks.py — root forward pre-hook +ctx.cuda_graph_stream = torch.cuda.Stream() +ctx.cuda_graph_pool = torch.cuda.graph_pool_handle() + +# CudaGraphRunner passes them to make_graphed_callables: +make_graphed_callables( + modules, + sample_args, + pool=ctx.cuda_graph_pool, + capture_stream=ctx.cuda_graph_stream, +) ``` -All modules share the same `TracePoolAllocator` — slot assignments are fixed -by `plan()` regardless of which modules are graphed. - -### Nesting limitation - -A parent FSDP module that contains other FSDP modules as children **cannot** -use `enable_cuda_graph=True`. Only leaf FSDP modules (those without FSDP -children) are eligible. This is enforced in `_init_fsdp_state` with a -`RuntimeError`. - -The reason: CUDA graph capture runs the module's forward without FSDP hooks -(the hooks are popped). For a parent module, the forward calls child -FSDP-module forwards, which may themselves have CUDA graph logic or require -FSDP unshard/reshard. Limiting CUDA graph to leaf modules avoids this -complexity. - -## 7. Capture stream invariants - -During CUDA graph capture, `FSDPCudaGraphRunner.capture_forward` disables -side-stream collectives and sets `cuda_graph_active=True` so that hooks -defer reshard and skip post-backward cleanup while the graph is being built. - -The root context provides: - -```python -@dataclass -class _FSDPRootContext: - cuda_graph_active: bool = False - """True only during CUDA graph capture (inside FSDPCudaGraphRunner). - Suppresses side-stream vs default-stream mismatches and defers reshard.""" - - @property - def cuda_graph_compatible(self) -> bool: - """Return True when both side-stream collectives are disabled.""" - return (not self.enable_unshard_prefetch) and (not self.enable_async_reduce_grad) -``` +## 7. Parameter gradients -``FSDPModule.cuda_graph_compatible`` adds an additional check that the -allocator is in ``"optimized"`` phase: - -```python -@property -def cuda_graph_compatible(self) -> bool: - ctx = self._fsdp_root_context - if not isinstance(ctx.bucket_allocator, TracePoolAllocator): - return False - if ctx.bucket_allocator.phase != "optimized": - return False - return ctx.cuda_graph_compatible -``` - -During replay, ``cuda_graph_active`` is **not** set — the graph callable -runs forward+backward atomically inside ``graphed(*flat)``, and the standard -hook lifecycle handles reshard/reduce_grad normally after the graph returns. - -## 8. Hook changes - -### 8.1 Forward pre-hook — capture trigger - -```python -# hooks.py — _register_forward_pre_hook -def unshard_param_groups(fsdp_module, args, kwargs): - # ... unshard as normal ... - if fsdp_module._fsdp_state.enable_cuda_graph and ( - not hasattr(fsdp_module, "_fsdp_cg_runner") - ): - cg_runner = FSDPCudaGraphRunner(fsdp_module) - cg_runner.capture_forward(*args, **kwargs) - cg_runner.install() - fsdp_module._fsdp_cg_runner = cg_runner -``` - -Capture happens exactly once — the first time the module is called in -an optimized forward pass. Subsequent calls skip because `_fsdp_cg_runner` -already exists. - -### 8.2 Defer post-forward reshard - -```python -# hooks.py — _register_forward_hook → reshard_param_groups -def reshard_param_groups(module, *unused): - ctx = module._fsdp_root_context - if ctx.backward_phase and id(module) == ctx.backward_module: - return - if ctx.cuda_graph_active: - return # deferred; cleanup runs after capture returns - module.reshard() -``` - -### 8.3 Suppress post_backward during capture - -```python -# hooks.py — _register_backward_hook → post_backward -def post_backward(module): - ctx = module._fsdp_root_context - if ctx.cuda_graph_active: - return # hooks are popped during capture; nothing to do - ctx.backward_done_modules.add(id(module)) - ctx._advance_backward_module() - module.reshard() - ... -``` - -### 8.4 Final callback — skips during capture - -```python -# hooks.py — _post_backward_final_callback -def _post_backward_final_callback(root_state, root_module): - ctx = root_module._fsdp_root_context - if ctx.cuda_graph_active: - return # capture manages cleanup; hooks are already popped - stream = ctx.rs_stream - for module in reversed(ctx.forward_order): - ... - ctx.cuda_graph_active = False # not reached during capture — clears - # any stale state at end of replay -``` - -``cuda_graph_active`` is **only needed during capture** (when -``FSDPCudaGraphRunner.capture_forward`` calls ``make_graphed_callables``). -The runner pops hooks before capture and restores them after, so normal -post-forward / post-backward cleanup runs correctly outside the graph region. - -During replay, the graph callable runs forward+backward atomically inside -``graphed(*flat)``. After it returns, the standard hook lifecycle handles -reshard and reduce_grad normally — no ``cuda_graph_active`` guard is required. - -### 8.5 Root forward pre-hook — restore slots - -```python -# hooks.py — _register_root_forward_pre_hook -def root_forward_pre_hook(_hook_module, args): - ctx = fsdp_module._fsdp_root_context - ... - if isinstance(ba, TracePoolAllocator) and ba.phase == "optimized": - ba.disable_flexible_mode() - ba.reset_cursor() - if any(getattr(m._fsdp_state, "enable_cuda_graph", False) - for m in ctx.forward_order): - ba.restore_slots() -``` - -## 9. Allocator interface - -The allocator requires three methods to support CUDA graph replay: - -```python -class TracePoolAllocator: - _captured_slot_state: List[bool] # snapshot of slot.in_use at capture end - - def snapshot_slots(self): - """Freeze current slot in_use state after all graphs are captured.""" - self._captured_slot_state = [s.in_use for s in self._slots] - - def restore_slots(self): - """Restore slot state to capture-time snapshot before each replay.""" - for i, in_use in enumerate(self._captured_slot_state): - self._slots[i].in_use = in_use -``` +`make_graphed_callables` includes module parameters in the autograd surface +(via `Graphed.apply(*(user_args + module_params))`). Parameter gradients are +computed during backward capture and returned by `Graphed.backward`. Autograd +sets `param.grad`; FSDP's post-backward hook (`reduce_grad`) consumes them via +`param.main_grad`. No manual gradient buffer management needed. -`snapshot_slots()` is called in the post-backward callback after `plan()` -(trace→optimized transition). `restore_slots()` is called in the root -forward pre-hook before each replay to ensure the same slot availability -as during capture. +## 8. te-graph-runtime local modifications -## 10. Slot state management +Vendored at `te_graph_runtime/`. Key local changes from upstream: -After the trace phase and `plan()`, `snapshot_slots()` freezes the -allocator's slot `in_use` state. Before each forward pass during replay, -`restore_slots()` resets slots to that snapshot. This ensures: +- **Non-tensor `sample_kwargs`**: `None`-safe guards at 6 access points in + warmup, backward capture, and `Graphed.forward`. +- **Positional arg replay**: `functionalized` reconstructs capture-time arg + order from both `user_args` and `user_kwargs`. +- **Memory cleanup**: `gc.collect()` + `torch.cuda.empty_cache()` between + warmup and capture. +- **Stream management**: optional `capture_stream` parameter; warmup uses + separate throwaway stream for `torch.compile` compatibility. +- **`param.grad` cleanup**: cleared after each backward hook iteration to + prevent gradient accumulation memory leaks. -- Flexible-mode allocs between microbatches see correct slot availability. -- Non-graphed modules during replay advance `in_use` flags correctly. -- The next replay starts from the same clean state. +See `te_graph_runtime/README.md` for details and upstream attribution. -## 11. Risks +## 9. Known issues & fixes -| Risk | Mitigation | -|------|-----------| -| Side-stream all-gather invisible to graph → corrupt param buffers | `cuda_graph_compatible` assertion + runner disables side streams during capture | -| Post-forward reshard frees buffers too early | `cuda_graph_active` guard in `reshard_param_groups` defers reshard during capture | -| AccumulateGrad nodes carry stale stream → `cudaErrorStreamCaptureImplicit` | `make_graphed_callables` handles warmup + capture on the same stream, keeping `.grad` tensors alive through warmup | -| Non-deterministic FP8/RNG across warmup iterations | `make_graphed_callables` runs 3 warmup iterations before capture — FP8 scales settle, RNG advances | -| Slot `in_use` flags stale after replay | `restore_slots()` resets to capture-end snapshot before each forward | -| Capture triggered for nested FSDP modules | `_init_fsdp_state` rejects `enable_cuda_graph=True` on modules with FSDP children | -| Graph capture OOM (pool too large) | Pool is pre-allocated by `plan()`; CUDA graph mempool uses same pattern | +| Issue | Root cause | Fix | Date | +|-------|-----------|-----|------| +| Convergence degradation | Missing `register_generator_state` on graphs — dropout masks baked | `register_generator_state` on both fwd/bwd graphs | 2026-06 | +| `cudaErrorStreamCaptureInvalidated` with `torch.compile` | Shared warmup/capture stream; compile recompilation breaks capture | Separate throwaway warmup stream (workaround; ideal: keep shared stream and fix compile guard mismatch) | 2026-06 | +| `cudaErrorStreamCaptureUnsupported` during compile | `autocast(enabled=False)` triggered guard mismatch and recompilation inside graph | `autocast(cache_enabled=False)` preserves bf16 autocast | 2026-06 | +| OOM during warmup | `warmup_outputs` held tensor refs across `gc.collect` + `empty_cache` | `param.grad = None` prevents gradient accumulation across warmup iterations | 2026-06 | +| Non-tensor kwargs crash `.requires_grad` | `tree_flatten` passes `None` into `static_input_surface` | 6 `is not None` / `isinstance` guards in te-graph-runtime | 2026-06 | +| Positional `hidden_states` missing in replay | `kwargs_keys` validation rejected positional args | `functionalized` checks both `user_args` and `user_kwargs` | 2026-06 | -## 12. Files +## 10. Files | File | Role | |------|------| -| `cuda_graph_runner.py` | `FSDPCudaGraphRunner` — capture via `make_graphed_callables`, install/uninstall forward patch | -| `fsdp_module.py` | `_FSDPRootContext.cuda_graph_active` flag, `cuda_graph_compatible` property, nested FSDP check in `_init_fsdp_state` | -| `hooks.py` | Capture trigger in forward pre-hook, reshard deferral in forward/backward hooks, slot restore in root forward pre-hook, cleanup in final callback | -| `allocator.py` | `TracePoolAllocator.snapshot_slots()` / `restore_slots()` for stable slot state | -| `fully_shard.py` | Accept `enable_cuda_graph` kwarg; passes to `_enable_cuda_graph()` | +| `cuda_graph_runner.py` | `CudaGraphRunner` — sample arg recording, batch capture orchestration | +| `hooks.py` | Capture trigger, hook save/restore, shared pool+stream creation | +| `fsdp_module.py` | `_FSDPRootContext.cuda_graph_runner` / `cuda_graph_stream` / `cuda_graph_pool` | +| `te_graph_runtime/` | Vendored `make_graphed_callables` with local modifications | +| `design/cuda_graph_design.md` | This document | -## 13. No user-visible knobs +## 11. No user-visible knobs | What happens | User action | |---|---| -| Trace collection | Automatic (MB0 forward) | +| Trace collection | Automatic (MB0) | | Pool planning | Automatic (MB0 post-backward) | -| Slot snapshot | Automatic (after plan) | -| Graph capture | Automatic (first optimized MB forward, via `FSDPCudaGraphRunner`) | -| Graph replay | Automatic (subsequent MBs, via patched forward) | -| Slot state management | Automatic (snapshot/restore around replay) | -| Side-stream suppression | Automatic (`cuda_graph_active` flag in hooks) | -| Flexible mode toggling | Automatic (existing hook lifecycle) | +| Sample arg recording | Automatic (MB1 forward pre-hook) | +| Graph capture | Automatic (MB1 post-backward, via `make_graphed_callables`) | +| Graph replay | Automatic (MB2+, via `Graphed.apply`) | +| RNG state | Handled by `make_graphed_callables` internally | +| Slot state | Automatic (snapshot/restore) | diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/dp_buffer.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/dp_buffer.py index d49085bd7f7..23e1f6a13d6 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/dp_buffer.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/dp_buffer.py @@ -586,6 +586,7 @@ def is_unsharded(self) -> bool: def unshard( self, bind_params: bool = False, + stream: Optional[torch.cuda.Stream] = None, ) -> torch.Tensor: """All-gather the full buffer from all shards and bind parameter storage. @@ -602,15 +603,14 @@ def unshard( sm = self.buffer_index.shard_meta shard_buffer = self.data[sm.local_data_index : sm.local_data_index + sm.size] - torch.distributed.all_gather_into_tensor( - output_tensor=full_buffer, - input_tensor=shard_buffer, - group=self.dp_group, - ) - if full_buffer.is_cuda: - # Temporary all-gather buckets may be released from another stream before - # the collective finishes; record the producer stream for allocator safety. - full_buffer.record_stream(torch.cuda.current_stream()) + stream = stream or torch.cuda.current_stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + torch.distributed.all_gather_into_tensor( + output_tensor=full_buffer, + input_tensor=shard_buffer, + group=self.dp_group, + ) if bind_params: self._bind_buffer_to_params(full_buffer) @@ -674,7 +674,10 @@ def fetch_buffer(self, *, as_shard: bool = False) -> torch.Tensor: @torch.no_grad() def reduce_grad( self, + *, overwrite_grad: bool = False, + grad_comm_dtype: Optional[torch.dtype] = None, + stream: Optional[torch.cuda.Stream] = None, ): """Reduce gradients into the optimizer-facing local shard. @@ -689,7 +692,7 @@ def reduce_grad( if self.sharding_strategy in ("no_shard", "optim"): overwrite_grad = True - grad_comm_dtype = self.mp_policy.grad_comm_dtype or self.dtype + grad_comm_dtype = grad_comm_dtype or self.mp_policy.grad_comm_dtype or self.dtype if self.gradient_scaling_factor in (None, 1.0): op = torch.distributed.ReduceOp.SUM @@ -705,15 +708,18 @@ def reduce_grad( local_grad_shard = self.data[sm.local_data_index : sm.local_data_index + sm.size] if not self.is_distributed and self.sharding_strategy == "no_shard": - comm_input = ( - self.data if grad_comm_dtype == self.dtype else self.data.to(grad_comm_dtype) - ) - if prescale: - comm_input.mul_(self.gradient_scaling_factor) - torch.distributed.all_reduce(comm_input, group=self.dp_group, op=op) - if grad_comm_dtype != self.dtype: - self.data.copy_(comm_input.to(self.dtype)) - return + stream = stream or torch.cuda.current_stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + comm_input = ( + self.data if grad_comm_dtype == self.dtype else self.data.to(grad_comm_dtype) + ) + if prescale: + comm_input.mul_(self.gradient_scaling_factor) + torch.distributed.all_reduce(comm_input, group=self.dp_group, op=op) + if grad_comm_dtype != self.dtype: + self.data.copy_(comm_input.to(self.dtype)) + return if self.is_distributed: # ZeRO-2/3 (optim_grads/optim_grads_params): ``self.data`` is the @@ -722,9 +728,6 @@ def reduce_grad( # accumulated into ``local_grad_shard`` for gradient accumulation. input_buffer = self.fetch_buffer() output_offset = sm.bucket_data_index - if input_buffer.is_cuda: - # Keep temporary reduce-scatter buffers tied to the stream that uses them. - input_buffer.record_stream(torch.cuda.current_stream()) else: # ZeRO-1 (optim): ``self.data`` is the replicated full grad # accumulation buffer. The optimizer consumes only this rank's @@ -733,23 +736,26 @@ def reduce_grad( input_buffer = self.fetch_buffer() output_offset = sm.local_data_index - comm_input = input_buffer.to(grad_comm_dtype) - if prescale: - comm_input.mul_(self.gradient_scaling_factor) - reduced_grad_shard = comm_input[output_offset : output_offset + sm.size] + stream = stream or torch.cuda.current_stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + comm_input = input_buffer.to(grad_comm_dtype) + if prescale: + comm_input.mul_(self.gradient_scaling_factor) + reduced_grad_shard = comm_input[output_offset : output_offset + sm.size] - torch.distributed.reduce_scatter_tensor( - output=reduced_grad_shard, input=comm_input, group=self.dp_group, op=op - ) + torch.distributed.reduce_scatter_tensor( + output=reduced_grad_shard, input=comm_input, group=self.dp_group, op=op + ) - # If the reduced shard is already in the local grad buffer, skip copy/accumulation. - if local_grad_shard.data_ptr() == reduced_grad_shard.data_ptr(): - return + # If the reduced shard is already in the local grad buffer, skip copy/accumulation. + if local_grad_shard.data_ptr() == reduced_grad_shard.data_ptr(): + return - if overwrite_grad: - local_grad_shard.copy_(reduced_grad_shard) - else: - local_grad_shard += reduced_grad_shard + if overwrite_grad: + local_grad_shard.copy_(reduced_grad_shard) + else: + local_grad_shard += reduced_grad_shard def check_all_fsdp_buffers(module) -> bool: diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/fsdp_module.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/fsdp_module.py index 5beb2418367..3491cafa88d 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/fsdp_module.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/fsdp_module.py @@ -150,13 +150,11 @@ class _FSDPRootContext: window indicates a bug.""" cuda_graph_pool: Optional[Any] = None - f"""Shared CUDA graph memory pool handle for CUDA graph capture. + """Shared CUDA graph memory pool handle for CUDA graph capture.""" - Obtained via ``torch.cuda.graph_pool_handle()``. Multiple - ``torch.cuda.CUDAGraph`` objects created with this handle share - the same backing memory pool, allowing the CUDA driver to reuse - graph memory across FSDP modules and reduce total GPU memory - consumption.""" + cuda_graph_runner: Optional[Any] = None + """``CudaGraphRunner`` instance. Created lazily on the first + optimized forward pre-hook and reused across micro-batches.""" backward_module: Optional[int] = None """``id(module)`` of the FSDP module whose backward is pending next. @@ -491,7 +489,7 @@ def _materialize_meta_module( After materialization, full parameters are broadcast from DP rank 0 before DTensor wrapping so every rank shards the same initialized value. """ - materialization_device = torch.cuda.current_device() + materialization_device = f"cuda:{torch.cuda.current_device()}" for name, m in self.named_modules(): if m in ignored_modules: continue @@ -519,17 +517,20 @@ def _materialize_meta_module( raise ValueError( f"Module {name} contains meta parameters but cannot reset them" ) - # Move materialized parameters to the same target device (e.g., GPU) m.to(materialization_device) if mesh is not None and mesh.size() > 1: dp_group = mesh.get_group() src_rank = torch.distributed.get_global_rank(dp_group, 0) - for param in self.parameters(): - if param.is_meta or isinstance(param, DTensor): + for name, m in self.named_modules(): + if m in ignored_modules: continue - torch.distributed.broadcast(param.data, src=src_rank, group=dp_group) + for param in m.parameters(): + if param.is_meta or isinstance(param, DTensor): + continue + param.data = param.data.to(materialization_device) + torch.distributed.broadcast(param.data, src=src_rank, group=dp_group) def _init_fsdp_state( self, @@ -656,13 +657,6 @@ def unshard(self, async_op: bool = False, bwd_pass: bool = False): ctx = self._fsdp_root_context stream = ctx.ag_stream if async_op else torch.cuda.current_stream() - if async_op: - # Synchronize ag_stream with current_stream to guarantee that main-stream - # writes to parameter data are visible before the all-gather kernel reads them - # on ag_stream. Without this barrier, stale or partially-written parameter - # shards may be gathered, causing convergence divergence. - stream.wait_stream(torch.cuda.current_stream()) - # Unshard this module and optionally prefetch next modules in the forward/backward pass if async_op: prefetch_modules = ctx.get_prefetch_next_modules(self, bwd_pass=bwd_pass) @@ -686,8 +680,7 @@ def unshard(self, async_op: bool = False, bwd_pass: bool = False): f"NaN detected in dist param for parameter {name}" ) - with torch.cuda.stream(stream): - param_group.unshard(bwd_pass=bwd_pass) + param_group.unshard(bwd_pass=bwd_pass, stream=stream) # Record event to track when unshard is done for this module if async_op: @@ -785,9 +778,14 @@ def reduce_grad(self, async_op: bool = False): # the Python-side ``setattr(param, "grad_added_to_main_grad", True)`` that # accompanies the eager backward is captured away. We record the per-param # flag during the trace micro-batch and restore it here. + zero_targets = [] + copy_srcs = [] + copy_dsts = [] + for name, param in zip(param_names, param_group.params): grad_added = getattr(param, "grad_added_to_main_grad", False) recorded = getattr(param, "_mfsdp_recorded_te_wgrad", False) + if grad_added or recorded: if param.grad is not None: del param.grad @@ -806,18 +804,23 @@ def reduce_grad(self, async_op: bool = False): param_main_grad is None or param_main_grad.data_ptr() != main_grad.data_ptr() ): - main_grad.zero_() + zero_targets.append(main_grad.view(-1)) else: main_grad = param.get_main_grad() - main_grad.copy_(param.grad.detach()) + copy_srcs.append(param.grad.detach().view(-1)) + copy_dsts.append(main_grad.view(-1)) del param.grad + # 2 kernel launches total (instead of N) + if zero_targets: + torch._foreach_zero_(zero_targets) + if copy_dsts: + torch._foreach_copy_(copy_dsts, copy_srcs) + if async_op: # ---- Overlapped path ---- # Switch to rs_stream for the reduce-scatter kernel - stream.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(stream): - param_group.reduce_grad() + param_group.reduce_grad(stream=stream) else: # ---- Non-overlapped path ---- # Reduce gradients immediately and release grad buffer @@ -825,28 +828,27 @@ def reduce_grad(self, async_op: bool = False): param_group.release_grad_buffer() # Install reduced gradients to distributed parameters - with torch.cuda.stream(stream): - for name, param, dist_param, dist_grad in zip( - param_names, - param_group.params, - param_group.dist_params, - param_group.dist_grads, - ): - if not param.requires_grad: - continue - if param_group.mp_policy.use_decoupled_grad: - setattr(dist_param, "decoupled_grad", dist_grad) - if dist_param.grad is not None: - del dist_param.grad - else: - assert ( - dist_grad is None or dist_param.dtype == dist_grad.dtype - ), ( - f"{name} Dist param dtype {dist_param.dtype} does not match dist grad dtype {dist_grad.dtype}" - ) - setattr(dist_param, "grad", dist_grad) - if hasattr(dist_param, "decoupled_grad"): - dist_param.decoupled_grad = None + for name, param, dist_param, dist_grad in zip( + param_names, + param_group.params, + param_group.dist_params, + param_group.dist_grads, + ): + if not param.requires_grad: + continue + if param_group.mp_policy.use_decoupled_grad: + setattr(dist_param, "decoupled_grad", dist_grad) + if dist_param.grad is not None: + del dist_param.grad + else: + assert ( + dist_grad is None or dist_param.dtype == dist_grad.dtype + ), ( + f"{name} Dist param dtype {dist_param.dtype} does not match dist grad dtype {dist_grad.dtype}" + ) + setattr(dist_param, "grad", dist_grad) + if hasattr(dist_param, "decoupled_grad"): + dist_param.decoupled_grad = None if async_op: event = stream.record_event() diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/fully_shard.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/fully_shard.py index 4fb05118845..09106235307 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/fully_shard.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/fully_shard.py @@ -39,6 +39,7 @@ from .utils import _init_default_fully_shard_mesh __all__ = ["FSDPModule", "fully_shard"] +_fsdp_class_cache = {} # module-level cache def fully_shard( @@ -99,7 +100,9 @@ def fully_shard( mp_policy = MixedPrecisionPolicy() cls = module.__class__ - new_cls = type(f"FSDP{cls.__name__}", (FSDPModule, cls), {}) + if cls not in _fsdp_class_cache: + _fsdp_class_cache[cls] = type(f"FSDP{cls.__name__}", (FSDPModule, cls), {}) + new_cls = _fsdp_class_cache[cls] module.__class__ = new_cls use_trace_pool = ( diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/hooks.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/hooks.py index 6e31bbab9b1..9620ade353f 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/hooks.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/hooks.py @@ -24,7 +24,7 @@ from torch.utils._pytree import tree_flatten, tree_map, tree_unflatten from .allocator import TracePoolAllocator -from .cuda_graph_runner import FSDPCudaGraphRunner +from .cuda_graph_runner import CudaGraphRunner from .fsdp_module import FSDPModule, _FSDPState from .utils import RegisterFSDPBackwardFunction @@ -55,6 +55,7 @@ def _find_fsdp_target(hook_module: nn.Module) -> Optional[FSDPModule]: return None +@torch.compiler.disable def mfsdp_forward_pre_hook(hook_module: nn.Module, args: Any, kwargs: Any): """Pre-forward hook for FSDP modules and fine-grained sub-modules. @@ -100,36 +101,24 @@ def mfsdp_forward_pre_hook(hook_module: nn.Module, args: Any, kwargs: Any): for param_group in target._fsdp_param_groups: param_group._maybe_free_grad_data() - # ---- CUDA graph capture (FSDPModule targets only) --------------------- - # Fine-grained hooks fire on sub-modules whose forward args differ from - # the FSDPModule's; CG capture is not meaningful there. + # ---- CUDA graph: record sample args (first optimized micro-batch) ----- + # Actual capture happens in mfsdp_post_backward_final_callback via + # ctx.cuda_graph_runner.capture_and_install(). if ( isinstance(hook_module, FSDPModule) and target._fsdp_state.enable_cuda_graph - and (not hasattr(target, "_fsdp_cg_runner")) + and (not getattr(target, "_fsdp_cg_installed", False)) and not ctx.backward_phase and target.cuda_graph_compatible ): - if torch.distributed.get_rank() == 0: - logger.debug( - "Capturing CUDA graph for module %s (id=%s)", - target._fsdp_module_name, - id(target), - ) - cg_runner = FSDPCudaGraphRunner( - target, graph_pool=ctx.cuda_graph_pool - ) - cg_runner.capture_forward(*args, **kwargs) - cg_runner.install() - target._fsdp_cg_runner = cg_runner - if torch.distributed.get_rank() == 0: - logger.debug( - "Captured CUDA graph for module %s (id=%s)", - target._fsdp_module_name, - id(target), + if ctx.cuda_graph_runner is None: + ctx.cuda_graph_runner = CudaGraphRunner( + graph_pool=ctx.cuda_graph_pool, ) + ctx.cuda_graph_runner.record_module(target, args, kwargs) +@torch.compiler.disable def mfsdp_post_forward_hook(module: nn.Module, *unused): """Post-forward hook: reshard parameters. @@ -186,11 +175,22 @@ def _register_forward_hook(module: FSDPModule): module._mfsdp_forward_hook = module.register_forward_hook(mfsdp_post_forward_hook) + +def _maybe_capture_cuda_graphs(ctx, root_module) -> None: + """Trigger batch CUDA graph capture via ``ctx.cuda_graph_runner``.""" + if ctx.cuda_graph_runner is not None: + with torch.enable_grad(): + ctx.cuda_graph_runner.capture_and_install( + root_module, capture_stream=ctx.cuda_graph_stream, + ) + + # --------------------------------------------------------------------------- # Internal: backward hook helpers # --------------------------------------------------------------------------- +@torch.compiler.disable def mfsdp_pre_backward_setup( hook_module: nn.Module, grads: Any = None, skip_final_callback: bool = False ): @@ -222,6 +222,7 @@ def mfsdp_pre_backward_setup( target._fsdp_pre_backward_done = True +@torch.compiler.disable def mfsdp_post_backward_hook(module: nn.Module): """Post-backward hook: reshard parameters and reduce gradients. @@ -252,6 +253,7 @@ def mfsdp_post_backward_hook(module: nn.Module): ctx._advance_backward_module() +@torch.compiler.disable def mfsdp_post_backward_final_callback(root_module: nn.Module): """Finalise the backward pass: drain skipped modules, reset state, clear fine-grained flags, and (on the first micro-batch) transition @@ -322,6 +324,31 @@ def mfsdp_post_backward_final_callback(root_module: nn.Module): f"Unexpected bucket allocator phase: {bucket_alloc.phase}" ) + # ---- CUDA graph: batch capture (after first optimized forward+backward) -- + _maybe_capture_cuda_graphs(ctx, root_module) + + + +def _maybe_capture_cuda_graphs(ctx, root_module) -> None: + """Trigger batch CUDA graph capture via ``ctx.cuda_graph_runner``. + + Guarded by asserts: TracePoolAllocator must be in use and in + ``"optimized"`` phase (stable buffer addresses required for CG). + """ + if ctx.cuda_graph_runner is not None: + allocator = ctx.bucket_allocator + assert isinstance(allocator, TracePoolAllocator), ( + "CUDA graph capture requires TracePoolAllocator" + ) + assert allocator.phase == "optimized", ( + f"CUDA graph capture requires allocator phase='optimized', " + f"got '{allocator.phase}'" + ) + with torch.enable_grad(), torch.cuda.amp.autocast(enabled=False): + ctx.cuda_graph_runner.capture_and_install( + root_module, capture_stream=ctx.cuda_graph_stream, + ) + # --------------------------------------------------------------------------- # Internal: backward hook helpers diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/mixed_precision.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/mixed_precision.py index 022c3efbbb1..79feb233e03 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/mixed_precision.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/mixed_precision.py @@ -409,7 +409,9 @@ def main_params_dtype_for_param(self, tensor: torch.Tensor) -> Optional[torch.dt no ``main_weight_buffer`` is allocated — the optimizer mutates the model-weight buffer directly. Set to ``torch.float32`` when the model uses quantized weights (FP8 / NVFP4) so the optimizer works on - high-precision copies. + high-precision copies. When this equals the model-weight dtype and + the sharding layout matches, ``_init_buffers`` skips the separate + buffer for the same reason (see ``ParamGroup._init_buffers``). """ return self.main_params_dtype diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/param_group.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/param_group.py index 7dc9ba26777..56bc4e2055e 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/param_group.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/param_group.py @@ -149,6 +149,8 @@ def _init_buffers(self) -> None: Buffer creation logic: - model_weight_buffer: always created; replicated unless "optim_grads_params" - main_weight_buffer: created if mp_policy.main_params_dtype is specified + AND it differs from the model-weight dtype or requires a different + sharding layout; otherwise the optimizer mutates model_weight_buffer - main_grad_buffer: created if requires_grad """ s = self.sharding_strategy @@ -172,9 +174,19 @@ def _init_buffers(self) -> None: tbuf.set_item(i, self.mp_policy.get_param_data(p, transpose=True)) self.transpose_weight_buffer = tbuf - # Create main weight buffer for mixed precision + # Create main weight buffer for mixed precision. Skip the redundant + # copy when the optimizer dtype matches the model-weight dtype AND the + # sharding layout is identical — in that case the optimizer mutates + # ``model_weight_buffer`` directly via the dist_param views (which the + # code below already binds to ``model_weight_buffer`` when + # ``main_weight_buffer`` is None). Quantized params (FP8/NVFP4) always + # need a separate main buffer because their model-weight dtype (uint8) + # differs from the optimizer dtype (fp32), so the dtype guard below + # already prevents skipping them. main_params_dtype = self.mp_policy.main_params_dtype_for_param(self.params[0]) - if main_params_dtype is not None: + if main_params_dtype is not None and ( + main_params_dtype != model_weight_dtype or shard_main_weights != shard_weights + ): mbuf = self._create_buffer(main_params_dtype, shard_main_weights, "main_weight") mbuf.init_data(torch.empty(mbuf.data_size, dtype=mbuf.dtype, device=self.device)) for i, p in enumerate(self.params): @@ -207,7 +219,12 @@ def _init_buffers(self) -> None: # Create distributed parameter views self._init_dist_params() - def unshard(self, bwd_pass: bool = False, bind_params: bool = True): + def unshard( + self, + bwd_pass: bool = False, + bind_params: bool = True, + stream: Optional[torch.cuda.Stream] = None, + ) -> None: """ Unshard model weights by all-gathering from sharded buffer. @@ -220,7 +237,7 @@ def unshard(self, bwd_pass: bool = False, bind_params: bool = True): self.model_weight_buffer, self.transpose_weight_buffer, bwd_pass=bwd_pass ): if weight_buffer is not None: - weight_buffer.unshard(bind_params=bind_params) + weight_buffer.unshard(bind_params=bind_params, stream=stream) self.mp_policy.post_unshard(self.params, bwd_pass=bwd_pass) @@ -255,7 +272,7 @@ def copy_main_weights_to_model_weights(self): self.transpose_weight_buffer, ) - def reduce_grad(self): + def reduce_grad(self, stream: Optional[torch.cuda.Stream] = None): """ Reduce gradients across DP ranks. @@ -269,6 +286,7 @@ def reduce_grad(self): # accumulating — no stale data from uninitialised or zeroed buffers. self.main_grad_buffer.reduce_grad( overwrite_grad=self._grad_buffer_is_fresh, + stream=stream, ) self._grad_buffer_is_fresh = False diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/te_graph_runtime/README.md b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/te_graph_runtime/README.md new file mode 100644 index 00000000000..dc3e4a15b52 --- /dev/null +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/te_graph_runtime/README.md @@ -0,0 +1,57 @@ +# te-graph-runtime (vendored) + +Vendored from [buptzyb/te-graph-runtime](https://github.com/buptzyb/te-graph-runtime). + +## Acknowledgements + +- [@buptzyb](https://github.com/buptzyb) (Robin Zhang) — original `te-graph-runtime` implementation, `capture_time_hooks` support +- [@shjwudp](https://github.com/shjwudp) (Jianbin Chang) — MFSDP v2 CUDA graph integration, local modifications + +TE-compatible CUDA graph callable runtime with `capture_time_hooks` support +— hooks that run outside CUDA graph capture (for FSDP unshard / reshard) +and are not replayed. + +## Why vendored + +The MFSDP v2 CUDA graph runner depends on `capture_time_hooks` (from TE PR #2831) +and the CUDA graph parameter-gradient lifetime fix (from TE PR #2937), which +are not yet available in upstream `torch.cuda.make_graphed_callables`. + +## API + +```python +from .te_graph_runtime import make_graphed_callables +``` + +## Local modifications + +### Non-tensor `sample_kwargs` support +Frameworks pass mixed tensor/non-tensor kwargs (e.g. `attention_mask=None`, +`img_shapes=[[1,64,64]]`). `tree_flatten` passes `None` entries into +`static_input_surface`, which crashed `.requires_grad` and `.data_ptr()`. + +- Added `is not None` / `isinstance(arg, torch.Tensor)` guards at 6 access points + in warmup, backward capture, and `Graphed.forward`. +- `functionalized` reconstructs capture-time arg order from both `user_kwargs` + (by name) and `user_args` (by position) to handle positional args during replay. + +### Memory cleanup between warmup and capture +Two rounds of `gc.collect()` + `torch.cuda.empty_cache()` between warmup and +forward capture release cached blocks from the CUDA caching allocator, giving +the graph pool a clean, unfragmented state. + +### Stream management +- `capture_stream` parameter accepted by `make_graphed_callables` and + `_make_graphed_callables` for shared-pool serialization. +- **Warmup and capture share the same CUDA stream for activation reuse** — + intermediate tensors allocated during warmup stay at the same addresses, + so capture reuses them instead of freeing and reallocating. Saves + significant GPU memory vs. a throwaway stream. +- A throwaway warmup stream is used only as a workaround for `torch.compile` + compatibility (see `cuda_graph_design.md` §9). The ideal state removes + this workaround once the underlying compile guard mismatch is fixed. + +## Prefer pip install + +If the fixes above are upstreamed, prefer `pip install te-graph-runtime` over +this vendored copy. See `cuda_graph_runner.py` for the fallback import logic. diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/te_graph_runtime/__init__.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/te_graph_runtime/__init__.py new file mode 100644 index 00000000000..c321557ec33 --- /dev/null +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/te_graph_runtime/__init__.py @@ -0,0 +1,20 @@ +"""TE-compatible CUDA graph callable runtime. + +Vendored from https://github.com/buptzyb/te-graph-runtime +Prefer ``pip install te-graph-runtime`` instead of the vendored copy. +""" + +from .graph import ( + UPSTREAM_TE_COMMIT, + UPSTREAM_TE_GRAPH_PATH, + UPSTREAM_TE_VERSION, + make_graphed_callables, +) + +__all__ = [ + "UPSTREAM_TE_COMMIT", + "UPSTREAM_TE_GRAPH_PATH", + "UPSTREAM_TE_VERSION", + "make_graphed_callables", +] + diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/te_graph_runtime/graph.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/te_graph_runtime/graph.py new file mode 100644 index 00000000000..416a8351fb1 --- /dev/null +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/te_graph_runtime/graph.py @@ -0,0 +1,1962 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Standalone TE-compatible CUDA graph callable runtime.""" +from __future__ import annotations + +from collections.abc import Iterable, Sequence +import contextlib +import gc +from math import ceil, prod +from typing import Any, Callable, Dict, List, Optional, Tuple, TypeVar, Union +import warnings + +UPSTREAM_TE_VERSION = "v2.16" +UPSTREAM_TE_COMMIT = "4220403e831d29e93868f7793693ea83f6b8b05b" +UPSTREAM_TE_GRAPH_PATH = "transformer_engine/pytorch/graph.py" + +__all__ = [ + "UPSTREAM_TE_COMMIT", + "UPSTREAM_TE_GRAPH_PATH", + "UPSTREAM_TE_VERSION", + "make_graphed_callables", +] + +_torch = None +torch = None +_tree_flatten = None +_tree_unflatten = None +_graph_pool_handle = None +_TE_AVAILABLE = None +_TE_IMPORT_ERROR = None + + +class _UnavailableTEType: + """Placeholder used when TransformerEngine is not installed.""" + + +DelayedScaling = _UnavailableTEType +Recipe = Any +dist_group_type = Any +TransformerEngineBaseModule = _UnavailableTEType +BasicOperation = _UnavailableTEType +Sequential = _UnavailableTEType +OperationFuser = _UnavailableTEType + + +class _FP8StateStub: + skip_fp8_weight_update_tensor = None + + +class _FP8GlobalStateManagerStub: + quantization_state = _FP8StateStub() + + @staticmethod + def is_first_fp8_module() -> bool: + return False + + @staticmethod + def reduce_and_update_fp8_tensors(*args, **kwargs) -> None: + return None + + @staticmethod + def is_fp8_enabled() -> bool: + return False + + @staticmethod + def get_fp8_recipe() -> None: + return None + + @staticmethod + def get_fp8_group() -> None: + return None + + @staticmethod + def add_fp8_tensors_to_global_buffer(*args, **kwargs) -> None: + return None + + +FP8GlobalStateManager = _FP8GlobalStateManagerStub + + +@contextlib.contextmanager +def _null_autocast(*args, **kwargs): + yield + + +autocast = _null_autocast + + +def get_default_fp8_recipe(): + raise RuntimeError( + "FP8 graph capture requires transformer_engine. Install te-graph-runtime[te] " + "or disable FP8/TE-specific options." + ) + + +def _require_torch(): + """Import torch lazily so package import does not require torch initialization.""" + global _torch, torch, _tree_flatten, _tree_unflatten, _graph_pool_handle + if _torch is None: + import torch as imported_torch + from torch.utils._pytree import tree_flatten, tree_unflatten + from torch._C import _graph_pool_handle as imported_graph_pool_handle + + _torch = imported_torch + torch = imported_torch + _tree_flatten = tree_flatten + _tree_unflatten = tree_unflatten + _graph_pool_handle = imported_graph_pool_handle + return _torch + + +def _load_optional_te() -> bool: + """Load TransformerEngine internals when available, without delegating graphing.""" + global _TE_AVAILABLE, _TE_IMPORT_ERROR + global DelayedScaling, Recipe, dist_group_type + global autocast, FP8GlobalStateManager, get_default_fp8_recipe + global get_all_rng_states, graph_safe_rng_available + global TransformerEngineBaseModule, BasicOperation, Sequential, OperationFuser + + if _TE_AVAILABLE is not None: + return _TE_AVAILABLE + try: + from transformer_engine.common.recipe import DelayedScaling as te_DelayedScaling + from transformer_engine.common.recipe import Recipe as te_Recipe + from transformer_engine.pytorch.constants import dist_group_type as te_dist_group_type + from transformer_engine.pytorch.quantization import ( + autocast as te_autocast, + FP8GlobalStateManager as te_FP8GlobalStateManager, + get_default_fp8_recipe as te_get_default_fp8_recipe, + ) + from transformer_engine.pytorch.distributed import ( + get_all_rng_states as te_get_all_rng_states, + graph_safe_rng_available as te_graph_safe_rng_available, + ) + from transformer_engine.pytorch.module.base import ( + TransformerEngineBaseModule as te_TransformerEngineBaseModule, + ) + from transformer_engine.pytorch.ops.op import BasicOperation as te_BasicOperation + from transformer_engine.pytorch.ops import Sequential as te_Sequential + from transformer_engine.pytorch.ops.fuser import OperationFuser as te_OperationFuser + except Exception as exc: # pragma: no cover - exact import failure is environment-specific + _TE_AVAILABLE = False + _TE_IMPORT_ERROR = exc + return False + + DelayedScaling = te_DelayedScaling + Recipe = te_Recipe + dist_group_type = te_dist_group_type + autocast = te_autocast + FP8GlobalStateManager = te_FP8GlobalStateManager + get_default_fp8_recipe = te_get_default_fp8_recipe + get_all_rng_states = te_get_all_rng_states + graph_safe_rng_available = te_graph_safe_rng_available + TransformerEngineBaseModule = te_TransformerEngineBaseModule + BasicOperation = te_BasicOperation + Sequential = te_Sequential + OperationFuser = te_OperationFuser + _TE_AVAILABLE = True + _TE_IMPORT_ERROR = None + return True + + +def _prepare_runtime() -> bool: + _require_torch() + return _load_optional_te() + + +def get_all_rng_states() -> Dict[Any, Any]: + return {} + + +def graph_safe_rng_available() -> bool: + _torch_mod = _require_torch() + return ( + hasattr(_torch_mod.cuda.CUDAGraph, "register_generator_state") + and hasattr(_torch_mod.Generator, "graphsafe_set_state") + and hasattr(_torch_mod.Generator, "graphsafe_get_state") + and hasattr(_torch_mod.Generator, "clone_state") + ) + + +def _te_required_error(feature: str) -> RuntimeError: + detail = f" Original import error: {_TE_IMPORT_ERROR}" if _TE_IMPORT_ERROR else "" + return RuntimeError( + f"{feature} requires transformer_engine internals compatible with {UPSTREAM_TE_VERSION}." + f" Install te-graph-runtime[te] or disable TE-specific graph options.{detail}" + ) + + +def _torch_dtype_to_np_typestr(dtype): + _torch_mod = _require_torch() + mapping = { + _torch_mod.float16: " 0 else 0, False), + "version": 3, + } + + +def make_weak_ref(x): + """Return a tensor-like weak reference so CUDA graph pool memory can be reused.""" + _torch_mod = _require_torch() + + def convert_to_torch_tensor(tensor): + if isinstance(tensor, _torch_mod.Tensor): + return tensor + old_ptr = tensor.data_ptr() + new_tensor = _torch_mod.as_tensor(tensor).view(tensor.dtype) + if old_ptr != new_tensor.data_ptr(): + raise RuntimeError("Data pointer mismatch after converting to torch.Tensor") + return new_tensor + + if isinstance(x, _torch_mod.Tensor): + return convert_to_torch_tensor(_WeakRefTensor(x.data_ptr(), x.dtype, x.shape)) if x.is_cuda else x + if isinstance(x, tuple): + return tuple(make_weak_ref(i) for i in x) + if isinstance(x, list): + return [make_weak_ref(i) for i in x] + if isinstance(x, dict): + return {k: make_weak_ref(v) for k, v in x.items()} + if isinstance(x, (int, float, bool)) or x is None: + return x + raise TypeError( + f"Invalid type {type(x).__name__} to make weak ref. Valid types are: " + "torch.Tensor, tuple, list, dict, int, float, bool, and None." + ) + + + + +_IS_GRAPH_CAPTURING = False + +_T = TypeVar("_T") +SingleOrTuple = Union[_T, Tuple[_T, ...]] + +_CAPTURE_TIME_HOOK_KEYS = ( + "forward_pre_hooks", + "forward_pre_hooks_with_kwargs", + "forward_hooks", + "forward_hooks_with_kwargs", + "backward_pre_hooks", + "backward_hooks", +) + + +def _empty_capture_time_hooks() -> Dict[str, Dict[Any, Any]]: + return {key: {} for key in _CAPTURE_TIME_HOOK_KEYS} + + +def _canonicalize_capture_time_hooks( + num_callables: int, + capture_time_hooks: Optional[List[Optional[Dict[str, Dict]]]], +) -> List[Dict[str, Dict[Any, Any]]]: + if capture_time_hooks is None: + return [_empty_capture_time_hooks() for _ in range(num_callables)] + if len(capture_time_hooks) != num_callables: + raise ValueError( + f"capture_time_hooks has {len(capture_time_hooks)} entries but there are " + f"{num_callables} callables" + ) + + canonicalized = [] + for callable_idx, hooks in enumerate(capture_time_hooks): + if hooks is None: + canonicalized.append(_empty_capture_time_hooks()) + continue + if not isinstance(hooks, dict): + raise TypeError( + "capture_time_hooks entries must be dicts or None, " + f"but entry {callable_idx} has type {type(hooks).__name__}" + ) + unknown_keys = sorted(set(hooks) - set(_CAPTURE_TIME_HOOK_KEYS)) + if unknown_keys: + raise ValueError( + f"Unknown capture_time_hooks keys for callable {callable_idx}: {unknown_keys}. " + f"Supported keys are {list(_CAPTURE_TIME_HOOK_KEYS)}" + ) + + callable_hooks = _empty_capture_time_hooks() + for key in _CAPTURE_TIME_HOOK_KEYS: + value = hooks.get(key, {}) + if value is None: + value = {} + if not isinstance(value, dict): + raise TypeError( + f"capture_time_hooks[{callable_idx!r}][{key!r}] must be a dict, " + f"but got {type(value).__name__}" + ) + callable_hooks[key] = dict(value) + canonicalized.append(callable_hooks) + return canonicalized + + +def _check_capture_time_hook_return(value: Any, hook_name: str, detail: str) -> None: + if value is not None: + raise RuntimeError( + f"capture_time_hooks {hook_name} must not return a value ({detail} must not be " + "modified via hook return)" + ) + + +def set_capture_start() -> None: + """Record beginning of `make_graphed_callables`.""" + global _IS_GRAPH_CAPTURING + _IS_GRAPH_CAPTURING = True + + +def set_capture_end() -> None: + """Record end of `make_graphed_callables`.""" + global _IS_GRAPH_CAPTURING + _IS_GRAPH_CAPTURING = False + + +def is_graph_capturing() -> bool: + """Return whether within `make_graphed_callables`.""" + return _IS_GRAPH_CAPTURING + + +def graph_pool_handle(): + """ + Returns an opaque token representing the id of a graph memory pool. + """ + _require_torch() + return _graph_pool_handle() + + +@contextlib.contextmanager +def _none_grad_context_wrapper(inputs): + """ + Wrapper to set the gradients of the inputs to None, + in case the backward pass makes grad accumulations. + """ + original_input_grads = [] + for input_tensor in inputs: + original_input_grads.append(input_tensor.grad) + input_tensor.grad = None + yield + for input_tensor, original_grad in zip(inputs, original_input_grads): + input_tensor.grad = original_grad + + +@contextlib.contextmanager +def _graph_context_wrapper(*args, **kwargs): + """Wrapper around `torch.cuda.graph`. + + This wrapper is a temporary workaround for a PyTorch bug: + automatic garbage collection can destroy a graph while another + graph is being captured, resulting in a CUDA error. See + https://github.com/pytorch/pytorch/pull/161037. + + """ + gc_is_enabled = gc.isenabled() + if gc_is_enabled: + gc.disable() + with torch.cuda.graph(*args, **kwargs): + yield + if gc_is_enabled: + gc.enable() + + +def _make_graphed_callables( + callables: SingleOrTuple[Callable], + sample_args: SingleOrTuple[Tuple[torch.Tensor, ...]], + num_warmup_iters: int = 3, + allow_unused_input: bool = False, + cache_quantized_params: bool = False, + sample_kwargs: Optional[SingleOrTuple[Dict[str, Any]]] = None, + _order: Optional[List[int]] = None, + _num_layers_per_chunk: Optional[List[int]] = None, + pool: Optional[Tuple[int, ...]] = None, + retain_graph_in_backward: bool = False, + _reuse_graph_input_output_buffers: bool = False, + _clone_param_grads_on_return: bool = True, + pre_warmup_hook: Optional[Callable] = None, + post_warmup_hook: Optional[Callable] = None, + capture_time_hooks: Optional[List[Optional[Dict[str, Dict]]]] = None, + capture_stream: Optional[torch.cuda.Stream] = None, +) -> SingleOrTuple[Callable]: + """ + Helper method for `make_graphed_callables` + """ + + if torch.is_autocast_enabled() and torch.is_autocast_cache_enabled(): + raise RuntimeError( + "make_graphed_callables does not support the autocast " + "caching. Please set `cache_enabled=False`." + ) + + # Default is to pass no kwargs to callables + if sample_kwargs is None: + if isinstance(callables, tuple): + sample_kwargs = tuple({} for _ in range(len(sample_args))) + else: + sample_kwargs = {} + + # Canonicalize args as tuples + just_one_callable = False + if not isinstance(callables, tuple): + just_one_callable = True + callables = (callables,) + sample_args = (sample_args,) + sample_kwargs = (sample_kwargs,) + + capture_time_hooks = _canonicalize_capture_time_hooks(len(callables), capture_time_hooks) + + # Check training/inference + is_training = all(c.training for c in callables) + if not is_training and any(c.training for c in callables): + raise RuntimeError( + "make_graphed_callables only supports when modules are all in training or all in" + " inference mode." + ) + + # Check sizes of args + _order_without_wgrad = None + delay_wgrad_compute = False + if _order is None: + if len(sample_args) != len(callables): + raise ValueError( + "Expected sample_args to have the same length as callables, " + f"but got {len(sample_args)} sample_args for {len(callables)} callables" + ) + if len(sample_kwargs) != len(callables): + raise ValueError( + "Expected sample_kwargs to have the same length as callables, " + f"but got {len(sample_kwargs)} sample_kwargs for {len(callables)} callables" + ) + else: + # Custom logic for interleaved pipeline parallelism + # Note: This is tightly coupled with the Megatron-core + # implementation of interleaved pipeline parallelism at + # https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/core/pipeline_parallel/schedules.py. + # Note: The model is assumed to consist of layers + # (corresponding to callables) that are grouped into + # model chunks. _num_layers_per_chunk is a list of integers + # that indicates the number of layers in each model chunk. + # _order is a list of chunk indices (1-indexed) that + # indicates the order in which the layers are evaluated. + # Positive values indicate forward passes and negative + # values indicate backward passes. Each + # entry in sample_args corresponds to one of the forward + # passes. + _order_without_wgrad = [] + for c_id in _order: + if ceil(c_id) != c_id: + delay_wgrad_compute = True + continue + _order_without_wgrad.append(c_id) + num_model_chunks = max(_order_without_wgrad) + num_microbatches = len(_order_without_wgrad) // num_model_chunks // 2 + if num_model_chunks * num_microbatches * 2 != len(_order_without_wgrad): + raise ValueError( + f"Pipeline-parallel order dimension mismatch: num_model_chunks ({num_model_chunks})" + f" * num_microbatches ({num_microbatches}) * 2 =" + f" {num_model_chunks * num_microbatches * 2}, but len(_order_without_wgrad) =" + f" {len(_order_without_wgrad)}" + ) + + # When delay_wgrad_compute is enabled, each layer is treated as a model chunk, which + # allows for fine-grained graph capture order. + if delay_wgrad_compute: + if _num_layers_per_chunk is None: + raise ValueError( + "'_num_layers_per_chunk' must be provided when delay_wgrad_compute is True." + ) + for num_layers in _num_layers_per_chunk: + if num_layers != 1: + raise ValueError( + "Each model chunk must have only one layer when delay_wgrad_compute is" + f" True, but got {num_layers} layers." + ) + + # Determine number of layers in each model chunk. + if _num_layers_per_chunk is None: + if not ( + len(sample_args) * 2 >= len(_order_without_wgrad) + and (len(sample_args) * 2 % len(_order_without_wgrad) == 0) + ): + raise ValueError( + f"{len(sample_args)} * 2 >= {len(_order_without_wgrad)} and" + f" {len(sample_args)} * 2 % {len(_order_without_wgrad)} == 0" + ) + num_layers = len(sample_args) // num_model_chunks // num_microbatches + _num_layers_per_chunk = [num_layers] * num_model_chunks + else: + if not ( + isinstance(_num_layers_per_chunk, int) + or len(_num_layers_per_chunk) == num_model_chunks + ): + raise ValueError( + "If _num_layers_per_chunk is provided, it must be an integer or a list of" + f" {num_model_chunks} integers, but got {_num_layers_per_chunk}." + ) + if isinstance(_num_layers_per_chunk, int): + _num_layers_per_chunk = [_num_layers_per_chunk] * num_model_chunks + total_num_layers = sum(_num_layers_per_chunk) + if len(callables) != total_num_layers: + raise ValueError( + f"Callables should have ({total_num_layers}) " + + f"entries when order input is provided but got {len(callables)}." + ) + if len(sample_args) != total_num_layers * num_microbatches: + raise ValueError( + f"Expected {total_num_layers * num_microbatches} " + + f"args tuple, but got {len(sample_args)}." + ) + + # Calculate the starting index of each chunk in callables for future use. + _prefix_num_layers = [0] + for m_chunk in range(num_model_chunks): + num_layers = _num_layers_per_chunk[m_chunk] + _prefix_num_layers.append(_prefix_num_layers[-1] + num_layers) + + if len(sample_kwargs) != len(sample_args): + raise ValueError( + "Pipeline-parallel schedule requires sample_kwargs and sample_args to have " + f"the same length, but got {len(sample_kwargs)} sample_kwargs " + f"for {len(sample_args)} sample_args" + ) + + # Check reuse graph conditions and reorganize sample_args and sample_kwargs. + # Note: When capturing a graph, we hold onto the args and kwargs so we have static buffers + # when the graph is replayed. If two model chunk microbatches have no overlap between their + # forward and backward, then we can reduce memory usage by reusing the same static buffers. + if _reuse_graph_input_output_buffers: + if _order is None: + raise ValueError( + "`_order` must be provided when `_reuse_graph_input_output_buffers` is True." + ) + if not is_training: + raise RuntimeError( + "`_reuse_graph_input_output_buffers` is only available in training mode." + ) + if isinstance(sample_args, tuple): + sample_args = list(sample_args) + if isinstance(sample_kwargs, tuple): + sample_kwargs = list(sample_kwargs) + + # Reorganize args and kwargs for input tensor reuse. + # fwd_sample_qs is keyed by model chunk index. The value is a queue of tuples. + # Each tuple contains the sample key signature and its fwd_idx. When we finish a backward + # chunk, we pop the corresponding fwd_idx and push to the consumed_sample_q. + # consumed_sample_q is keyed by the sample key signature. The value is a queue of the + # fwd_idx whose backward has been called so that we can reuse the same static buffers. + # In this way, we can reuse the same static input buffers for the non-overlapping samples + # with the same input signature. + fwd_sample_qs = {} + consumed_sample_q = {} + fwd_idx = [0] * num_model_chunks + for c_id in _order: + m_chunk = abs(ceil(c_id)) - 1 + + if c_id > 0: + sample_start_idx = (_prefix_num_layers[m_chunk] * num_microbatches) + ( + fwd_idx[m_chunk] * _num_layers_per_chunk[m_chunk] + ) + fwd_sample_idx = [ + sample_start_idx + i for i in range(_num_layers_per_chunk[m_chunk]) + ] + if m_chunk not in fwd_sample_qs: + fwd_sample_qs[m_chunk] = [] + for per_callable_fwd_idx in fwd_sample_idx: + sample_args_keys = tuple( + (t.shape, t.dtype, t.layout) for t in sample_args[per_callable_fwd_idx] + ) + sample_kwargs_keys = tuple( + (k, v.shape, v.dtype, v.layout) + for k, v in sorted(sample_kwargs[per_callable_fwd_idx].items()) + ) + sample_keys = sample_args_keys + sample_kwargs_keys + + fwd_sample_qs[m_chunk].append((sample_keys, per_callable_fwd_idx)) + if consumed_sample_q.get(sample_keys, []): + reuse_fwd_idx = consumed_sample_q[sample_keys].pop(0) + sample_args[per_callable_fwd_idx] = sample_args[reuse_fwd_idx] + sample_kwargs[per_callable_fwd_idx] = sample_kwargs[reuse_fwd_idx] + fwd_idx[m_chunk] += 1 + elif ceil(c_id) != c_id: + continue + else: + num_consumed_samples = min( + len(fwd_sample_qs[m_chunk]), _num_layers_per_chunk[m_chunk] + ) + for sample_keys, per_callable_fwd_idx in fwd_sample_qs[m_chunk][ + :num_consumed_samples + ]: + if sample_keys not in consumed_sample_q: + consumed_sample_q[sample_keys] = [] + consumed_sample_q[sample_keys].append(per_callable_fwd_idx) + fwd_sample_qs[m_chunk] = fwd_sample_qs[m_chunk][num_consumed_samples:] + + if cache_quantized_params: + # Initialize flag that controls FP8 weight updates + qstate = FP8GlobalStateManager.quantization_state + if qstate.skip_fp8_weight_update_tensor is None: + qstate.skip_fp8_weight_update_tensor = torch.empty( + 1, dtype=torch.float32, device="cuda" + ) + qstate.skip_fp8_weight_update_tensor.fill_(False) + + # Check callables + for c in callables: + if isinstance(c, torch.nn.Module): + if not ( + len(c._backward_hooks) == 0 + and len(c._backward_pre_hooks) == 0 + and len(c._forward_hooks) == 0 + and len(c._forward_pre_hooks) == 0 + ): + raise RuntimeError( + "Modules must not have hooks registered at the time they are passed. " + + "However, registering hooks on modules after passing them " + + "through make_graphed_callables is allowed. If you need hooks during " + + "capture, pass them with capture_time_hooks so they run outside CUDA " + + "graph capture and are not replayed." + ) + if not all(b.requires_grad is False for b in c.buffers()): + raise RuntimeError( + "In any :class:`~torch.nn.Module` passed to " + + ":func:`~make_graphed_callables`, only parameters may be trainable. " + + "All buffers must have ``requires_grad=False``." + ) + + # Flatten callable arguments + per_callable_kwargs_keys = [list(kwargs.keys()) for kwargs in sample_kwargs] + flatten_sample_args = [] + for args, kwargs, kwargs_keys in zip(sample_args, sample_kwargs, per_callable_kwargs_keys): + flatten_arg, _ = _tree_flatten(args) + flatten_kwarg, _ = _tree_flatten([kwargs[key] for key in kwargs_keys]) + flatten_sample_args.append(tuple(flatten_arg + flatten_kwarg)) + if not all(isinstance(arg, torch.Tensor) for arg in flatten_arg): + raise TypeError( + "In the beta API, sample_args " + + "for each callable must contain only Tensors. Other types are not allowed." + ) + + # If a callable is an nn.Module, its graph's full input surface is the args the user explicitly + # passes to forward (ie, its sample_args) AND the module's parameter attributes. + # Note: These per_callable_* variables are not actually + # per-callable, but per-forward-pass (see description of _order). + # The names are kept for consistency with + # PyTorch make_graphed_callables. + per_callable_len_user_args = [len(args) for args in flatten_sample_args] + if _order is None: + per_callable_module_params = [ + tuple(c.parameters()) if isinstance(c, torch.nn.Module) else () for c in callables + ] + per_callable_static_input_surfaces = [ + flatten_sample_args[i] + per_callable_module_params[i] for i in range(len(callables)) + ] + else: + per_callable_module_params = [] + for m_chunk in range(num_model_chunks): + for _ in range(num_microbatches): + for l_no in range(_num_layers_per_chunk[m_chunk]): + per_callable_module_params.append( + tuple(callables[_prefix_num_layers[m_chunk] + l_no].parameters()) + if isinstance( + callables[_prefix_num_layers[m_chunk] + l_no], + torch.nn.Module, + ) + else () + ) + if len(per_callable_module_params) != len(flatten_sample_args): + raise ValueError( + "Pipeline-parallel dimension mismatch: " + f"per_callable_module_params has {len(per_callable_module_params)} entries, " + f"but flatten_sample_args has {len(flatten_sample_args)} entries" + ) + per_callable_static_input_surfaces = [ + flatten_sample_args[i] + per_callable_module_params[i] + for i in range(len(flatten_sample_args)) + ] + + fwd_graphs = [torch.cuda.CUDAGraph() for _ in range(len(flatten_sample_args))] + bwd_graphs = [torch.cuda.CUDAGraph() for _ in range(len(flatten_sample_args))] + bwd_dw_graphs = [torch.cuda.CUDAGraph() for _ in range(len(flatten_sample_args))] + graph_callables = [None for _ in range(len(flatten_sample_args))] + + def _returned_param_grad_clone_slots(static_grad_inputs, module_params): + """Snapshot static grad slots that need clones before Graphed.backward returns.""" + if not _clone_param_grads_on_return: + return (False,) * len(static_grad_inputs) + module_param_start = len(static_grad_inputs) - len(module_params) + return tuple( + idx >= module_param_start + and not getattr(module_params[idx - module_param_start], "skip_backward_post_hook", False) + for idx in range(len(static_grad_inputs)) + ) + + # For cases with multiple active RNG states, e.g. TP. + if graph_safe_rng_available(): + for _, state in get_all_rng_states().items(): + for fwd_graph, bwd_graph, bwd_dw_graph in zip(fwd_graphs, bwd_graphs, bwd_dw_graphs): + fwd_graph.register_generator_state(state) + bwd_graph.register_generator_state(state) + bwd_dw_graph.register_generator_state(state) + + mempool = graph_pool_handle() if pool is None else pool + + # Warmup + # Hopefully prevents cudnn benchmarking and other lazy-initialization cuda work + # from ending up in any captures. + torch.cuda.synchronize() + + # Get warmup func and func_idx. + warmup_func_idx = [] + warmup_func = [] + if _order is None: + for func_idx, func in enumerate(callables): + warmup_func_idx.append(func_idx) + warmup_func.append(func) + else: + fwd_idx = [0] * num_model_chunks + for c_id in _order: + if c_id > 0: + m_chunk = c_id - 1 + for l_no in range(_num_layers_per_chunk[m_chunk]): + func = callables[_prefix_num_layers[m_chunk] + l_no] + func_idx = (_prefix_num_layers[m_chunk] * num_microbatches) + ( + fwd_idx[m_chunk] * _num_layers_per_chunk[m_chunk] + l_no + ) + warmup_func_idx.append(func_idx) + warmup_func.append(func) + fwd_idx[m_chunk] += 1 + if len(warmup_func) != len(sample_args): + raise ValueError(f"Warmup runs {len(warmup_func)} don't match args {len(sample_args)}.") + if len(warmup_func_idx) != len(set(warmup_func_idx)): + raise RuntimeError( + f"Warmup runs {len(warmup_func)} but only {len(set(warmup_func_idx))} are unique." + ) + + # Filter the TE modules that cudagraph can access. + visited_te_modules = {} + need_bwd_dw_graph = {} + + def _call_capture_time_forward_pre_hooks(callable_idx, func, args, kwargs) -> None: + hooks = capture_time_hooks[callable_idx] + with_kwargs = hooks["forward_pre_hooks_with_kwargs"] + for hook_id, hook in hooks["forward_pre_hooks"].items(): + if hook_id in with_kwargs: + _check_capture_time_hook_return( + hook(func, args, kwargs), + "forward_pre_hooks", + "args/kwargs", + ) + else: + _check_capture_time_hook_return( + hook(func, args), + "forward_pre_hooks", + "args", + ) + + def _call_capture_time_forward_hooks(callable_idx, func, args, kwargs, outputs) -> None: + hooks = capture_time_hooks[callable_idx] + with_kwargs = hooks["forward_hooks_with_kwargs"] + for hook_id, hook in hooks["forward_hooks"].items(): + if hook_id in with_kwargs: + _check_capture_time_hook_return( + hook(func, args, kwargs, outputs), + "forward_hooks", + "output", + ) + else: + _check_capture_time_hook_return( + hook(func, args, outputs), + "forward_hooks", + "output", + ) + + def _call_capture_time_backward_pre_hooks(callable_idx, func, grad_outputs) -> None: + for hook in capture_time_hooks[callable_idx]["backward_pre_hooks"].values(): + _check_capture_time_hook_return( + hook(func, grad_outputs), + "backward_pre_hooks", + "grad_output", + ) + + def _call_capture_time_backward_hooks(callable_idx, func, grad_inputs, grad_outputs) -> None: + for hook in capture_time_hooks[callable_idx]["backward_hooks"].values(): + _check_capture_time_hook_return( + hook(func, grad_inputs, grad_outputs), + "backward_hooks", + "grad_input", + ) + + def _make_grad_outputs(outputs): + return tuple( + torch.empty_like(o) if o is not None and o.requires_grad else None for o in outputs + ) + + def _run_warmup_forward(func_idx, func, callable_idx, register_discovery_hooks=True): + args = sample_args[func_idx] + kwargs = sample_kwargs[func_idx] + + def hook_fn(module, inputs, outputs, func_idx=func_idx): # pylint: disable=unused-argument + modules = set() + if isinstance(module, TransformerEngineBaseModule): + modules.add(module) + # If forward is called on a BasicOperation directly the hook will run. + elif isinstance(module, BasicOperation): + modules.add(module) + elif hasattr(module, "need_backward_dw") and hasattr(module, "backward_dw"): + modules.add(module) + # If forward is called on a te.ops.Sequential it is not called on its constituent ops. + elif isinstance(module, Sequential): + if module._module_groups is None: + raise RuntimeError("module._module_groups should have been initialized by warmup") + for module_group in module._module_groups: + if isinstance(module_group, OperationFuser): + for basic_op in module_group._basic_ops: + modules.add(basic_op) + if modules: + if func_idx not in visited_te_modules: + visited_te_modules[func_idx] = modules + else: + visited_te_modules[func_idx].update(modules) + + _call_capture_time_forward_pre_hooks(callable_idx, func, args, kwargs) + hooks = [] + if register_discovery_hooks and isinstance(func, torch.nn.Module): + for module in func.modules(): + hooks.append(module.register_forward_hook(hook_fn)) + outputs = func(*args, **kwargs) + for hook in hooks: + hook.remove() + _call_capture_time_forward_hooks(callable_idx, func, args, kwargs, outputs) + flatten_outputs, _ = _tree_flatten(outputs) + return flatten_outputs + + def _run_warmup_backward(func_idx, func, outputs, warmup_iter, callable_idx) -> None: + static_input_surface = per_callable_static_input_surfaces[func_idx] + inputs = tuple(i for i in static_input_surface if i is not None and i.requires_grad) + outputs_requiring_grad = tuple(o for o in outputs if o is not None and o.requires_grad) + grad_outputs = _make_grad_outputs(outputs) + + _call_capture_time_backward_pre_hooks(callable_idx, func, grad_outputs) + with _none_grad_context_wrapper(inputs): + torch.autograd.backward( + outputs_requiring_grad, + grad_tensors=tuple(o for o in grad_outputs if o is not None), + ) + grad_inputs = tuple(input.grad for input in inputs) + _call_capture_time_backward_hooks(callable_idx, func, grad_inputs, grad_outputs) + + # Filter module params that get None grad from grad_inputs and remove them + # from static_input_surface. This is to ensure that the backward hooks + # registered to these params are not wrongly triggered. + num_required_grad_sample_args = sum( + isinstance(arg, torch.Tensor) and arg.requires_grad + for arg in flatten_sample_args[func_idx] + ) + required_grad_input_idx = [] + for i, arg in enumerate(static_input_surface): + if isinstance(arg, torch.Tensor) and arg.requires_grad: + required_grad_input_idx.append(i) + module_params_with_grad = [] + for grad_inputs_idx, inputs_idx in enumerate(required_grad_input_idx): + if grad_inputs[grad_inputs_idx] is None and grad_inputs_idx < num_required_grad_sample_args: + if not allow_unused_input: + raise RuntimeError( + "The input tensor requires grad, but the grad is None after backward pass." + ) + elif grad_inputs[grad_inputs_idx] is not None and grad_inputs_idx >= num_required_grad_sample_args: + module_params_with_grad.append(static_input_surface[inputs_idx]) + if len(module_params_with_grad) != len(per_callable_module_params[func_idx]): + if warmup_iter != 0: + raise RuntimeError( + "no-grad params should only be used as inputs in the first warmup" + f" iteration, but found in iteration {warmup_iter}" + ) + per_callable_module_params[func_idx] = tuple(module_params_with_grad) + static_input_surface = flatten_sample_args[func_idx] + tuple(module_params_with_grad) + per_callable_static_input_surfaces[func_idx] = static_input_surface + + # Run wgrad. This is essential for some TE modules when they have + # delay_wgrad_compute enabled. + need_backward_dw = False + for module in visited_te_modules.get(func_idx, set()): + if hasattr(module, "need_backward_dw") and module.need_backward_dw(): + need_backward_dw = True + module.backward_dw() + need_bwd_dw_graph[func_idx] = need_backward_dw + + def _run_warmup_iteration(warmup_iter, register_discovery_hooks): + if _order is None: + warmup_outputs = [] + for func_idx, func in zip(warmup_func_idx, warmup_func): + outputs = _run_warmup_forward( + func_idx, + func, + func_idx, + register_discovery_hooks=register_discovery_hooks, + ) + warmup_outputs.append((func_idx, func, outputs)) + if is_training: + for func_idx, func, outputs in reversed(warmup_outputs): + _run_warmup_backward(func_idx, func, outputs, warmup_iter, func_idx) + return + + per_fwd_outputs = {} + fwd_idx = [0] * num_model_chunks + bwd_idx = [0] * num_model_chunks + for c_id in _order: + if c_id > 0: + m_chunk = c_id - 1 + for l_no in range(_num_layers_per_chunk[m_chunk]): + callable_idx = _prefix_num_layers[m_chunk] + l_no + per_callable_fwd_idx = (_prefix_num_layers[m_chunk] * num_microbatches) + ( + fwd_idx[m_chunk] * _num_layers_per_chunk[m_chunk] + l_no + ) + func = callables[callable_idx] + outputs = _run_warmup_forward( + per_callable_fwd_idx, + func, + callable_idx, + register_discovery_hooks=register_discovery_hooks, + ) + per_fwd_outputs[per_callable_fwd_idx] = outputs + fwd_idx[m_chunk] += 1 + elif ceil(c_id) == c_id: + if is_training: + m_chunk = -c_id - 1 + for l_no in reversed(range(_num_layers_per_chunk[m_chunk])): + callable_idx = _prefix_num_layers[m_chunk] + l_no + per_callable_bwd_idx = ( + _prefix_num_layers[m_chunk] * num_microbatches + ) + (bwd_idx[m_chunk] * _num_layers_per_chunk[m_chunk] + l_no) + func = callables[callable_idx] + outputs = per_fwd_outputs[per_callable_bwd_idx] + _run_warmup_backward( + per_callable_bwd_idx, + func, + outputs, + warmup_iter, + callable_idx, + ) + bwd_idx[m_chunk] += 1 + + # Run warmup on the same stream as capture so workspace buffers + # stay in the same CUDA context and don't need re-allocation. + capture_stream = capture_stream or torch.cuda.Stream() + with torch.cuda.stream(capture_stream): + if pre_warmup_hook is not None: + pre_warmup_hook() + + for warmup_iter in range(num_warmup_iters): + _run_warmup_iteration(warmup_iter, register_discovery_hooks=True) + + # TE discovery temporarily registers forward hooks, and Dynamo guards + # compiled modules on hook state. Capture runs after those hooks are + # removed, so warm the capture-equivalent specialization as well. + compiled_callables = any( + getattr(func, "_compiled_call_impl", None) is not None + or hasattr(getattr(func, "forward", None), "_torchdynamo_orig_callable") + for func in callables + ) + if num_warmup_iters > 0 and compiled_callables: + _run_warmup_iteration( + num_warmup_iters, + register_discovery_hooks=False, + ) + + if post_warmup_hook is not None: + post_warmup_hook() + torch.cuda.synchronize() + + import gc + gc.collect() + torch.cuda.empty_cache() + gc.collect() + torch.cuda.empty_cache() + + # All captures here share a mempool. To avoid replays corrupting each other's memory, + # the safest approach is to capture all passes in the same order they'll run: + # fwd 1, fwd 2, ... fwd N, then bwd N, bwd N-1, ... bwd 1. + + if _order is not None: # pylint: disable=too-many-nested-blocks + per_callable_static_outputs = [None] * len(flatten_sample_args) + per_callable_output_unflatten_spec = [None] * len(flatten_sample_args) + per_callable_static_grad_outputs = [None] * len(flatten_sample_args) + per_callable_static_grad_inputs = [None] * len(flatten_sample_args) + per_callable_returned_param_grad_clone_slots = [None] * len(flatten_sample_args) + fwd_idx = [0] * num_model_chunks + bwd_idx = [0] * num_model_chunks + static_grad_outputs_dict = {} + wgrad_validation_list = [None] * len(_order) + previous_chunk_last_callable_bwd_idx = None + for i, c_id in enumerate(_order): + if c_id > 0: + if not isinstance(c_id, int): + raise TypeError( + f"Forward order value must be an integer, but got {type(c_id).__name__}." + ) + # Capture forward graph for model chunk c_id, microbatch fwd_idx[c_id-1] + m_chunk = c_id - 1 + for l_no in range(_num_layers_per_chunk[m_chunk]): + callable_idx = _prefix_num_layers[m_chunk] + l_no + func = callables[callable_idx] + per_callable_fwd_idx = (_prefix_num_layers[m_chunk] * num_microbatches) + ( + fwd_idx[m_chunk] * _num_layers_per_chunk[m_chunk] + l_no + ) + args = sample_args[per_callable_fwd_idx] + kwargs = sample_kwargs[per_callable_fwd_idx] + fwd_graph = fwd_graphs[per_callable_fwd_idx] + _call_capture_time_forward_pre_hooks(callable_idx, func, args, kwargs) + with _graph_context_wrapper(fwd_graph, pool=mempool, stream=capture_stream): + outputs = func(*args, **kwargs) + _call_capture_time_forward_hooks(callable_idx, func, args, kwargs, outputs) + flatten_outputs, spec = _tree_flatten(outputs) + per_callable_static_outputs[per_callable_fwd_idx] = tuple(flatten_outputs) + per_callable_output_unflatten_spec[per_callable_fwd_idx] = spec + graph_callables[per_callable_fwd_idx] = func + fwd_idx[m_chunk] += 1 + else: + # Capture backward graph for model chunk c_id, microbatch bwd_idx[-c_id-1] + m_chunk = -ceil(c_id) - 1 + previous_per_callable_bwd_idx = None + for l_no in list(reversed(range(_num_layers_per_chunk[m_chunk]))): + callable_idx = _prefix_num_layers[m_chunk] + l_no + per_callable_bwd_idx = (_prefix_num_layers[m_chunk] * num_microbatches) + ( + bwd_idx[m_chunk] * _num_layers_per_chunk[m_chunk] + l_no + ) + if ceil(c_id) == c_id and need_bwd_dw_graph.get(per_callable_bwd_idx, False): + # Check if bwd graph has corresponding wgrad graph: + # Number of dgrad backward graphs should be equal to number of + # wgrad backward graphs. + # Note: For MCore, the validation rule is more strict (the next backward + # of dgrad graph must be corresponding wgrad graph). + if wgrad_validation_list[i] is None: + same_bwd_c_id_list = [i] + num_wgrad_c_id = 0 + for idx in range(i + 1, len(_order)): + if _order[idx] > 0: + continue + if _order[idx] == c_id: + same_bwd_c_id_list.append(idx) + if _order[idx] + 0.5 == c_id: + num_wgrad_c_id += 1 + if len(same_bwd_c_id_list) == num_wgrad_c_id: + for same_c_id_idx in same_bwd_c_id_list: + wgrad_validation_list[same_c_id_idx] = True + break + if len(same_bwd_c_id_list) < num_wgrad_c_id: + # It's impossible to have more wgrad than dgrad. + wgrad_validation_list[i] = False + break + if wgrad_validation_list[i] is None: + wgrad_validation_list[i] = False + if not wgrad_validation_list[i]: + raise RuntimeError( + f"Number of wgrad graph({num_wgrad_c_id}) doesn't match number " + f"of dgrad graphs ({len(same_bwd_c_id_list)}) for chunk {c_id}." + ) + elif ceil(c_id) != c_id: + per_callable_bwd_idx -= _num_layers_per_chunk[m_chunk] + if not is_training: + raise RuntimeError("Only training mode supports backward_dw.") + # If no one module needs the backward_dw, the bwd_dw_graph will be empty. + # So skip capturing it. For backward_dw, the order value is c_id - 0.5 to indicate + # the specific order of backward_dw. + if ceil(c_id) - c_id != 0.5: + raise ValueError( + "The order diff of wgrad and dgrad must be 0.5, " + f"get {ceil(c_id) - c_id}." + ) + if not need_bwd_dw_graph.get(per_callable_bwd_idx, False): + raise RuntimeError( + "No module needs wgrad computation but get float in order" + ) + bwd_dw_graph = bwd_dw_graphs[per_callable_bwd_idx] + with _graph_context_wrapper(bwd_dw_graph, pool=mempool, stream=capture_stream): + for module in visited_te_modules[per_callable_bwd_idx]: + if ( + hasattr(module, "need_backward_dw") + and module.need_backward_dw() + ): + module.backward_dw() + continue + + static_input_surface = per_callable_static_input_surfaces[per_callable_bwd_idx] + static_outputs = per_callable_static_outputs[per_callable_bwd_idx] + bwd_graph = bwd_graphs[per_callable_bwd_idx] + # For now, assumes all static_outputs require grad + if _reuse_graph_input_output_buffers: + # Note for _reuse_graph_input_output_buffers: grad output is only used + # within backward, so we can reuse the same static buffers every time. + static_grad_outputs_keys = tuple( + (o.shape, o.dtype, o.layout) + for o in static_outputs + if o is not None and o.requires_grad + ) + if static_grad_outputs_keys in static_grad_outputs_dict: + static_grad_outputs = static_grad_outputs_dict[static_grad_outputs_keys] + else: + static_grad_outputs = tuple( + torch.empty_like(o) if o is not None and o.requires_grad else None + for o in static_outputs + ) + static_grad_outputs_dict[static_grad_outputs_keys] = static_grad_outputs + else: + static_grad_outputs = tuple( + torch.empty_like(o) if o is not None and o.requires_grad else None + for o in static_outputs + ) + if is_training: + func = graph_callables[per_callable_bwd_idx] + _call_capture_time_backward_pre_hooks( + callable_idx, + func, + static_grad_outputs, + ) + inputs = tuple(i for i in static_input_surface if i is not None and i.requires_grad) + with _none_grad_context_wrapper(inputs), _graph_context_wrapper( + bwd_graph, pool=mempool, stream=capture_stream + ): + torch.autograd.backward( + tuple( + o for o in static_outputs if o is not None and o.requires_grad + ), + grad_tensors=tuple(o for o in static_grad_outputs if o is not None), + retain_graph=retain_graph_in_backward, + ) + grad_inputs = tuple(input.grad for input in inputs) + _call_capture_time_backward_hooks( + callable_idx, + func, + grad_inputs, + static_grad_outputs, + ) + + # Constructs a tuple suitable for returning from Graphed.backward: + # Pads out the actually-needed grads with Nones in gradient slots for inputs + # that don't require grad. I couldn't think of a one-liner for this pattern. + static_grad_inputs = [] + grad_idx = 0 + for arg in static_input_surface: + if is_training and isinstance(arg, torch.Tensor) and arg.requires_grad: + static_grad_inputs.append(grad_inputs[grad_idx]) + grad_idx += 1 + else: + static_grad_inputs.append(None) # type: ignore[arg-type] + static_grad_inputs = tuple(static_grad_inputs) # type: ignore[assignment] + + per_callable_static_grad_outputs[per_callable_bwd_idx] = static_grad_outputs + per_callable_static_grad_inputs[per_callable_bwd_idx] = static_grad_inputs + returned_param_grad_clone_slots = _returned_param_grad_clone_slots( + static_grad_inputs, + per_callable_module_params[per_callable_bwd_idx], + ) + per_callable_returned_param_grad_clone_slots[per_callable_bwd_idx] = ( + returned_param_grad_clone_slots + ) + + # Weak ref the static outputs and static grad inputs that are no longer needed + # in the following steps. These two type of tensors are both in cudagraph + # mempool, so we just deallocate them and let PyTorch's memory allocator + # reuse them elsewhere. + if _reuse_graph_input_output_buffers: + # Weak ref the static outputs of the forward pass of this backward. It's + # no longer needed after the corresponding backward graph is built up. + per_callable_static_outputs[per_callable_bwd_idx] = make_weak_ref( + static_outputs + ) + + # Parameter grads can be weak-refed immediately only when Graphed.backward + # will clone them before returning them to autograd users. + static_grad_inputs = per_callable_static_grad_inputs[per_callable_bwd_idx] + per_callable_static_grad_inputs[per_callable_bwd_idx] = tuple( + ( + make_weak_ref(grad_input) + if returned_param_grad_clone_slots[idx] and grad_input is not None + else grad_input + ) + for idx, grad_input in enumerate(static_grad_inputs) + ) + + # Weak ref the static grad inputs of the previous backward pass within the + # same chunk. + if previous_per_callable_bwd_idx is not None: + idx = previous_per_callable_bwd_idx + per_callable_static_grad_inputs[idx] = make_weak_ref( + per_callable_static_grad_inputs[idx] + ) + previous_per_callable_bwd_idx = per_callable_bwd_idx + + # Weak ref the static grad inputs of the previous chunk's last backward + # pass. + # Note: After a chunk's backward pass, we assume Mcore will send the grad + # input to another pipeline parallel rank and that the communication is + # finished before the end of the next chunk's backward pass. + if l_no == 0: + if previous_chunk_last_callable_bwd_idx is not None: + idx = previous_chunk_last_callable_bwd_idx + per_callable_static_grad_inputs[idx] = make_weak_ref( + per_callable_static_grad_inputs[idx] + ) + previous_chunk_last_callable_bwd_idx = per_callable_bwd_idx + if ceil(c_id) == c_id: + bwd_idx[m_chunk] += 1 + else: + # Capture forward graphs + per_callable_static_outputs = [] + per_callable_output_unflatten_spec = [] + for func_idx, (func, args, kwargs, fwd_graph) in enumerate( + zip(callables, sample_args, sample_kwargs, fwd_graphs) + ): + _call_capture_time_forward_pre_hooks(func_idx, func, args, kwargs) + with _graph_context_wrapper(fwd_graph, pool=mempool, stream=capture_stream): + outputs = func(*args, **kwargs) + _call_capture_time_forward_hooks(func_idx, func, args, kwargs, outputs) + graph_callables[func_idx] = func + + flatten_outputs, spec = _tree_flatten(outputs) + per_callable_static_outputs.append(tuple(flatten_outputs)) + per_callable_output_unflatten_spec.append(spec) + + # Capture backward graphs in reverse order + per_callable_static_grad_outputs = [] + per_callable_static_grad_inputs = [] + per_callable_returned_param_grad_clone_slots = [] + for static_input_surface, static_outputs, bwd_graph, bwd_dw_graph, bwd_idx in zip( + reversed(per_callable_static_input_surfaces), + reversed(per_callable_static_outputs), + reversed(bwd_graphs), + reversed(bwd_dw_graphs), + reversed(range(len(per_callable_static_input_surfaces))), + ): + # For now, assumes all static_outputs require grad + static_grad_outputs = tuple( + torch.empty_like(o) if o is not None and o.requires_grad else None + for o in static_outputs + ) + if is_training: + func = graph_callables[bwd_idx] + _call_capture_time_backward_pre_hooks(bwd_idx, func, static_grad_outputs) + inputs = tuple(i for i in static_input_surface if i is not None and i.requires_grad) + with _none_grad_context_wrapper(inputs), _graph_context_wrapper( + bwd_graph, pool=mempool + ): + torch.autograd.backward( + tuple(o for o in static_outputs if o is not None and o.requires_grad), + grad_tensors=tuple(o for o in static_grad_outputs if o is not None), + retain_graph=retain_graph_in_backward, + ) + grad_inputs = tuple(input.grad for input in inputs) + _call_capture_time_backward_hooks(bwd_idx, func, grad_inputs, static_grad_outputs) + + if need_bwd_dw_graph.get(bwd_idx, False): + with _graph_context_wrapper(bwd_dw_graph, pool=mempool, stream=capture_stream): + for module in visited_te_modules[bwd_idx]: + if hasattr(module, "need_backward_dw") and module.need_backward_dw(): + module.backward_dw() + # Constructs a tuple suitable for returning from Graphed.backward: + # Pads out the actually-needed grads with Nones in gradient slots for inputs that + # don't require grad. I couldn't think of a slick one-liner for this pattern. + static_grad_inputs = [] + grad_idx = 0 + for arg in static_input_surface: + if is_training and isinstance(arg, torch.Tensor) and arg.requires_grad: + static_grad_inputs.append(grad_inputs[grad_idx]) + grad_idx += 1 + else: + static_grad_inputs.append(None) # type: ignore[arg-type] + static_grad_inputs = tuple(static_grad_inputs) # type: ignore[assignment] + + per_callable_static_grad_outputs.append(static_grad_outputs) + per_callable_static_grad_inputs.append(static_grad_inputs) + per_callable_returned_param_grad_clone_slots.append( + _returned_param_grad_clone_slots( + static_grad_inputs, + per_callable_module_params[bwd_idx], + ) + ) + + # Reverse the most recent per-callable lists. + per_callable_static_grad_outputs = list(reversed(per_callable_static_grad_outputs)) + per_callable_static_grad_inputs = list(reversed(per_callable_static_grad_inputs)) + per_callable_returned_param_grad_clone_slots = list( + reversed(per_callable_returned_param_grad_clone_slots) + ) + # Now for every per_callable list, per_callable_*[i] holds the stuff for the ith callable. + + def make_graphed_autograd_function( + fwd_graph, + bwd_graph, + module_params, + kwargs_keys, + len_user_args, + output_unflatten_spec, + static_input_surface, + static_outputs, + static_grad_outputs, + static_grad_inputs, + returned_param_grad_clone_slots, + ): + class Graphed(torch.autograd.Function): + """Autograd function for graph replay.""" + + @staticmethod + def forward(ctx, skip_fp8_weight_update, cuda_graph_stream, cuda_graph_event, *inputs): + # pylint: disable=missing-function-docstring + + # Set flag for whether to update FP8 weight updates + ctx.is_first_module = FP8GlobalStateManager.is_first_fp8_module() + if ctx.is_first_module and skip_fp8_weight_update is not None: + FP8GlobalStateManager.quantization_state.skip_fp8_weight_update_tensor.fill_( + skip_fp8_weight_update + ) + ctx.cuda_graph_stream = cuda_graph_stream + ctx.cuda_graph_event = cuda_graph_event + # Copy values from new tensors into static tensors + for i in range(len_user_args): + if ( + isinstance(static_input_surface[i], torch.Tensor) + and inputs[i] is not None + and static_input_surface[i].data_ptr() != inputs[i].data_ptr() + ): + static_input_surface[i].copy_(inputs[i]) + + # Replay forward graph + if cuda_graph_stream != torch.cuda.current_stream(): + cuda_graph_stream.wait_stream(torch.cuda.current_stream()) + with cuda_graph_stream: + fwd_graph.replay() + if cuda_graph_event is not None: + torch.cuda.current_stream().wait_event(cuda_graph_event) + else: + torch.cuda.current_stream().wait_stream(cuda_graph_stream) + else: + fwd_graph.replay() + if not isinstance(static_outputs, tuple): + raise TypeError( + "Expected static_outputs to be a tuple, but got" + f" {type(static_outputs).__name__}" + ) + return tuple(o.detach() if o is not None else o for o in static_outputs) + + @staticmethod + @torch.autograd.function.once_differentiable + def backward(ctx, *grads): + # pylint: disable=missing-function-docstring + + # Replay backward graph + if len(grads) != len(static_grad_outputs): + raise ValueError( + "Backward graph grad dimension mismatch: " + f"received {len(grads)} grads, " + f"but expected {len(static_grad_outputs)} static_grad_outputs" + ) + for g, grad in zip(static_grad_outputs, grads): + if g is not None: + # don't copy if autograd gods have been kind and the + # incoming grad is already in the right place + if g.data_ptr() != grad.data_ptr(): + g.copy_(grad) + if ctx.cuda_graph_stream != torch.cuda.current_stream(): + ctx.cuda_graph_stream.wait_stream(torch.cuda.current_stream()) + with ctx.cuda_graph_stream: + bwd_graph.replay() + if ctx.cuda_graph_event is not None: + torch.cuda.current_stream().wait_event(ctx.cuda_graph_event) + else: + torch.cuda.current_stream().wait_stream(ctx.cuda_graph_stream) + else: + bwd_graph.replay() + + # Update FP8 scale factors if needed + if ctx.is_first_module: + FP8GlobalStateManager.reduce_and_update_fp8_tensors(forward=False) + + # Input args that didn't require grad expect a None gradient. + if not isinstance(static_grad_inputs, tuple): + raise TypeError( + "Expected static_grad_inputs to be a tuple, but got" + f" {type(static_grad_inputs).__name__}" + ) + grad_inputs = [] + for idx, grad_input in enumerate(static_grad_inputs): + if grad_input is None: + grad_inputs.append(None) + elif returned_param_grad_clone_slots[idx]: + # Returned parameter grads may be installed directly as param.grad. + # Clone to avoid exposing CUDA graph static buffers to autograd users. + grad_inputs.append(grad_input.detach().clone()) + else: + grad_inputs.append(grad_input.detach()) + return (None, None, None) + tuple(grad_inputs) + + def functionalized(*user_args, **user_kwargs): + + # Decide whether to update FP8 weights + skip_fp8_weight_update = None + if cache_quantized_params: + if "is_first_microbatch" not in user_kwargs or not isinstance( + user_kwargs["is_first_microbatch"], bool + ): + raise ValueError( + "`is_first_microbatch` boolean kwarg must be provided for FP8 weight" + " caching." + ) + + skip_fp8_weight_update = not user_kwargs["is_first_microbatch"] + + # The cuda_graph_stream and cuda_graph_event are used in the TE CUDA graph replay. + # When replaying the graph in the cuda graph stream, the graph replay could overlap + # with the work on main stream. + # When cuda_graph_event is given, it should be an external event recorded + # in the cuda graph and is used to sync-back to the main stream. + # If cuda_graph_event is not given, it will be None and the graph replay will block + # the main stream until it is finished. + if "cuda_graph_stream" in user_kwargs: + cuda_graph_stream = user_kwargs["cuda_graph_stream"] + user_kwargs.pop("cuda_graph_stream") + else: + cuda_graph_stream = torch.cuda.current_stream() + if "cuda_graph_event" in user_kwargs: + cuda_graph_event = user_kwargs["cuda_graph_event"] + user_kwargs.pop("cuda_graph_event") + else: + cuda_graph_event = None + # Runs the autograd function with inputs == all inputs to + # the graph that might require grad (explicit user args + + # module parameters) + # Assumes module params didn't change since capture. + # Reconstruct the same flattened arg order as capture time. + # User may pass some recorded kwargs as positional args, so + # check user_args first (by position), then user_kwargs. + user_pos_args = list(user_args) + kwarg_values = [] + for key in kwargs_keys: + if key in user_kwargs: + kwarg_values.append(user_kwargs[key]) + elif user_pos_args: + kwarg_values.append(user_pos_args.pop(0)) + # else: key was a default not passed — skip (not a tensor) + flatten_user_kwargs, _ = _tree_flatten(kwarg_values) + func_args = tuple(flatten_user_kwargs) + module_params + out = Graphed.apply( + skip_fp8_weight_update, cuda_graph_stream, cuda_graph_event, *func_args + ) + return _tree_unflatten(out, output_unflatten_spec) + + return functionalized + + def make_graphed_attribute_functions(graph_idx): + # Get te modules for current graph + te_modules = visited_te_modules.get(graph_idx, set()) + + # Attach backward_dw as an attribute to the graphed callable. + def backward_dw(): + if need_bwd_dw_graph.get(graph_idx, False): + bwd_dw_graphs[graph_idx].replay() + + # Trigger the grad accumulation hook for wgrad graphs. + for module in te_modules: + if ( + hasattr(module, "_trigger_wgrad_accumulation_and_reduce_hooks") + and module.need_backward_dw() + ): + module._trigger_wgrad_accumulation_and_reduce_hooks() + + # Attach reset as an attribute to the graphed callable. + def reset(): + fwd_graphs[graph_idx].reset() + bwd_graphs[graph_idx].reset() + bwd_dw_graphs[graph_idx].reset() + + return backward_dw, reset + + # Put together the final graphed callables + ret = [] + for i in range(len(sample_args)): + graphed = make_graphed_autograd_function( + fwd_graphs[i], + bwd_graphs[i], + per_callable_module_params[i], + per_callable_kwargs_keys[i], + per_callable_len_user_args[i], + per_callable_output_unflatten_spec[i], + per_callable_static_input_surfaces[i], + per_callable_static_outputs[i], + per_callable_static_grad_outputs[i], + per_callable_static_grad_inputs[i], + per_callable_returned_param_grad_clone_slots[i], + ) + + func = graph_callables[i] + te_modules = visited_te_modules.get(i, set()) + if isinstance(func, torch.nn.Module): + + def make_graphed_forward(func, graph_training_state, graphed, orig_fwd, te_modules): + def new_fwd(*user_args, **user_kwargs): + # If the module's training-or-eval state matches what we graphed, + # run the graph, otherwise run the original forward method + if func.training == graph_training_state: + # Set the FP8 group from global amax reduction. + if FP8GlobalStateManager.is_fp8_enabled(): + fp8_recipe = FP8GlobalStateManager.get_fp8_recipe() + for m in func.modules(): + if m not in te_modules: + # Only Set the FP8 meta for the modules included by forward + continue + if isinstance(m, TransformerEngineBaseModule): + from transformer_engine.pytorch.attention.dot_product_attention import ( + DotProductAttention, + ) + + if ( + isinstance(m, DotProductAttention) + and not fp8_recipe.fp8_mha + and not fp8_recipe.fp8_dpa + ): + # Don't need to update FP8 meta for non-FP8 DPA + continue + m.fp8_meta["fp8_group"] = FP8GlobalStateManager.get_fp8_group() + m.fp8_meta["recipe"] = FP8GlobalStateManager.get_fp8_recipe() + FP8GlobalStateManager.add_fp8_tensors_to_global_buffer( + m.fp8_meta, + ) + elif isinstance(m, BasicOperation): + for mode in ("forward", "backward"): + if m.num_quantizers(mode): + m._fp8_metas[mode][ + "fp8_group" + ] = FP8GlobalStateManager.get_fp8_group() + m._fp8_metas[mode][ + "recipe" + ] = FP8GlobalStateManager.get_fp8_recipe() + FP8GlobalStateManager.add_fp8_tensors_to_global_buffer( + m._fp8_metas[mode], + ) + return graphed(*user_args, **user_kwargs) + return orig_fwd(*user_args, **user_kwargs) + + return new_fwd + + forward = make_graphed_forward(func, func.training, graphed, func.forward, te_modules) + if _order is None: + func.forward = forward + ret.append(func) + else: + ret.append(forward) + else: + ret.append(graphed) + + backward_dw_func, reset_func = make_graphed_attribute_functions(i) + setattr(ret[-1], "backward_dw", backward_dw_func) + setattr(ret[-1], "reset", reset_func) + + if just_one_callable: + return ret[0] + + return tuple(ret) + + +def save_fp8_tensors( + modules: Iterable[torch.nn.Module], + recipe: Optional[Recipe], +) -> Optional[List[Any]]: + """ + Returns the FP8 tensors for all modules + with adjusted amax history sizes. + """ + + if not isinstance(recipe, DelayedScaling): + return None + + fp8_tensors = [] + for module in modules: + for m in module.modules(): + module_tensors = None + if isinstance(m, TransformerEngineBaseModule): + if m.primary_weights_in_fp8: + m.adjust_amax_history_length(recipe.amax_history_len) + module_tensors = m.get_fp8_meta_tensors() + elif isinstance(m, BasicOperation): + m.reset_recipe_state(recipe=recipe) + module_tensors = m._save_fp8_metas() + fp8_tensors.append(module_tensors) + return fp8_tensors + + +def restore_fp8_tensors( + modules: Iterable[torch.nn.Module], + fp8_tensors: Optional[List[Any]], +) -> None: + """Restore FP8 tensors.""" + + if fp8_tensors is None: + return + + for module in modules: + for m in module.modules(): + module_tensors = fp8_tensors.pop(0) + if isinstance(m, TransformerEngineBaseModule): + m.reset_fp8_meta_tensors(module_tensors) + elif isinstance(m, BasicOperation): + m._load_fp8_metas(module_tensors) + if len(fp8_tensors) != 0: + raise RuntimeError( + f"Got FP8 state for {len(fp8_tensors)} more modules than expected. " + "There is probably a discrepancy with `save_fp8_tensors`." + ) + + +def make_graphed_callables( + modules: SingleOrTuple[Callable], + sample_args: SingleOrTuple[Tuple[torch.Tensor, ...]], + num_warmup_iters: int = 3, + allow_unused_input: bool = False, + sample_kwargs: Optional[SingleOrTuple[Dict[str, Any]]] = None, + fp8_enabled: Optional[SingleOrTuple[bool]] = None, + fp8_calibrating: Optional[bool] = None, + fp8_recipe: Optional[Recipe] = None, + fp8_group: Optional[dist_group_type] = None, + fp8_weight_caching: Optional[bool] = None, + enabled: Optional[SingleOrTuple[bool]] = None, + calibrating: Optional[bool] = None, + recipe: Optional[Recipe] = None, + amax_reduction_group: Optional[dist_group_type] = None, + cache_quantized_params: Optional[bool] = None, + _order: Optional[List[int]] = None, + _num_layers_per_chunk: Optional[List[int]] = None, + pool: Optional[Tuple[int, ...]] = None, + retain_graph_in_backward: bool = False, + _reuse_graph_input_output_buffers: bool = False, + _clone_param_grads_on_return: bool = True, + pre_warmup_hook: Optional[Callable] = None, + post_warmup_hook: Optional[Callable] = None, + capture_time_hooks: Optional[List[Optional[Dict[str, Dict]]]] = None, + capture_stream: Optional[torch.cuda.Stream] = None, +) -> Union[Callable, Tuple[Callable, ...]]: + """ + Make CUDA graph version of Transformer Engine modules + + A variation of PyTorch's `make_graphed_callables` utility function + with support for Transformer Engine modules and FP8. Please see + the + original PyTorch implementation for more documentation. + + .. warning:: + + Arguments 'fp8_enabled', 'fp8_calibrating', 'fp8_recipe', 'fp8_group', and 'fp8_weight_caching' are deprecated. + Use arguments 'enabled', 'calibrating', 'recipe', 'amax_reduction_group', and 'cache_quantized_params' instead. + + Graphing parameters + ------------------- + modules: (tuple of) callable + Callable or callables to graph. + sample_args: (tuple of) tuple of torch.Tensor + Positional arguments to callable(s). + num_warmup_iters: int, default = 3 + Number of warmup iterations. + allow_unused_input: bool, default = False + Whether to handle case where callable inputs + and outputs are disconnected in compute graph. + sample_kwargs: (tuple of) dict, optional + Keyword arguments to callable(s) + pool: (tuple of) int, default = None, optional + An instance returned from function `torch.cuda.graph_pool_handle` that hints + this graph may share memory with the indicated pool. + retain_graph_in_backward: bool, default = False + Whether to set retain_graph=True in backward graph capture. + _reuse_graph_input_output_buffers: bool, default = False + Reduce memory usage by reusing input/output data buffers between + graphs. Only supported with Mcore interleaved pipeline parallelism, i.e. + when `_order` is provided. All callables in `modules` are assumed to have + inputs and outputs with the same dtype and shape. + _clone_param_grads_on_return: bool, default = True + Clone parameter gradients before returning them from CUDA graph replay. + Disabling this avoids the extra clone/copy and may improve performance, + but returned parameter gradients will alias CUDA graph static gradient + buffers. These tensors no longer have standard PyTorch returned-gradient + lifetime semantics: a later replay of the same graph, or reused-buffer + replay of another callable, may overwrite retained hook or `.grad` + tensors. Only disable this when the caller consumes returned parameter + gradients before any such overwrite can occur. + pre_warmup_hook: callable, default = None + A hook function that will be called once before all warmup iterations + (not once per callable). + post_warmup_hook: callable, default = None + A hook function that will be called once after all warmup iterations + (not once per callable). + capture_time_hooks: list of dict, optional + Per-callable hooks invoked during warmup and graph capture, but + intentionally executed outside CUDA graph capture so they are not + recorded into the graph and are not replayed. Each hook must return + ``None``. Each list entry corresponds to one callable and may include + these keys: + ``"forward_pre_hooks"`` maps hook IDs to hooks with signature + ``hook(module, args)`` or ``hook(module, args, kwargs)`` when the ID + is present in ``"forward_pre_hooks_with_kwargs"``; + ``"forward_hooks"`` maps hook IDs to hooks with signature + ``hook(module, args, output)`` or ``hook(module, args, kwargs, output)`` + when the ID is present in ``"forward_hooks_with_kwargs"``; + ``"backward_pre_hooks"`` maps hook IDs to + ``hook(module, grad_output)``; + ``"backward_hooks"`` maps hook IDs to + ``hook(module, grad_input, grad_output)``. + + Quantization parameters + ----------------------- + enabled: (tuple of) bool, default = False + whether or not to enable low precision quantization (FP8/FP4). + If tuple, the length must match the number of modules. + calibrating: bool, default = False + calibration mode allows collecting statistics such as amax and scale + data of quantized tensors even when executing without quantization enabled. + This is useful for saving an inference ready checkpoint while training + using a higher precision. + recipe: recipe.Recipe, default = None + recipe used for low precision quantization. + amax_reduction_group: torch._C._distributed_c10d.ProcessGroup, default = None + distributed group over which amaxes for the quantized tensors + are reduced at the end of each training step. + cache_quantized_params: bool, default = False + Whether or not to cache quantized weights across microbatches. if set to `True`, + the `is_first_microbatch` boolean argument must be passed into the forward + method for TransformerEngine modules. When storing primary weights in low precision + using TE's `quantized_model_init` API and using an quantization aware optimizer, + this arg must be set to `False` if calculating weight transposes' outside TE, e.g., + in the optimizer step. + + """ + + te_available = _prepare_runtime() + + # Handle deprecated args. If old kwargs are set, they are prioritized with warning. + if fp8_enabled is not None: + if enabled is not None: + raise ValueError( + "make_graphed_callables has deprecated `fp8_enabled` kwarg " + "in favor of `enabled`, but both kwargs are set." + ) + warnings.warn( + "make_graphed_callables has deprecated `fp8_enabled` kwarg in favor of `enabled`. " + "`fp8_enabled` will be removed in a future release.", + category=DeprecationWarning, + stacklevel=2, + ) + enabled = fp8_enabled + if enabled is None: + enabled = False + + if fp8_calibrating is not None: + if calibrating is not None: + raise ValueError( + "make_graphed_callables has deprecated `fp8_calibrating` kwarg " + "in favor of `calibrating`, but both kwargs are set." + ) + warnings.warn( + "make_graphed_callables has deprecated `fp8_calibrating` kwarg in favor of " + "`calibrating`. `fp8_calibrating` will be removed in a future release.", + category=DeprecationWarning, + stacklevel=2, + ) + calibrating = fp8_calibrating + if calibrating is None: + calibrating = False + + if fp8_recipe is not None: + if recipe is None: + warnings.warn( + "make_graphed_callables has deprecated `fp8_recipe` kwarg in favor of " + "`recipe`. `fp8_recipe` will be removed in a future release.", + category=DeprecationWarning, + stacklevel=2, + ) + else: + raise ValueError( + "make_graphed_callables has deprecated `fp8_recipe` kwarg " + "in favor of `recipe`, but both kwargs are set." + ) + recipe = fp8_recipe + + if fp8_group is not None: + if amax_reduction_group is None: + warnings.warn( + "make_graphed_callables has deprecated `fp8_group` kwarg in favor of " + "`amax_reduction_group`. `fp8_group` will be removed in a future release.", + category=DeprecationWarning, + stacklevel=2, + ) + else: + raise ValueError( + "make_graphed_callables has deprecated `fp8_group` kwarg " + "in favor of `amax_reduction_group`, but both kwargs are set." + ) + amax_reduction_group = fp8_group + + if fp8_weight_caching is not None: + if cache_quantized_params is not None: + raise ValueError( + "make_graphed_callables has deprecated `fp8_weight_caching` kwarg " + "in favor of `cache_quantized_params`, but both kwargs are set." + ) + warnings.warn( + "make_graphed_callables has deprecated `fp8_weight_caching` kwarg in favor of " + "`cache_quantized_params`. `fp8_weight_caching` will be removed in a future release.", + category=DeprecationWarning, + stacklevel=2, + ) + cache_quantized_params = fp8_weight_caching + if cache_quantized_params is None: + cache_quantized_params = False + + set_capture_start() + + # Handle single module. + just_one_callable = False + if not isinstance(modules, tuple): + just_one_callable = True + modules = (modules,) + + if not isinstance(enabled, tuple): + if not isinstance(enabled, bool): + raise TypeError( + f"enabled must be a bool or a tuple of bools, but got {type(enabled).__name__}" + ) + enabled = (enabled,) * len(modules) + else: + if len(enabled) != len(modules): + raise ValueError( + f"enabled length ({len(enabled)}) must match modules length ({len(modules)})" + ) + if not te_available and ( + any(enabled) + or calibrating + or recipe is not None + or amax_reduction_group is not None + or cache_quantized_params + ): + raise _te_required_error("FP8/TE-specific graph capture") + if any(enabled) and recipe is None: + recipe = get_default_fp8_recipe() + elif not any(enabled): + recipe = None + module_uses_fp8 = dict(zip((id(m) for m in modules), enabled)) + + # Store FP8 tensors to reset later. + saved_fp8_tensors = save_fp8_tensors(modules, recipe=recipe) + + # FP8 wrapper. + old_call_funcs = {} + + def wrap_autocast(block): + block_cls = type(block) + if block_cls in old_call_funcs: + return + + old_call_funcs[block_cls] = block_cls.__call__ + + # Wrap the original call function of the module class. + def call_func(self, *args, **kwargs): + with autocast( + enabled=module_uses_fp8.get(id(self), False), + calibrating=calibrating, + recipe=recipe, + amax_reduction_group=amax_reduction_group, + _graph=True, + ): + outputs = old_call_funcs[block_cls](self, *args, **kwargs) + return outputs + + block_cls.__call__ = call_func + + forward_funcs = [] + for module in modules: + if not isinstance(module, torch.nn.Module): + raise TypeError(f"Graphing for {type(module)} is not supported.") + wrap_autocast(module) + forward_funcs.append(module) + + if just_one_callable: + forward_funcs = forward_funcs[0] + else: + forward_funcs = tuple(forward_funcs) + + # Save RNG state. + if graph_safe_rng_available(): + generators = [ + torch.cuda.default_generators[torch.cuda.current_device()], + *get_all_rng_states().values(), + ] + original_rng_states = [state.get_state() for state in generators] + else: + original_rng_states = torch.cuda.get_rng_state() + + graphed_callables = _make_graphed_callables( + forward_funcs, + sample_args, + num_warmup_iters=num_warmup_iters, + allow_unused_input=allow_unused_input, + cache_quantized_params=cache_quantized_params, + sample_kwargs=sample_kwargs, + _order=_order, + _num_layers_per_chunk=_num_layers_per_chunk, + pool=pool, + retain_graph_in_backward=retain_graph_in_backward, + _reuse_graph_input_output_buffers=_reuse_graph_input_output_buffers, + _clone_param_grads_on_return=_clone_param_grads_on_return, + pre_warmup_hook=pre_warmup_hook, + post_warmup_hook=post_warmup_hook, + capture_time_hooks=capture_time_hooks, + capture_stream=capture_stream, + ) + + # Ensures warmup does not affect numerics for ops such as dropout. + if graph_safe_rng_available(): + for gen, state in zip(generators, original_rng_states): + gen.set_state(state) + else: + torch.cuda.set_rng_state(original_rng_states) + + # Remove FP8 wrapper. + for module_cls, old_call in old_call_funcs.items(): + module_cls.__call__ = old_call + + # Restore FP8 state. + restore_fp8_tensors(modules, saved_fp8_tensors) + + set_capture_end() + return graphed_callables diff --git a/tests/unit_tests/distributed/megatron_fsdp/v2/test_cuda_graph_runner.py b/tests/unit_tests/distributed/megatron_fsdp/v2/test_cuda_graph_runner.py new file mode 100644 index 00000000000..de9595c6d8f --- /dev/null +++ b/tests/unit_tests/distributed/megatron_fsdp/v2/test_cuda_graph_runner.py @@ -0,0 +1,94 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for Megatron-FSDP v2 CUDA graph capture hooks.""" + +from types import SimpleNamespace +from unittest.mock import patch + +import torch + +from megatron.core.distributed.fsdp.src.megatron_fsdp.v2.cuda_graph_runner import ( + CudaGraphRunner, + _make_bwd_post_hook, + _prepare_compiled_modules_for_capture, + _restore_compiled_modules_after_capture_failure, +) + + +def test_capture_backward_post_hook_clears_only_unsharded_parameter_grads(): + """Warmup grads must not survive outside the CUDA graph private pool.""" + full_param = torch.nn.Parameter(torch.ones(4)) + full_param.grad = torch.full_like(full_param, 2) + full_param.main_grad = torch.full_like(full_param, 3) + + dist_param = torch.nn.Parameter(torch.ones(2)) + dist_param.grad = torch.full_like(dist_param, 4) + + reshard_calls = [] + module = SimpleNamespace( + _fsdp_param_groups=[ + SimpleNamespace(params=[full_param], dist_params=[dist_param]) + ], + reshard=lambda: reshard_calls.append(True), + ) + + _make_bwd_post_hook(module)(module, (), ()) + + assert full_param.grad is None + assert torch.equal(full_param.main_grad, torch.full_like(full_param, 3)) + assert torch.equal(dist_param.grad, torch.full_like(dist_param, 4)) + assert reshard_calls == [True] + + +def test_module_compile_is_converted_to_compiled_forward_for_capture(): + module = torch.nn.Linear(2, 2) + original_forward = module.forward + compiled_call_impl = object() + compiled_forward = object() + module._compiled_call_impl = compiled_call_impl + + with patch.object(torch, "compile", return_value=compiled_forward) as compile_mock: + saved = _prepare_compiled_modules_for_capture([module]) + + compile_mock.assert_called_once_with( + original_forward, + dynamic=False, + options={"triton.cudagraphs": False}, + ) + assert module._compiled_call_impl is None + assert module.forward is compiled_forward + + _restore_compiled_modules_after_capture_failure(saved) + assert module._compiled_call_impl is compiled_call_impl + assert module.forward == original_forward + + +def test_module_compile_is_normalized_when_first_forward_is_recorded(): + module = torch.nn.Linear(2, 2) + module._compiled_call_impl = object() + + def compiled_forward(input): + return input + + runner = CudaGraphRunner(graph_pool=None) + sample = torch.ones(1, 2) + + with patch.object(torch, "compile", return_value=compiled_forward): + runner.record_module(module, (sample,), {}) + + assert module._compiled_call_impl is None + assert module.forward is compiled_forward + assert len(runner._compiled_module_state) == 1 + diff --git a/tests/unit_tests/distributed/megatron_fsdp/v2/test_dist_muon.py b/tests/unit_tests/distributed/megatron_fsdp/v2/test_dist_muon.py index 9378ee4f062..f5f9c0df3d1 100644 --- a/tests/unit_tests/distributed/megatron_fsdp/v2/test_dist_muon.py +++ b/tests/unit_tests/distributed/megatron_fsdp/v2/test_dist_muon.py @@ -118,7 +118,7 @@ def dist_env(): def _read_full_weight(pg, idx, shape): """All-gather this param's full master weight and return it (reshaped, cloned).""" - mbuf = pg.main_weight_buffer + mbuf = pg.main_weight_buffer or pg.model_weight_buffer full = mbuf.unshard(bind_params=False) start, end = mbuf.buffer_index._get_item_global_range(idx) out = full[start:end].clone().view(shape) From 21256aac3805a566681c627e02c5b71967037e12 Mon Sep 17 00:00:00 2001 From: jianbinc Date: Thu, 2 Jul 2026 17:47:27 +0800 Subject: [PATCH 27/28] Add real-data convergence checks to FSDP examples 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. --- examples/megatron_fsdp/README.md | 12 +- examples/megatron_fsdp/fsdp_toy.py | 131 +++++++++++++++++++-- examples/qwenimage_mfsdp/test_qwenimage.py | 80 ++++++++++++- 3 files changed, 199 insertions(+), 24 deletions(-) diff --git a/examples/megatron_fsdp/README.md b/examples/megatron_fsdp/README.md index ad166e65ee6..c8bf9efc9bf 100644 --- a/examples/megatron_fsdp/README.md +++ b/examples/megatron_fsdp/README.md @@ -9,7 +9,7 @@ Example scripts for training and checkpoint conversion using [Megatron-FSDP](../ Standalone toy example (not Megatron-LM) demonstrating Megatron-FSDP v2 usage: - Basic model wrapping with `fully_shard()` -- CUDA graph capture (`--cuda-graph` / `--no-cuda-graph`) +- CUDA graph capture (`--cuda-graph`, off by default) > **Experimental** — CUDA graph support is experimental and may change. - Activation checkpointing (`--activation-checkpoint`) - Distributed checkpointing @@ -19,15 +19,10 @@ torchrun --nproc_per_node=2 examples/megatron_fsdp/fsdp_toy.py \ --model-dim 512 --n-layers 2 --batch-size 4 \ --use-megatron-fsdp -# With CUDA graph +# With CUDA graph (Megatron-FSDP only) torchrun --nproc_per_node=2 examples/megatron_fsdp/fsdp_toy.py \ --model-dim 512 --n-layers 2 --batch-size 4 \ --use-megatron-fsdp --cuda-graph - -# Without CUDA graph -torchrun --nproc_per_node=2 examples/megatron_fsdp/fsdp_toy.py \ - --model-dim 512 --n-layers 2 --batch-size 4 \ - --use-megatron-fsdp --no-cuda-graph ``` | Flag | Default | Description | @@ -35,8 +30,7 @@ torchrun --nproc_per_node=2 examples/megatron_fsdp/fsdp_toy.py \ | `--model-dim` | `1024` | Hidden dimension size | | `--n-layers` | `3` | Number of transformer layers | | `--use-megatron-fsdp` | off | Use Megatron-FSDP v2 instead of PyTorch FSDP2 | -| `--cuda-graph` | on | Enable CUDA graph capture on transformer layers | -| `--no-cuda-graph` | — | Disable CUDA graph capture | +| `--cuda-graph` | off | Enable CUDA graph capture on transformer layers (Megatron-FSDP only) | | `--activation-checkpoint` | off | Enable activation checkpointing | ### `qwen3-30b-a3b.gbs128_mbs4_seq4096_n2_mfsdp2_mxfp8_wandb.sh` diff --git a/examples/megatron_fsdp/fsdp_toy.py b/examples/megatron_fsdp/fsdp_toy.py index 946f06b84d0..7c825eff7d7 100644 --- a/examples/megatron_fsdp/fsdp_toy.py +++ b/examples/megatron_fsdp/fsdp_toy.py @@ -68,6 +68,60 @@ def enable_activation_checkpointing(self): layer._use_activation_checkpointing = True +# ----------------------- +# Synthetic supervised data (teacher-student) +# ----------------------- + +def set_seed(seed: int) -> None: + """Seed CPU + CUDA RNGs so all ranks initialize identically.""" + torch.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + + +class TeacherStudentData: + """Deterministic teacher-student regression task for convergence checks. + + A frozen, randomly-initialized ``ToyModel`` acts as the teacher. Inputs are + drawn from a per-(step, rank) seeded generator so every rank sees a distinct + but reproducible stream (proper data parallelism), and targets are the + teacher's outputs. The student (the FSDP model) learns to match the teacher, + so the MSE loss must decrease monotonically toward zero. + """ + + def __init__( + self, + dim: int, + n_layers: int, + seed: int, + device: str = "cuda", + dtype: torch.dtype = torch.bfloat16, + ): + self.dim = dim + self.device = device + self.dtype = dtype + self.seed = seed + # Build the teacher identically on all ranks (seed differs from student + # init so the student does not start at the solution). + set_seed(seed) + self.teacher = ToyModel(dim=dim, n_layers=n_layers).to(device=device, dtype=dtype) + self.teacher.eval() + for p in self.teacher.parameters(): + p.requires_grad_(False) + + @torch.no_grad() + def sample( + self, batch_size: int, seq_len: int, step: int, rank: int + ) -> Tuple[torch.Tensor, torch.Tensor]: + gen = torch.Generator(device=self.device) + gen.manual_seed(self.seed + 1_000_003 * step + 9_176 * rank) + x = torch.randn( + batch_size, seq_len, self.dim, + device=self.device, dtype=self.dtype, generator=gen, + ) + y = self.teacher(x) + return x, y + + # ----------------------- # Distributed init / mesh # ----------------------- @@ -88,7 +142,7 @@ def build_fsdp_model( n_layers: int, use_megatron_fsdp: bool, use_activation_checkpointing: bool = False, - enable_cuda_graph: bool = True, + enable_cuda_graph: bool = False, enable_trace_pool: bool = False, ) -> Tuple["FSDPModule", torch.distributed.device_mesh.DeviceMesh]: if use_megatron_fsdp: @@ -246,20 +300,29 @@ def train( model: nn.Module, optimizer: torch.optim.Optimizer, start_step: int = 0, + data: "TeacherStudentData" = None, ) -> None: rank = dist.get_rank() model.train() step = start_step t_start = time.time() + initial_loss = None + final_loss = None for epoch in range(args.epochs): for _ in range(args.steps_per_epoch): - # Dummy data - x = torch.randn(args.batch_size, args.seq_len, args.model_dim, - device="cuda", dtype=torch.bfloat16) - y = model(x) - loss = y.sum() / (args.batch_size * args.seq_len) + if data is not None: + # Deterministic teacher-student regression: loss must decrease. + x, target = data.sample(args.batch_size, args.seq_len, step, rank) + y = model(x) + loss = torch.nn.functional.mse_loss(y.float(), target.float()) + else: + # Dummy data (no convergence signal). + x = torch.randn(args.batch_size, args.seq_len, args.model_dim, + device="cuda", dtype=torch.bfloat16) + y = model(x) + loss = y.sum() / (args.batch_size * args.seq_len) loss.backward() if args.use_megatron_fsdp and args.release_memory_pool: model.release_memory_pool() @@ -267,16 +330,25 @@ def train( optimizer.zero_grad() model.zero_grad() + # Average loss across DP ranks so the value is identical everywhere + # (safe for a consistent convergence assert on all ranks). + global_loss = loss.detach() + dist.all_reduce(global_loss, op=dist.ReduceOp.AVG) + global_loss = global_loss.item() + if initial_loss is None: + initial_loss = global_loss + final_loss = global_loss + if step % args.log_interval == 0 and rank == 0: t_now = time.time() elapsed = t_now - t_start - it_s = elapsed / max(step - start_step + 1, 1) + s_per_step = elapsed / max(step - start_step + 1, 1) alloc = _fmt_bytes(torch.cuda.memory_allocated()) max_reserved = _fmt_bytes(torch.cuda.max_memory_reserved()) print( - f"[rank0] epoch={epoch} step={step} loss={loss.item():.4f} " + f"[rank0] epoch={epoch} step={step} loss={global_loss:.4e} " f"alloc={alloc} max_reserved={max_reserved} " - f"it={elapsed:.1f}s ({it_s * 1000:.0f}ms/it)" + f"elapsed={elapsed:.1f}s ({s_per_step * 1000:.0f}ms/step)" ) if args.ckpt_dir and step % args.ckpt_interval == 0 and step > 0: @@ -288,6 +360,20 @@ def train( if args.ckpt_dir: save_checkpoint(model, optimizer, step, args.ckpt_dir) + # Convergence verification (teacher-student mode only). + if data is not None and initial_loss is not None: + ratio = final_loss / max(initial_loss, 1e-12) + if rank == 0: + print( + f"[rank0] convergence: initial_loss={initial_loss:.4e} " + f"final_loss={final_loss:.4e} ratio={ratio:.3f} " + f"(threshold={args.convergence_threshold})" + ) + assert final_loss < initial_loss * args.convergence_threshold, ( + f"Convergence check failed: final_loss={final_loss:.4e} is not < " + f"initial_loss={initial_loss:.4e} * {args.convergence_threshold}" + ) + # ----------------------- # __main__ entry @@ -309,8 +395,9 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--log-interval", type=int, default=5) parser.add_argument("--use-megatron-fsdp", action="store_true", help="Use Megatron-FSDP instead of PyTorch FSDP2") parser.add_argument("--activation-checkpoint", action="store_true", help="Enable activation checkpointing on transformer layers") - parser.add_argument("--no-cuda-graph", action="store_false", dest="cuda_graph", - default=True, help="Disable CUDA graph capture") + parser.add_argument("--cuda-graph", action="store_true", default=False, + help="Enable CUDA graph capture (Megatron-FSDP only). Off by default " + "to keep the Megatron-FSDP and PyTorch FSDP2 paths aligned.") parser.add_argument("--use-trace-pool", action="store_true", default=False, help="Use TracePoolAllocator for stable buffer addresses") parser.add_argument("--release-memory-pool", action="store_true", default=False, @@ -319,6 +406,13 @@ def parse_args() -> argparse.Namespace: "Only takes effect with --use-megatron-fsdp.") parser.add_argument("--record-memory-history", type=str, default=None, metavar="DIR", help="Enable CUDA memory recording, dump snapshot to this directory") + parser.add_argument("--use-real-data", action="store_true", default=False, + help="Train a deterministic teacher-student regression task with MSE " + "loss to verify convergence, instead of the dummy sum-loss.") + parser.add_argument("--seed", type=int, default=1234, + help="Random seed for model init and teacher-student data.") + parser.add_argument("--convergence-threshold", type=float, default=0.5, + help="With --use-real-data, assert final_loss < initial_loss * this.") return parser.parse_args() @@ -331,6 +425,9 @@ def main() -> None: stacks="all", ) + # Seed before model construction so all ranks share identical initial weights. + set_seed(args.seed) + model, _ = build_fsdp_model( dim=args.model_dim, n_layers=args.n_layers, @@ -341,11 +438,21 @@ def main() -> None: ) optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr) + data = None + if args.use_real_data: + # Teacher seed differs from student init so the student does not start + # at the solution; both are identical across ranks. + data = TeacherStudentData( + dim=args.model_dim, + n_layers=args.n_layers, + seed=args.seed + 10000, + ) + start_step = 0 if args.ckpt_dir: start_step = load_checkpoint_if_available(model, optimizer, args.ckpt_dir) - train(args, model, optimizer, start_step=start_step) + train(args, model, optimizer, start_step=start_step, data=data) if args.record_memory_history: rank = dist.get_rank() diff --git a/examples/qwenimage_mfsdp/test_qwenimage.py b/examples/qwenimage_mfsdp/test_qwenimage.py index ee1c46fd293..c173972595b 100644 --- a/examples/qwenimage_mfsdp/test_qwenimage.py +++ b/examples/qwenimage_mfsdp/test_qwenimage.py @@ -16,6 +16,7 @@ - Per-step GPU memory tracking - NVML real GPU memory logging - Numerical verification probe (--verify) + - Convergence overfit check (--real-data) - Nsight profiling support (--cuda_profiler_capture) Usage (per node): @@ -268,6 +269,12 @@ def parse(): p.add_argument("--gradient_checkpointing", action="store_true") p.add_argument("--verify", action="store_true", help="Cross-rank gloss + global grad-norm probe; timing INVALID for this run.") + p.add_argument("--real-data", action="store_true", + help="Overfit a fixed flow-matching batch (target velocity v=x1-x0) so the " + "loss decreases; verifies convergence instead of the dummy L2 loss. " + "Adds a cross-rank loss all-reduce, so timing is not meaningful.") + p.add_argument("--convergence-threshold", type=float, default=0.5, + help="With --real-data, assert final_loss < initial_loss * this.") p.add_argument("--cuda_profiler_capture", action="store_true", help="Bracket bench steps with cudaProfilerStart/Stop for Nsight capture ranges.") p.add_argument("--record-memory-history", type=str, default=None, metavar="DIR", @@ -434,6 +441,33 @@ def make_batch(args, device, dtype, gen): return hidden, timestep, prompt, img_shapes, [txt_len] * B +def make_fixed_flow_matching_batch(args, device, dtype, gen): + """Deterministic flow-matching batch to overfit for convergence checks. + + Builds a clean latent ``x1``, noise ``x0`` and per-sample time ``t in [0, 1]``. + The model input is the interpolant ``xt = (1 - t) * x0 + t * x1`` and the + regression target is the velocity ``v = x1 - x0``. Reusing this fixed batch + every step lets the model overfit, so ``MSE(model(xt, t), v)`` must decrease. + + The returned ``timestep`` is ``t * 1000`` to match the ``timestep / 1000`` + convention used at the model call site. + """ + B = args.batch_size + H, W = args.height // 16, args.width // 16 + seq = H * W + txt_len = args.instruction_seq_len + qwen25vl_vision_tokens( + args.height, args.width, args.vl_patch_size, args.vl_merge_size) + x1 = torch.randn(B, seq, 64, device=device, dtype=dtype, generator=gen) + x0 = torch.randn(B, seq, 64, device=device, dtype=dtype, generator=gen) + t = torch.rand(B, device=device, dtype=dtype, generator=gen) + t_b = t.view(B, 1, 1) + xt = (1.0 - t_b) * x0 + t_b * x1 + target = x1 - x0 + prompt = torch.randn(B, txt_len, 3584, device=device, dtype=dtype, generator=gen) + img_shapes = [[(1, H, W)]] * B + return xt, t * 1000.0, prompt, img_shapes, [txt_len] * B, target + + # ----------------------------- main ----------------------------- def main(): args = parse() @@ -508,6 +542,11 @@ def main(): _mem_log("after_model_init", rank=rank) gen = torch.Generator(device=device).manual_seed(args.seed + rank) + fixed_batch = None + real_initial_loss = None + real_final_loss = None + if args.real_data: + fixed_batch = make_fixed_flow_matching_batch(args, device, dtype, gen) step_times = [] cuda_profiler_active = False oom_occurred = False @@ -521,8 +560,12 @@ def main(): cuda_profiler_active = True if rank == 0: print("[nsys] cudaProfilerStart") - hs, ts, pe, img_shapes, txt_lens = make_batch(args, device, dtype, gen) - hs.requires_grad_(True) + if args.real_data: + xt, ts, pe, img_shapes, txt_lens, target = fixed_batch + hs = xt.detach().clone().requires_grad_(True) + else: + hs, ts, pe, img_shapes, txt_lens = make_batch(args, device, dtype, gen) + hs.requires_grad_(True) torch.cuda.synchronize(); dist.barrier() t0 = time.perf_counter() @@ -534,7 +577,10 @@ def main(): img_shapes=img_shapes, txt_seq_lens=txt_lens, return_dict=False) pred = (out[0] if isinstance(out, tuple) else out)[:, :hs.size(1)] pred = _NvtxFwdEndBwdStart.apply(pred) - loss = pred.float().pow(2).mean() + if args.real_data: + loss = torch.nn.functional.mse_loss(pred.float(), target.float()) + else: + loss = pred.float().pow(2).mean() loss.backward() t_bwd = time.perf_counter() @@ -560,11 +606,26 @@ def main(): dt = time.perf_counter() - t0 if step >= args.warmup_steps: step_times.append(dt) + + # Cross-rank mean loss for a consistent convergence signal + # (identical on all ranks; added after timing). + gl = None + if args.real_data: + g = loss.detach().float() + dist.all_reduce(g, op=dist.ReduceOp.AVG) + gl = g.item() + if real_initial_loss is None: + real_initial_loss = gl + real_final_loss = gl + if rank == 0: tag = "warmup" if step < args.warmup_steps else "bench " if args.verify: print(f"[{args.backend}] {tag} step {step:3d} | VERIFY (timing invalid) | " f"gloss={gloss:.6f} | gnorm={gnorm:.4f} | n_grad={n}") + elif gl is not None: + print(f"[{args.backend}] {tag} step {step:3d} | {dt*1000:8.2f} ms | " + f"loss={gl:.4e}") else: print(f"[{args.backend}] {tag} step {step:3d} | {dt*1000:8.2f} ms | " f"loss={loss.item():.4f}") @@ -592,6 +653,19 @@ def main(): print(f"\n[{args.backend}] avg step (n={len(step_times)}): {avg:.2f} ms | " f"peak mem: {peak.item():.2f} GB") + # Convergence verification (--real-data overfit only). ``real_final_loss`` is + # the cross-rank mean, so this assert is consistent across ranks. + if args.real_data and real_initial_loss is not None: + ratio = real_final_loss / max(real_initial_loss, 1e-12) + if rank == 0: + print(f"[{args.backend}] convergence: initial_loss={real_initial_loss:.4e} " + f"final_loss={real_final_loss:.4e} ratio={ratio:.3f} " + f"(threshold={args.convergence_threshold})") + assert real_final_loss < real_initial_loss * args.convergence_threshold, ( + f"Convergence check failed: final_loss={real_final_loss:.4e} is not < " + f"initial_loss={real_initial_loss:.4e} * {args.convergence_threshold}" + ) + _mem_log("final", rank=rank) if mem_mgr is not None: From 9d733b7fd0db152e5362ae1e8555053088baf2e2 Mon Sep 17 00:00:00 2001 From: jianbinc Date: Thu, 2 Jul 2026 23:12:54 +0800 Subject: [PATCH 28/28] handle split_dtensor --- .../fsdp/src/megatron_fsdp/uneven_dtensor.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/uneven_dtensor.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/uneven_dtensor.py index d1ae1c97101..fc1135deabf 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/uneven_dtensor.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/uneven_dtensor.py @@ -489,11 +489,13 @@ def split_dtensor( for size in split_size_or_sections: split_points.append(split_points[-1] + size) - # One collective call — result reused for all splits below. - assert hasattr(dtensor._local_tensor, "__create_chunk_list__"), ( - "DTensor local tensor is missing chunk metadata." - ) - chunk_meta = dtensor._local_tensor.__create_chunk_list__()[0] + # Chunk metadata is reused for all splits below. Use the cached closure + # when present (fast path, no collective); otherwise compute it once via a + # single collective. + if hasattr(dtensor._local_tensor, "__create_chunk_list__"): + chunk_meta = dtensor._local_tensor.__create_chunk_list__()[0] + else: + chunk_meta = gather_and_compute_chunk_metadata(dtensor) chunk_slice = slice(chunk_meta.offsets[dim], chunk_meta.offsets[dim] + chunk_meta.sizes[dim]) local_offset = chunk_meta.offsets[dim] local_tensor = dtensor.to_local()