Value-model (critic) support for the tinker-compat API via multi-LoRA (SAO, arXiv:2607.07508)

[qywu]

Motivation

SAO (Single-rollout Asynchronous Optimization, arXiv:2607.07508) — the RL method behind GLM-5.2/5.3's agentic pipeline — replaces GRPO group sampling with one rollout per prompt, stabilized by a trained value model (token-level GAE, critic updated K× per policy step, frozen-attention critic). Its DIS clipping is already expressible with our existing policy_loss/cispo losses; the missing piece in xorl is any notion of a value function: no value head, no value loss, no way to get V(s_t) out of a forward.

Our multi-LoRA server is an unusually good fit for the rest of SAO:

• the critic is just a second LoRA session sharing the frozen base model (no PPO memory doubling);
• "update the critic 2× per policy step" is client-side orchestration over the existing tinker-compat API;
• the paper's frozen-attention critic maps onto per-substrate target_modules (follow-up: per-session masking).

Design (v1)

The value head is a LoraLinear(hidden_size, 1) with a zero, frozen base weight, attached as model.value_head. A scalar head is rank-1 by definition, so the LoRA factorization V(s) = B·A·h·(α/r) loses no expressivity, and — because its params are literally named value_head.lora_A / value_head.lora_B — the entire multi-adapter machinery picks it up with no manager surgery: per-session copies, per-session optimizer state, tensor layouts, rank slicing, deterministic init (B=0 ⇒ V≡0 at init), gradient ownership compile, and save_state/load_state round-trips.

The head is consumed exactly like lm_head: the loss reads its straight-through folded weight (_get_effective_lm_head_weight_for), so gradients reach the adapter factors through the same DIRECT_OUTPUT_PROJECTION ownership lane lm-head LoRA already uses (direct = module is lm_head extends to module is value_head).

Two new server losses, dispatched through the existing elif chain in ModelRunner._compute_micro_batch_loss:

• value_loss: masked squared error of V(h_t) against per-token returns (optional PPO-style value clipping vs old_values), returning raw masked sums per the TokenPartial reducer contract; per-token values in the logprobs per-token channel, per-token squared errors in elementwise_loss.
• value_prediction: forward-only; returns per-token values (use with the existing no-grad forward op so the client can compute GAE).

No new API endpoints and no request-schema changes: a critic is created with the existing create_lora_training_client, trained with forward_backward(..., loss_fn="value_loss") where datums carry target_tokens + weights (action mask) + returns in loss_fn_inputs, and queried with forward(..., loss_fn="value_prediction").

Presence asymmetry: policy-loss steps never touch the value head, so value_head.* params are declared GradientPresencePolicy.AUTHORIZED_ZERO (everything else stays REQUIRED_IF_ACTIVE).

Client-side loop (target shape)

policy = svc.create_lora_training_client(base, rank=32, model_id="policy")
critic = svc.create_lora_training_client(base, rank=64, model_id="critic")

values = critic.forward(datums, loss_fn="value_prediction").result()   # V(s_t) per token
adv, ret = compute_skip_observation_gae(rewards, values, action_mask, gamma, lam)
for _ in range(K):                                                      # faster value update
    critic.forward_backward(with_returns(datums, ret), loss_fn="value_loss").result()
    critic.optim_step(critic_adam).result()
policy.forward_backward(with_adv(datums, adv, rollout_logprobs), loss_fn="policy_loss").result()
policy.optim_step(policy_adam).result()

v1 restrictions (validated at server startup)

enable_value_head=true requires:

• enable_lora=true (full-weights mode: follow-up),
• pipeline_parallel_size == 1 (the head is not in any PP terminal objective yet; PP already rejects unknown objectives),
• lm_head not in the LoRA targets (train_unembed=false or an explicit target list without lm_head): a value_loss backward produces no lm-head-LoRA grads and would trip REQUIRED_IF_ACTIVE,
• plain LoRA (no QLoRA / exact-contract lm-head lanes / fsdp_sharded_lm_head_loss / lm-head TP).

Work items

PR 1 (server, this repo) — see linked draft PR:

• ops/loss/value_loss.py: value_loss_function + value_prediction_function (+ LOSS_REGISTRY)
• model_builder.build_training_model(enable_value_head=...): create model.value_head after LoRA injection; zero+freeze base weight (re-zeroed post-load in the runner since it is absent from base checkpoints)
• torch_parallelize: value head in its own FSDP unit whose forward never runs (factors stay sharded DTensors; the stay-gathered grouping broke layout validation after forward-only ops — see PR Add scalar LoRA value-head (critic) support for SAO-style RL #85)
• ModelRunner: dispatch branches, _LOSS_EXCLUDE_KEYS entries (returns, old_values, …), effective-weight helper, direct classification for value_head, AUTHORIZED_ZERO presence, startup validation, post-load base-weight zeroing
• packing: returns/old_values join CAUSAL_TARGET_ALIGNED_FIELDS (HF-format shift; tinker-format flows through the generic seq-field loop)
• sampler export: exclude value_head.* from save_weights_for_sampler PEFT adapters (SGLang must never see it); keep it in save_state training checkpoints
• ServerArguments.enable_value_head + lora-config export + validation
• xorl.rl.advantages.compute_skip_observation_gae reference implementation (pure Python, paper Eq. 4–5)
• CPU tests: loss math/masking/reducer/clipping/grad-flow, packer alignment, GAE

Follow-ups (separate PRs):

• xorl-client: returns/old_values in Datum._KEY_TO_DTYPE, ship the GAE helper, docs + example loop
• GPU validation on 4×H100: critic session end-to-end (register → value_loss → optim_step → value_prediction → save/load_state), policy+critic interleaving, eviction/re-registration with value head
• Dedicated values field in LossFnOutput (wire + client) instead of reusing the logprobs per-token channel
• Per-session target-module masking (paper's frozen-attention critic without constraining the policy session)
• PP support (terminal objective + output_fqns), full-weights mode
• docs/server-training: value-model training guide; SAO recipe example
• Explained-variance tracking across a training run (the paper's key critic diagnostic)

Risks

• LoRA-rank critic capacity is the open scientific question; the paper's own ablations show critic quality gates single-rollout stability. Needs an empirical run (watch explained variance).
• The packer maps advantages == 0.0 → IGNORE_INDEX; returns deliberately does not reuse that convention — masking comes from weights/target_tokens only.

[qywu]

Status (2026-08-24): everything except PP support and full-weights mode is now implemented and validated.

PR #85 (now ready for review, was draft pending GPU validation) grew to cover:

• the dedicated state_values wire field (no more logprobs-channel reuse),
• per-session frozen_module_patterns — the paper's frozen-attention critic, verified to keep attention factors at exactly zero delta on disk while MLP + value head train,
• xorl.rl.explained_variance derived from the value-loss moment metrics (which finalize to true global means),
• a docs guide (server-training/value-model) and a runnable SAO example (examples/server/sao_critic/),
• a pre-existing tinker-compat fix: session normalization now accepts the SDK's default dropout: 0.0 (previously every create_lora_training_client call with a LoRA config was rejected).

GPU validation ran on 2×H100 and caught two real bugs, both fixed in PR #85's second commit:

1. the stay-gathered [norm, lm_head, value_head] FSDP grouping broke adapter layout validation after forward-only ops (FSDP2 reshard() cannot restore a group materialized by a no-grad pass) → the head now lives in its own never-unsharded unit and the loss consumes the folded delta via full_tensor();
2. per-rank factor-shard gradients were incomplete (~14% off analytic references) because each rank can only differentiate its own shard → the upstream weight gradient is now all-reduce-summed before reaching the factors; staged numerators now match analytic references exactly.

Validation artifacts: tests/server/runner/test_value_head_adapter_gpu.py (real-FSDP certification with analytic gradient references, incl. forward-only regression, AUTHORIZED_ZERO, rank slicing, frozen-session steps) and a live Qwen3-8B E2E (fresh critic predicts exact zeros → value loss 0.158→0.071 over 8 steps → client GAE over live state_values → sampler export excludes the head / save_state keeps it).

Client: togethercomputer/xorl-client#15 (state_values, dtype hints, GAE/EV helpers, frozen_module_patterns kwarg). Merge it with PR #85; the server tolerates old clients. The submodule pointer bump lands after that PR merges.

Still open (last two checkboxes): PP support (terminal objective + output_fqns) and full-weights mode — both intentionally out of scope for the LoRA-first v1; the server rejects those configs with clear errors. The remaining scientific question (LoRA-rank critic capacity) needs a real training run — watch explained_variance.

[qywu]

Real training results for the last open question (LoRA-rank critic capacity), run against a live server on the PR #85 branch — Qwen3-8B, 2×H100, rank-16 critics, tinker-compat client end-to-end.

A — critic capacity (frozen-attention vs full-adapter critic). Trajectories are natural sentences whose 0/1 terminal reward is determined by an early topic word; returns are per-token. Held-out explained variance (48 unseen trajectories):

step	0	5	10	15	20	30	59
frozen-attention	−1.06	−0.10	+0.81	+0.99	+1.00	+1.00	+1.00
full-adapter	−1.05	+0.04	+0.82	+0.97	+0.99	+1.00	+1.00

Value loss fell 0.2657 → 0.00004 (÷6,600) in 60 steps at lr 1e-3; no train/holdout gap. EV = −1 at init is itself a correctness check (a zero-predicting head on 0/1 returns gives exactly 1 − 0.5/0.25) — both critics landed there, confirming the zero-base head contract over the wire. The frozen_module_patterns critic matches the full critic while training fewer parameters.

B — full SAO loop (critic value_prediction → skip-observation GAE → K=2 value_loss steps → on-policy importance_sampling step, 40 rounds, 16 trajectories/round). The policy's held-out reward gap (mean action-token logprob on reward-1 minus reward-0 trajectories) went −0.135 → +0.066 (+0.20 nats), crossing zero — the policy now prefers rewarded trajectories it never trained on. The critic stayed calibrated the whole time (holdout EV ≥ 0.994 every round) and advantages stayed small (mean |A| ≈ 0.002–0.03), the signature of a working baseline.

Caveats: deliberately easy value function (answers does the machinery learn real values on a real model — yes, decisively — but cannot separate frozen vs full the way the paper's hard agentic tasks do), single seed, synthetic reward in B. The 128k-trajectory / MoE-backbone regime is still untested.

Also fixed while validating: the docs/example now use the exact wire metric keys (is_value_error_sq_mean:mean etc.), and the cpu (server) CI failure on PR #85 (a schema-shape pin that didn't know frozen_module_patterns) is fixed and green.

[qywu]

Harder settings + genericity + Wordle vs GRPO (all on the PR #85 branch; full curves: see the session's results page).

Hard critic (compositional XOR value under label noise). p = 0.9 if (topic is math) XOR (X>Y) else 0.1, rewards sampled Bernoulli(p), multi-turn trajectories with masked observation spans. Both rank-16 critics escape the mean baseline and partially compute the nonlinear value (held-out EV vs true p peaks ≈ 0.61–0.64 around step 70–100), then overfit the label noise (train loss 0.026 < Bayes floor 0.09 while held-out EV decays); frozen-attention holds up slightly better late. Practical: critic early-stopping and data scale matter, exactly per the paper's value-pretraining section.

Reward flip (paper Fig. 5 in miniature), critic-GAE vs running-mean arm, identical loss/data. The critic re-tracks instantly (EV −1.07 at the flip → +0.91 in 4 rounds → +0.998 in 8, K=2), and its policy re-aligns within ~4 rounds. The running-mean arm diverges (+127 nats!), lags the flip by ~12 rounds, then diverges again — the adaptation-lag result, amplified by the shared unclipped loss that the critic arm survives (small fitted-critic advantages are natural step-size control).

MoE genericity. Same harness on Qwen3-30B-A3B (128 routed experts, EP=2, expert-LoRA frozen-attention critic): fresh-critic zeros over the wire, loss 0.2657 → 0.00001, held-out EV +0.998 in 40 steps, policy coexistence, 37,248-key sampler export with zero value_head keys. This run surfaced and fixed a real bug (now in PR #85): the value-head weight gradient must be all-reduce-summed over its shard mesh, not WORLD — under EP with one data shard, a WORLD sum overcounts by the EP factor.

Wordle, critic single-rollout vs GRPO (no wordle test existed in xorl-client, so this is a tinker-cookbook-style one built on our API: live SGLang multi-turn rollouts, per-step adapter hot-reload, identical clipped policy loss in both arms, only the advantage estimator differs; Qwen3-8B base, 80 steps ≈ 640 rollouts/arm). Learning-inconclusive at this micro-scale (real from-scratch Wordle needs thousands of episodes) but mechanistically sharp: GRPO structurally skipped most updates (zero within-group reward variance — the paper's exact motivation for a critic) and its few updates taught it to emit the prompt's example word ('crane', eval 0.083 ≈ baseline); the critic arm updates on every rollout but sparse terminal rewards make most advantages small negatives, and over 640 updates the policy slowly degrades — SAO counters this with wide DIS clipping (ε_high = 5.0), large batches, and value pretraining. Two diagnostic runs are part of the finding: the deterministic xorl sampler needs per-request sampling_seeds or GRPO groups are byte-identical (zero updates ever), and unclipped IS + whitened sparse advantages collapses the policy outright.

Infra validated along the way: endpoint allowlist (XORL_OUTBOUND_ENDPOINT_ALLOWLIST), dynamic LoRA loading into SGLang every step (~160 adapter reloads without a hiccup), and concurrent multi-server operation (fixed-port collisions on 29500/5556 are worth a launcher follow-up: auto-pick free ports).

[qywu]

K₃ train-inference mismatch check (relates to the mismatch write-up at kiddyboots216.github.io/mismatch), measured on our dense Qwen3-8B pair: exact-contract trainer (ce_mode: bi_fused, fa4, Class-B RoPE — the resolved numerical program confirmed in logs) against SGLang with --rl-on-policy-target xorl (deterministic inference + exact dense numerics confirmed in its banner), with --enable-lora (csgmv backend) for runtime adapters. K₃ from compute_kl_stats on sampler-generated tokens replayed through importance_sampling:

adapter	prompts	K₃	max abs log-ratio
zero-delta	raw	3.3e-4	0.19
zero-delta	chat template	2.1e-4	0.12
trained (3 steps)	raw	1.7e-4	0.10
trained (3 steps)	chat template	9.6e-3	1.54

Reading: not a 0-mismatch reproduction in the runtime-multi-LoRA serving configuration. The base-model path holds ~1e-4 — roughly 100× below the write-up's mismatched-XoRL baseline (K₃ ≈ 0.017) but not zero — and the zero-delta A/B localizes the residual to the LoRA-kernel serving path: once adapter deltas are nonzero, K₃ can reach ~1e-2 with ±1.5-nat token ratios (the trainer consumes the canonically FOLDED weight, bit-matched to a merged export, while SGLang's csgmv kernels apply the delta unmerged at runtime — structurally different arithmetic). Sampling-transform replay was verified clean (meta_info shows temp 1.0 / top-p 1.0 / top-k disabled), and prompt templating is irrelevant to alignment since both engines consume the same raw ids.

Consequences: (1) DIS/TIS/CISPO-style clipping is genuinely load-bearing for runtime-LoRA RL on this stack — consistent with what our Wordle micro-runs showed; (2) a true 0-mismatch reproduction of the write-up's Wordle numbers needs either the merged/full-weight sync path or whatever multi-LoRA layout contract the write-up's runs used — worth pinning down before the big Qwen3.5-35B-A3B reproduction.