feat(grpo): surface mean TIS weight as a returned metric#10
Open
dongseokmotif wants to merge 2 commits into
Open
feat(grpo): surface mean TIS weight as a returned metric#10dongseokmotif wants to merge 2 commits into
dongseokmotif wants to merge 2 commits into
Conversation
The fused GRPO loss applies the vLLM importance-sampling (TIS) correction (per_token_loss *= vllm_is_ratio) but never returned its magnitude, so a consumer could only stub it as 0 -- indistinguishable from "TIS off". Append the mean applied vllm_is_ratio (over valid tokens) to the returned metrics, LAST and only when TIS is active, so existing positional consumers (metrics[0]=KL when beta!=0, then clip_fraction) and the chunked-loss base-class per-chunk aggregation are unaffected, and backward's *grad_metrics already absorbs the extra output. Mirror it in the torch reference so the vllm_is_ratio correctness test validates it. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
On the force_on_policy_tis path the kernel already forms the pre-clamp ratio rho = pi_theta/pi_vllm = exp(curr - gen). Compute monitor-only mismatch diagnostics from it (masked per-token means, normalized like the KL/clip metrics) and append them to the returned metrics (LAST, force path only): - vllm_kl = E[rho - log rho - 1] (K3 KL(pi_vllm||pi_theta), a.k.a. gen_kl_error) - E[rho^2] (chi^2 divergence = this - 1; raw 2nd moment logged so it survives per-microbatch sum aggregation) - tis_trunc_frac = fraction of tokens outside [tis_min, tis_max] Detects the rollout(vLLM block-fp8) vs training(bf16) probability mismatch (SGA bias/variance failure modes). Never enter the loss; backward unaffected. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem
The fused GRPO loss applies the vLLM importance-sampling (TIS) correction (
per_token_loss *= vllm_is_ratio) but never returned its magnitude. A consumer could only stubsampling_importance_ratioas0-- indistinguishable from "TIS off", which is misleading when debugging whether TIS is actually active.Change
grpo_loss.py: append the mean appliedvllm_is_ratio(over valid tokens) to the returnedmetrics, LAST and only when TIS is active. Existing positional consumers (metrics[0]=KL whenbeta!=0, then clip_fraction) and the chunked-loss base-class per-chunk aggregation are unaffected;backward's*grad_metricsalready absorbs the extra output.test_grpo_loss.py: mirror the same metric in the torch reference sotest_correctness_with_vllm_is_ratiovalidates it (lengths stay matched).Note
Could not run the chunked GRPO tests on this host (no CUDA/Triton) -- validated by mirroring the existing KL/clip metric pattern + the torch reference, and CPU-simulated the consumer parse. Needs a GPU/CI run to confirm numerics.
Pairs with MotifTechnologies/NeMoRL-torchtitan#91, which reads this metric and surfaces it as
train/sampling_importance_ratio(dual-compatible: works with or without this change).🤖 Generated with Claude Code