Repository navigation
Conversation
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>
9931c9e to
c7b1b01
Compare
|
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), Step 1 with the current exp(sum) weight: 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. |
|
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
The train/inference mismatch metrics are unchanged between the two runs, so the 100× change in |
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=Truebundles two behaviors into asingle flag:
exp(mean(log(π_curr/π_prev)))— numerically stableexp(sum(prev - gen))— collapses for longsequences
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 by10–50x with high variance across steps.
This PR adds a new
is_correction_levelconfig field toClippedPGLossConfigthat 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 (currentbehavior)
"sequence_mean"—exp(mean(prev - gen)), geometric mean, numerically stablefor long sequences
"token"—exp(prev - gen)per token, token-level correctionWhen unset (
None), behavior is derived fromsequence_level_importance_ratiosfor full backward compatibility.
Issues
None — discovered during GRPO training with Qwen3.5-35B-A3B (MoE, ~9k avg
response tokens) where
sampling_importance_ratioshowed mean ~0.4, std ~0.27,with some steps collapsing to 0.02.
Usage
Before your PR is "Ready for review"
Pre checks:
/NVIDIA-NeMo/RL/blob/main/docs/testing.md for how to run tests
/NVIDIA-NeMo/RL/blob/main/docs/documentation.md for how to write, build and test
the docs.
Additional Information
that exp(mean) produces IS weights closer to 1.0 than exp(sum) and that the loss
value matches hand-calculated expectations