diff --git a/miles/backends/megatron_utils/model.py b/miles/backends/megatron_utils/model.py index 3c50770aa0..04ead1b65f 100644 --- a/miles/backends/megatron_utils/model.py +++ b/miles/backends/megatron_utils/model.py @@ -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) diff --git a/miles/backends/training_utils/ess_lr.py b/miles/backends/training_utils/ess_lr.py new file mode 100644 index 0000000000..ea732a8dee --- /dev/null +++ b/miles/backends/training_utils/ess_lr.py @@ -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 diff --git a/miles/backends/training_utils/loss.py b/miles/backends/training_utils/loss.py index bf1eaf75a2..32d3674bab 100644 --- a/miles/backends/training_utils/loss.py +++ b/miles/backends/training_utils/loss.py @@ -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 diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index ae88f6078f..8087215ac3 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -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) diff --git a/tests/fast/backends/training_utils/test_ess_lr.py b/tests/fast/backends/training_utils/test_ess_lr.py new file mode 100644 index 0000000000..511162af42 --- /dev/null +++ b/tests/fast/backends/training_utils/test_ess_lr.py @@ -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