# FSDP2 Testing Guide This document contains test commands and scenarios used to validate the FSDP2 migration. ## Quick Test Commands ### Fast Test (3 iterations, block_size=4000) ```bash cd /home/victor/.cursor/worktrees/neon__SSH__victor-h100_/5gOyo/sunoGPT CUDA_VISIBLE_DEVICES=7 /home/victor/anaconda3/envs/suno_env_fa2_fa3_pt29/bin/python -u train.py \ --debug_val_only=False \ --master_port=12905 \ --out_dir="/app2/suno/checkpoints" \ --data_dir="/app2/suno/data/auk_v0" \ --train_metas_filename="metas_v8_tr_mini.jsonl" \ --val_metas_filename="metas_v8_val.jsonl" \ --step_save_iters=2_000 \ --eval_interval=1_000 \ --eval_iters=5 \ --batch_store_size=1 \ --n_layer=2 \ --n_head=32 \ --d_head=128 \ --learning_rate=5e-5 \ --max_iters=3 \ --warmup_iters=1 \ --block_size=4000 \ --batch_size=1 \ --fsdp=True \ --grad_checkpointing=True \ --compile=False \ --wandb_log=False \ --preload_optimizer=False \ --allow_skip=True ``` ### Test with Compile Enabled ```bash CUDA_VISIBLE_DEVICES=7 /home/victor/anaconda3/envs/suno_env_fa2_fa3_pt29/bin/python -u train.py \ --debug_val_only=False \ --master_port=12906 \ --out_dir="/app2/suno/checkpoints" \ --data_dir="/app2/suno/data/auk_v0" \ --train_metas_filename="metas_v8_tr_mini.jsonl" \ --val_metas_filename="metas_v8_val.jsonl" \ --step_save_iters=2_000 \ --eval_interval=1_000 \ --eval_iters=5 \ --batch_store_size=1 \ --n_layer=2 \ --n_head=32 \ --d_head=128 \ --learning_rate=5e-5 \ --max_iters=3 \ --warmup_iters=1 \ --block_size=4000 \ --batch_size=1 \ --fsdp=True \ --grad_checkpointing=True \ --compile=True \ --wandb_log=False \ --preload_optimizer=False \ --allow_skip=True ``` ### Test Checkpoint Loading ```bash # First save a checkpoint (use command above) # Then load it: CUDA_VISIBLE_DEVICES=7 /home/victor/anaconda3/envs/suno_env_fa2_fa3_pt29/bin/python -u train.py \ --debug_val_only=False \ --master_port=12907 \ --out_dir="/app2/suno/checkpoints" \ --data_dir="/app2/suno/data/auk_v0" \ --train_metas_filename="metas_v8_tr_mini.jsonl" \ --val_metas_filename="metas_v8_val.jsonl" \ --step_save_iters=2_000 \ --eval_interval=1_000 \ --eval_iters=5 \ --batch_store_size=1 \ --n_layer=2 \ --n_head=32 \ --d_head=128 \ --learning_rate=5e-5 \ --max_iters=5 \ --warmup_iters=1 \ --block_size=4000 \ --batch_size=1 \ --fsdp=True \ --grad_checkpointing=True \ --compile=False \ --wandb_log=False \ --preload_checkpoint=/app2/suno/checkpoints/YYYY-MM-DD_HH-MM-SS/last_ckpt.pt \ --preload_optimizer=True \ --allow_skip=True ``` ## SLURM Multi-GPU Tests ### Speed Test Script Located at: `/home/victor/neon/sunoGPT/slurm_scripts/gpt_speed_test.sh` ```bash cd /home/victor/neon/sunoGPT # Test without compile # (Edit script to set --compile=False) sbatch slurm_scripts/gpt_speed_test.sh # Test with compile # (Edit script to set --compile=True) sbatch slurm_scripts/gpt_speed_test.sh ``` Check job status: ```bash squeue -u $USER | grep gpt ``` Monitor output: ```bash tail -f /home/victor/neon/sunoGPT/slurm-JOBID.out ``` ## Test Scenarios Covered ### 1. Basic FSDP2 Training - ✅ FSDP2 wrapping with `fully_shard()` - ✅ Per-block sharding - ✅ Mixed precision (bfloat16) - ✅ DeviceMesh creation - ✅ Unified gradient clipping **Expected Output:** ``` [TIMESTAMP]: wrapping model in FSDP2... [TIMESTAMP]: GPU memory usage for model: 4.41GiB(5.55%) [TIMESTAMP]: training... ``` ### 2. Activation Checkpointing + FSDP2 - ✅ AC applied BEFORE FSDP (correct order) - ✅ Memory reduction verified **Expected Output:** ``` [TIMESTAMP]: wrapping model in FSDP2... [TIMESTAMP]: applying fsdp activation checkpointing... [TIMESTAMP]: GPU memory usage for model: 4.41GiB(5.55%) ``` ### 3. Compile + AC + FSDP2 - ✅ Correct order: AC → Compile → FSDP - ✅ Per-block compilation - ✅ Flash attention excluded from compile - ✅ No OOM with default mode (no fullgraph) **Expected Output:** ``` [TIMESTAMP]: wrapping model in FSDP2... [TIMESTAMP]: applying fsdp activation checkpointing... [TIMESTAMP]: compiling each transformer block... [TIMESTAMP]: model compilation complete [TIMESTAMP]: GPU memory usage for model: 4.41GiB(5.55%) ``` ### 4. Checkpoint Saving - ✅ Single-file `.pt` format maintained - ✅ FSDP2-compatible state dict APIs - ✅ Saves: `best_ckpt.pt`, `last_ckpt.pt`, `last_ckpt_infer.pt` **Expected Output:** ``` [TIMESTAMP]: saving checkpoint to /app2/suno/checkpoints/... [TIMESTAMP]: saving took XX.X seconds. ``` ### 5. Checkpoint Loading - ✅ Model state loaded correctly - ✅ Optimizer state loaded with `full_state_dict=True` - ✅ Training resumes from correct iteration **Expected Output:** ``` [TIMESTAMP]: loading model state_dict on gpu 0 [TIMESTAMP]: wrapping model in FSDP2... [TIMESTAMP]: training... ``` ## Performance Metrics ### Small Model (2 layers, 363M params) - **block_size=4000, 3 iters:** - Memory: 42.42GiB (53.40%) - Time: ~8-10s per iteration - **block_size=32000, 25 iters:** - Without compile: 35.34GiB (44.49%) - With compile: 53.51GiB (67.36%) ### Large Model (32 layers, 8xH100) - See WandB runs: - `fa2_fa3_pt29_baseline_compile` (compile=False) - `fa2_fa3_pt29_fsdp2_with_compile` (compile=True) ## Common Issues & Solutions ### Issue: OOM during compilation **Symptom:** CUDA OOM when using `--compile=True` with `fullgraph=True` **Solution:** Removed `fullgraph=True`, using default mode instead ### Issue: Checkpoint loading fails with DTensor error **Symptom:** `RuntimeError: got mixed torch.Tensor and DTensor` **Solution:** Added `full_state_dict=True` option when loading optimizer state ### Issue: AC applied after FSDP **Symptom:** Model breaks or inefficient **Solution:** Reordered to: AC → Compile → FSDP (following torchtitan) ## Key Files Modified 1. **sunoGPT/train.py** - FSDP2 wrapping logic (lines 477-541) - Correct ordering: AC → Compile → FSDP 2. **sunoGPT/utils/fsdp_policies.py** - Updated to `MixedPrecisionPolicy` 3. **sunoGPT/modules/base.py** - Flash attention compile exclusion - Fixed optimizer weight decay for FSDP2 4. **sunoGPT/utils/helpers.py** - FSDP2-compatible checkpoint save/load - Backward compatible with FSDP1 ## Verification Checklist - [ ] Model initializes without errors - [ ] FSDP2 wrapping completes successfully - [ ] Training runs for multiple iterations - [ ] Loss decreases (sanity check) - [ ] Checkpoint saves successfully - [ ] Checkpoint loads successfully - [ ] Training resumes from correct iteration - [ ] Memory usage is reasonable - [ ] Multi-GPU training works on SLURM - [ ] Compile mode works without OOM ## References - TorchTitan FSDP2 implementation: `/tmp/torchtitan/` - FSDP2 documentation: https://pytorch.org/docs/stable/fsdp.html - PR branch: `fsdp2-migration`