diff --git a/configs/README.md b/configs/README.md index 23a9f1811c3..b5fd52779ea 100644 --- a/configs/README.md +++ b/configs/README.md @@ -10,6 +10,9 @@ Configurations for running prime-rl. `minimax-m2.5` (swe), `nemotron-3-super` (swe), `deepseek-v4-flash` (sft + a standalone vLLM serving pre-flight). For the tutorialized equivalents, see [`examples/advanced/`](../examples/advanced). +- **[`benchmark/`](benchmark)** — SFT throughput benchmarks matched to torchtitan runs + (`qwen3-235b-a22b`, 4 nodes at seq 16k), with the torchtitan counterpart and the measured + numbers in its README. - **`ci/`** — integration and nightly configs used by CI. - **`debug/`** — throwaway configs for developing the framework itself: `algo/` (per-algorithm smokes), `fake/` (fake-data trainer/SFT smokes), `multi-env/` diff --git a/configs/benchmark/README.md b/configs/benchmark/README.md new file mode 100644 index 00000000000..3a7442edc1c --- /dev/null +++ b/configs/benchmark/README.md @@ -0,0 +1,95 @@ +# Benchmark configs + +Throughput benchmarks of the prime-rl SFT trainer against [torchtitan](https://github.com/pytorch/torchtitan), +with both stacks configured identically. Each config here has a torchtitan counterpart described below. +These runs measure step time only; they write no checkpoints. + +## Qwen3-235B-A22B, 4 nodes, seq 16k + +[`qwen3-235b-a22b/sft-4n-16k.toml`](qwen3-235b-a22b/sft-4n-16k.toml) + +| | prime-rl | torchtitan | +|---|---|---| +| hardware | 4 nodes x 8 B300, FSDP over 32 GPUs, EP 8 | same | +| model | `Qwen/Qwen3-235B-A22B-Thinking-2507`, HF weights | Qwen3 235B-A22B, random init | +| tokens per step | 32 x 16384, one packed sample per GPU (`PrimeIntellect/Reverse-Text-SFT`) | 32 x 16384, one packed sample per GPU (`c4_test`) | +| routing | `debug.force_balanced_routing` (round-robin) | `debug.moe_force_load_balance` (round-robin) | +| optimizer | AdamW lr 1e-5, fp32 master weights and states, bf16 compute, fp32 reduce | `default_adamw`, same precisions | +| attention | FA4 varlen | FA4 varlen (`attn_backend="varlen"`) | +| activation checkpointing | full (recompute the block, keep `aten::topk`) | `TopKSaveAC` = selective AC saving only `aten.topk` | +| compile | `fullgraph=true` per block | `fullgraph=True` per block | +| steps | 20, steady state = steps 5 to 20 | same | + +MFU uses the same formula in both stacks: `6 x active params (21.568B) + 6 x layers x heads x (qk_head_dim + v_head_dim) x seq` +with the real head dim of 128, i.e. 280.8 GFLOP/token at 16384, over a 2.25 PFLOP/s peak per B300. Note that this charges dense +causal attention while packed short documents make attention nearly free, so absolute MFU is roughly 2x overstated for both stacks; +the comparison is unaffected. + +### Results (2026-10-02, back to back on the same nodes) + +| run | s/step | tok/s/GPU | MFU | peak GiB | +|---|---|---|---|---| +| torchtitan `9e159aed7` | 5.04 | 3253 | 40.7% | 144 | +| prime-rl, this config | 4.39 | 3732 | 46.6% | 141 | + +prime-rl's number requires #3753 (EP experts stay in the block's FSDP unit) and #3754 (FA4 custom op); without them the block cannot +compile with `fullgraph=true` and the same config runs at 7.0 s/step. The Mega MoE dispatch (#3651) brings the same setup to 4.10 s/step +with full AC and 3.93 s/step with selective AC; see `docs/performance.md` once it lands. + +### Running it + +```bash +uv run sft @ configs/benchmark/qwen3-235b-a22b/sft-4n-16k.toml +``` + +Real routing: drop `[model.debug]`. Measured at 6.11 s/step with the HF checkpoint's routing (prime-rl only; torchtitan has no +trained router to compare against). + +### torchtitan counterpart + +torchtitan `9e159aed7` with these additions to `torchtitan/models/qwen3/config_registry.py`, launched with +`torchrun --nnodes=4 --nproc-per-node=8 -m torchtitan.train --module qwen3 --config qwen3_235b_a22b_primerl_bench_4n_s16k_fullac_bal` +(torch 2.15 nightly, `PYTORCH_ALLOC_CONF=expandable_segments:True`): + +```python +def qwen3_235b_a22b_primerl_bench_4n_s16k_fullac_bal(seq_len: int | None = 16384) -> Trainer.Config: + model_config = model_registry("235B-A22B", seq_len=seq_len, attn_backend="varlen") + config = Trainer.Config( + loss=ChunkedLossWrapper.Config(loss_fn=CrossEntropyLoss.Config(global_vocab_size=decoder_vocab_size(model_config))), + hf_assets_path="./assets/hf/Qwen3-235B-A22B", + metrics=MetricsProcessor.Config(log_freq=1), + model=model_config, + dataloader=GrainDataLoader.Config(dataset=ConcatThenSplitPackingConfig(dataset=DATASETS["c4_test"]), shuffle=False), + optimizer=default_adamw(lr=1e-5), + lr_scheduler=LRSchedulersContainer.Config(warmup_steps=1), + training=TrainingConfig( + num_tokens_per_microbatch_per_dp_rank=model_config.max_context_length, + max_context_length=model_config.max_context_length, + steps=20, + dtype="float32", + mixed_precision_param="bfloat16", + mixed_precision_reduce="float32", + disable_cuda_graphs=True, + ), + parallelism=ParallelismConfig(data_parallel_shard_degree=-1, expert_parallel_degree=8), + compile=CompileConfig(components=["model", "loss"]), + checkpointer=None, + activation_checkpoint=TopKSaveAC.Config(), + ) + config.debug = DebugConfig(moe_force_load_balance=True) + return config + + +class TopKSaveAC(SelectiveAC): + """Recompute the whole block except the router top-k: the counterpart of prime-rl's full AC. + torchtitan's FullAC recomputes the router, and a re-routed token breaks the EP dispatch.""" + + @dataclass(kw_only=True, slots=True) + class Config(SelectiveAC.Config): + force_recompute_mm_shapes_by_fqns: list[str] = field(default_factory=list) + + def get_save_ops(self) -> set: + return {torch.ops.aten.topk.default} +``` + +torchtitan's selective AC OOMs at this size, and its full AC crashes under EP, hence `TopKSaveAC`. diff --git a/configs/benchmark/qwen3-235b-a22b/sft-4n-16k.toml b/configs/benchmark/qwen3-235b-a22b/sft-4n-16k.toml new file mode 100644 index 00000000000..7d379fefda5 --- /dev/null +++ b/configs/benchmark/qwen3-235b-a22b/sft-4n-16k.toml @@ -0,0 +1,64 @@ +# Qwen3-235B-A22B SFT throughput benchmark, matched to the torchtitan run described in +# configs/benchmark/README.md: 4 nodes x 8 GPUs (FSDP over 32 GPUs, ep=8), seq 16384, one packed +# 16384-token sample per GPU per step (32 x 16384 tokens per step), full activation checkpointing, +# AdamW with fp32 master weights and fp32 optimizer states, bf16 compute, fp32 reduce, torch.compile +# with fullgraph, 20 steps. Routing is forced round-robin balanced because torchtitan trains from +# random init, whose natural routing is not comparable; real routing is ~40% slower on both stacks. +# +# uv run sft @ configs/benchmark/qwen3-235b-a22b/sft-4n-16k.toml +# +# There is no [ckpt] block: this run measures step time, it does not produce a usable checkpoint. +# `dist_timeout_seconds` covers the first run, which builds the HF -> prime-rl conversion cache. + +dist_timeout_seconds = 14400 +max_steps = 20 + +[env_vars] +HF_HOME = "/home/huggingface" + +# The renderer registry lists Qwen/Qwen3-235B-A22B but not the 2507 Thinking variant, which uses +# the same Qwen3 chat template. +[renderer] +name = "qwen3" + +[model] +name = "Qwen/Qwen3-235B-A22B-Thinking-2507" +impl = "custom" +seq_len = 16384 +ep = 8 +optim_cpu_offload = false +ac_offloading = "None" +# fullgraph needs the router inside the block's FSDP unit, which the fp32 router is not. +moe_router_dtype = "bfloat16" + +# Selective AC does not fit at seq 16384 on 4 nodes with fp32 AdamW states and the torch dispatch. +[model.ac] +mode = "full" +freq = 1 + +[model.compile] +fullgraph = true + +[model.moe.dispatch] +type = "torch" + +[data] +name = "PrimeIntellect/Reverse-Text-SFT" +seq_len = 16384 +batch_size = 32 +micro_batch_size = 1 + +[optim] +type = "adamw" +lr = 1e-5 + +[deployment] +type = "multi_node" +num_train_nodes = 4 + +[slurm] +job_name = "bench-qwen235b-sft-4n-16k" +partition = "all" + +[model.debug] +force_balanced_routing = true