Add NVFP4 four-over-six converter

ablation_arms.py

@@ -0,0 +1,168 @@
+"""Six-arm Qwen3-30B-A3B recipe-ablation configs (TitanRL + torchao).
+
+Untracked launcher module (selected via ``--module ablation_arms``),
+replicating the six-configuration recipe ablation from
+lmsys.org/blog/2026-07-29-mxfp8-nvfp4-rl at reduced scale on 2x4 GB200:
+
+1. arm_bf16       — BF16 training + BF16 rollout.
+2. arm_mxfp8_e2e  — end-to-end MXFP8 experts training + MXFP8 rollout.
+3. arm_mxfp8_hp   — MXFP8 rollout/forward, high-precision backward.
+4. arm_mxfp8_deq  — MXFP8 rollout/forward, dequantized backward.
+5. arm_nvfp4_hp   — per-token NVFP4 (4over6 row-scaled) rollout/forward,
+                    high-precision backward.
+6. arm_nvfp4_deq  — same forward, dequantized backward.
+
+Shared workload across all arms (the fixed-comparison contract): GRPO on
+dapo-math-17k, 8 prompts x 8 samples per step, 10 synchronous steps,
+max 2048 response tokens (blog: 8192 — scaled down to fit the window),
+AIME2025 greedy validation (16 samples) at start/end, AdamW lr 1e-6
+warmup 2 + linear decay, seeds pinned by the launch script.
+
+Low-precision arms additionally share (per the blog's recipe contract):
+routed-expert grouped GEMMs only, last 8 of 48 layers kept BF16, and
+weight decay 0 on the quantized expert weights (0.1 elsewhere).
+"""
+
+from torchtitan.components.optimizer import OptimizersContainer, ParamGroupConfig
+from torchtitan.components.quantization import (
+    MXFP8GroupedExpertsConverter,
+    NVFP4FourOverSixGroupedExpertsConverter,
+)
+from torchtitan.components.quantization.nvfp4 import nvfp4_bf16_first_last_fqns
+from torchtitan.experiments.rl.components.batcher import BatchConfig, Batcher
+from torchtitan.experiments.rl.controller import Controller, ValidationConfig
+from torchtitan.experiments.rl.environment import TokenEnv
+from torchtitan.experiments.rl.examples.alphabet_sort.config_registry import (
+    rl_grpo_qwen3_30b_a3b_varlen,
+)
+from torchtitan.experiments.rl.examples.dapo_math.data import AIME2025Dataset
+from torchtitan.experiments.rl.examples.dapo_math.rollouter import DapoMathRollouter
+from torchtitan.experiments.rl.rollout.advantage import AdvantageEstimator
+from torchtitan.models.qwen3 import model_registry
+
+_N_LAYERS = 48  # Qwen3-30B-A3B decoder layers
+_LAST_BF16 = 8  # blog: last 15% of layers kept BF16 -> 8 of 48
+_MAX_RESPONSE_TOKENS = 2048
+_MAX_TOTAL_TOKENS = 4096
+_NUM_VALIDATION_SAMPLES = 16
+
+# Allow-list of quantized layers (trailing dot so "layers.1." won't match
+# "layers.10"); shared by the converters of every low-precision arm.
+_QUANTIZED_LAYER_FQNS = nvfp4_bf16_first_last_fqns(_N_LAYERS, 0, _LAST_BF16)
+
+# Same layer window as a param-FQN regex: layers 0-39 routed-expert weights
+# (w1_EFD/w2_EDF/w3_EFD live at layers.N.moe.routed_experts.inner_experts,
+# verified via a meta-device build of the 30B-A3B spec). ParamGroupConfig
+# uses re.search on RAW named_parameters() names, which may carry wrapper
+# segments (e.g. _checkpoint_wrapped_module), so allow anything between the
+# layer index and routed_experts; router.gate stays in the default group.
+_QUANTIZED_EXPERT_PARAM_REGEX = r"layers\.([0-9]|[1-3][0-9])\..*routed_experts\."
+
+_ADAMW_COMMON = {"lr": 1e-6, "betas": (0.9, 0.95), "eps": 1e-8}
+
+
+def _lp_optimizer() -> OptimizersContainer.Config:
+    """AdamW with weight decay 0 on quantized expert weights, 0.1 elsewhere."""
+    return OptimizersContainer.Config(
+        param_groups=[
+            ParamGroupConfig(
+                pattern=_QUANTIZED_EXPERT_PARAM_REGEX,
+                optimizer_name="AdamW",
+                optimizer_kwargs={**_ADAMW_COMMON, "weight_decay": 0.0},
+            ),
+            ParamGroupConfig(
+                pattern=r".*",
+                optimizer_name="AdamW",
+                optimizer_kwargs={**_ADAMW_COMMON, "weight_decay": 0.1},
+            ),
+        ]
+    )
+
+
+def _ablation_base(converters: list | None) -> Controller.Config:
+    """Shared workload; arms differ ONLY in converters (+ LP weight decay)."""
+    config = rl_grpo_qwen3_30b_a3b_varlen()
+    config.model_spec = model_registry(
+        "30B-A3B",
+        attn_backend="varlen",
+        converters=converters or [],
+    )
+    config.rollouter = DapoMathRollouter.Config(
+        validation_dataset=AIME2025Dataset.Config(
+            num_samples=_NUM_VALIDATION_SAMPLES,
+        ),
+        token_env=TokenEnv.Config(
+            max_rollout_tokens=_MAX_TOTAL_TOKENS,
+            max_num_turns=1,
+        ),
+        advantage=AdvantageEstimator.Config(should_std_normalize=True),
+    )
+    config.async_loop.num_training_steps = 10
+    config.async_loop.num_prompts_per_train_step = 8
+    config.async_loop.num_samples_per_prompt = 8
+    config.async_loop.target_offpolicy_steps = 0
+    config.async_loop.window_fraction = None
+    # Keep zero-variance groups (their advantages are 0 via the +eps guard):
+    # DAPO-style dynamic resampling would otherwise make per-step rollout
+    # counts — and arm wall time — unpredictable across the six arms.
+    config.async_loop.training_sample_builder.drop_zero_std_reward_groups = False
+    config.async_loop.validation = ValidationConfig(
+        num_samples=_NUM_VALIDATION_SAMPLES,
+    )
+    config.async_loop.batcher = Batcher.Config(
+        batch=BatchConfig(local_batch_size=1, seq_len=_MAX_TOTAL_TOKENS),
+    )
+    config.generator.sampling.max_tokens = _MAX_RESPONSE_TOKENS
+    config.trainer.lr_scheduler.total_steps = 10
+    # Save updated weights at the final step for follow-up runs, but weights
+    # only: the base config's full-state save (optimizer moments included)
+    # would not fit 6 arms on this scratch volume.
+    config.trainer.checkpoint.last_save_model_only = True
+    if converters:
+        config.trainer.optimizer = _lp_optimizer()
+    return config
+
+
+def _nvfp4_converter(backward_override: str):
+    return NVFP4FourOverSixGroupedExpertsConverter.Config(
+        fqns=_QUANTIZED_LAYER_FQNS,
+        row_scaled_activation=True,
+        err_mode="mse",
+        e4m3_scale_bound=256,
+        weight_block="1x16",
+        backward_override=backward_override,
+        pad_multiple=128,
+    )
+
+
+def _mxfp8_converter(backward_override: str | None):
+    return MXFP8GroupedExpertsConverter.Config(
+        fqns=_QUANTIZED_LAYER_FQNS,
+        recipe_name="mxfp8_rceil",
+        backward_override=backward_override,
+        pad_multiple=128,
+    )
+
+
+def arm_bf16() -> Controller.Config:
+    return _ablation_base(None)
+
+
+def arm_mxfp8_e2e() -> Controller.Config:
+    return _ablation_base([_mxfp8_converter(None)])
+
+
+def arm_mxfp8_hp() -> Controller.Config:
+    return _ablation_base([_mxfp8_converter("high_precision")])
+
+
+def arm_mxfp8_deq() -> Controller.Config:
+    return _ablation_base([_mxfp8_converter("dequantized")])
+
+
+def arm_nvfp4_hp() -> Controller.Config:
+    return _ablation_base([_nvfp4_converter("high_precision")])
+
+
+def arm_nvfp4_deq() -> Controller.Config:
+    return _ablation_base([_nvfp4_converter("dequantized")])

nvfp4_4over6_results/README.md

@@ -0,0 +1,111 @@
+# Qwen3-30B-A3B GRPO recipe ablation — MXFP8 / NVFP4-4over6 experts
+
+Reduced-scale replication of the six-configuration RL quantization ablation
+from [lmsys.org/blog/2026-07-29-mxfp8-nvfp4-rl](https://lmsys.org/blog/2026-07-29-mxfp8-nvfp4-rl)
+on TitanRL (`torchtitan.experiments.rl`) with the torchao-backed
+`MXFP8GroupedExpertsConverter` / `NVFP4FourOverSixGroupedExpertsConverter`
+from this branch. Run 2026-08-25/26 on one GB200 node (4 GPUs).
+
+![training curves](training_curves.png)
+
+## Contract (identical across all six arms)
+
+- **Model**: Qwen3-30B-A3B (48 layers), HF checkpoint init.
+- **Topology**: 1 node x 4 GB200 — trainer TP2/EP2 on GPUs 0-1, one vLLM
+  generator TP2/EP2 on GPUs 2-3, synchronous on-policy loop
+  (`target_offpolicy_steps=0`), weights refit each step via torchstore
+  (BF16 masters; low-precision arms re-quantize dynamically each forward).
+- **Workload**: GRPO on dapo-math-17k, 10 training steps x 64 samples
+  (8 prompts x 8 samples), max 2048 response tokens / 4096 total
+  (blog uses 8192 — scaled down to fit the node window), std-normalized
+  advantages, zero-variance reward groups **kept** (advantage 0) so per-step
+  sample counts stay fixed across arms.
+- **Optimizer**: AdamW lr 1e-6, warmup 2 + linear decay over 10 steps,
+  weight decay 0 on quantized expert weights / 0.1 elsewhere
+  (BF16 arm: single default group).
+- **Quantized scope** (low-precision arms): routed-expert w1/w2/w3 grouped
+  GEMMs of layers 0-39 only. Last 8 of 48 layers, attention, router,
+  embeddings, lm_head, norms, and the KV cache stay BF16.
+- **Validation**: AIME2025, 16 samples, greedy, at step 0 and step 10.
+- **Seeds**: trainer and generator seeded 42; each arm starts from the same
+  HF checkpoint (the launcher refuses to reuse an existing output folder).
+
+### Arms
+
+| arm | experts forward (train + rollout) | experts backward |
+|---|---|---|
+| `arm_bf16` | BF16 | BF16 |
+| `arm_mxfp8_e2e` | MXFP8 (rceil) | MXFP8 quantized |
+| `arm_mxfp8_hp` | MXFP8 (rceil) | `high_precision` (BF16 on original inputs) |
+| `arm_mxfp8_deq` | MXFP8 (rceil) | `dequantized` (BF16 on dequantized fprop operands) |
+| `arm_nvfp4_hp` | NVFP4 4over6, row-scaled activations (MSE, bound 256, 1x16 weight blocks) | `high_precision` |
+| `arm_nvfp4_deq` | same | `dequantized` |
+
+## Results (10 steps each; means over steps 1-10)
+
+| arm | loss mean | grad-norm mean | reward mean | reward @10 | logprob diff mean | clip-frac mean |
+|---|---|---|---|---|---|---|
+| `arm_bf16` | 0.0111 | 0.123 | 0.098 | 0.172 | -0.00070 | 0.32% |
+| `arm_mxfp8_e2e` | 0.0074 | 0.110 | 0.073 | 0.125 | -0.00095 | 0.58% |
+| `arm_mxfp8_hp` | 0.0118 | 0.123 | 0.083 | 0.156 | -0.00088 | 0.60% |
+| `arm_mxfp8_deq` | 0.0032 | 0.112 | 0.075 | 0.172 | -0.00104 | 0.59% |
+| `arm_nvfp4_hp` | 0.0080 | 0.114 | 0.073 | 0.156 | -0.00273 | 1.96% |
+| `arm_nvfp4_deq` | 0.0027 | 0.099 | 0.064 | 0.156 | -0.00280 | 1.99% |
+
+(Full per-step data in `training_metrics.csv`, aggregates in `summary.csv`;
+regenerate both plus the chart with `make_results.py <outputs_dir> <results_dir>`.)
+
+### Findings
+
+1. **No divergence in any arm.** Loss oscillates within +-0.04 of zero,
+   gradient norms stay in 0.046-0.182, entropy curves are superimposable,
+   zero NaNs. Ten steps at lr 1e-6 is a short horizon — this is a
+   stability/parity smoke, not a convergence claim.
+2. **Quantization is visible exactly where the blog says it should be**: in
+   the rollout-vs-trainer logprob difference (BF16 floor ~0.0007 nats,
+   MXFP8 ~0.0010, NVFP4 ~0.0027) and the PPO ratio clipped fraction
+   (~0.3% / ~0.6% / ~2%). Backward mode (`high_precision` vs `dequantized`)
+   has no visible effect on either — the mismatch is a forward/rollout
+   property. Loss/reward/grad-norm/entropy show no precision ordering.
+3. **Step-3 zero-grad no-ops** in `arm_mxfp8_e2e` and `arm_nvfp4_deq`
+   (grad norm exactly 0) are all-zero-reward rollout batches: with
+   zero-variance groups kept, a batch where no sample earns reward has
+   all-zero advantages by construction. Not a numerics fault.
+4. **AIME2025 at this scale is noise**: 0 or 1 of 16 problems solved
+   pre/post across arms, with the +-1 flips uncorrelated with precision.
+   Expected for 10 updates at lr 1e-6.
+5. **Wall time** (launch to final save): BF16 ~80 min, MXFP8 ~91 min,
+   NVFP4 ~113 min; plus ~25 min closing validation + teardown. The first
+   arm on each node pays a ~525 s cold NFS checkpoint load (warm arms
+   load in ~10 s) — the BF16 figure includes one. Final model-only fp32 DCP save (115G) ~1060 s to NFS.
+   Closing 16-sample greedy validation: BF16 342 s, MXFP8 ~430 s,
+   NVFP4 ~525 s.
+6. **NVFP4 rollout needs the fused row-scaled forward**
+   (`FOUR_OVER_SIX_GROUPED_ROW_SCALED_FUSED_BF16_OUT=1`, set by the
+   launcher): the per-group loop path decodes at ~41 tok/s and validates in
+   2908 s vs ~525 s fused. Fused == loop + one extra BF16 rounding on the
+   GEMM output. With it, NVFP4 rollout decodes at ~250 tok/s at 64
+   concurrent requests vs ~400 tok/s BF16 — the fused path closes the gap
+   from ~10x to ~1.6x, it does not reach parity. The remaining NVFP4
+   rollout tax is per-forward re-quantization of expert weights; a
+   quantize-once-per-refit weight cache is the natural (numerics-neutral)
+   follow-up.
+
+## Files
+
+- `training_curves.png` — 2x3 small multiples, all six arms.
+- `training_metrics.csv` — long-form per-step scalars (arm, metric, step, value).
+- `summary.csv` — per-arm aggregates (source of the table above).
+- `make_results.py` — regenerates all three from the arms' tfevents.
+- `run_ablation_arm.sh` — launches one arm in the vLLM 26.08 container
+  (topology, paths, seeds, metrics sinks; CUDA preflight; GPU-mem sampler).
+- `run_wave.sh` — runs arms sequentially on a node, one log per arm.
+- `../ablation_arms.py` (repo root) — the six arm configs; workload knobs
+  live there so every arm shares them, selected via `--module ablation_arms
+  --config <arm>`.
+
+The launch scripts are records of the exact runs: paths, container image,
+and IMEX device flags are specific to the GB200 cluster they ran on.
+Per-arm outputs (tfevents, structured JSONL logs, rollout samples, and the
+step-10 fp32 model-only DCP checkpoints, 115G each) live outside the repo
+on the cluster scratch volume.