feat: ESS-guided learning-rate scaling for off-policy/async RL (default off)

[ZhentingWang]

What

Adds ESS-guided LR scaling (VCPO-style, arXiv:2602.17616) as a native, opt-in training feature. In fully-async / off-policy RL a few stale trajectories can dominate the update (heavy-tailed importance weights), collapsing the effective sample size (ESS). When that happens we shrink the effective LR by sqrt(ESS / B) for that update. Gated behind --use-ess-lr (default off → bit-exact legacy).

Algorithm

Per trajectory i, sequence-level (length-normalized = geometric-mean IS):

m_i      = mean_t (train_logp_t − rollout_logp_t)   over response (mask==1) tokens
w_i      = exp(m_i)                                   # geom-mean IS weight (NOT length-coupled exp(sum))
ESS      = (Σ_i w_i)^2 / (Σ_i w_i^2) ;  rho = ESS / B  ∈ (0, 1]
lr_scale = sqrt(rho), clamped to [--ess-lr-floor (0.1), 1.0]

rho → 1 ⇒ on-policy, weights even, no shrink. rho → 0 ⇒ a few stale trajectories dominate ⇒ shrink LR. The per-token quantity train_logp − rollout_logp is exactly the existing train_inference_mismatch signal, so ESS-LR is a complement to TIS that adds an adaptive decay at the LR level.

How it's wired

• Compute (backends/training_utils/ess_lr.py, new): ess_lr_compute() runs once per rollout from compute_advantages_and_returns() (over the full batch, before the rollout's optimizer-step loop) and stashes {scale, rho} in module state. On any missing input (e.g. --use-rollout-logprobs, the critic path, a non-last PP stage) it resets to a no-op scale=1 for that rollout rather than carrying a stale value; errors propagate rather than silently disabling.
• Apply (backends/megatron_utils/model.py): the scale is broadcast across the PP group (it is only computed on the last PP stage) and temporarily applied to the param-group LRs around optimizer.step(), then restored before the scheduler advances. Policy/actor only (loss_type == "policy_loss"); the critic is never scaled.
• loss.py: one import + one call — no logic.
• CLI (utils/arguments.py): --use-ess-lr, --ess-lr-floor.
• Metrics: train/lr_scale, train/rho_ess via the standard extra_metrics channel.

Distributed correctness

• CP: a trajectory's tokens are split across context-parallel ranks; we sum the per-token log-IS over each trajectory's local chunk and all-reduce over the CP group to recover the full-trajectory m_i. CP mask-chunking mirrors the existing advantage-whitening path.
• DP: Σw / Σw² / B are all-reduced over the intra-DP group so ESS is global; skipped when dp_size == 1 (local sums are already global, and avoids a None-group → WORLD fallback).
• PP: scale computed on the last stage, broadcast to all stages before the step. All ranks enter/skip the collectives on the same condition (use_ess_lr and loss_type == policy_loss), so no rank divergence.

Default-off / bit-exactness

When --use-ess-lr is unset, ess_lr_compute() returns immediately and the optimizer block is skipped entirely (no broadcast, no LR mutation) — the legacy path is byte-for-byte unchanged.

Scope / limitations

• Megatron backend only — enforced in miles_validate_args (--use-ess-lr asserts train_backend == "megatron"); the FSDP step does not consume the scale.
• Scale is per-rollout (reused across that rollout's optimizer steps), computed from the pre-training log-probs.

Tests

tests/fast/backends/training_utils/test_ess_lr.py (fast, no distributed): pure ESS math (equal-weights → 1, skew → sqrt(rho), floor clamp, cap ≤ 1), enabled-path state update (mock parallel state, w=[1,5] → rho=0.692), missing-input reset, and the default-off no-op.