Skip to content

Add PPO and SAO value-model training - #3699

Closed
nxphi47 wants to merge 10 commits into
mainfrom
feat/ppo
Closed

nxphi47 wants to merge 10 commits into
mainfrom
feat/ppo

Conversation

@nxphi47

@nxphi47 nxphi47 commented Sep 27, 2026 •

Copy link
Copy Markdown
Contributor

Summary

Add PPO and SAO with a separately trained token-value critic. The value head is a linear projection of final hidden states sized from lm_head.in_features; no model-name guard restricts the workstream to Qwen3. GRPO keeps its existing path when [value] is absent.

Data flow

  1. TrainSink converts the admitted traces in a completed group to original TrainingSamples and calls PPOAlgorithm.score_samples once. It sends one batched_token_ids HTTP request containing every trainable branch in that group.
  2. The critic packs complete token sequences up to its own value.model.seq_len, resets positions at each sequence boundary, balances packed bins across its DP replicas, and runs equal numbers of CP-sharded forwards on each replica. It returns values and bootstrap values in original branch order. PPO/SAO calculate skip-observation GAE and annotate the original samples.
  3. The orchestrator packs those samples twice: once for the policy's sequence length, CP, and DP, and once for the critic's. Independent ZMQ senders publish one grid item to each respective data rank. The critic receives and trains first, including masked dummy microbatches when its grid needs DP padding; the orchestrator waits for /status before sending the policy grid.
  4. The policy uses clipped PPO or SAO's direct-ratio mask loss. The critic minimizes masked value error with optional head-only warmup, JSONL pretraining, attention freezing, LoRA, and policy-backbone snapshot/retuning. Optimizers and checkpoints are separate.

Resources and supported layouts

  • Single-node 30B math comparisons allocate 2 policy, 2 critic, and 4 inference H200 GPUs; GRPO uses 2 policy and 6 inference GPUs. The full 30B policy and critic use CPU optimizer offload in this layout.
  • Multi-node SLURM can reserve deployment.num_value_nodes dedicated critic nodes or deployment.num_value_gpus_per_train_node GPUs on every policy node. Shared nodes start separate torchrun groups on disjoint GPUs. At launch, the binding node probes free rendezvous ports on the target host address plus loopback, excluding configured broadcast, service, and ZMQ ports; a shared node chooses both ports together, and failed probes retry. The probe no longer binds all network interfaces. The regular inference router waits for every worker's /v1/models response before discovering models, so slow workers are not registered as unknown. Its health-probe timeout is 30 seconds to tolerate busy long-context generation. The launcher connects ZMQ and HTTP hosts and propagates failures. Shared, dedicated, and no-critic rendered scripts pass bash -n. A live two-node SAO run passed batched scoring, two critic updates, a full policy optimizer step, and updated-weight broadcast.
  • Policy and critic derive DP from their own GPU count and CP degree. They may use different DP, CP, EP, and GPU counts when backbone sync is off. Snapshot sync currently requires matching CP/EP and GPU counts. The critic batch transport is ZMQ. The value service accepts token IDs, so VLM scoring needs a multimodal request path. Compatible text backbones have no Qwen3-only gate; only Qwen3 MoE has received a distributed real-model smoke so far.

Validation

  • Focused orchestrator/config/value tests: 70 passed, including batched versus per-sample score/GAE parity, HTTP batch and legacy request shapes, critic packing, masked dummies, multi-node DP configuration, and rendezvous port selection.
  • Four-GPU dummy-model runs passed at critic (CP, DP) = (1, 4), (2, 2), (4, 1), including sparse and bursty scoring, unequal policy/critic DP ZMQ traffic, and one critic update from received data.
  • A small real Qwen3 MoE critic with FSDP passed mixed-length packed versus individual score parity and a masked-dummy critic update at CP=2/DP=2 and CP=4/DP=1 (Ulysses). These use random small weights, not the 30B pretrained checkpoint.
  • Earlier short local 30B Qwen3 MoE math PPO/SAO and hosted R2E-Gym PPO/SAO updates completed, and a tiny LoRA critic sync/resume passed. Port selection skips occupied and reserved sockets; loopback and nonloopback direct probes and the shared, dedicated, and no-critic SLURM dry runs pass. A rendered router function retried a delayed dummy backend and launched only when both /v1/models endpoints were ready. Ruff and diff checks pass. Non-Qwen3 text backbones, VLM scoring, and long-run resume remain unverified.

Ongoing 16K math comparison (2026-09-27 08:00 UTC)

Jobs PPO 1579, SAO 1580, and GRPO 1581 were running toward 300 steps at this check. Confirmed completed policy steps: 15 PPO, 12 SAO, 13 GRPO. They use 16,384-token context, a 15,000-token completion cap, batch 32, and group size 8. PPO and SAO use the paper's length-adaptive policy GAE (alpha=1.5) and value_lambda=1. Their latest completed-step effective correctness was 0.28125 PPO, 0.09375 SAO, and 0.28125 GRPO; raw all correctness was 0.1387, 0.1706, and 0.1174. These cohorts have different admission rules, so these figures do not establish relative model quality. The runs started before the multi-node change and have no AIME evaluation. No live job was changed for this PR.

Two-node 32K SAO launch (2026-09-27 22:09 UTC)

At the user's request, jobs 1579 and 1580 were cancelled after local tests and the two-node launch checks passed. Job 1618 used a standalone SAO variant of math32k_sl1node_off8.toml: 400 steps, batch 64, group 8, 31K completion cap, one eight-GPU inference node and one node split 4 policy/4 critic GPUs. It loaded both 30B models but failed when critic rendezvous took the policy's NCCL broadcast port. Job 1621 then passed startup and policy weight broadcast but had 30/30 AIME eval errors because aime25 requires a network policy unsupported by the copied subprocess runtime. Its 128 long concurrent rollouts also made both vLLM API servers miss health checks; the router marked them unhealthy and new math requests failed. Job 1621 was stopped before any scored batch or optimizer update. Replacement job 1626 is running on nodes 036/024 with automatic rendezvous port discovery (29561 and 29957), Prime runtime for AIME, and 32 concurrent rollouts. AIME step-0 eval finished with reward 0.70 and no errors. Math batch 1 scored with reward 0.265625 and 17/64 trainable samples; critic step 1 loss was 0.06884 in 227 seconds. Batch 2 reward was 0.40625 with 64/64 trainable samples; critic step 2 loss was 0.07975 in 773.6 seconds. Policy step 1 completed with loss -0.0011 and v1 broadcast applied in 48.7 seconds; policy step 2 has started. Both replicas occasionally missed the running router's five-second probes under load but continued generating; the new 30-second timeout applies only to future launches. The earlier GRPO job was not touched.

Comment thread src/prime_rl/utils/ports.py Fixed
Comment thread tests/unit/test_configs.py Fixed
@nxphi47
nxphi47 marked this pull request as ready for review September 28, 2026 18:22
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>
@faresobeid faresobeid closed this Nov 3, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants