From e91bbed8116121fa99dddafdfdc1b353dd0b2794 Mon Sep 17 00:00:00 2001 From: yinjjiew Date: Tue, 12 May 2026 04:30:21 +0000 Subject: [PATCH] Fix top-k OPD imports for OpenClaw combine Keep the top-k OPD helper modules importable from openclaw-combine and avoid enabling incompatible full recompute for Qwen3-VL deepstack bridge providers. --- openclaw-combine/hint_opd_loss.py | 648 ++++++++++++++++++ openclaw-combine/hint_opd_select_loss.py | 537 +++++++++++++++ .../run_qwen3_4b_openclaw_topk_select.sh | 2 +- .../backends/megatron_utils/model_provider.py | 14 +- 4 files changed, 1199 insertions(+), 2 deletions(-) create mode 100644 openclaw-combine/hint_opd_loss.py create mode 100644 openclaw-combine/hint_opd_select_loss.py diff --git a/openclaw-combine/hint_opd_loss.py b/openclaw-combine/hint_opd_loss.py new file mode 100644 index 0000000..9276f8e --- /dev/null +++ b/openclaw-combine/hint_opd_loss.py @@ -0,0 +1,648 @@ +"""Combined GRPO + top-K On-Policy-Distillation loss (verl-aligned). + +``` +loss = w_rl * grpo_pg_loss + w_opd * topk_opd_loss +``` + +with weights driven by env vars (defaults: ``w_rl=0``, ``w_opd=1``). + +Top-K OPD design (aligned with verl's ``compute_topk_opd_loss``) +---------------------------------------------------------------- +For each response position ``t`` and each candidate ``v in S_t``: + +* Subset selection (``--distill-subset-mode``): + - ``student``: S_t = top-K(pi_old) + - ``teacher``: S_t = top-K(pi_T) (slime ships this aligned to the student + via the prm_teacher gather pass) + - ``overlap``: S_t = top-K(pi_old) ∩ top-K(pi_T); per-token |S_t| varies + and rows with empty intersection are masked out. + + The subset has TWO purposes ONLY: + (1) per-token filter for which (t, v) pairs enter the loss; + (2) the support over which the IS weight ``w_v`` sums to 1. + The subset is NEVER used to renormalize log-probabilities. + +* Three GLOBAL log-probs per (t, v in S_t): + ``ell_old(v) = log pi_old(v|s_t)`` (detached, from old_actor forward) + ``ell_cur(v) = log pi_theta(v|s_t)`` (autograd-connected, current actor) + ``ell_T(v) = log pi_T(v|s_t)`` (detached, from teacher forward) + +* Detached IS weight (subset-normalized -- the ONLY use of subset normalization): + ``w_v = softmax(ell_old over S_t)``, sum_{v in S_t} w_v = 1. + +* Detached advantage (GLOBAL log-ratio, weighted by w_v): + ``A_v = (ell_T(v) - ell_old(v)) * w_v``. + Sign: A_v > 0 iff teacher places more GLOBAL mass on v than (old) student. + +* PPO ratio (GLOBAL, gradient through ell_cur only): + ``rho_v = exp(ell_cur(v) - ell_old(v))``. + Under on-policy training rho_v == 1 by value, but its gradient + ``∇_θ ell_cur(v)`` is non-zero. + +* PPO clipped surrogate per (t, v): + ``L_v = max(-A_v * rho_v, -A_v * clip(rho_v, 1-eps_lo, 1+eps_hi))``. + +* Per-token aggregation: SUM over S_t (NOT mean -- since sum_v w_v = 1, the + sum is naturally bounded to O(|ell_T - ell_old|) and concentrates gradient + pressure on the head of S_t): + ``L_t = sum_{v in S_t} L_v``. + +* Trajectory aggregation: ``sum_of_sample_mean`` (matches slime conventions). + +Why GLOBAL ratios and GLOBAL advantage (vs subset-renormalized) +--------------------------------------------------------------- +A subset-renormalized ratio ``q_bar_new / q_bar_old`` is invariant to any +global rescaling of the student's distribution that preserves within-S_t +ordering. The student can satisfy "match teacher on subset" by collapsing +``pi_theta(S_t)`` to near zero -- nothing in the surrogate prevents it. +Using the GLOBAL ratio ``pi_theta(v) / pi_old(v)`` keeps the IS correction +honest: any global mass leaving S_t shrinks rho_v proportionally. + +Likewise, subset-renormalizing the teacher's log-probs would replace +``log pi_T(v)`` with ``log pi_T(v | v in S_t)``, which discards the +absolute confidence of the teacher (and the magnitude of disagreement +when the teacher prefers some v* outside S_t). + +Memory note +----------- +Computing the GLOBAL log-prob ``ell_cur(v) = raw[v] - global_lse(raw)`` +with autograd materializes the local softmax of the student's logits in +backward (one full-vocab pass per sample). This is the unavoidable cost +of a faithful global PPO ratio. +""" + +from __future__ import annotations + +import os +from argparse import Namespace +from collections.abc import Callable + +import torch +import torch.distributed as dist +from megatron.core import mpu + +from slime.backends.megatron_utils.loss import get_log_probs_and_entropy, get_responses +from slime.utils.ppo_utils import compute_approx_kl, compute_policy_loss + + +# --------------------------------------------------------------------------- +# TP-aware gather of raw logits at given GLOBAL vocab indices (autograd). +# --------------------------------------------------------------------------- + + +class _VocabParallelGatherRawLogits(torch.autograd.Function): + """Single-pass raw-logit gather at K global vocab indices, TP-sharded. + + Forward: + logits: ``[R, V_local]`` (vocab dim is TP-sharded). + idx_global: ``[R, K]`` global vocab ids in ``[0, V)``. + + Returns ``raw ∈ [R, K]`` where ``raw[r,k] = logits_global[r, idx[r,k]]``, + reconstructed across TP via a single all_reduce(SUM) on the masked + per-rank gather (off-shard entries contribute 0). + + Backward: + ``∂L/∂logits[r,v] = sum_k g[r,k] · 1{v == idx[r,k] && in_shard}`` + Implemented via scatter_add on the rank's local logits region. + + Memory: saves only ``[R, K]`` index data + the ``in_shard`` mask; the + raw-logit slice itself is rederived from the saved ``logits`` ref via + a tiny gather. For training-time use the activation cost is O(R*K) + instead of O(R*V_local). + """ + + @staticmethod + def forward(ctx, logits, idx_global, tp_group, tp_world, tp_rank): + V_local = logits.size(-1) + shard_lo = tp_rank * V_local + + in_shard = (idx_global >= shard_lo) & (idx_global < shard_lo + V_local) + idx_local = (idx_global - shard_lo).clamp(min=0, max=V_local - 1) + gathered = torch.gather(logits, dim=-1, index=idx_local) # [R, K] + gathered_masked = torch.where(in_shard, gathered, torch.zeros_like(gathered)) + if tp_world > 1: + dist.all_reduce(gathered_masked, op=dist.ReduceOp.SUM, group=tp_group) + + ctx.save_for_backward(idx_local, in_shard) + ctx.logits_shape = logits.shape + ctx.logits_dtype = logits.dtype + ctx.logits_device = logits.device + ctx.tp_world = tp_world + return gathered_masked # [R, K] + + @staticmethod + def backward(ctx, grad_out): + idx_local, in_shard = ctx.saved_tensors + # Mask the incoming grad to only this rank's shard, then scatter into the + # rank-local logits gradient. Off-shard contributions are zeroed. + masked_grad = torch.where(in_shard, grad_out, torch.zeros_like(grad_out)) + grad_input = torch.zeros( + ctx.logits_shape, dtype=ctx.logits_dtype, device=ctx.logits_device + ) + grad_input.scatter_add_(dim=-1, index=idx_local, src=masked_grad.to(ctx.logits_dtype)) + return grad_input, None, None, None, None + + +def _gather_student_raw_logits_at_indices( + logits_chunk: torch.Tensor, + indices: torch.Tensor, + tp_group, +) -> torch.Tensor: + """Gather raw global logits at ``indices`` (autograd-connected).""" + tp_world = dist.get_world_size(group=tp_group) if dist.is_initialized() else 1 + tp_rank = dist.get_rank(group=tp_group) if dist.is_initialized() else 0 + return _VocabParallelGatherRawLogits.apply( + logits_chunk, indices, tp_group, tp_world, tp_rank + ) + + +# --------------------------------------------------------------------------- +# TP-aware global LSE over the full vocab (autograd). +# --------------------------------------------------------------------------- + + +class _VocabParallelGlobalLSE(torch.autograd.Function): + """Numerically-stable global LSE over the FULL vocab dim, TP-aware. + + Forward: + logits: ``[R, V_local]`` (vocab dim is TP-sharded). + Returns ``lse ∈ [R, 1]`` where + ``lse[r] = log sum_{v=0..V-1} exp(logits_global[r, v])``. + + Backward: + ``∂L/∂logits[r, v] = grad_out[r, 0] * softmax_global(r, v)``. + Locally, ``softmax_global(r, v) = exp(logits[r, v] - lse[r])`` for + v in this rank's shard. Each TP rank scatters its own local-shard + softmax, which together reconstruct the global softmax gradient. + + Memory: backward materializes ``[R, V_local]`` softmax probabilities. + This is unavoidable for a faithful global ratio. + """ + + @staticmethod + def forward(ctx, logits, tp_group, tp_world): + logits_f = logits.float() + row_max = logits_f.max(dim=-1, keepdim=True).values # [R, 1] + if tp_world > 1: + dist.all_reduce(row_max, op=dist.ReduceOp.MAX, group=tp_group) + shifted = logits_f - row_max + sum_exp = shifted.exp().sum(dim=-1, keepdim=True) # [R, 1] local + if tp_world > 1: + dist.all_reduce(sum_exp, op=dist.ReduceOp.SUM, group=tp_group) + lse = row_max + sum_exp.clamp_min(1e-30).log() # [R, 1] + + ctx.save_for_backward(logits, lse) + return lse + + @staticmethod + def backward(ctx, grad_out): + # grad_out: [R, 1]. d(lse)/d(logits[v]) = global softmax(v). + # Local-shard softmax = exp(logits - global_lse). Multiplying by + # grad_out broadcasts [R, 1] across the local vocab dim. + logits, lse = ctx.saved_tensors + softmax_local = (logits.float() - lse).exp() # [R, V_local] + grad_in = grad_out * softmax_local + return grad_in.to(logits.dtype), None, None + + +def _vocab_parallel_global_lse( + logits_chunk: torch.Tensor, + tp_group, +) -> torch.Tensor: + """Global LSE over the full vocab, TP-aware (autograd-connected). + + Returns ``[R, 1]`` in float32. Use to convert raw logits to global + log-probabilities: ``log pi(v) = logits[v] - global_lse``. + """ + tp_world = dist.get_world_size(group=tp_group) if dist.is_initialized() else 1 + return _VocabParallelGlobalLSE.apply(logits_chunk, tp_group, tp_world) + + +# --------------------------------------------------------------------------- +# Knobs +# --------------------------------------------------------------------------- + + +def _w_rl() -> float: + return float(os.environ.get("HINT_OPD_W_RL", "0.0")) + + +def _w_opd() -> float: + return float(os.environ.get("HINT_OPD_W_OPD", "1.0")) + + +def _eps_clip_lo(args: Namespace) -> float: + v = os.environ.get("HINT_OPD_PPO_CLIP_EPS_LO", "") + return float(v) if v else float(args.eps_clip) + + +def _eps_clip_hi(args: Namespace) -> float: + v = os.environ.get("HINT_OPD_PPO_CLIP_EPS_HI", "") + return float(v) if v else float(args.eps_clip_high) + + +def _adv_diff_clip() -> float | None: + """Magnitude clamp on the per-candidate teacher-vs-old log-ratio. + + Applied as ``diff = (ell_T - ell_old).clamp(-t, t)`` before forming + ``A_v = diff * w_v``. Bounds advantage magnitude when the teacher and + old student disagree wildly on a candidate (common early in training, + on rare-token candidates, or when distillation crosses tokenizers). + + Env var: ``HINT_OPD_ADV_DIFF_CLIP``. + Returns: + positive float t -> clamp diff to [-t, t] + None (env var unset/empty -> defaults to 2.0) + non-positive (env var <= 0) -> disabled, returns None + """ + raw = os.environ.get("HINT_OPD_ADV_DIFF_CLIP", "") + if raw == "": + return 2.0 + val = float(raw) + return val if val > 0.0 else None + + +# --------------------------------------------------------------------------- +# verl-aligned PPO surrogate over S_t for one sample. +# --------------------------------------------------------------------------- + + +_NEG_INF = float("-inf") +# verl-style numerical guard on the log-ratio before exp(). Prevents +# exp(huge) NaNs on rare tail candidates early in training. +_PPO_KL_CLAMP = 20.0 + + +def _local_lse(values: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: + """Numerically-stable lse over the True positions of ``mask``. + + Both tensors are ``[R, K]``. Returns ``[R, 1]``. If the entire row's + mask is all-False we return 0 (the row will be masked out downstream). + """ + masked_vals = torch.where(mask, values, values.new_full((), _NEG_INF)) + row_max = masked_vals.max(dim=-1, keepdim=True).values + # If a row is all -inf (no valid k), replace its max with 0 so the + # log/exp stays finite; the row's loss will be zeroed by row_valid. + row_max = torch.where(torch.isfinite(row_max), row_max, torch.zeros_like(row_max)) + shifted = masked_vals - row_max + exped = torch.where(mask, shifted.exp(), torch.zeros_like(shifted)) + sum_exp = exped.sum(dim=-1, keepdim=True) + sum_exp = torch.clamp(sum_exp, min=1e-30) + return row_max + sum_exp.log() + + +def _local_softmax(values: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: + """Softmax over the True positions of ``mask``. + + Both tensors are ``[R, K]``. Returns ``[R, K]``: off-mask entries are 0 + and rows where the mask is all-False return all-zero (the row is + masked out downstream by ``row_valid``). + + Used for the IS weight ``w_v = softmax(ell_old over S_t)``. + """ + lse = _local_lse(values, mask) # [R, 1] + out = (values - lse).exp() + out = torch.where(mask, out, torch.zeros_like(out)) + return out + + +def _opd_one_sample( + logits_chunk: torch.Tensor, + *, + student_indices: torch.Tensor, # [R, Kq] long, GLOBAL vocab ids + student_old_lp: torch.Tensor, # [R, Kq] log pi_old(v) for v in student_indices + teacher_indices: torch.Tensor, # [R, Kp] long + teacher_lp: torch.Tensor, # [R, Kp] log pi_T(v) for v in teacher_indices + eps_lo: float, + eps_hi: float, + diff_clip: float | None, + tp_group, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """verl-aligned top-K OPD surrogate for one sample. + + Pipeline (notation matches the file docstring): + 1. Build ``S_t = student_indices ∩ teacher_indices``, expressed as + a [R, Kq] boolean mask aligned with ``student_indices``. For the + ``student`` / ``teacher`` modes Kq == Kp and the two index lists + are identical, so the mask is all-True. + 2. Pull GLOBAL teacher log-probs ``ell_T(v)`` at ``student_indices`` + via per-row index match. (The teacher gather pass already gave + us global log pi_T values; we just reorder them.) + 3. Compute GLOBAL student-current log-probs: + ``ell_cur(v) = raw_logits(v) - global_lse(raw_logits)`` + via the autograd-connected gather + global LSE helpers. Both + halves contribute gradient to the student's logits. + 4. Detached IS weight on S_t: + ``w_v = softmax(ell_old over S_t)``, sum_v w_v = 1. + 5. Detached advantage (with optional magnitude clamp on the diff): + ``diff_v = clamp(ell_T(v) - ell_old(v), -t, +t)`` (t = diff_clip, + skipped when diff_clip is None), + ``A_v = diff_v * w_v``. + 6. PPO surrogate using the GLOBAL log-ratio: + ``ppo_kl = ell_old - ell_cur`` (clamped for numerical safety), + ``L_v = max(-A_v * rho_v, -A_v * clip(rho_v, 1-eps_lo, 1+eps_hi))`` + via ``compute_policy_loss(ppo_kl, A)``. + 7. Per-token aggregation: SUM over S_t (NOT mean -- the IS weight + already normalises within S_t, so summing concentrates gradient + pressure on the high-w_v candidates without dividing it away). + + Returns: + per_token_pg [R] OPD surrogate per token (sum over S_t) + per_token_clip [R] fraction of S_t entries that hit the PPO clip + per_token_diff [R] w-weighted |ell_T - ell_old| per token (monitor) + row_valid [R] bool, True iff |S_t| >= 1 + """ + R = student_indices.size(0) + if R == 0: + z = student_indices.new_zeros((0,), dtype=torch.float32) + b = student_indices.new_zeros((0,), dtype=torch.bool) + return z, z, z, b + + # 1) Build the S_t mask on the student-indices axis. + if torch.equal(student_indices, teacher_indices): + # Fast path: student/teacher modes ship identical index sets. + sub_mask = torch.ones_like(student_indices, dtype=torch.bool) + eq = None # not needed + else: + # eq: [R, Kq, Kp] booleans, at most one True per (r, k). + eq = student_indices.unsqueeze(-1) == teacher_indices.unsqueeze(-2) + sub_mask = eq.any(dim=-1) # [R, Kq] + row_valid = sub_mask.any(dim=-1) # [R] + mask_f = sub_mask.float() # [R, Kq] + + # 2) Align teacher GLOBAL log-probs to student_indices ordering. + if eq is None: + teacher_lp_aligned = teacher_lp.float() + else: + # Weighted sum picks the unique matched teacher_lp per (r, k); + # off-mask entries collapse to 0 (zeroed out by mask_f anyway). + eq_f = eq.float() + teacher_lp_aligned = (eq_f * teacher_lp.unsqueeze(-2).float()).sum(dim=-1) + + # 3) GLOBAL student-current log-probs at student_indices (autograd). + # raw_at_K - global_lse(raw) = log pi_theta(v). + student_new_raw = _gather_student_raw_logits_at_indices( + logits_chunk, student_indices, tp_group + ).float() # [R, Kq] + global_lse = _vocab_parallel_global_lse(logits_chunk, tp_group) # [R, 1] + ell_cur = student_new_raw - global_lse # [R, Kq] + + ell_old = student_old_lp.float() # [R, Kq] + ell_T = teacher_lp_aligned # [R, Kq] + + # 4) Detached IS weight w_v = softmax(ell_old | S_t). + # Built from ell_old alone, so the diff_clip below does NOT change w. + w = _local_softmax(ell_old, sub_mask).detach() # [R, Kq] + + # 5) Detached advantage A_v = (ell_T - ell_old) * w_v. + # Optionally clamp the per-candidate teacher-vs-old log-ratio to + # [-diff_clip, +diff_clip]. Bounds advantage magnitude when the + # teacher and old student disagree wildly on a candidate (rare-token + # candidates, early training, cross-tokenizer distillation). + diff = (ell_T - ell_old).detach() # [R, Kq] + if diff_clip is not None: + diff = diff.clamp(min=-diff_clip, max=diff_clip) + advantage = (diff * w).detach() # [R, Kq] + + # 6) PPO surrogate with the GLOBAL log-ratio. + # compute_policy_loss expects ppo_kl = log(p_old / p_new) and returns + # -min(rho*A, clip(rho)*A), exactly our L_v. + ppo_kl = (ell_old - ell_cur).clamp(min=-_PPO_KL_CLAMP, max=_PPO_KL_CLAMP) + pg, clip = compute_policy_loss(ppo_kl, advantage, eps_lo, eps_hi) + + # 7) SUM over S_t (mask out non-subset entries). + per_token_pg = (pg * mask_f).sum(dim=-1) # [R] + # clipfrac: fraction of S_t entries that were clipped (not summed). + n_per_token = mask_f.sum(dim=-1).clamp(min=1.0) + per_token_clip = (clip * mask_f).sum(dim=-1) / n_per_token # [R] + # diff monitor: w-weighted |ell_T - ell_old| -- matches the loss's + # importance weighting so it's a calibrated "mean teacher gap" signal. + per_token_diff = (diff.abs() * w).sum(dim=-1) # [R] + + # Zero rows with empty subset so they don't pollute the trajectory mean. + row_valid_f = row_valid.float() + return ( + per_token_pg * row_valid_f, + per_token_clip * row_valid_f, + per_token_diff * row_valid_f, + row_valid, + ) + + +# --------------------------------------------------------------------------- +# Slime entry point +# --------------------------------------------------------------------------- + + +def hint_opd_loss_function( + args: Namespace, + batch: dict, + logits: torch.Tensor, + sum_of_sample_mean: Callable[[torch.Tensor], torch.Tensor], +) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: + """Slime ``--custom-loss-function-path`` entry point.""" + response_lengths = batch["response_lengths"] + total_lengths = batch["total_lengths"] + max_seq_lens = batch.get("max_seq_lens", None) + + w_rl = _w_rl() + w_opd = _w_opd() + eps_lo = _eps_clip_lo(args) + eps_hi = _eps_clip_hi(args) + diff_clip = _adv_diff_clip() + need_entropy_for_loss = args.entropy_coef != 0.0 + + # ---- forward: per-token student log-probs (and optional entropy) ---- + _, log_probs_and_entropy = get_log_probs_and_entropy( + logits, + args=args, + unconcat_tokens=batch["unconcat_tokens"], + total_lengths=total_lengths, + response_lengths=response_lengths, + with_entropy=need_entropy_for_loss, + max_seq_lens=max_seq_lens, + ) + new_log_probs_per_sample = log_probs_and_entropy["log_probs"] + new_log_probs = torch.cat(new_log_probs_per_sample, dim=0) + + # Hard guarantee: every old log-prob we use is from the Megatron actor + # forward (batch["log_probs"]), NEVER from SGLang rollout. + assert not getattr(args, "use_rollout_logprobs", False), ( + "hint_opd loss requires --use-rollout-logprobs to be unset so old-policy " + "log-probs come from the Megatron old_actor forward." + ) + + # ----------------- GRPO branch (token-level PG, weight w_rl) ----------------- + grpo_pg_loss = torch.zeros((), device=logits.device, dtype=torch.float32) + grpo_pg_clipfrac = torch.zeros((), device=logits.device, dtype=torch.float32) + ppo_kl_mean_sampled = torch.zeros((), device=logits.device, dtype=torch.float32) + if w_rl != 0.0: + old_log_probs = torch.cat(batch["log_probs"], dim=0) + ppo_kl_sampled = old_log_probs - new_log_probs + grpo_advantages = torch.cat(batch["advantages"], dim=0) + pg_loss_tokens, pg_clipfrac_tokens = compute_policy_loss( + ppo_kl_sampled, grpo_advantages, eps_lo, eps_hi + ) + grpo_pg_loss = sum_of_sample_mean(pg_loss_tokens) + grpo_pg_clipfrac = sum_of_sample_mean(pg_clipfrac_tokens) + ppo_kl_mean_sampled = sum_of_sample_mean(ppo_kl_sampled) + + # ----------------- Top-K OPD branch (verl-aligned, weight w_opd) ----------------- + opd_loss = torch.zeros((), device=logits.device, dtype=torch.float32) + opd_clipfrac_scalar = torch.zeros((), device=logits.device, dtype=torch.float32) + teacher_student_logp_diff_mean: torch.Tensor | None = None + subset_size_mean: torch.Tensor | None = None + + student_topk_lp = batch.get("topk_log_probs") + student_topk_idx = batch.get("topk_indices") + teacher_topk_lp = batch.get("prm_teacher_topk_log_probs") + teacher_topk_idx = batch.get("prm_teacher_topk_indices") + + have_student_side = ( + student_topk_lp is not None + and student_topk_idx is not None + and len(student_topk_lp) > 0 + and len(student_topk_idx) > 0 + ) + have_teacher_side = ( + teacher_topk_lp is not None + and teacher_topk_idx is not None + and len(teacher_topk_lp) > 0 + and len(teacher_topk_idx) > 0 + ) + + if w_opd != 0.0: + if not (have_student_side and have_teacher_side): + raise RuntimeError( + "hint_opd_loss requires both ('topk_log_probs', 'topk_indices') " + "and ('prm_teacher_topk_log_probs', 'prm_teacher_topk_indices') in " + "the batch. Confirm --distill-topk > 0, --prm-teacher-load is set, " + "and the slime old_actor pass runs inside emit_topk_logprobs() so " + "student-side top-K data is shipped." + ) + + tp_group = mpu.get_tensor_model_parallel_group() + all_pg = [] + all_clip = [] + all_diff = [] + all_size = [] + for i, (logits_chunk, _tokens_chunk) in enumerate( + get_responses( + logits, + args=args, + unconcat_tokens=batch["unconcat_tokens"], + total_lengths=total_lengths, + response_lengths=response_lengths, + max_seq_lens=max_seq_lens, + ) + ): + s_idx = student_topk_idx[i].to(device=logits_chunk.device, dtype=torch.long) + s_lp = student_topk_lp[i].to(device=logits_chunk.device, dtype=torch.float32) + t_idx = teacher_topk_idx[i].to(device=logits_chunk.device, dtype=torch.long) + t_lp = teacher_topk_lp[i].to(device=logits_chunk.device, dtype=torch.float32) + assert s_idx.size(0) == logits_chunk.size(0), ( + f"student topk size mismatch: s_idx[0]={s_idx.size(0)} " + f"vs logits_chunk[0]={logits_chunk.size(0)}" + ) + assert t_idx.size(0) == logits_chunk.size(0), ( + f"teacher topk size mismatch: t_idx[0]={t_idx.size(0)} " + f"vs logits_chunk[0]={logits_chunk.size(0)}" + ) + + pg_t, clip_t, diff_t, valid_t = _opd_one_sample( + logits_chunk, + student_indices=s_idx, + student_old_lp=s_lp, + teacher_indices=t_idx, + teacher_lp=t_lp, + eps_lo=eps_lo, + eps_hi=eps_hi, + diff_clip=diff_clip, + tp_group=tp_group, + ) + all_pg.append(pg_t) + all_clip.append(clip_t) + all_diff.append(diff_t) + # Per-token subset size for monitoring. + if torch.equal(s_idx, t_idx): + size_t = s_idx.new_full((s_idx.size(0),), s_idx.size(-1), dtype=torch.float32) + else: + eq = s_idx.unsqueeze(-1) == t_idx.unsqueeze(-2) + size_t = eq.any(dim=-1).float().sum(dim=-1) + all_size.append(size_t * valid_t.float()) + + opd_pg_tokens = torch.cat(all_pg, dim=0) + opd_clip_tokens = torch.cat(all_clip, dim=0) + opd_diff_tokens = torch.cat(all_diff, dim=0) + opd_size_tokens = torch.cat(all_size, dim=0) + opd_loss = sum_of_sample_mean(opd_pg_tokens) + opd_clipfrac_scalar = sum_of_sample_mean(opd_clip_tokens) + teacher_student_logp_diff_mean = sum_of_sample_mean(opd_diff_tokens) + subset_size_mean = sum_of_sample_mean(opd_size_tokens) + + # ----------------- entropy term ----------------- + if need_entropy_for_loss: + entropy = torch.cat(log_probs_and_entropy["entropy"], dim=0) + entropy_loss = sum_of_sample_mean(entropy) + else: + with torch.no_grad(): + _, ent_data = get_log_probs_and_entropy( + logits, + args=args, + unconcat_tokens=batch["unconcat_tokens"], + total_lengths=total_lengths, + response_lengths=response_lengths, + with_entropy=True, + max_seq_lens=max_seq_lens, + ) + entropy_loss = sum_of_sample_mean(torch.cat(ent_data["entropy"], dim=0)) + + # ----------------- combine ----------------- + loss = w_rl * grpo_pg_loss + w_opd * opd_loss - args.entropy_coef * entropy_loss + + # ----------------- optional KL-to-ref regulariser ----------------- + kl_loss = torch.tensor(0.0, device=logits.device) + if args.use_kl_loss and batch.get("ref_log_probs") is not None: + ref_log_probs = torch.cat(batch["ref_log_probs"], dim=0) + kl = compute_approx_kl( + new_log_probs, ref_log_probs, kl_loss_type=args.kl_loss_type, + ) + kl_loss = sum_of_sample_mean(kl) + loss = loss + args.kl_loss_coef * kl_loss + + if new_log_probs.numel() == 0: + loss = loss + 0 * logits.sum() + + # ----------------- monitoring ----------------- + train_rollout_logprob_abs_diff = None + if "rollout_log_probs" in batch and batch["rollout_log_probs"]: + rollout_lp = torch.cat(batch["rollout_log_probs"], dim=0) + train_rollout_logprob_abs_diff = sum_of_sample_mean( + (new_log_probs.detach() - rollout_lp).abs() + ) + + reported: dict[str, torch.Tensor] = { + "loss": loss.clone().detach(), + "grpo_pg_loss": grpo_pg_loss.clone().detach(), + "opd_loss": opd_loss.clone().detach(), + "entropy_loss": entropy_loss.clone().detach(), + "grpo_pg_clipfrac": grpo_pg_clipfrac.clone().detach(), + "opd_pg_clipfrac": opd_clipfrac_scalar.clone().detach(), + "ppo_kl_sampled": ppo_kl_mean_sampled.clone().detach(), + "w_rl": torch.tensor(w_rl, device=loss.device), + "w_opd": torch.tensor(w_opd, device=loss.device), + } + if teacher_student_logp_diff_mean is not None: + reported["opd_teacher_student_logp_topk_abs_mean"] = ( + teacher_student_logp_diff_mean.clone().detach() + ) + if subset_size_mean is not None: + reported["opd_subset_size"] = subset_size_mean.clone().detach() + if train_rollout_logprob_abs_diff is not None: + reported["train_rollout_logprob_abs_diff"] = ( + train_rollout_logprob_abs_diff.clone().detach() + ) + if args.use_kl_loss: + reported["kl_loss"] = kl_loss.clone().detach() + + return loss, reported diff --git a/openclaw-combine/hint_opd_select_loss.py b/openclaw-combine/hint_opd_select_loss.py new file mode 100644 index 0000000..4d66494 --- /dev/null +++ b/openclaw-combine/hint_opd_select_loss.py @@ -0,0 +1,537 @@ +"""Hint-OPD top-K loss with K-candidate teacher supervision selection. + +Drop-in alternative to ``hint_opd_loss.hint_opd_loss_function`` for the +``hint_opd_hint_rollout_select`` path (``--hint-m K``, ``--hint-selection``). + +Supports all 9 combinations of ``--hint-selection`` × ``--distill-subset-mode``: + +============================== ================= ============ ============== +``--hint-selection`` ``student`` ``overlap`` ``teacher`` +============================== ================= ============ ============== +``shortest`` in-kernel k*=0 in-kernel k*=0 legacy (actor-side k*=0) +``token_optimal`` in-kernel k*(t) in-kernel k*(t) legacy (actor-side k*(t)) +``sequence_optimal`` in-kernel k* in-kernel k* legacy (actor-side k*) +============================== ================= ============ ============== + +For ``hint_opt_exp`` ``sequence_optimal`` is per-sample (one ``k*`` for all +response tokens of a prompt) -- the rollout module ships a single CoT per +sample with no PRM-step structure (``step_wise_step_token_spans`` is +absent), so the kernel's "sequence = whole response" fallback fires. +This is the natural granularity here: the prompt has one response, the +response IS the sequence. Contrast with the retool variant +(``hybrid_stepwise_topk_opd_select_loss``) where ``step_token_spans`` +IS shipped and ``sequence_optimal`` is per-PRM-step. + +Selection score (all three subset modes): + + O[k, t] = | S^q_t ∩ S^p_{t,k} | + + * ``token_optimal`` : k*(t) = argmax_k O[k, t] + * ``sequence_optimal`` : k* = argmax_k Σ_t O[k, t] (one k* per sample, + broadcast over all response tokens). This is + the natural granularity for the hint_opt_exp + single-CoT pipeline (no tool-calling steps). + If a future caller starts shipping + ``step_token_spans`` metadata the kernel will + transparently switch to per-step k* via the + helper ``_select_k_star_per_token``. + +Subset-mode semantics (orthogonal to selection): + + * ``--distill-subset-mode student`` : ``S_t = S^q_t``. Teacher log-probs + arrive in ``prm_teacher_topk_log_probs_cand`` GATHERED at the + student's top-K (the actor-side multi-cand gather pass populates + this). Indices on S^q are constant across k; the per-(k, t) + selection signal travels in + ``prm_teacher_native_topk_indices_cand`` (each candidate's own + native top-K). The loss feeds ``_opd_one_sample`` with + ``teacher_indices = student_topk_idx`` so its subset matches S^q. + + * ``--distill-subset-mode overlap`` : ``S_t = S^q_t ∩ S^p_{t, k*(t)}``. + Teacher log-probs / indices arrive at the candidate's native top-K + in the ``_cand`` keys. The kernel computes the overlap mask + internally inside ``_opd_one_sample``. + + * ``--distill-subset-mode teacher`` : ``S_t = S^p_{t, k*(t)}``. The + actor-side ``train_actor`` does selection AND the extra student-old + re-gather at the chosen S^p, then collapses the cand tensors into + legacy single-cand keys. The kernel here simply delegates to the + legacy single-cand ``hint_opd_loss_function`` for that case. + +The GRPO branch (``w_rl * grpo_pg``) and the KL-to-ref / entropy +branches are imported verbatim from ``hint_opd_loss``; only the OPD +branch differs. + +Environment variables (same as the topk-OPD baseline) +----------------------------------------------------- + HINT_OPD_W_RL weight on GRPO PG (default 0.0) + HINT_OPD_W_OPD weight on top-K OPD (default 1.0) + HINT_OPD_PPO_CLIP_EPS_LO override args.eps_clip (optional) + HINT_OPD_PPO_CLIP_EPS_HI override args.eps_clip_high (optional) + HINT_OPD_ADV_DIFF_CLIP clamp on (ell_T - ell_old) (default 2.0) +""" + +from __future__ import annotations + +from argparse import Namespace +from collections.abc import Callable + +import torch +from megatron.core import mpu + +from hint_opd_loss import ( + _adv_diff_clip, + _eps_clip_hi, + _eps_clip_lo, + _opd_one_sample, + _w_opd, + _w_rl, + hint_opd_loss_function, +) +from slime.backends.megatron_utils.loss import get_log_probs_and_entropy, get_responses +from slime.utils.ppo_utils import compute_approx_kl, compute_policy_loss + + +# --------------------------------------------------------------------------- +# Per-candidate / student top-K overlap on the response-token axis. +# --------------------------------------------------------------------------- + + +def _overlap_count_per_token( + student_idx: torch.Tensor, # [R, K_q] + teacher_idx: torch.Tensor, # [K, R, K_p] +) -> torch.Tensor: + """``O[k, t] = | S^q_t ∩ S^p_{t,k} |`` via broadcasted equality. + + Memory: ``[K, R, K_q, K_p]`` boolean intermediate. With K_q=K_p=20 and + R≤8192, K≤8 this is ~25 MB per sample — small. + """ + eq = student_idx.unsqueeze(0).unsqueeze(-1) == teacher_idx.unsqueeze(-2) + return eq.any(dim=-1).sum(dim=-1).to(torch.long) # [K, R] + + +def _select_k_star_per_token( + overlap_kr: torch.Tensor | None, # [K, R] or None for shortest + *, + hint_selection: str, + step_token_spans: list[list[int]] | None, # per-sample list of [t0, t1] + R: int, + device: torch.device, +) -> torch.Tensor: + """Compute ``k*(t) ∈ [0, K)`` per response token. + + * ``shortest`` : ``k* = 0`` for every token (the rollout module + orders candidates shortest-first, so candidate 0 is the shortest hint + that survived the min-token / dedup filter). ``overlap_kr`` may be + ``None`` since no overlap signal is required. + * ``token_optimal`` : argmax over candidates per token. + * ``sequence_optimal`` : argmax over candidates per PRM step + (broadcast over tokens in the step). When ``step_token_spans`` is + missing/empty -- the case for the hint_opt_exp single-CoT pipeline + where the "sequence" IS the whole response -- this collapses to + per-sample argmax (one ``k*`` for every token in the response). + """ + if hint_selection == "shortest": + return torch.zeros(R, dtype=torch.long, device=device) + assert overlap_kr is not None, ( + "_select_k_star_per_token: overlap_kr is required for " + f"hint_selection={hint_selection!r}." + ) + K, R_kr = overlap_kr.shape + assert R_kr == R, f"overlap_kr R mismatch: {R_kr} vs {R}" + if K == 1: + return torch.zeros(R, dtype=torch.long, device=device) + if hint_selection == "token_optimal": + return overlap_kr.argmax(dim=0) # [R] + if hint_selection == "sequence_optimal": + if not step_token_spans: + k_star_scalar = int(overlap_kr.sum(dim=-1).argmax().item()) + return torch.full((R,), k_star_scalar, dtype=torch.long, device=device) + out = torch.zeros(R, dtype=torch.long, device=device) + for span in step_token_spans: + t0, t1 = int(span[0]), int(span[1]) + t0 = max(0, min(t0, R)) + t1 = max(t0, min(t1, R)) + if t1 == t0: + continue + seg_score = overlap_kr[:, t0:t1].sum(dim=-1) # [K] + out[t0:t1] = int(seg_score.argmax().item()) + return out + raise ValueError( + f"Unknown --hint-selection: {hint_selection!r}. Expected " + "'shortest' / 'token_optimal' / 'sequence_optimal'." + ) + + +def _gather_along_K( + cand_tensor: torch.Tensor, # [K, R, *] + k_star_per_token: torch.Tensor, # [R] +) -> torch.Tensor: + """Slice ``cand_tensor[k*(t), t, ...]`` per token. Returns ``[R, *]``.""" + R = cand_tensor.size(1) + trailing = cand_tensor.shape[2:] + expand_shape = (1, R, *trailing) + view_shape = (1, R) + tuple(1 for _ in trailing) + gather_idx = k_star_per_token.view(view_shape).expand(expand_shape) + return torch.gather(cand_tensor, dim=0, index=gather_idx).squeeze(0) + + +# --------------------------------------------------------------------------- +# Slime entry point +# --------------------------------------------------------------------------- + + +def hint_opd_select_loss_function( + args: Namespace, + batch: dict, + logits: torch.Tensor, + sum_of_sample_mean: Callable[[torch.Tensor], torch.Tensor], +) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: + """``--custom-loss-function-path`` entry point for hint_opd_hint_select. + + Dispatches on ``args.hint_selection`` AND ``args.distill_subset_mode``. + All 9 combinations are supported: + + * ``subset_mode == "teacher"`` (any ``hint_selection``): the actor-side + ``train_actor`` already collapsed the cand teacher tensors into + legacy single-cand keys -- selection + extra student-old re-gather + happen there. ``_select_teacher_cand_per_sample`` short-circuits + ``shortest`` to ``k* = 0``. We delegate to the legacy single-cand + ``hint_opd_loss_function`` for all three teacher cells. + + * ``subset_mode in {student, overlap}`` (any ``hint_selection``): + runs the in-kernel selection branch below. Under multi-cand the + ``train_actor`` does NOT collapse the cand keys (only teacher mode + does), so the legacy single-cand kernel CANNOT be used here -- the + ``_cand`` keys are present but the legacy ``prm_teacher_topk_*`` + keys are absent. The in-kernel branch consumes the ``_cand`` keys + directly and slices at ``k*(t)`` (=0 for ``shortest``). + """ + hint_selection = getattr(args, "hint_selection", "shortest") + subset_mode = getattr(args, "distill_subset_mode", "student") + if hint_selection not in ("shortest", "token_optimal", "sequence_optimal"): + raise ValueError( + f"Unknown --hint-selection: {hint_selection!r}. Expected one of " + "'shortest', 'token_optimal', 'sequence_optimal'." + ) + if subset_mode == "teacher": + return hint_opd_loss_function(args, batch, logits, sum_of_sample_mean) + + response_lengths = batch["response_lengths"] + total_lengths = batch["total_lengths"] + max_seq_lens = batch.get("max_seq_lens", None) + + w_rl = _w_rl() + w_opd = _w_opd() + eps_lo = _eps_clip_lo(args) + eps_hi = _eps_clip_hi(args) + diff_clip = _adv_diff_clip() + need_entropy_for_loss = args.entropy_coef != 0.0 + + # ---- forward: per-token student log-probs (and optional entropy) ---- + _, log_probs_and_entropy = get_log_probs_and_entropy( + logits, + args=args, + unconcat_tokens=batch["unconcat_tokens"], + total_lengths=total_lengths, + response_lengths=response_lengths, + with_entropy=need_entropy_for_loss, + max_seq_lens=max_seq_lens, + ) + new_log_probs = torch.cat(log_probs_and_entropy["log_probs"], dim=0) + + assert not getattr(args, "use_rollout_logprobs", False), ( + "hint_opd_select_loss requires --use-rollout-logprobs to be unset so " + "old-policy log-probs come from the Megatron old_actor." + ) + + # ----------------- GRPO branch (token PG, weight w_rl) ----------------- + grpo_pg_loss = torch.zeros((), device=logits.device, dtype=torch.float32) + grpo_pg_clipfrac = torch.zeros((), device=logits.device, dtype=torch.float32) + ppo_kl_mean_sampled = torch.zeros((), device=logits.device, dtype=torch.float32) + if w_rl != 0.0: + old_log_probs = torch.cat(batch["log_probs"], dim=0) + ppo_kl_sampled = old_log_probs - new_log_probs + rl_advantages = torch.cat(batch["advantages"], dim=0) + pg_loss_tokens, pg_clipfrac_tokens = compute_policy_loss( + ppo_kl_sampled, rl_advantages, eps_lo, eps_hi + ) + grpo_pg_loss = sum_of_sample_mean(pg_loss_tokens) + grpo_pg_clipfrac = sum_of_sample_mean(pg_clipfrac_tokens) + ppo_kl_mean_sampled = sum_of_sample_mean(ppo_kl_sampled) + + # ----------------- K-candidate top-K OPD branch ----------------- + opd_loss = torch.zeros((), device=logits.device, dtype=torch.float32) + opd_clipfrac_scalar = torch.zeros((), device=logits.device, dtype=torch.float32) + teacher_student_logp_diff_mean: torch.Tensor | None = None + subset_size_mean: torch.Tensor | None = None + sel_overlap_mean: torch.Tensor | None = None + sel_k_star_mean: torch.Tensor | None = None + + student_topk_lp = batch.get("topk_log_probs") + student_topk_idx = batch.get("topk_indices") + teacher_topk_lp_cand = batch.get("prm_teacher_topk_log_probs_cand") + teacher_topk_idx_cand = batch.get("prm_teacher_topk_indices_cand") + teacher_native_idx_cand = batch.get("prm_teacher_native_topk_indices_cand") + step_spans_per_sample = batch.get("step_wise_step_token_spans") + + have_student = ( + student_topk_lp is not None + and student_topk_idx is not None + and len(student_topk_lp) > 0 + and len(student_topk_idx) > 0 + ) + have_teacher_cand = ( + teacher_topk_lp_cand is not None + and teacher_topk_idx_cand is not None + and len(teacher_topk_lp_cand) > 0 + and len(teacher_topk_idx_cand) > 0 + ) + + if w_opd != 0.0: + if not (have_student and have_teacher_cand): + raise RuntimeError( + "hint_opd_select_loss requires both student top-K " + "(topk_log_probs / topk_indices) and the candidate-axis " + "teacher top-K (prm_teacher_topk_log_probs_cand / " + "prm_teacher_topk_indices_cand) in the batch. Confirm " + "--distill-topk > 0, --hint-m > 0, --distill-subset-mode " + "in {student, overlap, teacher}, and that the rollout " + "function path is " + "hint_opd_hint_rollout_select.generate_rollout_with_hint_select." + ) + if ( + subset_mode == "student" + and hint_selection != "shortest" + and ( + teacher_native_idx_cand is None + or len(teacher_native_idx_cand) == 0 + ) + ): + raise RuntimeError( + "subset_mode=student with --hint-selection in " + "{token_optimal, sequence_optimal} requires " + "prm_teacher_native_topk_indices_cand (the per-candidate " + "selection signal) in the batch. Confirm slime's " + "gather_at_indices multi-cand path is engaged. " + "(shortest does not need it because k*=0 always.)" + ) + + tp_group = mpu.get_tensor_model_parallel_group() + all_pg = [] + all_clip = [] + all_diff = [] + all_size = [] + all_overlap_sel = [] + all_k_star = [] + for i, (logits_chunk, _tokens_chunk) in enumerate( + get_responses( + logits, + args=args, + unconcat_tokens=batch["unconcat_tokens"], + total_lengths=total_lengths, + response_lengths=response_lengths, + max_seq_lens=max_seq_lens, + ) + ): + s_idx = student_topk_idx[i].to(device=logits_chunk.device, dtype=torch.long) + s_lp = student_topk_lp[i].to(device=logits_chunk.device, dtype=torch.float32) + t_idx_cand = teacher_topk_idx_cand[i].to( + device=logits_chunk.device, dtype=torch.long + ) + t_lp_cand = teacher_topk_lp_cand[i].to( + device=logits_chunk.device, dtype=torch.float32 + ) + + R = logits_chunk.size(0) + assert s_idx.dim() == 2 and s_idx.size(0) == R, ( + f"student topk shape mismatch: s_idx={tuple(s_idx.shape)} " + f"vs R={R}" + ) + assert t_idx_cand.dim() == 3 and t_idx_cand.size(1) == R, ( + f"teacher topk_cand shape mismatch: t_idx_cand=" + f"{tuple(t_idx_cand.shape)} vs R={R}; expected [K, R, K_p]." + ) + assert t_lp_cand.shape == t_idx_cand.shape, ( + f"teacher logp/idx cand shape mismatch: " + f"lp={tuple(t_lp_cand.shape)} idx={tuple(t_idx_cand.shape)}" + ) + + # Selection signal: only computed when actually needed. + # * ``shortest`` : k*=0 always, no overlap signal needed. + # * ``student`` mode : t_idx_cand is a constant copy of S^q + # across k, so overlap with it would be + # trivially K_q; use each candidate's + # NATIVE top-K (emitted by slime under + # ``emit_native_topk_indices``). + # * ``overlap`` mode : use the candidate's own top-K + # (= t_idx_cand under this subset mode). + if hint_selection == "shortest": + overlap_kr = None + else: + if subset_mode == "student": + sel_idx_src = teacher_native_idx_cand[i].to( + device=logits_chunk.device, dtype=torch.long + ) + assert sel_idx_src.shape[1] == R, ( + f"native_topk shape mismatch: native=" + f"{tuple(sel_idx_src.shape)} vs R={R}; " + "expected [K, R, K_p]." + ) + else: + sel_idx_src = t_idx_cand + overlap_kr = _overlap_count_per_token(s_idx, sel_idx_src) # [K, R] + + spans_i = ( + step_spans_per_sample[i] + if step_spans_per_sample is not None and i < len(step_spans_per_sample) + else None + ) + k_star_per_token = _select_k_star_per_token( + overlap_kr, + hint_selection=hint_selection, + step_token_spans=spans_i, + R=R, + device=logits_chunk.device, + ) + + # Slice cand tensors at k*(t) per token. Under student mode the + # log-probs are at S^q (constant indices across k); under overlap + # they are at the candidate's own top-K. + t_lp_sel = _gather_along_K(t_lp_cand, k_star_per_token) + if subset_mode == "student": + # Force the kernel's S^p to S^q so the loss subset is S^q. + # (Indices in t_idx_cand are already constant across k, but + # this keeps the contract explicit.) + t_idx_sel = s_idx + else: + t_idx_sel = _gather_along_K(t_idx_cand, k_star_per_token) + + if overlap_kr is None: + # ``shortest``: report selected-overlap as 0 so the wandb + # panel is well-defined and visibly distinct from the + # optimal modes (where it's ~K_q for student mode). + overlap_sel_per_token = torch.zeros( + R, device=logits_chunk.device, dtype=torch.float32 + ) + else: + row_idx = torch.arange(R, device=overlap_kr.device) + overlap_sel_per_token = overlap_kr[k_star_per_token, row_idx].float() + + pg_t, clip_t, diff_t, valid_t = _opd_one_sample( + logits_chunk, + student_indices=s_idx, + student_old_lp=s_lp, + teacher_indices=t_idx_sel, + teacher_lp=t_lp_sel, + eps_lo=eps_lo, + eps_hi=eps_hi, + diff_clip=diff_clip, + tp_group=tp_group, + ) + all_pg.append(pg_t) + all_clip.append(clip_t) + all_diff.append(diff_t) + if torch.equal(s_idx, t_idx_sel): + size_t = s_idx.new_full( + (s_idx.size(0),), s_idx.size(-1), dtype=torch.float32 + ) + else: + eq = s_idx.unsqueeze(-1) == t_idx_sel.unsqueeze(-2) + size_t = eq.any(dim=-1).float().sum(dim=-1) + all_size.append(size_t * valid_t.float()) + all_overlap_sel.append(overlap_sel_per_token * valid_t.float()) + all_k_star.append(k_star_per_token.float() * valid_t.float()) + + opd_pg_tokens = torch.cat(all_pg, dim=0) + opd_clip_tokens = torch.cat(all_clip, dim=0) + opd_diff_tokens = torch.cat(all_diff, dim=0) + opd_size_tokens = torch.cat(all_size, dim=0) + opd_overlap_sel_tokens = torch.cat(all_overlap_sel, dim=0) + opd_k_star_tokens = torch.cat(all_k_star, dim=0) + opd_loss = sum_of_sample_mean(opd_pg_tokens) + opd_clipfrac_scalar = sum_of_sample_mean(opd_clip_tokens) + teacher_student_logp_diff_mean = sum_of_sample_mean(opd_diff_tokens) + subset_size_mean = sum_of_sample_mean(opd_size_tokens) + sel_overlap_mean = sum_of_sample_mean(opd_overlap_sel_tokens) + sel_k_star_mean = sum_of_sample_mean(opd_k_star_tokens) + + # ----------------- entropy term ----------------- + if need_entropy_for_loss: + entropy = torch.cat(log_probs_and_entropy["entropy"], dim=0) + entropy_loss = sum_of_sample_mean(entropy) + else: + with torch.no_grad(): + _, ent_data = get_log_probs_and_entropy( + logits, + args=args, + unconcat_tokens=batch["unconcat_tokens"], + total_lengths=total_lengths, + response_lengths=response_lengths, + with_entropy=True, + max_seq_lens=max_seq_lens, + ) + entropy_loss = sum_of_sample_mean(torch.cat(ent_data["entropy"], dim=0)) + + # ----------------- combine ----------------- + loss = w_rl * grpo_pg_loss + w_opd * opd_loss - args.entropy_coef * entropy_loss + + # ----------------- optional KL-to-ref regulariser ----------------- + kl_loss = torch.tensor(0.0, device=logits.device) + if args.use_kl_loss and batch.get("ref_log_probs") is not None: + ref_log_probs = torch.cat(batch["ref_log_probs"], dim=0) + kl = compute_approx_kl( + new_log_probs, ref_log_probs, kl_loss_type=args.kl_loss_type, + ) + kl_loss = sum_of_sample_mean(kl) + loss = loss + args.kl_loss_coef * kl_loss + + if new_log_probs.numel() == 0: + loss = loss + 0 * logits.sum() + + # ----------------- monitoring ----------------- + train_rollout_logprob_abs_diff = None + if "rollout_log_probs" in batch and batch["rollout_log_probs"]: + rollout_lp = torch.cat(batch["rollout_log_probs"], dim=0) + train_rollout_logprob_abs_diff = sum_of_sample_mean( + (new_log_probs.detach() - rollout_lp).abs() + ) + + # Reported-metric naming matches `hint_opd_loss_function` so wandb + # panels created for the single-cand baseline keep working. + reported: dict[str, torch.Tensor] = { + "loss": loss.clone().detach(), + "grpo_pg_loss": grpo_pg_loss.clone().detach(), + "opd_loss": opd_loss.clone().detach(), + "entropy_loss": entropy_loss.clone().detach(), + "grpo_pg_clipfrac": grpo_pg_clipfrac.clone().detach(), + "opd_pg_clipfrac": opd_clipfrac_scalar.clone().detach(), + "ppo_kl_sampled": ppo_kl_mean_sampled.clone().detach(), + "w_rl": torch.tensor(w_rl, device=loss.device), + "w_opd": torch.tensor(w_opd, device=loss.device), + } + if teacher_student_logp_diff_mean is not None: + reported["opd_teacher_student_logp_topk_abs_mean"] = ( + teacher_student_logp_diff_mean.clone().detach() + ) + if subset_size_mean is not None: + reported["opd_subset_size"] = subset_size_mean.clone().detach() + if sel_overlap_mean is not None: + reported["sel_overlap_at_k_star"] = sel_overlap_mean.clone().detach() + if sel_k_star_mean is not None: + reported["sel_k_star_mean"] = sel_k_star_mean.clone().detach() + if train_rollout_logprob_abs_diff is not None: + reported["train_rollout_logprob_abs_diff"] = ( + train_rollout_logprob_abs_diff.clone().detach() + ) + if args.use_kl_loss: + reported["kl_loss"] = kl_loss.clone().detach() + + # Embed selection-mode / subset-mode tags (constant ints) so the wandb + # dashboard can group runs without needing the run config. + mode_id = {"shortest": 0, "token_optimal": 1, "sequence_optimal": 2}[hint_selection] + subset_id = {"student": 0, "overlap": 1, "teacher": 2}.get(subset_mode, -1) + reported["hint_selection_mode_id"] = torch.tensor(mode_id, device=loss.device) + reported["distill_subset_mode_id"] = torch.tensor(subset_id, device=loss.device) + + return loss, reported diff --git a/openclaw-combine/run_qwen3_4b_openclaw_topk_select.sh b/openclaw-combine/run_qwen3_4b_openclaw_topk_select.sh index baf3e33..d548f64 100755 --- a/openclaw-combine/run_qwen3_4b_openclaw_topk_select.sh +++ b/openclaw-combine/run_qwen3_4b_openclaw_topk_select.sh @@ -261,7 +261,7 @@ ray start --head --node-ip-address "${MASTER_ADDR}" --num-gpus "${NUM_GPUS}" --d RUNTIME_ENV_JSON="{ \"env_vars\": { - \"PYTHONPATH\": \"${REPO_ROOT}/Megatron-LM:${SCRIPT_DIR}:${REPO_ROOT}/openclaw-opd:${REPO_ROOT}/hint_opt_exp:${SLIME_ROOT}\", + \"PYTHONPATH\": \"${REPO_ROOT}/Megatron-LM:${SCRIPT_DIR}:${REPO_ROOT}/openclaw-opd:${SLIME_ROOT}\", \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\", \"OPENCLAW_EVAL_MODE\": \"${OPENCLAW_EVAL_MODE}\", \"OPENCLAW_COMBINE_OPD_TEACHER_SOURCE\": \"${OPENCLAW_COMBINE_OPD_TEACHER_SOURCE}\", diff --git a/slime/slime/backends/megatron_utils/model_provider.py b/slime/slime/backends/megatron_utils/model_provider.py index 09971a5..326a1e6 100644 --- a/slime/slime/backends/megatron_utils/model_provider.py +++ b/slime/slime/backends/megatron_utils/model_provider.py @@ -109,7 +109,19 @@ def get_model_provider_func( # the most consequential one: without it, full activations stay resident # for every layer and large-context RL runs OOM. Forward it explicitly so # bridge runs are at least memory-comparable to raw runs. - if getattr(args, "recompute_granularity", None) is not None: + # Qwen3-VL bridge providers with deepstack visual features cannot use + # full-layer checkpointing: the upstream checkpoint path passes + # deepstack_visual_embeds as list[Tensor] into save_for_backward(). + skip_full_recompute = ( + getattr(args, "recompute_granularity", None) == "full" + and bool(getattr(provider, "deepstack_visual_indexes", None)) + ) + if skip_full_recompute: + print( + "Bridge provider: skipped full activation recompute for " + f"{type(provider).__name__} with deepstack visual features" + ) + elif getattr(args, "recompute_granularity", None) is not None: provider.recompute_granularity = args.recompute_granularity provider.recompute_method = args.recompute_method provider.recompute_num_layers = args.recompute_num_layers