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

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 #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.

Review the remaining PP support and full-weights-mode follow-ups, starting with ModelRunner, torch_parallelize, and the existing startup validation. Check the linked draft PR and completed CPU/GPU coverage before changing anything. Done means the remaining restrictions are addressed with matching tests and the value-model path still works end to end.