chore(configs): add the Qwen3-235B SFT benchmark config matched to torchtitan by MarioSieg · Pull Request #3765 · PrimeIntellect-ai/prime-rl

MarioSieg · 2026-10-03T00:16:08Z

Adds configs/benchmark/ with the Qwen3-235B-A22B SFT config used for the prime-rl vs torchtitan throughput comparison, so the number is reproducible by anyone with 4 nodes.

What's in it

configs/benchmark/qwen3-235b-a22b/sft-4n-16k.toml: 4 nodes x 8 GPUs, FSDP over 32 GPUs, EP 8, seq 16384, one packed 16384-token sample per GPU per step, full AC, AdamW with fp32 states, bf16 router, compile.fullgraph=true, round-robin balanced routing, 20 steps.
configs/benchmark/README.md: the setup table for both stacks, the shared MFU formula and its caveat, the measured numbers, and the torchtitan counterpart (the config_registry.py additions and the TopKSaveAC policy) with the launch command.
One bullet in configs/README.md.

Measured (2026-10-02, same 4 nodes, back to back)

run	s/step	tok/s/GPU	MFU
torchtitan 9e159aed7	5.04	3253	40.7%
prime-rl, this config	4.39	3732	46.6%

The prime-rl number needs #3753 and #3754 merged; on current main the same config cannot compile with fullgraph=true and runs at 7.0 s/step.

Mega MoE variants and a docs page with the full table come in a separate PR on top of #3651.

…rchtitan Adds configs/benchmark/ with the 4-node seq-16k Qwen3-235B-A22B SFT config used to benchmark the trainer against torchtitan, and a README with the matched torchtitan configuration, the shared MFU formula, the measured numbers, and how to run both sides.

Replaces the three switches `model.fsdp_cpu_offload`, `model.optim_cpu_offload` and `model.full_offload` with one field: ```toml [model] offload = "optimizer" # default; also "full", "none" ``` This applies to both the RL trainer and SFT. The `fsdp` switch is not migrated: the mode is removed outright (rationale below). The `[model.full_offload] numa_bind` knob is dropped: no recipe, deployment, template or doc sets it. Full offload now always pins each rank's CPUs to its GPU's NUMA node. That was already the default. Ranks whose GPU reports no NUMA node still skip binding with a warning. On hosts without GPU sysfs topology (`/sys/bus/pci/devices/<bus>/numa_node`), binding is no longer optional: full offload fails there instead of being switched off. Why: the three switches could only be used one at a time, enforced by `cpu_offload_mutual_exclusion`. Because `optim_cpu_offload` defaulted to `true`, enabling full offload also meant writing `optim_cpu_offload = false`, otherwise validation failed (see the old `configs/advanced/deepseek-v4-flash/sft.toml` comment). With one value the conflicting combinations cannot be expressed, so that validator is deleted. The default behaviour does not change: optimizer-state offload stays on. Why drop `fsdp`: the mode was the stock torch FSDP2 `CPUOffloadPolicy` with no in-tree config, example, or test user; Muon rejected it and it broke NCCL weight broadcast. It was kept only for long-context runs whose fp32 masters/moments no longer fit on GPU (the reason it was restored in #3052); the `full` mode added in #3234 keeps FP32 masters, moments, and gradients in host RAM, covering that memory pressure instead. Gone with it: the `CPUOffloadPolicy` selection in `setup_fsdp`, CPU-side model init in `setup_model`/`load_dcp_from_hf`, `_move_buffers_to_cuda`, and the fsdp leg of the Gloo special case. Changes: - `ModelConfig.offload: Literal["none", "optimizer", "full"] = "optimizer"` replaces the three fields. `OptimizerInBackwardOffloadConfig` (and with it `numa_bind`) and its bool-normalizing `BeforeValidator` are gone. - Validators in `TrainerConfig` and `SFTConfig`: - full offload accepts only AdamW or SignSGD and disables clipping, now reading `offload == "full"`; - the `Muon optimizer does not support FSDP CPU offload` validators are deleted with the mode. - Runtime: - `setup_optimizer(..., offload=...)` replaces `cpu_offload` + `full_offload_config`, and its "both enabled" check is gone. - `FullCPUOffloadOptimizer` drops its unused `offload_config` argument. - `setup_full_cpu_optimizer_offload()` takes no arguments and always binds. - `model.py`: FSDP always uses the default `OffloadPolicy()`; CPU-side model init in `setup_model`/`load_dcp_from_hf` and `_move_buffers_to_cuda` are gone. - Tomls: `glm-4.5-air/swe-budget.toml` and `ci/integration/reverse-text-sft/full-offload-resume.toml` → `offload = "full"`; `glm-5.3/sft/h200/16k-12-node.toml` → `offload = "none"`; `glm-5.3/sft/h200/base.toml` → `offload = "optimizer"`. Also updated `docs/scaling.md`, the `deepseek-v4-flash/sft.toml` comment, and `skills/training/start-run/SKILL.md`, which documented the removed fields. Breaking: the old keys now fail with `extra_forbidden`, and `offload = "fsdp"` fails the `Literal`. There is no shim. Interaction with #3849 (one trainer-common config base): #3849 moves the full-offload and Muon/FSDP validators from `TrainerConfig`/`SFTConfig` onto `BaseTrainerConfig`. This PR edits the full-offload validator's condition and deletes the Muon/FSDP one, so the conflict is textual only: keep #3849's single copy of the full-offload validators with `self.model.offload == "full"` and drop its Muon/FSDP copy. Whichever lands second takes a small rebase. Other open PRs that set or read the old keys need a rebase: - tomls: #3883, #3879, #3766, #3765, #3729, #3629, #3627, #3616, #3544 - code: #3885 (`optim_cpu_offload`), #3739 (`full_offload` validators), #3724, #3719, #3699, #3610, #3487 (`fsdp_cpu_offload` / `enable_gloo`) Test plan - `pytest tests/unit/orchestrator tests/unit/test_configs.py tests/unit/train/test_cpu_adam.py tests/unit/train/test_state_offload.py -m "not gpu"`: 336 passed, 1 skipped (Qwen3-VL snapshot not cached on this host), 4 deselected (GPU; this host has no GPU). The existing offload config tests are updated to the new field. `test_load_configs` covers the 4 edited tomls. - `cli(TrainerConfig|SFTConfig, ...)`: `--model.offload full|none` and the default (`optimizer`) resolve as expected; `optim_cpu_offload` and `offload = "fsdp"` are rejected. - `import prime_rl.trainer.rl.train, prime_rl.trainer.sft.train` works. - `ruff check` and `ruff format --check` on the changed Python files: clean. - Not run: GPU or multi-node training with `offload = "full"`.  --- > [!NOTE] > **Medium Risk** > Breaking TOML/API surface touches distributed training memory paths (FSDP offload, optimizer-in-backward, Gloo); behavior should match prior modes when mapped correctly, but mis-migrated configs will fail at parse time. > > **Overview** > **Breaking config change:** `model.fsdp_cpu_offload`, `model.optim_cpu_offload`, and `model.full_offload` are replaced by a single **`model.offload`** setting (`"optimizer"` default, or `"full"`, `"fsdp"`, `"none"`). RL and SFT share the same field; example TOMLs and scaling docs are updated accordingly. > > The mutual-exclusion validator between the old flags is removed because only one mode can be expressed. **`[model.full_offload]` / `numa_bind` is dropped** — `"full"` always runs NUMA binding and CPU optimizer thread setup via parameterless `setup_full_cpu_optimizer_offload()`. Runtime paths (`setup_optimizer`, FSDP CPU offload in `model.py`, Gloo when offload is `"full"` or `"fsdp"`) key off the new enum; `FullCPUOffloadOptimizer` no longer takes an offload config object. > > Default training behavior is unchanged: optimizer-state CPU offload remains the default (`"optimizer"`). Old TOML keys are not aliased and should fail validation. > > <sup>Reviewed by [Cursor Bugbot](https://cursor.com/bugbot) for commit 35da73a. Bugbot is set up for automated code reviews on this repo. Configure [here](https://www.cursor.com/dashboard/bugbot).</sup>  --------- Co-authored-by: sami jaghouar <sami.jaghouar@gmail.com>

… example - examples/advanced/qwen3-235b-a22b/sft-max-perf.toml: 4 x 8 B300, seq 16k, EP 8, full AC, fullgraph, fp32 AdamW, no offloading, balanced routing. Measured 4.39 s/step, 46.6% MFU. Ported from #3765 to the current schema (no `model.impl`, `offload = "none"`, no cluster-specific paths). - docs/performance-trainer.md: how MFU is computed, the full setup and numbers for the Qwen3-235B run, and the GLM-5.3 SFT before/after for the cuDNN + FlashMLA sparse attention backend (#3829). - Linked from docs nav, docs/scaling.md#benchmarking, and examples/README.md. Co-Authored-By: Mario Sieg <mario.sieg.64@gmail.com> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

…r perf example The defaults are tuned for memory (optimizer offload, AC offloading, compile without fullgraph). Document the settings that trade that memory back for throughput, with two examples. - docs/scaling.md: new "Max-Throughput Recipe" section after the memory-tight one. - examples/advanced/qwen3-235b-a22b/sft-max-perf.toml (+ README): 4 x 8 B300, seq 16k, EP 8, full AC, fullgraph, fp32 AdamW, no offloading, balanced routing. Ported from #3765 to the current schema. Co-Authored-By: Mario Sieg <mario.sieg.64@gmail.com> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

@MarioSieg

…r perf example (#4046) The trainer defaults are tuned for memory (optimizer CPU offload, AC offloading, compile without `fullgraph`) so large models fit on few GPUs. This documents how to trade that memory back for throughput, with an example to reproduce it. - **`docs/scaling.md`**: new "Max-Throughput Recipe" section after the memory-tight one: `offload = "none"`, `ac_offloading = "None"`, `fullgraph = true`, bf16 router, add nodes instead of CP, one full-length sequence per GPU. Points to the two examples below. - **`examples/advanced/qwen3-235b-a22b/sft-max-perf.toml`** (+ README): Qwen3-235B-A22B SFT on 4 x 8 B300 at 16k with that recipe (EP 8, full AC, fp32 AdamW, balanced routing, 20 steps), around 46% MFU. - The existing GLM-5.3 `16k-12-node.toml` overlay is referenced as the second example. Config ported from @MarioSieg's #3765 to the current schema (no `model.impl`, `offload = "none"`, no cluster-specific paths); supersedes it. Validated with `tests/unit/test_configs.py`. 🤖 Generated with [Claude Code](https://claude.com/claude-code)  --- > [!NOTE] > **Low Risk** > Documentation and example TOML only; no trainer or runtime code changes. > > **Overview** > Adds a **Max Trainer Throughput Recipe** to `docs/scaling.md` as the counterpart to the memory-tight section: when GPU memory allows, disable optimizer and activation offloading, enable `torch.compile` with `fullgraph`, use bf16 MoE routing, prefer scaling nodes over context parallelism, and keep `micro_batch_size = 1` per full-length sequence. The doc calls out RL’s fp32 router default as a caveat for `fullgraph`. > > Ships a new **`examples/advanced/qwen3-235b-a22b/`** reference run (`sft-max-perf.toml` + README): 20-step SFT on `Qwen3-235B-A22B` across 4×8 B300 at 16k with that recipe (~46% MFU), balanced routing for stable benchmarking, and no checkpoint block. **`examples/README.md`** lists the example; the scaling doc also points at the existing GLM-5.3 `16k-12-node.toml` overlay. > > <sup>Reviewed by [Cursor Bugbot](https://cursor.com/bugbot) for commit de9878c. Bugbot is set up for automated code reviews on this repo. Configure [here](https://www.cursor.com/dashboard/bugbot).</sup>  --------- Co-authored-by: Mario Sieg <mario.sieg.64@gmail.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

faresobeid mentioned this pull request Oct 10, 2026

refactor(configs)!: one model.offload field for CPU offloading #3988

Merged

samsja mentioned this pull request Oct 12, 2026

docs(scaling): max-throughput recipe + Qwen3-235B-A22B SFT max trainer perf example #4046

Merged

Navigation Menu

Repository navigation

chore(configs): add the Qwen3-235B SFT benchmark config matched to torchtitan - #3765

chore(configs): add the Qwen3-235B SFT benchmark config matched to torchtitan#3765
MarioSieg wants to merge 1 commit into
mainfrom
mario/optimize-benchmark-config

MarioSieg commented Oct 3, 2026

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Conversation

MarioSieg commented Oct 3, 2026

What's in it

Measured (2026-10-02, same 4 nodes, back to back)

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant