LLM360 · ZhentingWang · Jun 29, 2026
@@ -488,7 +488,35 @@ def forward_step(data_iterator: DataIterator, model: GPTModel, return_schedule_p
         # creates optimizer state on first step, so release inactive blocks here
         # before tiny state allocations fail with reserved-but-free memory.
         clear_memory()
+        # === ESS-guided LR scaling (no-op / bit-exact unless --use-ess-lr) ===
+        # The scale is computed on the last PP stage (where log-probs live), so
+        # broadcast it across the PP group, temporarily scale the param-group LRs
+        # for this step, and restore them before opt_param_scheduler advances.
+        # Only the policy/actor optimizer is scaled; the critic (loss_type ==
+        # "value_loss") must not inherit the policy's off-policy ESS scale.
+        _ess_saved_lrs = None
+        if getattr(args, "use_ess_lr", False) and getattr(args, "loss_type", "policy_loss") == "policy_loss":
+            from ..training_utils.ess_lr import _ESS_LR_STATE
+
+            _ess_buf = torch.tensor(
+                [float(_ESS_LR_STATE.get("scale", 1.0)), float(_ESS_LR_STATE.get("rho_ess", 1.0))],
+                device=torch.cuda.current_device(),
+            )
+            torch.distributed.broadcast(
+                _ess_buf,
+                src=mpu.get_pipeline_model_parallel_last_rank(),
+                group=mpu.get_pipeline_model_parallel_group(),
+            )
+            _ESS_LR_STATE["scale"] = float(_ess_buf[0].item())
+            _ESS_LR_STATE["rho_ess"] = float(_ess_buf[1].item())
+            if _ESS_LR_STATE["scale"] != 1.0:
+                _ess_saved_lrs = [g["lr"] for g in optimizer.param_groups]
+                for g in optimizer.param_groups:
+                    g["lr"] = g["lr"] * _ESS_LR_STATE["scale"]
         update_successful, grad_norm, num_zeros_in_grad = optimizer.step()
+        if _ess_saved_lrs is not None:
+            for g, _saved_lr in zip(optimizer.param_groups, _ess_saved_lrs, strict=False):
+                g["lr"] = _saved_lr
 
         # Update learning rate.
         assert update_successful
@@ -659,6 +687,14 @@ def train(
             if args.enable_mtp_training:
                 extra_metrics["mtp_loss"] = mtp_losses
 
+            if getattr(args, "use_ess_lr", False) and getattr(args, "loss_type", "policy_loss") == "policy_loss":
+                from ..training_utils.ess_lr import _ESS_LR_STATE
+
+                # _ESS_LR_STATE was just synced across the PP group in the
+                # optimizer.step() block above, so this rank holds the value.
+                extra_metrics["lr_scale"] = float(_ESS_LR_STATE.get("scale", 1.0))
+                extra_metrics["rho_ess"] = float(_ESS_LR_STATE.get("rho_ess", 1.0))
+
             for param_group_id, param_group in enumerate(optimizer.param_groups):
                 extra_metrics[f"lr-pg_{param_group_id}"] = opt_param_scheduler.get_lr(param_group)
 

@@ -0,0 +1,116 @@
+"""ESS-guided learning-rate scaling (VCPO-style; arXiv:2602.17616).
+
+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)``.
+
+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
+    ESS      = (sum_i w_i)^2 / (sum_i w_i^2) ; rho = ESS / B  in (0, 1]
+    lr_scale = sqrt(rho), clamped to [args.ess_lr_floor, 1.0]
+
+``ess_lr_compute`` runs once per ROLLOUT over the full batch (from
+``compute_advantages_and_returns``, before the rollout's optimizer-step loop) and
+stashes the scale in ``_ESS_LR_STATE``; the same scale is reused for every
+optimizer step in that rollout. ``megatron_utils/model.py`` broadcasts it across
+the PP group and applies it around ``optimizer.step()`` (only for the policy/actor
+loss, not the critic). Everything is a no-op (bit-exact legacy) unless
+``args.use_ess_lr`` is set.
+"""
+
+from argparse import Namespace
+
+import torch
+
+from miles.utils.types import RolloutBatch
+
+from .cp_utils import get_logits_and_tokens_offset_with_cp
+from .parallel import ParallelState
+
+# Step-level scale/rho, written by ess_lr_compute() on the last PP stage and
+# synced across the PP group at optimizer-step time. Read by model.py.
+_ESS_LR_STATE: dict[str, float] = {"scale": 1.0, "rho_ess": 1.0}
+
+
+def ess_lr_scale_from_sums(
+    sum_w: torch.Tensor, sum_w2: torch.Tensor, batch: torch.Tensor, floor: float
+) -> tuple[float, float]:
+    """Map the (already DP-reduced) ESS sums to ``(lr_scale, rho_ess)``. Pure / testable."""
+    rho = (sum_w * sum_w) / (batch * sum_w2 + 1e-8)  # ESS / B in (0, 1]
+    scale = float(torch.sqrt(rho.clamp_min(1e-8)).clamp(min=floor, max=1.0).item())
+    return scale, float(rho)
+
+
+def ess_lr_compute(args: Namespace, parallel_state: ParallelState, rollout_data: RolloutBatch) -> None:
+    """Compute the ESS-guided LR scale and stash it in ``_ESS_LR_STATE``.
+
+    No-op unless ``args.use_ess_lr``. Only the last pipeline stage holds the
+    train/rollout log-probs, so this early-returns elsewhere; the computed scale
+    is broadcast across the PP group at optimizer-step time (see model.py).
+    """
+    if not getattr(args, "use_ess_lr", False):
+        return
+    # Reset up front so a missing-input early-return (e.g. --use-rollout-logprobs
+    # where train log-probs are not recomputed, the critic path, or a non-last PP
+    # stage) yields a no-op scale=1 for THIS rollout instead of silently reusing
+    # the previous rollout's value.
+    _ESS_LR_STATE["scale"] = 1.0
+    _ESS_LR_STATE["rho_ess"] = 1.0
+    train_log_probs = rollout_data.get("log_probs")
+    rollout_log_probs = rollout_data.get("rollout_log_probs")
+    loss_masks = rollout_data.get("loss_masks")
+    if not train_log_probs or not rollout_log_probs or not loss_masks:
+        return  # not last pp stage, or rollout log-probs unavailable
+    total_lengths = rollout_data.get("total_lengths")
+    response_lengths = rollout_data.get("response_lengths")
+    max_seq_lens = rollout_data.get("max_seq_lens", None)
+    cp_size = parallel_state.cp.size
+    device = train_log_probs[0].device
+
+    local_num: list[torch.Tensor] = []  # CP-local-chunk sum of (train - rollout) log-prob over masked tokens
+    full_cnt: list[torch.Tensor] = []  # full-trajectory masked token count
+    for i in range(len(train_log_probs)):
+        d = train_log_probs[i].float() - rollout_log_probs[i].float()  # CP-local per-token log-IS
+        full_mask = loss_masks[i].to(device=device).float()  # FULL response mask
+        if cp_size == 1:
+            local_mask = full_mask
+        else:
+            # Mirror the CP mask-chunking used by advantage whitening in loss.py so
+            # the local mask aligns with the CP-local log-prob chunk `d`.
+            prompt_len = int(total_lengths[i]) - int(response_lengths[i])
+            max_seq_len = max_seq_lens[i] if max_seq_lens is not None else None
+            _, _, _, token_offsets = get_logits_and_tokens_offset_with_cp(
+                int(total_lengths[i]), int(response_lengths[i]), args.qkv_format, max_seq_len
+            )
+            (s0, e0), (s1, e1) = token_offsets[0], token_offsets[1]
+            res_s0, res_e0 = max(0, s0 - prompt_len), max(0, e0 - prompt_len)
+            res_s1, res_e1 = max(0, s1 - prompt_len), max(0, e1 - prompt_len)
+            parts = []
+            if res_e0 > res_s0:
+                parts.append(full_mask[res_s0:res_e0])
+            if res_e1 > res_s1:
+                parts.append(full_mask[res_s1:res_e1])
+            local_mask = torch.cat(parts) if parts else torch.zeros(0, device=device, dtype=full_mask.dtype)
+        n = min(d.numel(), local_mask.numel())
+        local_num.append((d[:n] * local_mask[:n]).sum())
+        full_cnt.append(full_mask.sum())
+
+    if not local_num:
+        return
+    local_num_t = torch.stack(local_num).float()  # [B_local]
+    full_cnt_t = torch.stack(full_cnt).float().clamp_min(1.0)
+    if cp_size > 1:
+        torch.distributed.all_reduce(local_num_t, group=parallel_state.cp.group)  # -> full-trajectory numerator
+    m = local_num_t / full_cnt_t  # per-traj mean log-IS (geom-mean exponent)
+    w = torch.exp(m.clamp(min=-30.0, max=30.0))  # geom-mean IS weight (clamp guards exp overflow)
+    stat = torch.stack([w.sum(), (w * w).sum(), torch.tensor(float(w.numel()), device=device)])
+    if parallel_state.intra_dp.size > 1:
+        # Skip when dp_size == 1: the local sums are already global, and
+        # intra_dp.group may be None (-> all_reduce would fall back to WORLD and
+        # deadlock, since only the last PP stage reaches this line).
+        torch.distributed.all_reduce(stat, group=parallel_state.intra_dp.group)  # ESS sums over DP
+    scale, rho = ess_lr_scale_from_sums(stat[0], stat[1], stat[2], float(args.ess_lr_floor))
+    _ESS_LR_STATE["scale"] = scale
+    _ESS_LR_STATE["rho_ess"] = rho
@@ -26,6 +26,7 @@
     get_logits_and_tokens_offset_with_cp,
     get_sum_of_sample_mean,
 )
+from .ess_lr import ess_lr_compute
 from .parallel import get_parallel_state
 
 
@@ -444,6 +445,8 @@ def compute_advantages_and_returns(args: Namespace, rollout_data: RolloutBatch)
             chunk_lengths = [chunk.size(0) for chunk in advantages]
             advantages = list(torch.split(whitened_advs_flat, chunk_lengths))
 
+    ess_lr_compute(args, parallel_state, rollout_data)
+
     rollout_data["advantages"] = advantages
     rollout_data["returns"] = returns
 

@@ -137,6 +137,22 @@ def add_train_arguments(parser):
                 default=False,
                 help="Whether to enable true-on-policy mode.",
             )
+            parser.add_argument(
+                "--use-ess-lr",
+                action="store_true",
+                default=False,
+                help=(
+                    "Enable ESS-guided learning-rate scaling (VCPO-style, arXiv:2602.17616). "
+                    "Shrinks the effective LR by sqrt(ESS/B) when off-policy importance weights "
+                    "are heavy-tailed (low effective sample size). Default off (bit-exact legacy)."
+                ),
+            )
+            parser.add_argument(
+                "--ess-lr-floor",
+                type=float,
+                default=0.1,
+                help="Lower clamp for the ESS LR scale; only used when --use-ess-lr.",
+            )
             parser.add_argument(
                 "--train-env-vars",
                 type=json.loads,
@@ -2099,6 +2115,12 @@ def miles_validate_args(args):
             args.use_dynamic_batch_size is False
         ), "Dynamic batch size is not supported for bshd format. Please specify --micro-batch-size instead."
 
+    if getattr(args, "use_ess_lr", False):
+        # ESS-LR scales the optimizer LR inside the megatron train step; the FSDP
+        # backend never consumes the scale, so enabling it there would silently
+        # no-op. Fail loudly instead.
+        assert args.train_backend == "megatron", "--use-ess-lr is only supported for the megatron backend."
+
     _maybe_apply_dumper_overrides(args)
 
 

@@ -0,0 +1,120 @@
+"""Fast unit tests for ESS-guided LR scaling (VCPO-style; arXiv:2602.17616).
+
+Covers the pure ESS->lr_scale math (``ess_lr_scale_from_sums``) and the
+default-off no-op path of ``ess_lr_compute``. The distributed all-reduce
+plumbing inside ``ess_lr_compute`` is exercised by integration runs, not here.
+"""
+
+import math
+import types
+from argparse import Namespace
+
+import pytest
+
+torch = pytest.importorskip("torch")
+
+from miles.backends.training_utils.ess_lr import (  # noqa: E402
+    _ESS_LR_STATE,
+    ess_lr_compute,
+    ess_lr_scale_from_sums,
+)
+
+
+def _sums(weights):
+    w = torch.tensor(weights, dtype=torch.float32)
+    return w.sum(), (w * w).sum(), torch.tensor(float(w.numel()))
+
+
+def test_equal_weights_give_scale_one():
+    # ESS == B  ->  rho == 1  ->  lr_scale == 1 (no shrink, fully on-policy).
+    sum_w, sum_w2, b = _sums([1.0, 1.0, 1.0, 1.0])
+    scale, rho = ess_lr_scale_from_sums(sum_w, sum_w2, b, floor=0.1)
+    assert rho == pytest.approx(1.0, abs=1e-6)
+    assert scale == pytest.approx(1.0, abs=1e-6)
+
+
+def test_skewed_weights_shrink_lr():
+    # One heavier weight collapses ESS below B  ->  rho < 1  ->  scale = sqrt(rho).
+    weights = [1.0, 1.0, 1.0, 5.0]
+    sum_w, sum_w2, b = _sums(weights)
+    scale, rho = ess_lr_scale_from_sums(sum_w, sum_w2, b, floor=0.1)
+    expected_rho = (8.0**2) / (4.0 * 28.0)  # 64 / 112
+    assert rho == pytest.approx(expected_rho, rel=1e-5)
+    assert scale == pytest.approx(math.sqrt(expected_rho), rel=1e-5)
+    assert 0.1 < scale < 1.0
+
+
+def test_floor_clamps_extreme_collapse():
+    # rho = 1/B = 0.005  ->  sqrt(rho) ~= 0.0707 < floor  ->  clamped to floor.
+    sum_w = torch.tensor(1.0)
+    sum_w2 = torch.tensor(1.0)
+    b = torch.tensor(200.0)
+    scale, rho = ess_lr_scale_from_sums(sum_w, sum_w2, b, floor=0.1)
+    assert rho == pytest.approx(0.005, rel=1e-3)
+    assert math.sqrt(rho) < 0.1
+    assert scale == pytest.approx(0.1, abs=1e-6)
+
+
+def test_scale_never_exceeds_one():
+    # rho can numerically exceed 1 with tiny denominators; scale must stay <= 1.
+    sum_w = torch.tensor(10.0)
+    sum_w2 = torch.tensor(1.0)
+    b = torch.tensor(1.0)
+    scale, _ = ess_lr_scale_from_sums(sum_w, sum_w2, b, floor=0.1)
+    assert scale <= 1.0
+
+
+def test_compute_is_noop_when_disabled():
+    # Default off must not touch _ESS_LR_STATE (bit-exact legacy path).
+    _ESS_LR_STATE["scale"] = 1.0
+    _ESS_LR_STATE["rho_ess"] = 1.0
+    args = Namespace(use_ess_lr=False)
+    rollout_data = {
+        "log_probs": [torch.zeros(4)],
+        "rollout_log_probs": [torch.zeros(4)],
+        "loss_masks": [torch.ones(4)],
+    }
+    ess_lr_compute(args, parallel_state=None, rollout_data=rollout_data)
+    assert _ESS_LR_STATE["scale"] == 1.0
+    assert _ESS_LR_STATE["rho_ess"] == 1.0
+
+
+def _fake_parallel_state(cp_size=1, dp_size=1):
+    # cp_size == 1 skips the CP all-reduce; dp_size == 1 skips the DP all-reduce,
+    # so the math runs locally on one rank with no process groups needed.
+    return types.SimpleNamespace(
+        cp=types.SimpleNamespace(size=cp_size, group=None),
+        intra_dp=types.SimpleNamespace(size=dp_size, group=None),
+    )
+
+
+def test_compute_updates_state_when_enabled():
+    # Single rank (cp=dp=1) -> no all-reduce. Two trajectories with IS weights
+    # w = [1, 5]  ->  rho = 36 / (2 * 26) = 0.6923, scale = sqrt(rho).
+    _ESS_LR_STATE["scale"] = 0.5  # poison: must be overwritten
+    _ESS_LR_STATE["rho_ess"] = 0.5
+    args = Namespace(use_ess_lr=True, ess_lr_floor=0.1, qkv_format="thd")
+    rollout_data = {
+        "log_probs": [torch.zeros(4), torch.full((4,), math.log(5.0))],
+        "rollout_log_probs": [torch.zeros(4), torch.zeros(4)],
+        "loss_masks": [torch.ones(4), torch.ones(4)],
+        "total_lengths": [4, 4],
+        "response_lengths": [4, 4],
+    }
+    ess_lr_compute(args, _fake_parallel_state(cp_size=1), rollout_data)
+    expected_rho = 36.0 / (2.0 * 26.0)
+    assert _ESS_LR_STATE["rho_ess"] == pytest.approx(expected_rho, rel=1e-4)
+    assert _ESS_LR_STATE["scale"] == pytest.approx(math.sqrt(expected_rho), rel=1e-4)
+
+
+def test_compute_resets_when_inputs_missing():
+    # Enabled but rollout_log_probs absent (e.g. the --use-rollout-logprobs path,
+    # the critic path, or a non-last PP stage): must reset to a no-op scale this
+    # rollout, not silently reuse the previous rollout's value.
+    _ESS_LR_STATE["scale"] = 0.42
+    _ESS_LR_STATE["rho_ess"] = 0.17
+    args = Namespace(use_ess_lr=True, ess_lr_floor=0.1, qkv_format="thd")
+    rollout_data = {"log_probs": [torch.zeros(4)], "loss_masks": [torch.ones(4)]}  # no rollout_log_probs
+    ess_lr_compute(args, _fake_parallel_state(cp_size=1), rollout_data)
+    assert _ESS_LR_STATE["scale"] == 1.0
+    assert _ESS_LR_STATE["rho_ess"] == 1.0