Add scalar LoRA value-head (critic) support for SAO-style RL

[qywu]

Part of #84 (design + full plan there). Adds value-model (critic) support to the tinker-compat API so SAO-style single-rollout RL (arXiv:2607.07508, the GLM-5.2/5.3 recipe) can run with a critic as a second LoRA session on the shared base model. GPU-validated: 2×H100 FSDP2 certification test with analytic gradient references, plus a live end-to-end Qwen3-8B server run. Client counterpart: togethercomputer/xorl-client#15.

Design in one paragraph

The value head is a LoraLinear(hidden_size, 1) attached as model.value_head with a zero, frozen base weight — a scalar head is rank-1, so the LoRA factorization V(s) = B·A·h·(α/r) loses nothing. Because its parameters are named value_head.lora_A/value_head.lora_B, the multi-adapter manager picks them up with essentially no manager changes: per-session copies, per-session optimizer, tensor layouts, rank slicing, deterministic init (B=0 ⇒ V≡0 at init), and save_state/load_state all compose. The head gets its own FSDP unit whose forward never runs, so its factors stay sharded DTensors at all times; the loss consumes the folded delta via full_tensor() through the existing direct-output-projection ownership lane, with the upstream weight gradient all-reduce-summed across ranks (_SumGradAcrossRanks) so each rank's factor-shard gradients are complete. No new endpoints; the API surface grows by two loss_fn names, two per-token datum fields, one response field, and one per-session LoRA option.

What's included

• ops/loss/value_loss.py — value_loss (masked squared error vs per-token returns, optional PPO-style clipping vs old_values, sum-composable moment metrics) and value_prediction (forward-only per-token V(s_t) for the no-grad forward op → client-side GAE). Raw-sum TokenPartial reducer contract.
• LossFnOutput.state_values — dedicated wire field for per-token values (value losses no longer reuse the logprobs channel).
• Per-session frozen_module_patterns — the paper's frozen-attention critic: matching adapter factors are skipped at gradient staging, declared AUTHORIZED_ZERO, and provably stay at zero delta, without constraining other sessions. Legacy session specs stay byte-identical when the option is absent.
• Model build / FSDP — enable_value_head on build_training_model; head created after LoRA injection in its own never-unsharded FSDP unit; base weight re-zeroed post-load.
• Runner — dispatch branches, _LOSS_EXCLUDE_KEYS, effective-weight helper (delta full_tensor + grad all-reduce), direct classification, AUTHORIZED_ZERO presence for value-head and session-frozen factors, startup validation.
• Packer — returns/old_values join CAUSAL_TARGET_ALIGNED_FIELDS (a 0.0 return does not mask the token).
• Checkpointing — sampler exports drop value_head.* (verified on disk: 504-key adapter, zero value-head keys); training saves keep them so critics resume.
• Session normalization — accepts the SDK's default dropout: 0.0 as a no-op (previously every create_lora_training_client call with an explicit LoRA config was rejected — pre-existing tinker-compat bug).
• xorl.rl — compute_skip_observation_gae (paper Eq. 4–5) and explained_variance (the paper's critic diagnostic, derived from the is_value_*/is_return_* moment metrics, which finalize to true global means).
• Docs + example — docs/server-training/value-model guide; examples/server/sao_critic/ runnable SAO loop.

GPU validation (what the draft was waiting on)

Certification test (tests/server/runner/test_value_head_adapter_gpu.py, 2×H100, real FSDP2 + NCCL, analytic references): ownership classifications (DIRECT_OUTPUT_PROJECTION + AUTHORIZED_ZERO), value-loss step with per-parameter staged-numerator equality for trunk and head, policy-style step on the same session (absent value-head grads accepted, exact-zero transport — no stale-numerator leakage), a forward-only op between training steps (regression for the layout-contract bug below), heterogeneous rank slicing, frozen-trunk session across two steps with factors provably unmoved, cross-session isolation.

Live E2E (Qwen3-8B, 2 GPUs, real server + client SDK): fresh critic predicts exact zeros over the wire → 8 value-loss steps drop the loss 0.158→0.071 and move V(s) toward returns → client GAE over live state_values → frozen-attention critic leaves q_proj.lora_B ≡ 0 on disk while MLP + head train → policy session (importance_sampling) coexists → sampler export excludes the head, save_state keeps it.

Validation surfaced and fixed two real integration bugs (details in the second commit): the stay-gathered FSDP grouping broke adapter layout validation after forward-only ops (FSDP2 reshard() cannot restore a group materialized by a no-grad pass — hence the own-unit design), and per-rank factor-shard gradients were incomplete without the upstream all-reduce (staged numerators were ~14% off analytic references before the fix; exact after).

Not in this PR (tracked in #84)

• PP support, full-weights mode; per-session masking beyond substring patterns.
• A larger-scale training run to assess LoRA-rank critic capacity (watch explained variance).
• Submodule pointer bump for xorl-client (merge Add value-model (critic) support: state_values, GAE/EV helpers, frozen modules xorl-client#15 first; the server tolerates old clients).

[broly-code-security-scanner]

Broly Security Scan

Note

✅ Clean scan
No vulnerabilities detected in this PR.

Note

Re-scan this PR anytime with /broly scan — useful after /broly undismiss, or to refresh findings without a new push.

Broly — SAST (zai-org/GLM-5.2) · Secrets · SCA · IaC · GH Actions · Base Images · Supply Chain Threats · Exploit Chains · Adversarial Verification

We're continuously improving Broly's accuracy and finding quality — your feedback is valuable. False positives, missed findings, bugs, and feature requests all welcome.

 

[kiddyboots216]

Hm this doesn't really conform to the existing adapter gradient contract design that is currently present...