Skip to content

feat: decouple IS correction level from GSPO policy ratio - #2269

Open
bo-ke wants to merge 1 commit into
NVIDIA-NeMo:mainfrom
bo-ke:feat/decouple-is-correction-level
Open

bo-ke wants to merge 1 commit into
NVIDIA-NeMo:mainfrom
bo-ke:feat/decouple-is-correction-level

Conversation

@bo-ke

@bo-ke bo-ke commented Apr 15, 2026

Copy link
Copy Markdown

What does this PR do?

Decouple IS correction granularity from the GSPO policy ratio so they can be
configured independently.

Currently sequence_level_importance_ratios=True bundles two behaviors into a
single flag:

  1. GSPO policy ratio: exp(mean(log(π_curr/π_prev))) — numerically stable
  2. IS correction weight: exp(sum(prev - gen)) — collapses for long
    sequences

For long sequences (~9k+ tokens), even a small per-token logprob mismatch between
training and inference backends (e.g., ~2% with vLLM + Megatron) accumulates via
exp(sum) to near-zero IS weights, effectively reducing the learning rate by
10–50x with high variance across steps.

This PR adds a new is_correction_level config field to ClippedPGLossConfig
that controls IS correction weight computation independently of the PPO ratio
(which remains controlled by sequence_level_importance_ratios).

Supported values:

  • "sequence" — exp(sum(prev - gen)), one weight per sequence (current
    behavior)
  • "sequence_mean" — exp(mean(prev - gen)), geometric mean, numerically stable
    for long sequences
  • "token" — exp(prev - gen) per token, token-level correction

When unset (None), behavior is derived from sequence_level_importance_ratios
for full backward compatibility.

Issues

None — discovered during GRPO training with Qwen3.5-35B-A3B (MoE, ~9k avg
response tokens) where sampling_importance_ratio showed mean ~0.4, std ~0.27,
with some steps collapsing to 0.02.

Usage

loss_fn:
  # GSPO policy ratio (unchanged)                                                
  sequence_level_importance_ratios: true
  token_level_loss: false                                                        
  # IS correction now independently configurable                               
  is_correction_level: "sequence_mean"  # or "token", or "sequence" (default)    
  use_importance_sampling_correction: true                                       
  truncated_importance_sampling_ratio: 2                                         
  truncated_importance_sampling_type: tis                                        

Before your PR is "Ready for review"

Pre checks:

  • Make sure you read and followed /NVIDIA-NeMo/RL/blob/main/CONTRIBUTING.md
  • Did you write any new necessary tests?
  • Did you run the unit tests and functional tests locally? Visit our
    /NVIDIA-NeMo/RL/blob/main/docs/testing.md for how to run tests
  • Did you add or update any necessary documentation? Visit our
    /NVIDIA-NeMo/RL/blob/main/docs/documentation.md for how to write, build and test
    the docs.

Additional Information

  • All 44 existing tests pass (backward compat verified)
  • New test test_clipped_pg_loss_gspo_is_correction_level_sequence_mean validates
    that exp(mean) produces IS weights closer to 1.0 than exp(sum) and that the loss
    value matches hand-calculated expectations
  • seq-mask-tis compatibility assertion updated to account for the new field

@bo-ke
bo-ke requested review from a team as code owners April 15, 2026 03:42
@copy-pr-bot

copy-pr-bot Bot commented Apr 15, 2026

Copy link
Copy Markdown

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

Add `is_correction_level` config field to `ClippedPGLossConfig` that
controls IS correction weight computation independently of the PPO ratio
(which is still controlled by `sequence_level_importance_ratios`).

Currently `sequence_level_importance_ratios=True` bundles two behaviors:
1. GSPO policy ratio: exp(mean(log(π_curr/π_prev))) — stable
2. IS correction: exp(sum(prev - gen)) — collapses for long sequences

For long sequences (~9k+ tokens), even a small per-token logprob
mismatch between training and inference backends accumulates via exp(sum)
to near-zero IS weights, effectively reducing the learning rate by 10-50x
with high variance across steps.

The new `is_correction_level` field supports:
- "sequence": exp(sum) — current behavior
- "sequence_mean": exp(mean) — geometric mean, numerically stable
- "token": per-token exp(prev - gen)

When unset (None), behavior is derived from
`sequence_level_importance_ratios` for full backward compatibility.

Signed-off-by: kebo01 <kebo01@baidu.com>
@bo-ke
bo-ke force-pushed the feat/decouple-is-correction-level branch from 9931c9e to c7b1b01 Compare April 15, 2026 03:46
@chtruong814 chtruong814 added the needs-follow-up Issue needs follow-up label Apr 17, 2026
@svcnvidia-nemo-ci svcnvidia-nemo-ci added waiting-on-maintainers Waiting on maintainers to respond and removed needs-follow-up Issue needs follow-up labels Apr 21, 2026
@zhiyuanma-agents

Copy link
Copy Markdown

Independent reproduction on a second model family, supporting this PR.

Setup: Nemotron 3.5 Super VL 120B-A12B (MoE, 12B active; early-access checkpoint), LoRA, async GRPO non-colocated (vLLM generation node + DTensor trainer node), 32 prompts × 8 rollouts per step, 40–60-turn tool-use episodes averaging ~8–10k generated tokens (context up to 128k), loss_fn.sequence_level_importance_ratios: true, use_importance_sampling_correction: true, truncated_importance_sampling_type: tis, ratio 2.0.

Step 1 with the current exp(sum) weight: train/sampling_importance_ratio 0.0596, versus 0.9999 at step 1 of the same recipe with token-level ratios. The policy ratio itself was healthy — train/probs_ratio 1.0001 (max 1.0026, min 0.9975), gen_kl_error 0.0027, token_mult_prob_error 1.024, is_oob_ratio 0.011 — yet grad_norm was 8.6e-5 and the loss −1.7e-4: the whole batch was down-weighted ~17× by the sequence-summed weight, not by any real train/inference mismatch. Switching the weight to exp(mean over tokens) — this PR's is_correction_level: sequence_mean — is the fix we applied locally; we will follow up here with the post-fix step-1 sampling_importance_ratio from the restarted run.

One addition worth considering: log the pre-clamp weight alongside the clamped one. With TIS the [0, 2] clamp hides how far the raw weight sits from 1, which is exactly the signal that would have flagged this on step 1.

@zhiyuanma-agents

Copy link
Copy Markdown

Follow-up with the post-fix numbers, as promised above. Same recipe and prompt pool, restarted with the length-normalised sequence weight (exp of the token mean, i.e. this PR's sequence_mean), step 1:

metric exp(sum) — before exp(mean) — after
train/sampling_importance_ratio 0.0596 0.9972
is_oob_ratio (TIS cap 2.0) 0.011 0.0
grad_norm 8.6e-5 8.6e-3 (100×)
probs_ratio (max / min) 1.0001 (1.0026 / 0.9975) 0.99999 (1.0031 / 0.99698)
gen_kl_error 0.0027 0.00277
token_mult_prob_error 1.024 1.0247

The train/inference mismatch metrics are unchanged between the two runs, so the 100× change in grad_norm is entirely the importance weight. One caveat: the loss magnitude at step 1 is also depressed by the recipe's LinearLR warm-up (start factor 1e-8 over 10 iterations), so grad_norm is the clean before/after signal here, not the loss value.

@svcnvidia-nemo-ci svcnvidia-nemo-ci added waiting-on-maintainers Waiting on maintainers to respond and removed waiting-on-maintainers Waiting on maintainers to respond labels Sep 22, 2026

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

community-request waiting-on-maintainers Waiting on maintainers to respond

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants