diff --git a/examples/megatron_fsdp/README.md b/examples/megatron_fsdp/README.md index cc37911c12d..0acbb166d49 100644 --- a/examples/megatron_fsdp/README.md +++ b/examples/megatron_fsdp/README.md @@ -4,6 +4,52 @@ 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`, off by default) + > **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 (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 +``` + +| 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` | 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` + +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 new file mode 100644 index 00000000000..7c825eff7d7 --- /dev/null +++ b/examples/megatron_fsdp/fsdp_toy.py @@ -0,0 +1,472 @@ +# 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 time +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 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): + """MLP-style block with expand→project, similar to transformer FFN.""" + + def __init__(self, dim: int, expansion: int = 4, dropout: float = 0.0): + super().__init__() + 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): + 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.dropout(self.down(torch.nn.functional.gelu(self.gate(x)) * self.up(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) + + def enable_activation_checkpointing(self): + for layer in self.layers: + 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 +# ----------------------- + +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, + use_activation_checkpointing: bool = False, + enable_cuda_graph: bool = False, + enable_trace_pool: bool = False, +) -> Tuple["FSDPModule", torch.distributed.device_mesh.DeviceMesh]: + if use_megatron_fsdp: + try: + from megatron_fsdp.v2 import FSDPModule, fully_shard + 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(device="cuda", dtype=torch.bfloat16) + + if use_activation_checkpointing: + model.enable_activation_checkpointing() + + 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, + ) + else: + sublayer_kwargs = {} + kwargs = {} + + 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) + + if use_megatron_fsdp: + model._log_parameter_groups() + + 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): + 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, + ) + + 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): + 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() + 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}") + + 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} ({time.time() - t0:.1f}s)") + + +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 + + 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 + + +# ----------------------- +# 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, + 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): + 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() + optimizer.step() + 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 + 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={global_loss:.4e} " + f"alloc={alloc} max_reserved={max_reserved} " + 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: + save_checkpoint(model, optimizer, step, args.ckpt_dir) + + step += 1 + + # Final checkpoint + 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 +# ----------------------- + +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("--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=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=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, + 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") + 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() + + +def main() -> None: + args = parse_args() + + if args.record_memory_history: + torch.cuda.memory._record_memory_history( + max_entries=100000, + 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, + 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) + + 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, data=data) + + 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() + + +if __name__ == "__main__": + main() 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}" 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..c173972595b --- /dev/null +++ b/examples/qwenimage_mfsdp/test_qwenimage.py @@ -0,0 +1,700 @@ +""" +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) + - Convergence overfit check (--real-data) + - 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("--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", + 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 + + +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() + 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) + 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 + 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") + 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() + + 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) + 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() + + 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) + + # 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}") + + 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") + + # 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: + 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/distributed_data_parallel_config.py b/megatron/core/distributed/distributed_data_parallel_config.py index 56ec9e89539..7666f05f816 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) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import Optional, Tuple import torch @@ -96,6 +96,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. @@ -146,6 +155,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 @@ -221,6 +239,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. + """ + megatron_fsdp_prefetch_recompute_forward_weights: bool = False """If set to True, Megatron-FSDP prefetches rowwise weights needed by activation recomputation during backward before prefetching backward transpose weights. diff --git a/megatron/core/distributed/fsdp/checkpoint.py b/megatron/core/distributed/fsdp/checkpoint.py new file mode 100644 index 00000000000..06667c4e711 --- /dev/null +++ b/megatron/core/distributed/fsdp/checkpoint.py @@ -0,0 +1,1580 @@ +# 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_mamba_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_chunk_meta_source +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 + +try: + from megatron.core.ssm.mamba_mixer import MambaMixer +except ImportError: + MambaMixer = None + +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", + "handle_mamba_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, MambaMixer, FP8, expert remapping — see + individual ``handle_*_v2`` functions). ``load_state_dict()`` uses PyTorch's + ``set_state_dict``. + """ + + 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 + 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'): + 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"``). + 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", + ) + + +# ------------------------------------------------------------------ +# 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", + ) + + +# ------------------------------------------------------------------ +# 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 +# ------------------------------------------------------------------ + + +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 + + 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/MambaMixer 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 + + 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) + 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) + _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 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 + + 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="."): + """ + 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 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, 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`` + 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. + + 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. + """ + 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 ---- + if model is not None: + _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_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 + 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. + 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 + 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 ---- + + # 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) + 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 and hasattr(lt, "__create_chunk_list__"): + 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 5cf48d293f8..1dc673891ff 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: @@ -41,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 @@ -50,12 +54,17 @@ MegatronFSDP, MixedPrecisionPolicy, ) + from megatron.core.distributed.fsdp.src.megatron_fsdp.v2 import FSDPModule 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 +77,6 @@ class FullyShardedDataParallel(_BaseDataParallel): _MODULE_TYPE_REGISTRY: Dict[str, set] = { "column": { "ColumnParallelLinear", - "LinearCrossEntropyModule", "TEColumnParallelLinear", "TELayerNormColumnParallelLinear", "TEColumnParallelGroupedLinear", @@ -104,6 +112,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__) @@ -178,6 +199,7 @@ def __init__( "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,13 +212,6 @@ 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 @@ -215,18 +230,237 @@ 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 + from megatron.core.distributed.fsdp.src.megatron_fsdp.v2.mixed_precision import ( + FullyShardFP8Policy, + MixedPrecisionPolicy, + FullyShardNVFP4Policy, + ) + 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, MambaLayer, TEGroupedMLP, SequentialMLP] + + 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 = MixedPrecisionPolicy( + 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 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 + 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 + + if fsdp_unit_modules is not None: + 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=mesh, + gradient_scaling_factor=grad_sf, + enable_cuda_graph=False, + **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 synchronize_param_gather(): + ctx = self.module._fsdp_root_context + torch.cuda.current_stream().wait_stream(ctx.ag_stream) + + 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.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'): @@ -448,8 +682,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() @@ -468,6 +714,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 938e17a5b3f..fb97d95a33c 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 @@ -38,7 +38,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', @@ -155,6 +159,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. + """ + megatron_fsdp_prefetch_recompute_forward_weights: bool = False """If set to True, Megatron-FSDP prefetches rowwise weights needed by activation recomputation during backward before prefetching backward transpose weights. 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 59b86e74feb..f914b429968 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py @@ -1436,6 +1436,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..fc1135deabf 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: @@ -93,36 +96,12 @@ 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 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, source=source) def validate_uneven_dtensor(dtensor: DTensor) -> None: @@ -249,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 @@ -414,26 +392,91 @@ 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. 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 +489,37 @@ 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) + # 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() - # 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 +533,253 @@ 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), source="split") 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, source="make_uneven") + 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. + + 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 + ) + + # 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 + + +# ------------------------------------------------------------------ +# 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__ + 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): + """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, 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)) + + 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 + dtensor._local_tensor._chunk_meta_source = source diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/utils.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/utils.py index f771c17c17d..4e856ab0801 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/utils.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/utils.py @@ -864,5 +864,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 new file mode 100644 index 00000000000..44624770de4 --- /dev/null +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/README.md @@ -0,0 +1,347 @@ +# Megatron FSDP v2 + +> ⚠️ 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 + +``` +v2/ +├── README.md # This file +├── __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 (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. + +## 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 + +### MixedPrecisionPolicy + +Controls parameter/gradient dtypes and communication precision. + +**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. 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. | + +**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, + FullyShardFP8Policy, FullyShardNVFP4Policy, +) + +# 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 + 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, + nvfp4=FullyShardNVFP4Policy(enabled=True), +) +fully_shard(model, mp_policy=mp_policy) +``` + +### 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 / grad sync | All-reduce no-shard grads or reduce-scatter ZeRO grads | + +### 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. + +**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 +for layer in model.layers: + fully_shard(layer, enable_cuda_graph=True) +fully_shard(model) # root without CUDA graph +``` + +**How it works:** + +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 +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 +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 +``` + +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: + +- `unshard()` — all-gather to full tensor +- `reshard()` — free temporary buffer +- `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 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 + +### 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 + +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: 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 | **Supported** | Like DDP: replicated weights, full-gradient all-reduce, replicated optimizer states. No param-gather overlap. | + +## 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. + +### Sharding Strategies + +`no_shard` (DDP-like), `optim` (ZeRO-1), `optim_grads` (ZeRO-2), and +`optim_grads_params` (ZeRO-3 equivalent) are implemented. `no_shard`, +`optim`, and `optim_grads` use replicated compute weights, so parameter gather +overlap (prefetch/unshard pipelining) is not applicable. + +### CUDA Graph + +- **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 + +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 + +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. + +```bash +python pretrain_gpt.py \ + --use-megatron-fsdp-v2 \ + --use-megatron-fsdp \ + ... +``` + +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_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()` +- CUDA graph capture (`--cuda-graph` / `--no-cuda-graph`) +- 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 + +# 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 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``. ``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 + +```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 + +# Run specific test files +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" +``` 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..222409d204a --- /dev/null +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/__init__.py @@ -0,0 +1,55 @@ +# 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. + +# 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, + 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, + MixedPrecisionPolicy, + FullyShardNVFP4Policy, +) +from .param_group import ParameterGroup + +__all__ = [ + "FSDPModule", + "fully_shard", + "FullyShardFP8Policy", + "MixedPrecisionPolicy", + "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/allocator.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/allocator.py new file mode 100644 index 00000000000..d5505bfe955 --- /dev/null +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/allocator.py @@ -0,0 +1,704 @@ +# 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 +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 + 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.""" + + 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 for CUDA graph-compatible training. + + 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** (first micro-batch) + + 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()``) + + Replays the trace to build per-key live intervals, then uses + **conflict-graph coloring** to assign slots: + + * 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. + + 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`` / ``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 ---------------------------------------------------- # + + @dataclasses.dataclass + class _SlotInfo: + """Metadata for a physical slot (backed by its own tensor).""" + + 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: + """A single alloc or free recorded during the trace phase.""" + + seq: int + op: str # "alloc" | "free" + key: AllocatorKey + + # -- Init ----------------------------------------------------------- # + + def __init__(self) -> None: + super().__init__() + self._phase: str = "trace" # "trace" | "optimized" + + # 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] = {} + self._active_keys: Set[AllocatorKey] = set() + + # Pool state — populated by plan(), used in optimized phase + 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] = {} + + # -- Public interface ------------------------------------------------ # + + 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 self._phase == "released": + self._auto_resume() + if self._phase != "optimized": + return self._trace_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, + ) -> None: + key = _resolve_key(key, param_group_id) + if self._phase == "released": + self._auto_resume() + if self._phase != "optimized": + self._trace_free(key) + else: + self._optimized_free(key) + + # -- Phase 1: trace -------------------------------------------------- # + + def _trace_allocate( + self, key: AllocatorKey, size: int, dtype: torch.dtype, device: torch.device + ) -> Bucket: + if key in self._active_keys: + return self._buckets[key] + + 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) + ) + 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: + 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) + self._active_keys.discard(key) + + # -- Phase 2: plan --------------------------------------------------- # + + def plan(self) -> int: + """Build the static key→slot plan from the recorded trace. + + 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 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 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: + if ev.key in alloc_stack and alloc_stack[ev.key]: + alloc_seq = alloc_stack[ev.key].pop(0) + intervals_per_key[ev.key].append((alloc_seq, ev.seq)) + + # 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 + + if not intervals_per_key: + self._phase = "optimized" + return 0 + + # 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._key_to_slot.clear() + self._key_to_view.clear() + + total_elems = 0 + 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.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( + self, + 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: + """Conflict-graph coloring + per-slot tensor allocation. + + 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. + + 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 + ) + ) + + # 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] + + return sum(slot_sizes) + + # -- Phase 3: optimized runtime ------------------------------------- # + + def _optimized_allocate( + self, key: AllocatorKey, size: int, dtype: torch.dtype, device: torch.device + ) -> Bucket: + slot_idx = self._key_to_slot[key] + slot = self._slots[slot_idx] + assert size <= slot.size, ( + f"requested {size} > slot capacity {slot.size} (key={key!r})" + ) + slot.in_use = True + self._active_keys.add(key) + # Return a view of the pre-allocated slot tensor + view = self._key_to_view[key] + return Bucket(data=view) + + def _optimized_free(self, key: AllocatorKey) -> None: + if key not in self._active_keys: + return + 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() + + 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 + + @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.""" + lines = [] + 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) + 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 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_str} " + f"{'IN_USE' if slot.in_use else 'free'} " + f"keys={keys_in_slot}" + ) + 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) + + +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. + + 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 _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 _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 _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/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 new file mode 100644 index 00000000000..5f3a7ba5329 --- /dev/null +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/cuda_graph_runner.py @@ -0,0 +1,383 @@ +# 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 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 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) + +# --------------------------------------------------------------------------- +# Hook save / restore +# --------------------------------------------------------------------------- + +_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 _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 + + +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) + + +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}, + ) + except Exception: + _restore_compiled_modules_after_capture_failure(saved) + raise + return saved + + +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. + + Created once by the root forward pre-hook and stored on + ``ctx.cuda_graph_runner``. + """ + + 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 + + # 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 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]) + ) + + 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, + ) + + 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 + + modules = self._modules_ordered + n = len(modules) + + if torch.distributed.is_initialized() and torch.distributed.get_rank() == 0: + logger.info("CudaGraphRunner: capturing %d modules", n) + + # 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), + ) + + # 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) + + 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) + + _mem_after = _mem_snapshot() + _peak_alloc = torch.cuda.max_memory_allocated() + + 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), + ) + + if not isinstance(graphed, tuple): + graphed = (graphed,) + + # 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 = [] + + if torch.distributed.is_initialized() and torch.distributed.get_rank() == 0: + logger.info("CudaGraphRunner: installed CUDA graphs on %d modules", n) + + +# --------------------------------------------------------------------------- +# capture_time_hooks (unshard / reshard outside graph, not replayed) +# --------------------------------------------------------------------------- + + +def _make_fwd_pre_hook(module): + def hook(mod, args, kwargs): + module.unshard() + return hook + + +def _make_fwd_post_hook(module): + def hook(mod, args, kwargs, output): + module.reshard() + return hook + + +def _make_bwd_pre_hook(module): + def hook(mod, grad_output): + module.unshard(bwd_pass=True) + return hook + + +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/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/design/allocator_design.md b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/design/allocator_design.md new file mode 100644 index 00000000000..e38129ee72f --- /dev/null +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/design/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/design/cuda_graph_design.md b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/design/cuda_graph_design.md new file mode 100644 index 00000000000..09fd399a19f --- /dev/null +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/design/cuda_graph_design.md @@ -0,0 +1,248 @@ +# 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. + +## 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 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 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 + +```python +fully_shard(module, enable_cuda_graph=True) +``` + +## 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 + +Four minimal hooks attached via `capture_time_hooks` to each module: + +| 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` | + +These run outside `torch.cuda.graph()` — unshard/reshard are never captured. + +## 4. Lifecycle + +``` +╔═══════════════════════════════════════════════════════════════╗ +║ Stage 1 — Microbatch 0 (trace) ║ +║ ║ +║ 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 (record sample args + eager) ║ +║ ║ +║ root_forward_pre_hook: ║ +║ Creates shared graph_pool_handle + capture_stream ║ +║ reset_cursor(); restore_slots() ║ +║ ║ +║ forward (per graphed FSDPModule): ║ +║ 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() ║ +║ ║ +║ forward (per graphed FSDPModule): ║ +║ forward_pre_hook: unshard; _fsdp_cg_installed → skip record ║ +║ module.forward (graphed) → Graphed.apply → fwd_graph.replay ║ +║ post_forward_hook: reshard ║ +║ ║ +║ backward: ║ +║ pre_backward: unshard (bwd_pass) ║ +║ Graphed.backward → bwd_graph.replay ║ +║ post_backward: reshard + reduce_grad ║ +╚═══════════════════════════════════════════════════════════════╝ +``` + +## 5. Hook strategy + +### 5.1 During capture (`make_graphed_callables`) + +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. + +### 5.2 During replay + +Real FSDP hooks are **restored** after capture. They fire normally around +`module.__call__`: + +- `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 + +## 6. Capture stream & shared pool + +```python +# 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, +) +``` + +## 7. Parameter gradients + +`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. + +## 8. te-graph-runtime local modifications + +Vendored at `te_graph_runtime/`. Key local changes from upstream: + +- **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. + +See `te_graph_runtime/README.md` for details and upstream attribution. + +## 9. Known issues & fixes + +| 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 | + +## 10. Files + +| File | Role | +|------|------| +| `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 | + +## 11. No user-visible knobs + +| What happens | User action | +|---|---| +| Trace collection | Automatic (MB0) | +| Pool planning | Automatic (MB0 post-backward) | +| 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/design/design.md b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/design/design.md new file mode 100644 index 00000000000..fee7f75b521 --- /dev/null +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/design/design.md @@ -0,0 +1,868 @@ +# Design: Megatron FSDP v2 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()`, `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 | + +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-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. + + # --- 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(): + # 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: + 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 +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. + +**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-Scatter 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 +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 --- + 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 +``` + +--- + +## 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`. + +### 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 + +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, 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 + + 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=input_buffer, group=self.dp_group + ) + 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` +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 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:** 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.** + +```python +allocator = TracePoolAllocator() +# … run one iteration (trace phase) … +pool_elems = allocator.plan() # returns total element count +# … 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) → 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 — each +slot owns its own tensor, and buckets are lightweight views into them. 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/design/lazy_grad_buffer_design.md b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/design/lazy_grad_buffer_design.md new file mode 100644 index 00000000000..d6cd56928a1 --- /dev/null +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/design/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/design/mcore_fsdp_checkpoint_design.md b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/design/mcore_fsdp_checkpoint_design.md new file mode 100644 index 00000000000..0011c0414c4 --- /dev/null +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/design/mcore_fsdp_checkpoint_design.md @@ -0,0 +1,1062 @@ +# 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 | +|---------------|--------------|-------|-----------|-------| +| 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) | + +--- + +## 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/` | (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`, `_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) + +| 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 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). | +| `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. | +| `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. | +| `_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 Megatron 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 Megatron 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 │ +│ _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 ✓) │ +└──────────────────────────┬───────────────────────────────────┘ + │ +┌──────────────────────────▼───────────────────────────────────┐ +│ 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 +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?** + +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)``). +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`. + +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. + +### 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) | 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) | +| `"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 + +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 (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()` | + +### 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 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 +``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` — 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 +- [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 (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`) + +- [ ] 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. **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 + +--- + +## 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/design/nvfp4_design.md b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/design/nvfp4_design.md new file mode 100644 index 00000000000..f364ce2508e --- /dev/null +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/design/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 (Megatron FSDP v2 FP8 path) + +- `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: + `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 `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()`. +- 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 `MixedPrecisionPolicy`** with `nvfp4` field: + +```python +@dataclass(frozen=True) +class MixedPrecisionPolicy: + ... + 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 = MixedPrecisionPolicy( + ... + 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) + │ + ▼ +[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) + │ + ▼ +[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/design/tp_support_design.md b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/design/tp_support_design.md new file mode 100644 index 00000000000..5f5e828e610 --- /dev/null +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/design/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/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..23e1f6a13d6 --- /dev/null +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/dp_buffer.py @@ -0,0 +1,791 @@ +# 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, _free_storage +from .mixed_precision import MixedPrecisionPolicy +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 → (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 (start, end) in the full unsharded buffer for the given item.""" + idx = self.item_index_map[item_id] + 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. + + 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: MixedPrecisionPolicy, + *, + 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.sharding_strategy = sharding_strategy + 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 + 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``. + + 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.""" + 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 = False, + stream: Optional[torch.cuda.Stream] = None, + ) -> 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 is returned directly. If a replicated buffer only has this + rank's updated shard, the shard is all-gathered into self.data first. + """ + 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] + 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) + + setattr(full_buffer, "_dirty", False) # mark the buffer as clean (unsharded) + + return full_buffer + + 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] + 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[start:end].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. + + Memory allocation always occurs on the default stream for deterministic + caching-allocator behaviour. + """ + 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: + assert self.data is not None, "DataParallelBuffer data not initialized" + 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, + *, + 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. + + 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. + """ + if self.sharding_strategy in ("no_shard", "optim"): + overwrite_grad = True + + 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 + 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 + + 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": + 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 + # 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 + 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.fetch_buffer() + output_offset = sm.local_data_index + + 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 + ) + + # 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 + + +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..3491cafa88d --- /dev/null +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/fsdp_module.py @@ -0,0 +1,1168 @@ +# 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 +import weakref +from contextlib import contextmanager, nullcontext +from dataclasses import dataclass, field +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, TracePoolAllocator +from .mixed_precision import MixedPrecisionPolicy +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 + self.enable_cuda_graph: bool = 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. + ``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 + """Shared CUDA graph memory pool handle for CUDA graph capture.""" + + 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. + 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") + + 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: + """ + 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 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 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], + ignored_params: Optional[set], + mp_policy: MixedPrecisionPolicy, + gradient_scaling_factor: Optional[float] = None, + sharding_strategy: str = "optim_grads_params", + ): + """ + 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, + gradient_scaling_factor=gradient_scaling_factor, + sharding_strategy=sharding_strategy, + ) + 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 + 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[start:end].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[MixedPrecisionPolicy] = 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 = f"cuda:{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" + ) + # 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 name, m in self.named_modules(): + if m in ignored_modules: + continue + 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, + 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. + + 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._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 = [ + 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. + + 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() + + # 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}" + ) + + param_group.unshard(bwd_pass=bwd_pass, stream=stream) + + # 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 gradient reduction + 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 + + # 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): + 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. + # + # 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. + 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 + # 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) + if ( + param_main_grad is None + or param_main_grad.data_ptr() != main_grad.data_ptr() + ): + zero_targets.append(main_grad.view(-1)) + else: + main_grad = param.get_main_grad() + 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 + param_group.reduce_grad(stream=stream) + 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: + 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() + 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 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 in ("no_shard", "optim") + for param_group in child._fsdp_param_groups + ): + # 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): + 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.""" + 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(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.""" + self.zero_grad() + + 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}}. + 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(): + 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} + 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 + ) + 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: + 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) + 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_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, + mp_policy: MixedPrecisionPolicy, + 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. + + 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, + sharding_strategy=sharding_strategy, + ) + ) + + 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..09106235307 --- /dev/null +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/fully_shard.py @@ -0,0 +1,162 @@ +# 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 +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_forward_hook, + _register_forward_pre_hook, +) +from .mixed_precision import MixedPrecisionPolicy +from .utils import _init_default_fully_shard_mesh + +__all__ = ["FSDPModule", "fully_shard"] +_fsdp_class_cache = {} # module-level cache + + +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[MixedPrecisionPolicy] = 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, + 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. + + 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 + + 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( + "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 = MixedPrecisionPolicy() + + cls = module.__class__ + 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 = ( + 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", + ) + bucket_allocator = TracePoolAllocator() if use_trace_pool else StorageFreeingBucketAllocator() + + module._init_named_param_groups( + mesh, + ignored_params, + mp_policy=mp_policy, + gradient_scaling_factor=gradient_scaling_factor, + sharding_strategy=sharding_strategy, + ) + module._init_fsdp_state( + 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() + + _register_forward_pre_hook( + 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, + 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() + + return 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 new file mode 100644 index 00000000000..9620ade353f --- /dev/null +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/hooks.py @@ -0,0 +1,592 @@ +# 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, Optional, 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 .cuda_graph_runner import CudaGraphRunner +from .fsdp_module import FSDPModule, _FSDPState +from .utils import RegisterFSDPBackwardFunction + +logger = logging.getLogger(__name__) + + +# --------------------------------------------------------------------------- +# 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). + + 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 + + +@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. + + 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 + + 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 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: 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 getattr(target, "_fsdp_cg_installed", False)) + and not ctx.backward_phase + and target.cuda_graph_compatible + ): + 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. + + 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 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( + mfsdp_forward_pre_hook, prepend=True, with_kwargs=True, + ) + else: + 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 _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 +): + """Pre-backward hook for FSDP modules and fine-grained sub-modules. + + 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 + + _pre_backward_setup(target, skip_final_callback=skip_final_callback) + 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. + + 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__}" + ) + 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() + + +@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 + the bucket allocator from trace to optimized plan. + + 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__}" + ) + 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) + + # ---- 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}" + ) + + # ---- 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 +# --------------------------------------------------------------------------- + + +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, + 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. + """ + 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 + + # 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, mfsdp_post_backward_hook), + with_kwargs=True, + ) + + +# --------------------------------------------------------------------------- +# Post-backward final callback +# --------------------------------------------------------------------------- + + +def _register_post_backward_final_callback( + state: _FSDPState, module: nn.Module +) -> None: + """ + Enqueue a *single* engine callback that fires after every module's + backward pass has completed. + + Registered once by the root FSDP module (avoids duplicates). + 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 + + state._post_backward_callback_queued = True + Variable._execution_engine.queue_callback( + functools.partial(mfsdp_post_backward_final_callback, 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..79feb233e03 --- /dev/null +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/mixed_precision.py @@ -0,0 +1,797 @@ +# 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 MixedPrecisionPolicy: + """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 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. 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 + + def main_grads_dtype_for_param(self, tensor: torch.Tensor) -> torch.dtype: + """Return the main-gradient dtype for a parameter group. + + 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 + + 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.""" + 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 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)] + + 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 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: + 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 + ) + 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: " + "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 + ] + ) + else: + fp8_params = [] + 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=not no_shard) + 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=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) + 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 + ) + + # 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: + """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..56bc4e2055e --- /dev/null +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/v2/param_group.py @@ -0,0 +1,465 @@ +# 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 MixedPrecisionPolicy +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: MixedPrecisionPolicy, + 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() + 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 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 + + # 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, + ) + + def _init_buffers(self) -> None: + """ + Initialize all buffers based on sharding strategy. + + 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 + 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. + 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. 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 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): + 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) + + 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]) + gbuf = self._create_buffer(main_grads_dtype, shard_grads, "main_grad") + self.main_grad_buffer = gbuf + + # Create distributed parameter views + self._init_dist_params() + + 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. + + 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 + ): + if weight_buffer is not None: + weight_buffer.unshard(bind_params=bind_params, stream=stream) + + self.mp_policy.post_unshard(self.params, bwd_pass=bwd_pass) + + def has_unsharded_weight_buffers(self, bwd_pass: bool = False) -> bool: + """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_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._ensure_buffers_on_gpu() + 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, stream: Optional[torch.cuda.Stream] = None): + """ + Reduce gradients across DP ranks. + + 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._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( + overwrite_grad=self._grad_buffer_is_fresh, + stream=stream, + ) + self._grad_buffer_is_fresh = False + + 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 _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. + + 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": grads accumulate replicated, optimizer consumes reduced shards + - "no_shard": replicated, no sharding (DDP-equivalent) + """ + self.dist_params = [] + self.dist_grads = [] # placeholder, populated in _init_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) + 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) + + 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 + + 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) + 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 + + 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.""" + if set_to_none: + for dist_param in self.dist_params: + if dist_param.grad is not None: + dist_param.grad = None + if hasattr(dist_param, "decoupled_grad"): + dist_param.decoupled_grad = None + 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/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/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/__init__.py b/megatron/core/optimizer/__init__.py index 27b675d1b8d..433ef01effa 100644 --- a/megatron/core/optimizer/__init__.py +++ b/megatron/core/optimizer/__init__.py @@ -972,6 +972,114 @@ 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 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'), + ) + + 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) + + 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) + + def get_megatron_optimizer( config: OptimizerConfig, model_chunks: List[MegatronModule], @@ -1022,7 +1130,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, @@ -1066,6 +1179,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/clip_grads.py b/megatron/core/optimizer/clip_grads.py index 3c5491d39a1..6e40dc89354 100644 --- a/megatron/core/optimizer/clip_grads.py +++ b/megatron/core/optimizer/clip_grads.py @@ -240,9 +240,12 @@ def count_zeros_fp32( # If the parameter is managed by Megatron FSDP, the gradient is # an FSDP-sharded DTensor, and we should use the local shard. use_megatron_fsdp = True - grad = getattr(param, grad_attr)._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 = getattr(param, grad_attr)._local_tensor + num_zeros = grad.numel() - torch.count_nonzero(grad) + total_num_zeros += num_zeros continue is_not_shared = param_is_not_shared(param) is_not_tp_duplicate = param_is_not_tensor_parallel_duplicate(param, tp_group=tp_group) diff --git a/megatron/core/optimizer/distrib_optimizer.py b/megatron/core/optimizer/distrib_optimizer.py index 0b6fb468c22..0de6ff174e8 100644 --- a/megatron/core/optimizer/distrib_optimizer.py +++ b/megatron/core/optimizer/distrib_optimizer.py @@ -701,8 +701,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. @@ -816,7 +816,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 = {} @@ -900,9 +908,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 @@ -1426,8 +1434,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 @@ -1488,12 +1499,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 @@ -1612,7 +1625,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/fully_shard_v2_muon.py b/megatron/core/optimizer/fully_shard_v2_muon.py new file mode 100644 index 00000000000..04472cc225c --- /dev/null +++ b/megatron/core/optimizer/fully_shard_v2_muon.py @@ -0,0 +1,735 @@ +# 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 collections import OrderedDict +from dataclasses import dataclass +from typing import List + +import torch +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 + +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 + + +@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] + 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) + + @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=momentum.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. + 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``. + 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__( + self, + params, + grads, + momentum_buffers, + *, + 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", + tp_mode: str = "duplicated", + ) -> None: + 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) + if not param_list: + raise ValueError("FullyShardV2Muon got no parameters to manage.") + assert all( + isinstance(param, DTensor) for param in param_list + ), "FullyShardV2Muon expects every param to be a DTensor (FSDP-v2 dist_param)." + 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 + self._weight_decay = weight_decay + 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 + + # 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))) + + self._shard_ranges = self._compute_shard_ranges() + self._roots = self._assign_roots() + + 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. + + 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) + """ + 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 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 _assign_roots(self): + """Assign each param an NS root rank, load-balancing NS work per dp_group. + + Steps: + 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): + 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)): + 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 + work = [] + for param_idx in param_indices: + shape = self._main_weights[param_idx].shape + ns_cost = shape.numel() * min(shape) if len(shape) == 2 else 0 + 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] + 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_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 + ) + (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, + ) + ) + + no_comm_packages = [ + MuonGradPackage( + param_ids=comm_free_params[start : start + package_size], + has_comm=False, + ) + 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): + """Run one Muon update over all managed 2D matrix params. + + 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 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 + gather_reqs = [[] for _ in range(len(self._packages))] + full_grads = {} + + # 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 package_idx, package in enumerate(self._packages): + package.update_momentum(self, coef) + gather_reqs[package_idx] = package.issue_gather(self, full_grads) + + # 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) + + for package_idx, package in enumerate(self._packages): + orths = package.orthogonalize(self, gather_reqs[package_idx], full_grads) + owned_orths.update(orths) + 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) + if weight_shard.numel() == 0: + continue + 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) + + +class FullyShardV2MuonOptimizer(MegatronOptimizer): + """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, config, model_chunks, **muon_hyperparams): + self.model_chunks = list(model_chunks) if model_chunks else [] + 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. The inner optimizer orders params by communication need. + 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 module_name, m in root.named_modules(): + if not isinstance(m, FSDPModule): + continue + 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( + dist_param, "is_embedding_or_output_parameter", False + ): + params.append(dist_param) + grads.append(dist_grad) + 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. --- + 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() + # 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: + 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): + 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 _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): + 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): + 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): + 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/megatron/core/optimizer/optimizer_config.py b/megatron/core/optimizer/optimizer_config.py index 1048220953a..3168cc99c0f 100644 --- a/megatron/core/optimizer/optimizer_config.py +++ b/megatron/core/optimizer/optimizer_config.py @@ -340,6 +340,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/core/pipeline_parallel/combined_1f1b.py b/megatron/core/pipeline_parallel/combined_1f1b.py index 81524363993..7c52c325569 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,123 @@ 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 + + def _release_tensor_storage(tensors): """Release tensor storage after all backward users are done.""" if tensors is None: @@ -32,6 +150,11 @@ def _release_tensor_storage(tensors): tensor.untyped_storage().resize_(0) +# --------------------------------------------------------------------------- +# Schedules +# --------------------------------------------------------------------------- + + def combined_1f1b_schedule_for_no_pipelining( forward_step_func, data_iterator, @@ -65,13 +188,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( @@ -90,6 +217,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 @@ -113,7 +243,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 @@ -130,7 +262,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 @@ -202,12 +336,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." @@ -300,7 +432,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. @@ -315,6 +449,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. @@ -346,8 +486,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() @@ -398,20 +542,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 @@ -512,7 +651,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/arguments.py b/megatron/training/arguments.py index f082674935c..7f59ad07958 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -1226,6 +1226,14 @@ 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 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 @@ -1251,11 +1259,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"] @@ -1871,7 +1883,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. " @@ -4129,6 +4146,25 @@ 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( + '--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', @@ -4167,6 +4203,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/megatron/training/checkpointing.py b/megatron/training/checkpointing.py index e6965dd74f7..6e6dd3a9f57 100644 --- a/megatron/training/checkpointing.py +++ b/megatron/training/checkpointing.py @@ -52,6 +52,10 @@ 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, + _load_torch_dist_into_megatron_fsdp_v2, + ) from megatron.core.transformer.fsdp_dtensor_checkpoint import ( handle_experts_in_state_dict, handle_fp8_extra_state_case, @@ -1193,6 +1197,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, @@ -1231,6 +1246,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 @@ -1279,6 +1297,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: @@ -1448,7 +1469,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: @@ -1456,6 +1477,18 @@ 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, + model=model, + v2_state_dict=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, @@ -1521,7 +1554,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 @@ -1604,6 +1642,7 @@ def _load_base_checkpoint( iteration, release, checkpointing_context=checkpointing_context, + model=model, ) elif ckpt_format == "torch": ckpt_type = CheckpointType.LEGACY @@ -1918,7 +1957,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 ) @@ -2027,6 +2066,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 @@ -2146,7 +2195,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. @@ -2200,6 +2254,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/megatron/training/config/training_config.py b/megatron/training/config/training_config.py index 0b26d0d8bb5..4b3cf34aec1 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/initialize.py b/megatron/training/initialize.py index 86eda9bd3e1..9e49a97ef0f 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/megatron/training/training.py b/megatron/training/training.py index d8c72f7c159..212a03ed64d 100644 --- a/megatron/training/training.py +++ b/megatron/training/training.py @@ -4266,6 +4266,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/megatron/training/utils/common_utils.py b/megatron/training/utils/common_utils.py index 7b5e3b46fe1..1284112d76a 100644 --- a/megatron/training/utils/common_utils.py +++ b/megatron/training/utils/common_utils.py @@ -201,7 +201,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/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_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 new file mode 100644 index 00000000000..f5f9c0df3d1 --- /dev/null +++ b/tests/unit_tests/distributed/megatron_fsdp/v2/test_dist_muon.py @@ -0,0 +1,348 @@ +# 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 test for the single-root distributed Muon optimizer. + +``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, 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``). + +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 +""" + +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 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, +) +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 ( + 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 ( + 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 +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 +# 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"), + pytest.param( + [(1500, 2048), (3000, 2048), (2048, 1200), (2048, 2600), (2048,)], id="large" + ), +] + + +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.""" + 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 _read_full_weight(pg, idx, shape): + """All-gather this param's full master weight and return it (reshaped, cloned).""" + 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) + mbuf.reshard() + 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) + + +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) +@pytest.mark.parametrize("shapes", SHAPE_SETS) +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). + 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=MixedPrecisionPolicy( + main_params_dtype=torch.float32, main_grads_dtype=torch.float32 + ), + 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 + # of a param is empty. + 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, + weight_decay=weight_decay, + num_ns_steps=num_ns_steps, + ) + + # 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] + + 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)) + + # 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) + + # 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 = _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) + + # 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"[{mesh_case}/{strategy}] param {i} shape {s} step {step}: max abs diff " + f"{(got - ref_w[i]).abs().max().item():.3e}" + ) + + torch.distributed.barrier() 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..0958d60bea6 --- /dev/null +++ b/tests/unit_tests/distributed/megatron_fsdp/v2/test_fully_shard.py @@ -0,0 +1,987 @@ +# 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 ( + MixedPrecisionPolicy, +) + +# ------------------------------------------------------------------ # +# 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" + + 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)], + ) + 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 = MixedPrecisionPolicy(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 = MixedPrecisionPolicy(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 = MixedPrecisionPolicy( + 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. 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 +# ------------------------------------------------------------------ # + + +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 + + +# ------------------------------------------------------------------ # +# 10. Checkpoint — get_state_dict and preprocess_state_dict_for_uneven_dtensor +# ------------------------------------------------------------------ # + + +class TestCheckpoint: + 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" + + 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) + + 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 + + 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..a8c6ab2522c --- /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", fsdp_trace_pool=True), + dict(data_parallel_sharding_strategy="optim_grads_params", fsdp_trace_pool=True), + 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..78c0cb7036e --- /dev/null +++ b/tests/unit_tests/distributed/megatron_fsdp/v2/test_mcore_nd_parallel.py @@ -0,0 +1,680 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +import copy +import time + +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 + +STRICT_LOSS_ATOL = 5e-3 +STRICT_PARAM_ATOL = 5e-3 +STRICT_PARAM_RTOL = 1e-3 + + +@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 _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): + """ + 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) + 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( + 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, + expert_model_parallel_size=EP, + 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 = [] + param_snapshots = [] + + # Training loop + for step in range(NUM_TRAINING_STEPS): + t0 = time.time() + 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() + if verify_replicated_weight_buffers: + TestMegatronFSDPE2E._assert_replicated_weight_buffers_match(model_chunks) + + # 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) + ) + + 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( + 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, + fsdp_trace_pool=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", + ), + 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): + 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") + + 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, + 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}" + ), + ) + + @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"): + """ + 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..6c400931b9d --- /dev/null +++ b/tests/unit_tests/distributed/megatron_fsdp/v2/test_param_group.py @@ -0,0 +1,304 @@ +# 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 ( + 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 + +# ------------------------------------------------------------------ # +# 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 + if torch.distributed.is_initialized(): + 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 [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] + + +# ------------------------------------------------------------------ # +# 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=MixedPrecisionPolicy(), + 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: 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": (True, 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 not in ("no_shard", "optim_grads_params"): + pytest.skip( + "This test currently covers no_shard and optim_grads_params, " + f"skipping {strategy}." + ) + + 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) + # -- main_grad_buffer -- + if pg.requires_grad: + assert pg.main_grad_buffer is not None + assert pg.main_grad_buffer.is_distributed == g_dist + assert pg.main_grad_buffer.data is None # lazy init + + 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 not in ("no_shard", "optim_grads_params"): + pytest.skip( + "This test currently covers no_shard and optim_grads_params, " + f"skipping {strategy}." + ) + + 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 + assert wbuf is not None + + shard_before = wbuf.data.clone() + unsharded = wbuf.unshard() + + if not w_dist: + # Non-distributed: unshard returns self.data directly, no comm + assert unsharded is wbuf.data + 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): + 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 + wbuf.reshard() + if w_dist: + assert wbuf._unsharded_buffer is None + assert torch.equal(wbuf.data, shard_before) + + 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 not in ("no_shard", "optim_grads_params"): + pytest.skip( + "This test currently covers no_shard and optim_grads_params, " + f"skipping {strategy}." + ) + + groups, _, dp_group, rank, ws, device = _build_groups(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 + continue + + 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: + # 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) + + ref_shard = Ref.reduce_scatter(full.clone(), dp_group) + + 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 + 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() 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.", )