Repository navigation
docs: performance.md with the Qwen3-235B numbers vs torchtitan and the max-MFU benchmark configs - #3766
Draft
MarioSieg wants to merge 19 commits into
Draft
docs: performance.md with the Qwen3-235B numbers vs torchtitan and the max-MFU benchmark configs#3766MarioSieg wants to merge 19 commits into
MarioSieg wants to merge 19 commits into
Conversation
Add a mega_moe MoE dispatch config and MegaMoeTokenDispatcher that runs DeepGEMM's fused dispatch + SwiGLU expert compute + combine kernels (forward and backward) over the expert-parallel group. Includes a benchmark against the bf16 TorchTokenDispatcher path.
…o/overlap-moe-kernel-integration Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Mega MoE is now a `MegaMoEExpertCompute` installed on `GroupedExperts.compute` by `configure_moe_runtime`, like the other expert compute backends, instead of a token dispatcher that injected a `fused` callback into `GroupedExperts.forward`. Because the kernel fuses dispatch and combine, it consumes undispatched tokens and their routing: `FusedDispatchExpertCompute` describes such a compute, and `FusedTokenDispatcher` hands the tokens to it without communicating. `GroupedExperts.forward(x, *routing)` forwards either `num_tokens_per_expert` or `top_scores, selected_experts_indices` to its compute. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
…e extra Mega MoE kernels come from prime-mega-moe, a closed-source DeepGEMM fork. Its wheel is built by hand and attached to the prime-kernels release; prime-rl pins the URL and never builds it from source. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
SonicMoE landed on main squashed (#3535) with an ExpertCompute.validate hook and GroupedExperts.set_compute. Mega MoE adopts both: MegaMoEExpertCompute.validate carries the expert-structure checks and configure_moe_runtime installs it through set_compute. uv.lock was regenerated with the previous deep_gemm wheel URL because the prime-mega-moe wheel is not uploaded yet; re-run uv lock once it is. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
…ck compiles with fullgraph The fused forward was a custom op wrapped in a Python autograd.Function that saved x.to(bfloat16).contiguous() for backward. When the block compiles with fullgraph under activation checkpointing, dynamo rejects that saved tensor because it aliases the autograd.Function's input inside the checkpoint subgraph. Make the backward a custom op too and register autograd on prime_rl::mega_moe_forward, saving the op's own inputs. Both ops take the caller's dtypes, run the kernels in bf16, and return gradients in the inputs' dtypes, so the expert dW dtype is unchanged.
…n and the max-MFU benchmark configs docs/performance.md reports the matched prime-rl vs torchtitan SFT benchmark on Qwen3-235B-A22B (4 nodes, seq 16k): torch dispatch 4.39 s/step vs torchtitan 5.04 s/step, Mega MoE 4.10 s/step with full AC and 3.93 s/step with selective AC, the shared MFU formula and its caveat, and the per-step trace breakdown. configs/benchmark/qwen3-235b-a22b/sft-4n-16k-mega-moe.toml is the highest-MFU measured config. configs/benchmark/qwen3-30b-a3b/sft-1n-16k-mega-moe.toml scales it to one node and is marked as not yet measured.
samsja
added a commit
that referenced
this pull request
Oct 10, 2026
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"`. <!-- CURSOR_SUMMARY --> --- > [!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> <!-- /CURSOR_SUMMARY --> --------- Co-authored-by: sami jaghouar <sami.jaghouar@gmail.com>
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Stacked on #3651 because the Mega MoE configs need the
mega_moedispatch. Companion to the torch-dispatch benchmark config in the main-based PR (mario/optimize-benchmark-config).What's in it
docs/performance.md(and a nav entry): the Qwen3-235B-A22B SFT table vs torchtitan from the 2026-10-02 back-to-back runs, the shared MFU formula with its dense-attention caveat, the real-routing number, the per-step trace breakdown, and reproduction commands.configs/benchmark/qwen3-235b-a22b/sft-4n-16k-mega-moe.toml: the highest-MFU measured config (Mega MoE, selective AC, fullgraph): 3.93 s/step, 52.0% MFU.--model.ac.mode fullgives the 4.10 s/step full-AC row.configs/benchmark/qwen3-30b-a3b/sft-1n-16k-mega-moe.toml: the same setup scaled to one 8-GPU node. Not measured yet; the config header and the docs say so. The existing 30B runs on disk predate the fullgraph work and were never tuned, so there is no honest max-MFU number for it until one run.Table
Needs #3753 and #3754 under #3651 for
fullgraph=trueto compile, and the Mega MoE autograd fix already on this branch.