feat(deepseek_v4): fp8 RL by garrett361 · Pull Request #3627 · PrimeIntellect-ai/prime-rl

garrett361 · 2026-09-22T13:00:50Z

Experimental FP8 RL branch for DeepSeek V4.

Builds on:

fix(trainer): match vLLM's FP8 weight cast and remove the DeepGEMM MoE wrapper overhead #3647
feat(trainer): run DeepSeek V4's grouped attention output projection in FP8, matching vLLM #3955
feat(trainer): match vLLM's FP8 KV cache for DeepSeek V4 #3956

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>

## What - `VLLM_CACHE_ROOT=/tmp/.vllm-cache-<uid>` and `FLASHINFER_WORKSPACE_BASE=/tmp/.flashinfer-cache-<uid>` are added to `DEFAULT_INFERENCE_ENV_VARS`. `TRITON_CACHE_DIR` got the same treatment in #3877. - The by-hand copies (`/tmp/.vllm-cache`, `/tmp/.flashinfer-cache`) are removed from the four GLM-4.5-Air recipes. Those paths were shared by every user on a node. - docs/scaling.md gets one paragraph that lists the node-local cache defaults. The paths are dot-prefixed so the SLURM start-up cleanup (`rm -rf /dev/shm/vllm-* /dev/shm/vllm_* /tmp/vllm-* /tmp/vllm_* /tmp/torch-* /tmp/torchelastic_*` in the multi_node_rl, multi_node_sft and inference sbatch templates) does not wipe them. These are the same names the recipes used before. Precedence is the same as for the other launcher defaults: `[env_vars]` or `[inference.env_vars]` override them; a value exported in the shell does not. ## Why - Without these, vLLM's AOT/Inductor compile cache lives in `~/.cache/vllm` and FlashInfer's JIT cache under `$HOME`, which is usually the shared filesystem. - Measured while testing #3877 (GLM-4.5-Air TP8, vLLM 0.31): a crashed start left a partial shared vLLM cache. On the next start, 4 TP ranks loaded those artifacts and the other 4 recompiled. The ranks diverged and hung in a collective. So one crash breaks every later start that uses that cache. - The original scaleswe recipe (#2862) moved this cache to `/tmp` because the NFS cache hung inference startup. - Cost: the first start on each new node has a cold vLLM cache. GLM-4.5-Air: 465 s cold vs 375 s with a warm vLLM cache, so about 1.5 min once per node. ## Overlap - #3627 (garrett361, draft) adds `VLLM_CACHE_ROOT`/`FLASHINFER_WORKSPACE_BASE` under `/tmp/garrett/` to new DeepSeek-V4 recipes. With this PR those lines can go. ## Test plan - `pytest tests/unit/orchestrator tests/unit/test_configs.py tests/unit/utils -m "not gpu"`: 365 passed (rebased on origin/main 41a8ed0). - Ran the cleanup globs in a scratch dir that held `.vllm-cache-<uid>`, `.flashinfer-cache-<uid>` and, as a control, `vllm-cache-<uid>`. Only the control matched. - `ruff check` and `ruff format --check` on `process.py`: clean. - `rl @ examples/advanced/glm-4.5-air/swe.toml --dry-run`: - The rendered `rl.sbatch` exports both vars only in the inference block (`/tmp/.vllm-cache-<uid>`, `/tmp/.flashinfer-cache-<uid>`). - With `--inference.env-vars '{"VLLM_CACHE_ROOT": "/scratch/vc"}'`, the sbatch exports `/scratch/vc`.  --- > [!NOTE] > **Low Risk** > Launcher default env vars for inference only; configs can still override, with no auth or training logic changes. > > **Overview** > Inference launches now default **vLLM** and **FlashInfer** compile/JIT caches to per-user paths under `/tmp` (`VLLM_CACHE_ROOT`, `FLASHINFER_WORKSPACE_BASE`), matching the existing Triton cache behavior and avoiding shared `$HOME` or NFS caches that can leave tensor-parallel ranks on mismatched artifacts and hang. > > The four GLM-4.5-Air example recipes drop redundant `[inference.env_vars]` blocks that duplicated those settings. **docs/scaling.md** documents the launcher defaults and notes overrides via `[env_vars]` or component `env_vars`. > > <sup>Reviewed by [Cursor Bugbot](https://cursor.com/bugbot) for commit fe8a34d. Bugbot is set up for automated code reviews on this repo. Configure [here](https://www.cursor.com/dashboard/bugbot).</sup>

Squash of #3627 at d35a408, with its stacked #3647, #3955 and #3956. Its DeepSeek V4 C128 and attn_sink patches predate #4040 on main, so main's versions are kept and the duplicates dropped. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

…FP8 weight scales Promotes the PRIME_DIAG_UE8M0_WEIGHTS diagnostic to a config field. For dequantized FP8 releases the weights already sit on the e4m3 grid under power-of-two block scales, so quantizing with those scales is exact where amax / 448 injects full e4m3 rounding error. Measured against a bf16 reference on 203k SWE tokens, combined with VLLM_USE_DEEP_GEMM_E8M0=1 it gives the lowest mismatch of every FP8 configuration tried (KL 3.5e-3 vs 5.6e-3 for fp32 scales, IPO masked fraction 0.0000 vs 0.0002) and removes the decode-step glitch tokens. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

…ne FP8 weights Round-to-nearest pins a weight that starts on the e4m3 grid until the trainer has moved it half a bin, about 1000 steps at lr 1e-6, so the served policy stays at step 0 while the trainer drifts. Rounding each weight to its far neighbour with probability equal to the fractional distance makes the served weight unbiased and lets it track sub-bin updates in expectation. On-grid input rounds to nearest, so the checkpoint quantizes to identical bytes. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

Ported from the pre-rebuild feat/ds-v4-fp8-rl tip (backup/ds-v4-fp8-rl-2026-10-08) with the schema migrations main needs: drop the removed trainer.model.impl, move group_size and algo under orchestrator.train, and pin trainer.model.freeze_moe_router = false so runs stay comparable with the recorded baselines now that main resolves the RL default to true. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

On current main, compiled DeepSeek V4 decoder blocks fail on the second step's backward inside activation-checkpoint recompute: Dynamo config overrides are thread-local, so the raised recompile limit is invisible on the autograd thread, and the layers' structural differences exhaust the default limit there. Passing the limit per region removes the crash but raises mismatch KL about 30x from the second step on (0.046 against 0.0015 on the reverse-text probe), so these configs disable compile until the compiled path is fixed. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

…read `model.recompile_limit` (default 64) now sets dynamo's recompile limit at model setup, replacing the import-time global, and is also passed to each compiled transformer block's `compile` call. Dynamo config is thread-local, so the global only reached the main thread. Activation checkpointing recompiles blocks during backward on the autograd thread, which saw PyTorch's default of 8. On DeepSeek V4 the shared decoder layer code object reaches about 18 cache entries (one per layer variant plus the dynamic-shape recompile), so recompute hit the limit, fell back to eager, and crashed at step 2. A per-region limit applies on any thread. The removed `cache_size_limit = 64` line is an alias of `recompile_limit`, so the effective default is unchanged at 64. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

…onfigs The FP8 trainer now has a grouped FP8 path for o_a_proj that matches vLLM's wo_a, so drop it from the FP8 ignore_patterns in swe-fp8-fp8.toml and rl_fp8_fp8.toml. The rl_fp8_fp8 probe also stops keeping vLLM's wo_a in bf16, so both sides run the projection in FP8. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

…ner feature on swe-fp8-full.toml runs scaleswe at 131k like swe-fp8-fp8.toml, with a compiled trainer (not fullgraph, the fp32 router is its own FSDP unit), the grouped attention output projection in FP8, the FP8 KV cache simulation and no optimizer CPU offload. It defaults to 12 nodes, 8 trainer plus 4 one-node inference replicas with 64 rollouts in flight per replica, and its header notes the 16-node setting. Memory without optimizer offload and compile at 131k are unproven, so the first launch is a validation run. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

…che paths The five SWE configs hardcoded one user's /tmp and home paths in [env_vars]. Top-level env vars are emitted as double-quoted shell exports in the RL sbatch, so $USER and $HOME expand on the node. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

…FP8 config swe-fp8-full.toml now runs with no CPU offloading of any kind. Activation offloading defaulted on and ran synchronously (use_streams is hardcoded off), so it only added CPU round trips. The first launch measured a steady-state peak of 92-116 GiB per GPU with it on, and keeping the checkpointed block inputs on the GPU should add roughly 11 GiB, so the peak needs rechecking. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

…eek V4 SWE FP8 config The pre-run command deleted every sandbox carrying the config's label, so a second concurrent run of swe-fp8-full.toml would delete the first run's live sandboxes. Orphaned sandboxes expire on the runtime's idle timeout instead. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

With activation offload off, cp 4 peaked at 133 GiB on the worst of 64 ranks (117 GiB with offload on) and OOMed in a step-7 backward. cp 8 halves each rank's tokens to 16k. ep 8 still divides the mesh, and each context-parallel ring is one node, so ring traffic stays on NVLink. Data-parallel workers drop from 16 to 8. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

…64-step off-policy limit The 12-node run (8 trainer, 4 inference) spent ~72% of each 233 s step waiting on rollouts. Eight replicas should roughly halve the step. Long multi-turn episodes then span more than 32 steps, so max_off_policy_steps rises to 64 to keep them, and max_inflight stays at 64 per replica (512). Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

garrett361 force-pushed the feat/ds-v4-fp8-rl branch 2 times, most recently from a400152 to c76a621 Compare September 22, 2026 14:01

mikasenghaas mentioned this pull request Sep 22, 2026

fix(inference): make online FP8 weight updates safe and unbiased #3628

Closed

garrett361 force-pushed the feat/ds-v4-fp8-rl branch from c76a621 to 4348fc7 Compare September 23, 2026 15:40

garrett361 force-pushed the feat/ds-v4-fp8-rl branch 6 times, most recently from f416ea0 to 4e3b18d Compare October 10, 2026 01:26

This was referenced Oct 10, 2026

chore(trainer)!: remove ac_offloading.pin_memory #3972

Merged

fix(launch): node-local vLLM and FlashInfer caches by default #3981

Merged

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

Merged

garrett361 force-pushed the feat/ds-v4-fp8-rl branch from 4e3b18d to d35a408 Compare October 12, 2026 15:05

nxphi47 mentioned this pull request Oct 13, 2026

feat(trainer): add FP8 QAT with shared stochastic rounding #4073

Open

samsja mentioned this pull request Oct 14, 2026

feat: MXFP8 RL fixes (vLLM MoE reload, 32-token padding) and DeepSeek-V4 MXFP8 recipe #4098

Merged

garrett361 and others added 11 commits October 14, 2026 20:09

garrett361 force-pushed the feat/ds-v4-fp8-rl branch from d35a408 to c57822b Compare October 15, 2026 02:44

Navigation Menu

Repository navigation

feat(deepseek_v4): fp8 RL - #3627

feat(deepseek_v4): fp8 RL#3627
garrett361 wants to merge 12 commits into
mainfrom
feat/ds-v4-fp8-rl

garrett361 commented Sep 22, 2026 •

edited

Loading

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Conversation

garrett361 commented Sep 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

garrett361 commented Sep 22, 2026 •

edited

Loading