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

docs/astro.config.mjs

@@ -50,6 +50,7 @@ export default defineConfig({
               ],
             },
             { label: "Inference: xorl-sglang", slug: "server-training/sglang" },
+            { label: "Value-Model (Critic) Training", slug: "server-training/value-model" },
             {
               label: "Client SDK (xorl-client)",
               collapsed: true,

docs/src/content/docs/server-training/value-model.md

@@ -0,0 +1,88 @@
+---
+title: "Value-Model (Critic) Training"
+---
+
+xorl supports training a **value model (critic)** alongside policy sessions for PPO/SAO-style RL — the recipe behind single-rollout asynchronous RL ([SAO, arXiv:2607.07508](https://arxiv.org/abs/2607.07508)), where one rollout per prompt replaces GRPO group sampling and a trained critic supplies the advantage baseline.
+
+## Design
+
+The critic is **just another LoRA session** on the shared base model, so it costs a LoRA adapter — not a second model. When the server is launched with `enable_value_head: true`, the model carries a scalar value head (`hidden_size → 1`) implemented as a LoRA module with a zero, frozen base weight: the value function lives entirely in per-session adapter factors, every session gets its own independent copy, and a fresh critic predicts exactly `V(s) = 0`.
+
+Two loss functions become available:
+
+| `loss_fn` | Op | Inputs (`loss_fn_inputs`) | Output |
+|---|---|---|---|
+| `value_prediction` | `forward` (no-grad) | `target_tokens`, `weights` | per-token `V(s_t)` in `LossFnOutput.state_values` |
+| `value_loss` | `forward_backward` | `target_tokens`, `weights`, `returns`, optional `old_values` | masked squared error; `state_values` + per-token errors in `elementwise_loss` |
+
+`value_loss` params (via `loss_fn_params`): `vf_coef` (default 1.0) and `clip_range` (default 0.0 = off; with `old_values`, applies the PPO clipped-value objective).
+
+Like `advantages` and `logprobs`, the `returns` / `old_values` fields are **target-aligned** per-token vectors. Unlike `advantages`, a `returns` value of exactly `0.0` does **not** mask the token — masking comes only from `weights` / `target_tokens`.
+
+## Server configuration
+
+```yaml
+enable_lora: true
+enable_value_head: true
+# lm_head must NOT be a LoRA target (a value_loss backward produces no
+# lm-head adapter gradients): use an explicit list, or train_unembed: false.
+lora_target_modules: ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]
+```
+
+Current restrictions: plain LoRA only (no QLoRA), `pipeline_parallel_size: 1`, no `fsdp_sharded_lm_head_loss` / lm-head tensor parallelism.
+
+The value head never reaches inference: `save_weights_for_sampler` adapters exclude it (SGLang has no such module), while `save_state` training checkpoints keep it so critic sessions resume.
+
+## The SAO training loop
+
+```python
+from xorl_client import ServiceClient, compute_skip_observation_gae, explained_variance
+
+svc = ServiceClient(base_url=SERVER)
+policy = svc.create_lora_training_client(BASE_MODEL, model_id="policy")
+critic = svc.create_lora_training_client(BASE_MODEL, model_id="critic")
+
+# Per completed rollout (single-rollout, asynchronous):
+values_out = critic.forward(datums, loss_fn="value_prediction").result()
+values = values_out.loss_fn_outputs[0].state_values.data     # V(s_t) per token
+
+advantages, returns = compute_skip_observation_gae(
+    rewards, values, action_mask, gamma=1.0, lam=0.95,
+)
+
+for _ in range(K):                                            # faster value update (K=2 in the paper)
+    fb = critic.forward_backward(with_returns(datums, returns), loss_fn="value_loss").result()
+    critic.optim_step(critic_adam).result()
+
+policy.forward_backward(with_advantages(datums, advantages, rollout_logprobs),
+                        loss_fn="policy_loss").result()       # DIS: IcePop masking + rollout logprobs
+policy.optim_step(policy_adam).result()
+```
+
+`compute_skip_observation_gae` implements the paper's skip-observation estimator (Eq. 4–5): the Bellman recursion chains across **action tokens only**, so critic noise never propagates through environment-feedback tokens the model did not generate.
+
+## Monitoring critic health
+
+Explained variance is the paper's key critic diagnostic (it should climb toward 1.0; near 0 the critic is no better than the mean return). Every `value_loss` step reports sum-composable moments that reduce to global means, from which:
+
+```python
+ev = explained_variance(
+    value_error_sq_mean=metrics["is_value_error_sq_mean:mean"],
+    return_mean=metrics["is_return_mean:mean"],
+    return_sq_mean=metrics["is_return_sq_mean:mean"],
+)
+```
+
+## Frozen-attention critic
+
+The paper's frozen-attention critic (its strongest ablation) is a per-session option — it constrains only the critic, not the policy sessions sharing the substrate:
+
+```python
+critic = svc.create_lora_training_client(
+    BASE_MODEL,
+    model_id="critic",
+    frozen_module_patterns=["q_proj", "k_proj", "v_proj", "o_proj"],
+)
+```
+
+Patterns are substring matches against adapter parameter names. Matching factors keep their zero-delta initialization for this session: they are skipped at gradient staging and their optimizer state never moves, so the critic trains only its MLP factors and the value head.

examples/server/sao_critic/README.md

@@ -0,0 +1,31 @@
+# SAO-style critic training (value model)
+
+Minimal single-rollout RL loop with a trained value model, following
+[SAO (arXiv:2607.07508)](https://arxiv.org/abs/2607.07508): one rollout per
+prompt, skip-observation GAE from a critic that shares the base model with the
+policy as a second LoRA session, and the critic updated K× per policy step.
+
+## Server
+
+Launch a LoRA training server with the value head enabled (lm_head must not be
+a LoRA target):
+
+```bash
+python -m xorl.server.launcher --mode auto \
+  --config examples/server/configs/lora/qwen3_8b_lora.yaml \
+  --server.enable_value_head true --api-port 8300
+```
+
+## Run
+
+```bash
+python examples/server/sao_critic/run_sao_loop.py \
+  --base-url http://127.0.0.1:8300 --model Qwen/Qwen3-8B --steps 20
+```
+
+The script uses a toy reward (no environment needed) so the loop mechanics —
+`value_prediction` → GAE → `value_loss` ×K → `policy_loss` — can be verified
+end-to-end. Explained variance of the critic is printed each step; it should
+climb toward 1.0. Swap in real rollouts (a SamplingClient against SGLang,
+rollout logprobs, and a real reward) to make this a production loop; see
+`docs/server-training/value-model`.

examples/server/sao_critic/run_sao_loop.py

@@ -0,0 +1,115 @@
+"""Minimal SAO-style single-rollout RL loop with a trained value model.
+
+Demonstrates the loop mechanics against a live xorl training server started
+with ``enable_value_head: true`` (see README.md). Rollouts are synthetic (a
+toy terminal reward on fixed token sequences) so the critic/policy plumbing
+can be verified without an environment:
+
+    per rollout:
+      V(s_t)            <- critic.forward(loss_fn="value_prediction")
+      A_t, R_t          <- compute_skip_observation_gae(...)
+      critic            <- K x forward_backward(loss_fn="value_loss") + optim_step
+      policy            <- forward_backward(loss_fn="policy_loss") + optim_step
+
+Usage:
+    python run_sao_loop.py --base-url http://127.0.0.1:8300 --model Qwen/Qwen3-8B --steps 20
+"""
+
+import argparse
+import random
+
+from xorl_client import ServiceClient, compute_skip_observation_gae, explained_variance
+from xorl_client.types.adam_params import AdamParams
+from xorl_client.types.datum import Datum
+from xorl_client.types.model_input import ModelInput
+
+
+def make_rollout(rng: random.Random, length: int = 12):
+    """A synthetic 'rollout': tokens, an action mask (first 3 tokens are
+    prompt), and a terminal reward correlated with the token pattern."""
+    tokens = [rng.randrange(100, 5000) for _ in range(length)]
+    action_mask = [0] * 3 + [1] * (length - 4)  # target-aligned, length-1
+    reward = 1.0 if tokens[-1] % 2 == 0 else 0.0
+    return tokens, action_mask, reward
+
+
+def to_datum(tokens, action_mask, extra):
+    return Datum(
+        model_input=ModelInput.from_ints(tokens[:-1]),
+        loss_fn_inputs={
+            "target_tokens": tokens[1:],
+            "weights": [float(m) for m in action_mask],
+            **extra,
+        },
+    )
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--base-url", required=True)
+    parser.add_argument("--model", required=True)
+    parser.add_argument("--steps", type=int, default=20)
+    parser.add_argument("--critic-updates", type=int, default=2, help="K: critic steps per policy step")
+    parser.add_argument("--gamma", type=float, default=1.0)
+    parser.add_argument("--lam", type=float, default=0.95)
+    parser.add_argument("--policy-lr", type=float, default=1e-5)
+    parser.add_argument("--critic-lr", type=float, default=5e-5)
+    args = parser.parse_args()
+
+    svc = ServiceClient(base_url=args.base_url)
+    policy = svc.create_lora_training_client(args.model, model_id="sao-policy")
+    critic = svc.create_lora_training_client(args.model, model_id="sao-critic")
+    policy_adam = AdamParams(learning_rate=args.policy_lr)
+    critic_adam = AdamParams(learning_rate=args.critic_lr)
+    rng = random.Random(0)
+
+    for step in range(args.steps):
+        tokens, action_mask, reward = make_rollout(rng)
+        base = to_datum(tokens, action_mask, {})
+
+        # 1) Critic predicts V(s_t) for the rollout.
+        pred = critic.forward([base], loss_fn="value_prediction").result()
+        values = list(pred.loss_fn_outputs[0].state_values.data)
+
+        # 2) Skip-observation GAE across action tokens (terminal reward).
+        rewards = [0.0] * (len(values) - 1) + [reward]
+        advantages, returns = compute_skip_observation_gae(rewards, values, action_mask, gamma=args.gamma, lam=args.lam)
+
+        # 3) Faster value update: K critic steps per policy step.
+        ev = float("nan")
+        for _ in range(args.critic_updates):
+            fb = critic.forward_backward(
+                [to_datum(tokens, action_mask, {"returns": returns})],
+                loss_fn="value_loss",
+            ).result()
+            critic.optim_step(critic_adam).result()
+            metrics = fb.metrics
+            ev = explained_variance(
+                value_error_sq_mean=metrics.get("is_value_error_sq_mean:mean", float("nan")),
+                return_mean=metrics.get("is_return_mean:mean", float("nan")),
+                return_sq_mean=metrics.get("is_return_sq_mean:mean", float("nan")),
+            )
+
+        # 4) Policy step. With real rollouts, ``logprobs`` are the sampler's
+        # behavior logprobs (DIS: the ratio is policy/rollout); the synthetic
+        # stand-in just exercises the wire format.
+        n = len(tokens) - 1
+        policy.forward_backward(
+            [
+                to_datum(
+                    tokens,
+                    action_mask,
+                    {"advantages": advantages, "logprobs": [-2.0] * n},
+                )
+            ],
+            loss_fn="policy_loss",
+        ).result()
+        policy.optim_step(policy_adam).result()
+
+        print(f"step {step:3d}  reward {reward:.1f}  critic EV {ev:+.3f}")
+
+    print("done")
+
+
+if __name__ == "__main__":
+    main()

src/xorl/distributed/torch_parallelize.py

@@ -567,6 +567,10 @@ def _experts_shard_placement_fn(param):
     if exact_dsv4_lm_head and getattr(parallel_state, "lm_head_tp_size", 1) != 8:
         raise RuntimeError("The exact DSV4-Flash lm head requires lm_head_tensor_parallel_size=8")
     if lm_head_mod is not None and (fsdp_sharded_lm_head_loss or exact_dsv4_lm_head):
+        if getattr(model, "value_head", None) is not None:
+            raise NotImplementedError(
+                "enable_value_head is not supported with fsdp_sharded_lm_head_loss or the exact lm-head lanes"
+            )
         if parallel_state.tp_enabled:
             raise NotImplementedError("fsdp_sharded_lm_head_loss is not supported with tensor parallelism.")
         if not parallel_state.cp_enabled and not lm_head_tp:
@@ -628,6 +632,15 @@ def _experts_shard_placement_fn(param):
             last_fsdp_kwargs["reshard_after_forward"] = False
             fully_shard(last_modules, **last_fsdp_kwargs)
 
+    # A scalar value head (critic) gets its own FSDP unit whose forward is
+    # never invoked: its parameters therefore remain sharded DTensors at all
+    # times (stable for adapter layout validation across forward-only ops),
+    # and the loss consumes the folded LoRA delta through the direct-DTensor
+    # lane (full_tensor of a Partial matmul), like direct lm-head factors.
+    value_head_mod = getattr(model, "value_head", None)
+    if value_head_mod is not None and not pp_enabled:
+        fully_shard(value_head_mod, **fsdp_kwargs)
+
     # shard root model
     # Collect all _skip_fsdp experts params so they're also ignored by the
     # root-level fully_shard (layer-level already ignores them above, but the

src/xorl/lora/utils.py

@@ -1118,6 +1118,7 @@ def save_lora_checkpoint(
     transpose_moe_lora_to_peft: bool = True,
     lora_export_format: str = "peft",
     preserve_lora_dtype: bool = False,
+    exclude_value_head: bool = False,
 ) -> str:
     """
     Save LoRA weights in PEFT-compatible format.
@@ -1150,6 +1151,10 @@ def save_lora_checkpoint(
         preserve_lora_dtype: Keep LoRA tensor dtypes in the safetensors file
             instead of exporting bf16 weights. Use this for training-resume
             checkpoints; keep the default bf16 export for inference adapters.
+        exclude_value_head: Drop ``value_head`` factors from the export. Use
+            for inference/sampler adapters — serving engines have no scalar
+            value head module and would reject the unknown target. Training
+            checkpoints must keep the factors so critic sessions resume.
 
     Returns:
         Path to saved checkpoint directory
@@ -1175,6 +1180,9 @@ def save_lora_checkpoint(
     else:
         lora_state_dict = slice_lora_state_dict_to_active_rank(model, lora_state_dict)
 
+    if exclude_value_head:
+        lora_state_dict = {key: value for key, value in lora_state_dict.items() if "value_head" not in key}
+
     if lora_export_format == "dsv4_expert_banks":
         from xorl.models.transformers.deepseek_v4.exact_contract import (  # noqa: PLC0415
             DSV4_FLASH_LOGICAL_FACTOR_COUNT,

src/xorl/ops/loss/__init__.py

@@ -19,6 +19,7 @@
 from xorl.ops.loss.opd_loss import OPDLossMetrics, opd_loss_function, opd_vocab_parallel_loss_function
 from xorl.ops.loss.policy_loss import policy_loss_function
 from xorl.ops.loss.reducers import Reducer, SequencePartial, TokenPartial
+from xorl.ops.loss.value_loss import value_loss_function, value_prediction_function
 from xorl.ops.loss.vocab_parallel_cross_entropy import vocab_parallel_cross_entropy
 
 
@@ -39,6 +40,8 @@
     "policy_loss": policy_loss_function,
     "drgrpo": drgrpo_loss_function,
     "opd_loss": opd_loss_function,
+    "value_loss": value_loss_function,
+    "value_prediction": value_prediction_function,
 }
 
 
@@ -72,5 +75,7 @@ def register_loss_function(name: str, fn: Callable) -> None:
     "opd_loss_function",
     "opd_vocab_parallel_loss_function",
     "policy_loss_function",
+    "value_loss_function",
+    "value_prediction_function",
     "vocab_parallel_cross_entropy",
 ]