Merge caia-wandb-fp8-b200 into lambda-moe-tflops

docs/bs_sweep_perf_plan.md

@@ -0,0 +1,112 @@
+# Perf optimization plan — building on the bf16+compile baseline
+
+Companion to [bs_sweep_results.md](bs_sweep_results.md). This is a quick triage of perf knobs we can pull next, ordered by effort, with **rough** estimated per-GPU tflops on 8×B200 starting from those baselines.
+
+## Where we are today
+
+Steady-state, 8×B200, FSDP-only, bf16 + `torch.compile(components=["model","loss"])`, seq_len=8192:
+
+| Model | bs | tflops/gpu | mfu | headroom to bf16 peak (~2.25 PF/s) | headroom to fp8 peak (~4.5 PF/s) |
+| --- | --- | --- | --- | --- | --- |
+| gpt_oss_20b | 10 | ~573 | 25.5% | 4.0× | 7.9× |
+| qwen3_30b (Qwen3-30B-A3B) | 12 | ~455 | 20.2% | 5.0× | 9.9× |
+
+Estimates below are **back-of-envelope only**, calibrated against the relative cost of GEMM vs comm/AC/launch in MoE LLMs at this scale. Treat the `+%` columns as "expected order of magnitude, real result needs to be measured."
+
+## 1. Free wins (no torchao, ~hours of work)
+
+### 1.1 Loosen activation checkpointing for Qwen3
+
+Right now the local `qwen3_30b` helper hard-codes `mode="full"` whenever compile is on. `full` AC recomputes ~all transformer block forward in backward (~+33% compute). gpt_oss_20b already uses `selective` at bs<full-ac threshold in upstream, which is why its mfu is ~5 pp higher.
+
+- **Mechanism**: switch to `mode="selective"` with `selective_ac_option="op"` (recompute only attention sdpa + a few cheap ops), keep `full` only as a fallback.
+- **Catch**: At qwen3 bs=12 we are at 94% HBM — selective AC will OOM. Must combine with either reduced bs (bs=8–10) **or** fp8 (item 2) to free memory for the activations we no longer recompute.
+- **Estimated lift (Qwen3 at bs=8, selective AC)**: 455 → **520–560** tflops/gpu (+15–22%), mfu 23–25%.
+- **Estimated lift (gpt_oss)**: already at selective in upstream; negligible.
+
+### 1.2 Inductor max-autotune for GEMMs
+
+- **Mechanism**: set `TORCHINDUCTOR_MAX_AUTOTUNE_GEMM=1` (or `--compile.mode=max-autotune` if we add it as a knob — currently CompileConfig only exposes `enable / components / backend`). Inductor will template-tune the matmul kernels for the actual shapes seen at runtime.
+- **Cost**: ~5–15 min added to first-step compile; cache persists in `inductor_cache/`.
+- **Estimated lift**: +3–7% on both models. gpt_oss: 573 → **590–615**. qwen3: 455 → **470–490**.
+
+### 1.3 Compile lr-scheduler / optimizer step
+
+The trainer currently only compiles `model` and `loss`. The fused-Adam-style optimizer step is many small kernels; on B200 it's launch-overhead-bound.
+
+- **Mechanism**: extend [torchtitan/config/configs.py](../torchtitan/config/configs.py) `CompileConfig.components` to accept `"optimizer"` and wrap `clip_grad_norm_` + `optimizer.step` in `torch.compile(fullgraph=True)`. Trainer wrappers already use `sl.log_trace_span("optim")` so the boundary is clean.
+- **Estimated lift**: +2–5%. gpt_oss: 573 → **585–600**. qwen3: 455 → **465–480**.
+- **Risk**: low. Optimizer compile is well-trodden in torchao examples.
+
+Stacking 1.1+1.2+1.3 (Qwen3 at bs=8, selective AC, max-autotune, compiled optim):
+- Qwen3-30B-A3B: 455 → **~570–610** tflops/gpu (+25–34%, mfu ~25–27%).
+- gpt_oss-20b: 573 → **~610–650** tflops/gpu (+6–13%).
+
+## 2. torchao fp8 (medium effort)
+
+`torchtitan/components/quantization/float8.py` already exposes `Float8LinearConverter` and `Float8GroupedExpertsConverter` — see [float8.md](../torchtitan/components/quantization/float8.md). The local `_qwen3_30b_base` even imports `Float8LinearConverter` but never plumbs it through to `model_registry(..., quantization=...)`. So this is mostly a config-wiring change.
+
+### 2.1 Float8LinearConverter on dense linears (attn proj + lm_head/embed routing)
+
+- **Mechanism**: rowwise dynamic scaling on attention Q/K/V/O and the shared dense MLP. Recipe: `recipe_name="rowwise"`, `filter_fqns=["output", "router.gate", "auto_filter_small_kn"]`, `model_compile_enabled=True`.
+- **Compute share affected**: ~15–25% of MoE training time (attention + lm_head; experts excluded).
+- **Estimated lift**: +5–10% overall. gpt_oss: 573 → **610–630**. qwen3: 455 → **485–500**.
+- **Risk**: low–medium. Rowwise fp8 has converged-loss runs in torchao for llama3-405B; need a short loss-vs-bf16 check on c4 per [.claude/CLAUDE.md](../torchtitan/CLAUDE.md) "Validating Numerics".
+
+### 2.2 Float8GroupedExpertsConverter on MoE experts (big lever)
+
+- **Mechanism**: stack `Float8GroupedExpertsConverter.Config(model_compile_enabled=True)` after the linear converter. Uses torchao scaled grouped GEMMs for the expert path. This is the optimization deepseek_v3's config registry already opts into.
+- **Compute share affected**: 55–70% of training time (experts dominate at these sizes).
+- **Estimated lift**: +25–40% on both models — but bigger relative win on qwen3 since its `active=3.4B` ratio is more expert-heavy.
+- **gpt_oss-20b**: 573 → **720–800** tflops/gpu, mfu 32–36%.
+- **Qwen3-30B-A3B**: 455 → **570–640** tflops/gpu, mfu 25–28%.
+- **Memory bonus**: fp8 weights for experts halve their resident memory → frees ~30–40 GiB on Qwen3, finally making bs=14 feasible and re-enabling looser AC.
+- **Risk**: medium. Grouped GEMM fp8 is newer than per-linear fp8; need stricter loss check.
+
+Stacking 2.1+2.2 with the 1.x wins (selective AC freed by fp8 memory savings):
+- **gpt_oss-20b**: 573 → **~820–900** tflops/gpu (+43–57%, mfu ~37–40%).
+- **Qwen3-30B-A3B (bs ≥ 12 now safe)**: 455 → **~680–760** tflops/gpu (+50–67%, mfu ~30–34%).
+
+## 3. DeepGEMM grouped fp8 (larger effort)
+
+torchao's grouped scaled GEMM is reasonable but not state-of-art. DeepSeek's [DeepGEMM](https://github.com/deepseek-ai/DeepGEMM) ships hand-tuned fp8 grouped GEMMs targeted at Hopper-class hardware and B200. Reported ~1.3–1.7× over best alternatives on the relevant shapes.
+
+- **Mechanism**: write a `DeepGEMMGroupedExpertsConverter` that mirrors `Float8GroupedExpertsConverter` but routes expert forward through `deep_gemm.m_grouped_gemm_fp8_fp8_bf16_nt_contiguous`. Lives in `torchtitan/experiments/` per [.claude/CLAUDE.md](../torchtitan/CLAUDE.md) experiments rule (third-party kernel dependency).
+- **Estimated incremental lift over 2.2**: +5–15% on expert compute, i.e. another +3–10% overall.
+- **gpt_oss-20b (on top of 2.x)**: 820 → **870–950** tflops/gpu.
+- **Qwen3-30B-A3B (on top of 2.x)**: 680 → **730–820** tflops/gpu, mfu ~32–36%.
+- **Risk**: higher. Custom expert-forward needs autograd backward via the same kernel path; need a numerics + perf microbench before integration. Also adds a non-PyTorch dep, so it must stay under `torchtitan/experiments/` (core principle #1 in CLAUDE.md).
+
+## 4. graph_trainer (aggressive, experimental)
+
+`torchtitan/experiments/graph_trainer/` already has [qwen3](../torchtitan/experiments/graph_trainer/qwen3/config_registry.py) and deepseek_v3 entries. It captures forward + loss + backward (and optionally optim) as a single FX graph and applies passes for fusion, comm overlap, and CUDAGraph capture.
+
+- **Mechanism**: switch `--module` from `qwen3` to `graph_trainer.qwen3` (and similarly for gpt_oss once an entry is added). Inherits SimpleFSDP, regional CUDAGraph, async expert comm.
+- **Estimated lift over the stack in (3)**: +5–15% additional, mainly from launch-overhead reduction (CUDAGraph capture on a B200 dominated by lots of small MoE kernels) and overlap of expert all-to-all with compute.
+- **Risk**: experimental. Not all parallelism combos work; loss must be re-validated.
+
+## Summary
+
+Estimated path from baseline to a stacked configuration (per-GPU tflops, mfu in parens, B200 dense bf16 peak ≈ 2.25 PF/s):
+
+| Stage | gpt_oss-20b | qwen3-30B-A3B | Confidence |
+| --- | --- | --- | --- |
+| Baseline (today) | 573 (25.5%) | 455 (20.2%) | measured |
+| + 1.1–1.3 free wins | 610–650 (+10%) | 570–610 (+30%) | high |
+| + 2.x torchao fp8 | 820–900 (+50%) | 680–760 (+60%) | medium |
+| + 3 DeepGEMM | 870–950 (+60%) | 730–820 (+70%) | low |
+| + 4 graph_trainer | 920–1050 (+70%) | 770–900 (+80%) | low |
+
+## Recommended order
+
+1. **First**: 1.1 selective AC + 1.2 max-autotune on Qwen3 (one-day; no third-party deps; validates the memory model for fp8 next).
+2. **Second**: 2.1 → 2.2 — wire `Float8LinearConverter` then `Float8GroupedExpertsConverter` through `_qwen3_30b_base` and the gpt_oss config. Land each independently with a loss-vs-bf16 check on c4.
+3. **Third**: 1.3 compiled optimizer step (small, orthogonal — can land any time after 2.x).
+4. **Fourth (experimental)**: DeepGEMM converter under `torchtitan/experiments/deepgemm_moe/` with a microbench harness.
+5. **Optional**: graph_trainer once 2.x and 3 are stable, as a side-by-side comparison rather than a replacement.
+
+## Sanity checks before each step
+
+Per [.claude/CLAUDE.md](../torchtitan/CLAUDE.md):
+- Non-computation changes (1.1, 1.3, partly 4): bitwise-identical loss & grad_norm under `--debug.seed=42 --debug.deterministic`.
+- Computation changes (1.2 with new autotune kernels, all of 2.x, 3, 4): short c4 run, compare loss curve vs bf16 baseline via `scripts/loss_compare.py`.

docs/bs_sweep_results.md

@@ -0,0 +1,81 @@
+# BF16 + torch.compile batch-size sweep — GPT-OSS-20B and Qwen3-30B-A3B
+
+Results of a `local_batch_size` sweep run on one 8×B200 node, using bf16 + `torch.compile` and FSDP-only sharding. All other knobs are held fixed so each row isolates the effect of changing batch size.
+
+## torchtitan pull
+
+- Upstream: https://github.com/pytorch/torchtitan.git
+- Branch: `main`
+- Commit at run time: [`1690e0ae`](https://github.com/pytorch/torchtitan/commit/1690e0aeafa0a631251fd522b870112343d8bf77) — *"[cpu-offloading] Encode last consumer as dep arg in ao.wait (#3333)"* (Michael Lazos, 2026-05-13)
+- Local uncommitted changes on top of that commit:
+  - [torchtitan/models/qwen3/config_registry.py](../torchtitan/models/qwen3/config_registry.py) — added a `qwen3_30b` config registry entry (Qwen3-30B-A3B, `seq_len=8192`, FSDP-only, AC=`full` when compile is on, `lr=3e-4`, `warmup_steps=600`, `steps=3000`).
+  - [torchtitan/trainer.py](../torchtitan/trainer.py) — added a `[GRAD-PROBE]` debug print at step 1 and every 50 steps that logs gradient norms for `attn_wq`, `attn_wo`, `expert_mlp1`, `expert_mlp2` on layer 0.
+
+## Common setup
+
+| Setting | Value |
+| --- | --- |
+| Hardware | 1 node, 8× NVIDIA B200 (192 GB), node `ml-16-b200-node-002` |
+| Launcher | `torchrun --nproc_per_node=8`, single-node slurm job |
+| Parallelism | `pp=1, dp_replicate=1, dp_shard=8, cp=1, tp=1, ep=1` (pure FSDP) |
+| Precision | `--training.dtype=bfloat16` |
+| Compile | `--compile.enable` (components: `model`, `loss`) |
+| Activation checkpoint | `full` for Qwen3 (set in the local helper when compile is on); gpt_oss uses the upstream `gpt_oss_20b` defaults |
+| Sequence length | 8192 |
+| Dataset | `c4` (HuggingFace) |
+| Loss | `ChunkedCELoss` |
+| Optimizer / LR | `lr=3e-4`, `warmup_steps=600`, `total_steps=3000` |
+| Alloc env | `PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True` |
+| Run length | 30-min dev slurm allocation per job — jobs are stopped well before 3000 steps; metrics below are steady-state mid-run |
+
+Sbatch scripts: [scripts/run_gpt_oss_20b.sbatch](../scripts/run_gpt_oss_20b.sbatch), [scripts/run_qwen_30b.sbatch](../scripts/run_qwen_30b.sbatch).
+
+`tps` is tokens/sec/GPU (rank-local). `tflops` is per-GPU. `mfu` is fraction of B200 peak bf16 FLOPs. `memory` is reserved CUDA memory on rank 0, with the % being the fraction of the 192 GB B200 HBM (as reported by torchtitan).
+
+## GPT-OSS-20B (20.9B params total / 4.2B active)
+
+Run via `--module gpt_oss --config gpt_oss_20b` (upstream config, no local changes).
+
+| local_bs | global_bs | job id | tps (tok/s/gpu) | tflops/gpu | mfu | memory | last logged step | notes |
+| --- | --- | --- | --- | --- | --- | --- | --- | --- |
+| 10 | 80 | 1424 | ~18,300 | ~573 | ~25.5% | 69.5 GiB (38.96%) | 190 | baseline, longest run |
+| 10 | 80 | 1425 | ~18,200 | ~570 | ~25.3% | 69.5 GiB (38.96%) | 60 | rerun of bs=10, matches 1424 |
+| 14 | 112 | 1428 | ~17,150 | ~537 | ~23.9% | 84.8 GiB (47.54%) | 50 | |
+| 20 | 160 | 1430 | ~16,740 | ~524 | ~23.3% | 104.5 GiB (58.60%) | 140 | largest bs tested |
+
+Observations
+- tps/GPU drops monotonically as local bs grows (18.3k → 17.2k → 16.7k), so larger micro-batches are slightly less compute-efficient on this stack even though global_bs more than doubles from bs=10 to bs=20.
+- Memory scales close to linearly: +~1.5 GiB per extra sample (bs=10→14: +15.3 GiB; bs=14→20: +19.7 GiB).
+- Even at bs=20 the model uses only ~59% of HBM, so plenty of headroom remains.
+
+## Qwen3-30B-A3B (30.5B params total / 3.4B active)
+
+Run via `--module qwen3 --config qwen3_30b` (local registry entry — see the diff above).
+
+| local_bs | global_bs | job id | tps (tok/s/gpu) | tflops/gpu | mfu | memory | last logged step | notes |
+| --- | --- | --- | --- | --- | --- | --- | --- | --- |
+| 8 | 64 | 1423 | ~11,865 | ~446 | ~19.8% | 124.0 GiB (69.56%) | 160 | |
+| 10 | 80 | 1426 | ~11,950 | ~449 | ~20.0% | 146.1 GiB (81.92%) | 70 | |
+| 12 | 96 | 1429 | ~12,120 | ~455 | ~20.2% | 168.1 GiB (94.27%) | 50 | near memory ceiling |
+| 14 | — | 1427 | — | — | — | 176.5 GiB (98.98%) at step 10 | OOM | mapping failed in CUDACachingAllocator with repeated `expandable_segments` failures across all 8 ranks; only step 1 completed normally |
+
+Observations
+- Unlike gpt_oss, Qwen3-30B-A3B *gains* tps as local bs grows from 8 → 12 (~11.9k → ~12.1k), suggesting bs=8 is slightly compute-starved for this MoE layout.
+- Memory grows ~22 GiB per +2 samples (8→10→12). bs=12 sits at 94.3% reserved, leaving little headroom for activation peaks.
+- bs=14 OOMs at the very start of training — peak reserved hit 98.98% before failure; this is the practical ceiling at `seq_len=8192` / FSDP-only on 8×B200.
+
+## Cross-model takeaways
+
+- gpt_oss_20b is roughly 1.5× faster per GPU than qwen3_30b at comparable bs (~17–18k vs ~12k tps), in line with the smaller total parameter count and lower expert capacity.
+- mfu is ~5 pp higher on gpt_oss (23–25%) than on qwen3 (~20%) across the swept range; both are well below the ~50% mfu typical of dense compute-bound training, reflecting MoE + activation-checkpointing overheads.
+- Largest safe local_bs in this configuration:
+  - gpt_oss_20b: ≥ 20 (not pushed to OOM here).
+  - qwen3_30b: 12 (bs=14 OOMs).
+
+## How the numbers were extracted
+
+Lines of the form
+```
+step:  N  loss: ...  grad_norm: ...  memory: X GiB(Y%)  tps: Z  tflops: T  mfu: M%
+```
+are emitted by [torchtitan/components/metrics.py](../torchtitan/components/metrics.py) every `metrics.log_freq` steps. The tables above use steady-state values reported after compile warmup (step ≥ 20 for short runs, step ≥ 60 for the longer ones); step-1 numbers are excluded since they include compilation cost (typically 7–8× slower than steady state).

scripts/run_gpt_oss_20b.sbatch

@@ -0,0 +1,72 @@
+#!/bin/bash
+#SBATCH --job-name=tt-gptoss-20b-base-compile-b8
+#SBATCH --nodes=1
+#SBATCH --ntasks-per-node=1
+#SBATCH --gres=gpu:8
+#SBATCH --cpus-per-task=64
+#SBATCH --time=00:30:00
+#SBATCH --output=/data/cc/torchtitan/logs/base-compile-oss-b8-%j.out
+#SBATCH --error=/data/cc/torchtitan/logs/base-compile-oss-b8-%j.err
+#SBATCH --partition=dev
+#SBATCH --qos=dev
+#SBATCH --account=ubuntu
+
+set -euo pipefail
+
+mkdir -p /data/cc/torchtitan/logs
+
+source /data/cc/torchtitan/.venv/bin/activate
+cd /data/cc/torchtitan
+
+: "${HF_TOKEN:?HF_TOKEN must be exported before sbatch (do NOT hardcode in this file)}"
+export HF_TOKEN
+
+export TORCHINDUCTOR_CACHE_DIR=/data/cc/torchtitan/inductor_cache
+export HF_HOME=/data/cc/torchtitan/.hf_cache
+export HF_DATASETS_CACHE=$HF_HOME/datasets
+export HUGGINGFACE_HUB_CACHE=$HF_HOME/hub
+export TRITON_CACHE_DIR=/data/cc/torchtitan/.triton_cache
+export XDG_CACHE_HOME=/data/cc/torchtitan/.xdg_cache
+mkdir -p "$TORCHINDUCTOR_CACHE_DIR" "$HF_HOME" "$HF_DATASETS_CACHE" "$HUGGINGFACE_HUB_CACHE" "$TRITON_CACHE_DIR" "$XDG_CACHE_HOME"
+
+export MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1)
+export MASTER_PORT=29500
+export NCCL_DEBUG=WARN
+export NCCL_IB_DISABLE=0
+export PYTHONFAULTHANDLER=1
+export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
+export TORCH_NCCL_HEARTBEAT_TIMEOUT_SEC=3600
+
+# --- WandB ----------------------------------------------------------------
+# torchtitan's WandBLogger reads these env vars (see components/metrics.py).
+# Logging is enabled only if the wandb package is importable; otherwise the
+# --metrics.enable_wandb flag is omitted to avoid a hard failure.
+export WANDB_PROJECT="mfu-b200"
+export WANDB_RUN_NAME="base-compile-oss-b8-${SLURM_JOB_ID}"
+export WANDB_TEAM="${WANDB_TEAM:-caia-costello-lambd}"
+WANDB_FLAG=""
+if python -c "import wandb" 2>/dev/null; then
+  WANDB_FLAG="--metrics.enable_wandb"
+fi
+# --------------------------------------------------------------------------
+
+echo "=== GPT-OSS-20B: BF16 + compile + bs=4 (main baseline) ==="
+echo "Job ID:  $SLURM_JOB_ID"
+echo "Nodes:   $SLURM_NNODES"
+echo "Master:  $MASTER_ADDR:$MASTER_PORT"
+echo "Python:  $(which python)"
+echo "Torch:   $(python -c 'import torch;print(torch.__version__)')"
+echo ""
+
+srun --export=ALL torchrun \
+    --nproc_per_node=8 \
+    --nnodes="$SLURM_NNODES" \
+    --rdzv_id="$SLURM_JOB_ID" \
+    --rdzv_backend=c10d \
+    --rdzv_endpoint="$MASTER_ADDR:$MASTER_PORT" \
+    -m torchtitan.train \
+        --module gpt_oss \
+        --config gpt_oss_20b \
+        --training.local_batch_size=8 \
+        --compile.enable \
+        $WANDB_FLAG