[ruiling-smartbear]

Status: ready for review

A high-learning-rate (1e-5) online pilot collapsed with correction disabled: seed 29 had 24 consecutive zero-reward/zero-gradient batches among 27 captured batches, and held-out avg@4 fell from 0.515625 to 0 by rollout 24. This was ten times the repository recipe's 1e-6, so it cannot be attributed to the opt-in correction. A controlled rerun with the same model, train/eval data (matching SHA256), seed, 100-rollout schedule and harness changed only the learning rate and output paths. At 1e-6, both disabled and cap=2 runs completed 100 rollouts / 200 optimizer updates with nonzero rewards throughout and no collapse. Final held-out avg@4 was 0.9375 without correction and 0.796875 with cap=2 on 16 arithmetic prompts × 4 samples. Thus this small single-seed fixture supports excessive learning rate as the cause of the pilot collapse, but does not demonstrate a correction benefit. The full 32-GPU Qwen3-4B/DAPO recipe, additional seeds and long-run convergence remain unverified. The implementation is ready for code review as a default-disabled option; no reward-improvement claim is made. Broader efficacy is a separate research question.

Summary

Replay-anchored AR GRPO currently samples from the rollout policy but only uses saved rollout log-probabilities for diagnostics. Add opt-in rollout_is_threshold to multiply the existing clipped token loss by detached min(exp(log_pi_old - log_pi_rollout), threshold) weights, with a [-20,20] log-ratio safety clamp following verl token TIS. For example, old_logp_source: replay plus rollout_is_threshold: 2.0 enables capped correction; the default None preserves the existing objective.

The multiplier uses frozen anchors, keeps existing reduction denominators, and reports active-token weight/clipping diagnostics and explicitly micro-local ESS. Invalid thresholds, rollout-anchor combinations and missing saved behavior log-probabilities fail explicitly.

Related Issue

Closes #556 (RFC with mathematical reproduction, real GPU measurements and bounded online validation).

Test Plan

Main sync (2026-10-23)

Synced main (6fd82cb) at cb9e237. 83 CPU cases pass, including independent gradient oracle, frozen anchors and disabled-path parity against current main.

External harness commands: PYTHONPATH=. python <external>/rollout-correction/verify_correction.py. SKIP=no-commit-to-branch pre-commit run --from-ref origin/main --to-ref HEAD and git diff --check origin/main HEAD pass. Independent review passed. GPU was not rerun for this sync; earlier GPU evidence below remains tied to its tested commits.

Verification harnesses were run outside the checkout and are not committed, per AGENTS.md. Commands ran from this branch's checkout; $PROBE below is /Users/gaoruilin/Desktop/unirl-experiments/rollout-correction.

• PYTHONPATH=. /tmp/unirl-rollout-correction-venv/bin/python "$PROBE/verify_correction.py" — PASS, 83 cases: independent loss/gradient reference, all three reductions, asymmetric clipping, ragged/empty rows, partial/all masks, repeated updates with frozen anchors, disabled bitwise comparison against base, equal-policy unit weights, invalid thresholds/mode, missing behavior early failure, Hydra and packed select/concat, microbatch gradient parity.

• PYTHONPATH=. /tmp/unirl-rollout-correction-venv/bin/python "$PROBE/verify_distributed.py" — PASS, two-rank CPU Gloo global-token gradient equivalence, including a fully masked rank.

• modal run "$PROBE/run_modal.py" — PASS, Modal job ap-QMqXI6bRWpwXgyfabjeruE; 1 x L40S, Qwen3-0.6B bf16 revision c1899de289a04d12100db370d81485cdf75e47ca, vLLM 0.28.0 / torch 2.13.0+cu130 / transformers 5.15.1. Eight synthetic prompts x four siblings x 64 tokens, temperature=1, top_p=1, top_k disabled, TP=1, eager vLLM and UniRL Qwen3ARStage SDPA replay. Mean absolute log-prob difference 0.0274566; ESS fraction 0.997892 at cap=2. Actual GRPO backward on eight rows/512 tokens, final attention q_proj trainable, controlled +/-1 advantages: reference gradient max absolute difference 0; two finite AdamW steps with unchanged anchors. Runtime GRPO SHA256 matched this commit. No training recipe or reward-service claim.

• SKIP=no-commit-to-branch pre-commit run --from-ref 54cc7b69698332b7c1b167b174263fe02ce81ea7 --to-ref HEAD — all applicable hooks passed.

• git diff --check 54cc7b69698332b7c1b167b174263fe02ce81ea7 HEAD — PASS.

• Additional supported-version and two-GPU TrainStack/FSDP checks passed; details below. Passed: bounded full ARTrainer/Ray + reward-service training and online engine weight synchronization at recipe learning rate on the 2-GPU fixture; see controlled online validation below. Multi-node and long-run convergence remain unverified.

• SKIP=no-commit-to-branch pre-commit run --all-files --show-diff-on-failure — PASS locally at current PR head 300170e after the main-branch merge. GitHub Lint run 37365528822 was cancelled after 15 minutes with no recorded job steps; the two preceding Lint runs at 72c935a passed. The GitHub check remains non-green and requires a maintainer/admin rerun; the PR author lacks permission to rerun it.

Additional validation (2026-10-05)

• modal run "$PROBE/run_pinned.py" — PASS, job ap-bmetFyaYu6NdUX5p5n8uil, one L40S. Reran actual vLLM capture, UniRL replay and backward/reference/optimizer checks with the repository-supported transformers 5.12.1, torch 2.13.0+cu130 and vLLM 0.28.0. Earlier 5.15.1 results are supplementary only. The pinned run reproduced the reported 2,048-token discrepancy and zero maximum reference-gradient difference.
• modal run "$PROBE/run_fsdp.py" — PASS, job ap-rn3WhxaTsXO2amsFaPm4du, two L40S. Executed real TrainStack.train_track and FSDPBackend on Qwen3-0.6B: full-parameter optimization, fp32 master weights / bf16 compute, activation checkpointing, full sharding, reshard-after-forward, micro_batch_size=1, token-mean reduction, two disjoint optimizer updates per batch, three recorded-rollout batches (six updates total). Each rank used four variable-length responses (40/48/56/64 tokens), partial masks and controlled +/-1 advantages. Batch 1 made rank 1 entirely masked. Twelve micro-backward anchor checks per rank passed; global loss agreed exactly between ranks, all losses/gradient norms were finite, parameters changed, and full model state gathering completed. Both anchor fields stayed equal to the pre-update snapshots at every micro. Data were previously captured real vLLM outputs; this does not claim fresh resampling at each batch.

The new test closes the multi-GPU CUDA TrainStack/FSDP smoke gap. Still unverified: the 32-GPU Qwen3-4B/DAPO recipe, multi-node training, multiple seeds and long-run convergence. UniRL's native vLLM IPC path currently restricts model type to Qwen3-MoE, so this dense Qwen3-0.6B test cannot establish that path. No repository code changed for the added validation; tested GRPO SHA256 remains 31531e0edc1301ce5f80012499ae89d6266715dfac871aaea266b5b2ef3d87c7.

Controlled online validation (2026-10-05)

modal run "$PROBE/run_e2e.py" --steps 100 --seed 29 --cap none and --cap 2 — PASS, Modal apps ap-d4wi4D40O5gyJYy4XaGIqN and ap-PSvAtGpfvWHi8rMSec9YjZ. Each used two L40S GPUs, Qwen3-0.6B, torch 2.13.0+cu130, transformers 5.12.1, SGLang 0.5.19, Ray 2.59.0 and math-verify 0.9.0. Actual ARTrainer generation, reward scoring, two FSDP optimizer updates per rollout and TensorWeightSync ran for 100 rollouts. The experiment holds the arithmetic fixture and seed 29 fixed, using the original recipe learning rate 1e-6 but smaller model/batch/response limits than the 32-GPU recipe.

Both runs had finite losses and gradients, 100 recorded rollout sample metrics and seven exact BF16 engine weight readbacks per rank (including changed parameter hashes); hashes agreed between ranks. Without correction: 78/100 batches had nonzero aggregated gradients, first/last 20 mean rewards 0.6875/0.896875, held-out avg@4 0.515625 initially and 0.9375 finally. With cap=2: 75/100 batches had nonzero aggregated gradients, rewards 0.7125/0.93125, held-out avg@4 0.515625 initially and 0.796875 finally. Neither run had a zero-reward/zero-gradient streak. Cap=2 clipped no tokens in this fixture; mean weight 1.000034 and minimum micro ESS fraction 0.998536. Mean response length was 154.1 versus 157.4 tokens; truncation ratios 0.0444 versus 0.0512 (disabled versus enabled). These small single-seed runs show stable integration at recipe learning rate, not efficacy or statistical superiority.

The stopped 1e-5 pilot used the same fixture and training configuration aside from learning rate and output paths. Its disabled seed-29 run dropped from reward 0.9375 (batch 0) to 0.3125 and 0.125, then had 24 consecutive zero-reward/zero-gradient batches starting at batch 3. The pilot was stopped; it was not a 100-batch pass.

Compatibility / Risk

Default-disabled; no recipe, checkpoint or packed-data schema changes. Token truncation trades bias for bounded weights. This does not repair tokenization, sampling-support or temperature mismatches, nor implement sequence IS/rejection sampling. Micro ESS is not global ESS. The measured backend discrepancy justifies optional support, not a claim that all recipes require it or that reward improves.

Reviewer Notes

AI-assisted implementation and validation. Full diff independently reviewed by a fresh-context gpt-6-sol reviewer under the local PR review skill; no supported blocking findings. Kept the single-use correction in GRPO and reused existing frozen packed anchor fields and loss reductions; no unrelated cleanup or committed harnesses. The current PR head 300170e has an exact-commit review record with Status: pass; merging main did not change the two-file PR diff.

Duplicate-work checks searched open issues and PRs for rollout correction, importance sampling and rollout. #430/#504/#505 focus on numerical parity, while #523 concerns anchor mode; this PR adds the optional loss multiplier and does not take over those efforts. Submitting for maintainer review of the full diff and RFC contract. The submitter approved proceeding with the existing reviewed commit; broader training efficacy is not claimed.

Checklist

• I reviewed the changed code and removed unrelated/generated artifacts.
• I updated tests, docs, and configs where needed, or explained why not.