Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
36 commits
Select commit Hold shift + click to select a range
31b968b
feat(mfsdp-v2): mixed-precision infrastructure and MXFP8 support (by …
Autumn1998 May 27, 2026
7c841b1
feat(mfsdp-v2): checkpoint infrastructure, NVFP4 support, and FSDP v2…
shjwudp May 27, 2026
646ea0a
docs: polish v2 README, add TP design doc, activation checkpointing s…
shjwudp May 29, 2026
f50b7bb
fix: checkpoint save in fsdp_toy, import naming, test cleanup
shjwudp May 29, 2026
8b4f6e8
update README
shjwudp May 29, 2026
041dcc2
feat(mfsdp-v2): ZeRO-1/2/3 sharding strategies, MixedPrecisionPolicy …
Autumn1998 May 31, 2026
2874c91
feat(mfsdp-v2): sync module states after checkpoint load, harden unev…
shjwudp May 31, 2026
4456557
Add Megatron FSDP v2 Muon optimizer
Autumn1998 Jun 3, 2026
9aedf56
Clean up FullyShardV2Muon step() and docstrings
Autumn1998 Jun 4, 2026
a2bf8a0
Rework distributed Muon UT: scales, nesterov, bit-exact reference
Autumn1998 Jun 4, 2026
a4d9231
docs(mfsdp-v2): update README sharding strategies and MixedPrecisionP…
shjwudp Jun 5, 2026
63d0f8c
feat(mfsdp-v2): add Qwen3-30B-A3B 2-node MXFP8 training example with …
shjwudp Jun 5, 2026
ae61ac0
Merge pull request #11 from Autumn1998/develop/mfsdp-refactor-main-st…
shjwudp Jun 5, 2026
13c954c
Refactor Muon toward DTensor-driven, ParameterGroup-free (WIP)
Autumn1998 Jun 8, 2026
fda44d8
Make distributed Muon copy-free via batched isend/irecv
Autumn1998 Jun 8, 2026
811cac7
Clean up Muon: trim comments + drop dead single-root plumbing
Autumn1998 Jun 8, 2026
2680a95
Merge pull request #12 from Autumn1998/feat/fsdp-v2-muon-refactor
shjwudp Jun 9, 2026
651aadb
Refine FSDP v2 Muon gather flow
Autumn1998 Jun 9, 2026
d0d8e70
Megatron FSDP v2 built-in CUDA graph, allocator, checkpoint
shjwudp Jun 9, 2026
2b7efd5
Support no-shard in FSDP v2
Autumn1998 Jun 9, 2026
9364459
Merge remote-tracking branch 'origin/mfsdp_refactor' into feat/fsdp-v…
Autumn1998 Jun 9, 2026
276cd35
Improve Muon NS and P2P overlap
Autumn1998 Jun 10, 2026
aa4092a
Merge pull request #13 from Autumn1998/feat/fsdp-v2-muon-refactor
shjwudp Jun 10, 2026
b70e14d
Fix Muon subgroup P2P peers
Autumn1998 Jun 10, 2026
1ff3e08
Add Muon momentum checkpoint support
Autumn1998 Jun 10, 2026
eb3ccc8
Merge pull request #14 from Autumn1998/feat/fsdp-v2-muon-refactor
shjwudp Jun 11, 2026
8c3bea7
Skip no-op unshard for clean replicated buffers
Autumn1998 Jun 11, 2026
6397e2f
Merge pull request #15 from Autumn1998/fix/fsdp-v2-no-shard-unshard
shjwudp Jun 11, 2026
421f50f
Initialize clean replicated FSDP weight buffers
Autumn1998 Jun 11, 2026
1b6ef54
Merge pull request #16 from Autumn1998/fix/fsdp-v2-no-shard-unshard
shjwudp Jun 12, 2026
8502038
feat(fsdp-v2): lazy grad buffer, CPU offload infra, allocator release…
shjwudp Jun 12, 2026
0436cbd
Integrate M-FSDP v2 with 1F1B pipeline overlap; add lazy grad buffers…
shjwudp Jun 23, 2026
35be059
CUDA graph capture via vendored te-graph-runtime with `capture_time_h…
shjwudp Jun 30, 2026
21256aa
Add real-data convergence checks to FSDP examples
shjwudp Jul 2, 2026
e3021f9
Merge remote-tracking branch 'nvidia/dev' into mfsdp_refactor_dev_ups…
shjwudp Jul 2, 2026
9d733b7
handle split_dtensor
shjwudp Jul 2, 2026
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
46 changes: 46 additions & 0 deletions examples/megatron_fsdp/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
Loading