feat(algorithms): add AR PPO with GAE and value loss

[yhl48]

Summary

Part 3/3 of #86: add PPO with token-level GAE and a clipped value objective for AR training.

Stacked on #256. Merge #256 first, then rebase this branch onto main so the PR reduces to the single PPO/trainer commit.

• Adds PPO / PPOConfig with clipped policy loss and clipped critic loss.
• Freezes old critic values with the same planned micro-batch geometry used by the training forward, then computes GAE only after those anchors are reassembled.
• Applies loss_mask consistently to both actor and critic optimization and metrics.
• Adds advantage_mode: grpo | gae to both synchronous and asynchronous AR trainers.
• Adds a current-schema Qwen3-4B-Base + DAPO + colocated SGLang PPO recipe.
• Excludes train-only value_head.* tensors from SGLang weight sync.

Related Issue

Part of #86. Depends on #256.

Test Plan

No PR-specific test files are included.

• SKIP=no-commit-to-branch pre-commit run --all-files --show-diff-on-failure — all hooks passed.
• python -m compileall -q unirl — passed.
• Hydra composition of ar/qwen3_ppo_4b_base_dapo_sglang with placeholder dataset/model paths — passed.
• Temporary, uncommitted correctness harness — passed:
  • exact multi-trajectory GAE values and terminal-reward placement;
  • masked tokens have zero actor and critic gradient contribution;
  • critic anchors follow planned micro geometry (3 + 2 tokens), rather than a mismatched full-batch forward;
  • synchronous and asynchronous advantage_mode wiring;
  • meta-init + sharded-load Qwen3 replay followed by the actual PPO.prepare_segment -> prepare_part -> compute_loss_and_backward path.
• GPU smoke on an H20 with a 9.8M-parameter tiny Qwen3 checkpoint — passed:
  • eager/meta checkpoint state maximum difference: 0.0;
  • first-update ratio_mean: 1.0;
  • finite combined PPO loss;
  • finite, nonzero value-head gradient (L2=2.7303159).
• Full Qwen3-4B-Base comparison on 8 x H20, BF16, SGLang 0.5.12.post1 — passed:
  • 8 SGLang engines loaded the real checkpoint and completed full-tensor sync on every rollout;
  • 3 rollouts, batch size 8, 2 samples/prompt, 32 generated tokens, deterministic top_k=1, one nonzero-gradient PPO update per rollout;
  • a temporary alternating reward was used only to force nonzero advantages and was not committed;
  • both SGLang and trainside completed all three optimizer steps without OOM, NaN, exception, or clipping.

| rollout | SGLang loss | trainside loss | SGLang ratio | trainside ratio | SGLang mean |Δlogp| | trainside mean |Δlogp| |
| --- | ---: | ---: | ---: | ---: | ---: | ---: |
| 1 | -0.2223 | -0.2238 | 1.0048 | 0.9981 | 0.0114 | 0.0110 |
| 2 | -0.2206 | -0.2209 | 0.9988 | 1.0026 | 0.0120 | 0.0096 |
| 3 | -0.2199 | -0.2191 | 0.9996 | 0.9994 | 0.0066 | 0.0080 |

The maximum absolute loss difference was 1.5e-3. SGLang's rollout-to-replay log-probability drift stayed in the same range as trainside's own replay drift, so this smoke showed no additional numerical-precision regression from SGLang sync/replay.

Compatibility / Risk

• Opt-in: advantage_mode=grpo and use_value_head=false remain the defaults.
• PPO recipes must enable the value head, use advantage_mode: gae, and omit critic tensors from rollout sync.
• feat(algorithms): add AR PPO with GAE and value loss #259 is intentionally still a draft while feat(models): add AR value head and GAE wiring #256 is unmerged.
• After feat(models): add AR value head and GAE wiring #256 merges, rebase this branch onto main before marking it ready.
• The full-model check is a three-update correctness smoke, not a long-horizon convergence run.

Reviewer Notes

Review the anchor timing first: old values are frozen per planned micro before any update, reassembled, and then used to compute GAE once for the part. AI assistance was used; the final diff was reviewed and validated against the current repository state.

Checklist

• I reviewed the changed code and removed unrelated/generated artifacts.
• I updated configs and documented why no PR-specific test files are included.

[leviking98z-rgb]

Updated the validation result with a real Qwen3-4B-Base end-to-end comparison on 8 x H20 (SGLang 0.5.12.post1, BF16). No test files were added.

The run used 3 rollouts, batch size 8, 2 samples/prompt, 32 generated tokens, deterministic top_k=1, one PPO optimizer step per rollout, and full-tensor sync every rollout. A temporary alternating reward forced nonzero advantages so that all three rounds exercised backward, optimizer update, and the following SGLang weight resync.

| rollout | SGLang loss | trainside loss | SGLang ratio | trainside ratio | SGLang mean |Δlogp| | trainside mean |Δlogp| | SGLang grad norm | trainside grad norm |
| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: |
| 1 | -0.2223 | -0.2238 | 1.0048 | 0.9981 | 0.0114 | 0.0110 | 16.7814 | 16.9356 |
| 2 | -0.2206 | -0.2209 | 0.9988 | 1.0026 | 0.0120 | 0.0096 | 16.6435 | 16.1213 |
| 3 | -0.2199 | -0.2191 | 0.9996 | 0.9994 | 0.0066 | 0.0080 | 16.5123 | 16.4853 |

Both paths completed all three updates without OOM, NaN, exception, or clipping. The maximum absolute loss difference was 1.5e-3; SGLang's rollout-to-replay log-probability drift stayed in the same range as trainside's own replay drift. This smoke therefore found no extra numerical-precision regression caused by SGLang weight sync/replay.

The earlier review fixes are preserved in cda25822:

• old critic values use the exact planned training micro geometry;
• GAE runs only after anchors are reassembled;
• loss_mask excludes tokens from both actor and critic loss/metrics;
• sync and async AR both receive advantage_mode;
• the example uses the current nested SGLang config and omits value_head.* from sync.

This is a short correctness comparison rather than a long-horizon convergence run. #259 remains draft until #256 is merged and the stack is rebased.

[leviking98z-rgb]

Restored the original commit attribution without changing the tested code:

• the original PPO implementation is authored and committed solely by yhl48 (f40dcd0b);
• post-review fixes are a separate commit (cda25822);
• the final tree is byte-for-byte identical to the previously tested c8f1afe2 tree (9101889b4271909739254f2c468f2197197b0e1b);
• no test files were added.

compileall and every pre-commit hook passed after the history-only rebuild. The full Qwen3-4B SGLang/trainside results in the preceding comment therefore remain valid.