From 4b123193213410cdeaedb0d3ac26f7afa78bf38a Mon Sep 17 00:00:00 2001 From: Yi-Fu Wu Date: Sun, 22 Feb 2026 13:37:39 -0800 Subject: [PATCH 1/5] add overlong filtering for gym rollout + configurable advantage clipping c1abc09b4b9fa1a75cbd10f5ba07c74442c5ff67 https://github.com/NVIDIA-NeMo/RL/issues/1949 Signed-off-by: Yi-Fu Wu Co-Authored-By: Jiaqi Zeng Signed-off-by: Arnav Kundu --- nemo_rl/algorithms/grpo.py | 54 +++++++++++++++++++++++++++++++++++--- 1 file changed, 51 insertions(+), 3 deletions(-) diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py index 0337c8e05b2..9052d30f088 100644 --- a/nemo_rl/algorithms/grpo.py +++ b/nemo_rl/algorithms/grpo.py @@ -144,6 +144,11 @@ class GRPOConfig(TypedDict): max_num_steps: int max_rollout_turns: int normalize_rewards: bool + # Clipping bounds for normalized advantages to prevent extreme values + # When set, advantages are clipped to [advantage_clip_low, advantage_clip_high] after normalization + # Default: null (no clipping) + advantage_clip_low: NotRequired[float | None] + advantage_clip_high: NotRequired[float | None] use_leave_one_out_baseline: bool val_period: int val_batch_size: int | None # None for NeMo-Gym compatibility @@ -1976,6 +1981,14 @@ def grpo_train( ) del baseline_for_log + # Clip advantages to prevent extreme values from small std normalization + clip_low = master_config["grpo"].get("advantage_clip_low") + clip_high = master_config["grpo"].get("advantage_clip_high") + if clip_low is not None: + train_data["advantages"] = train_data["advantages"].clamp(min=clip_low) + if clip_high is not None: + train_data["advantages"] = train_data["advantages"].clamp(max=clip_high) + memory_tracker.snapshot_start_of_stage("Policy train", dir()) print("▶ Preparing for training...", flush=True) with timer.time("training_prep"): @@ -2956,9 +2969,36 @@ def async_grpo_train( # Prepare training data (same as sync version) with timer.time("data_processing"): - add_grpo_token_loss_masks_and_generation_logprobs( - repeated_batch["message_log"] - ) + # Apply overlong filtering - mask out truncated sequences from loss computation + with timer.time("overlong_filter"): + use_overlong_filtering = master_config["grpo"]["overlong_filtering"] + if use_overlong_filtering: + loss_multiplier = repeated_batch["loss_multiplier"].clone() + truncated = repeated_batch["truncated"] + + if isinstance(truncated, list): + truncated = torch.tensor(truncated, dtype=torch.bool) + + loss_multiplier[truncated] = 0 + repeated_batch["loss_multiplier"] = loss_multiplier + + # Add loss mask to each message + # Only unmask assistant messages that were actually generated (have generation_logprobs), + # not assistant messages that were part of the prompt history + for i, message_log in enumerate(repeated_batch["message_log"]): + for j, message in enumerate(message_log): + if message["role"] == "assistant" and "generation_logprobs" in message: + message["token_loss_mask"] = torch.ones_like( + message["token_ids"] + ) + else: + message["token_loss_mask"] = torch.zeros_like( + message["token_ids"] + ) + if "generation_logprobs" not in message: + message["generation_logprobs"] = torch.zeros_like( + message["token_ids"], dtype=torch.float32 + ) # Convert to flat format for training flat_messages, input_lengths = batched_message_log_to_flat_message( @@ -3086,6 +3126,14 @@ def async_grpo_train( f" 📊 Advantages stats: min={advantages.min():.4f}, max={advantages.max():.4f}, mean={advantages.mean():.4f}, std={advantages.std():.4f}" ) + # Clip advantages to prevent extreme values from small std normalization + clip_low = master_config["grpo"].get("advantage_clip_low") + clip_high = master_config["grpo"].get("advantage_clip_high") + if clip_low is not None: + train_data["advantages"] = train_data["advantages"].clamp(min=clip_low) + if clip_high is not None: + train_data["advantages"] = train_data["advantages"].clamp(max=clip_high) + print("▶ Preparing for training...") with timer.time("training_prep"): policy.prepare_for_training() From 92838c058484666182a5e439bab6f35295d2597e Mon Sep 17 00:00:00 2001 From: Arnav Kundu Date: Wed, 3 Jun 2026 14:35:15 -0700 Subject: [PATCH 2/5] fixed incompatible interfaces to Master Config Signed-off-by: Arnav Kundu --- nemo_rl/algorithms/grpo.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py index 9052d30f088..b7b25316191 100644 --- a/nemo_rl/algorithms/grpo.py +++ b/nemo_rl/algorithms/grpo.py @@ -1982,8 +1982,8 @@ def grpo_train( del baseline_for_log # Clip advantages to prevent extreme values from small std normalization - clip_low = master_config["grpo"].get("advantage_clip_low") - clip_high = master_config["grpo"].get("advantage_clip_high") + clip_low = master_config.grpo.get("advantage_clip_low", None) + clip_high = master_config.grpo.get("advantage_clip_high", None) if clip_low is not None: train_data["advantages"] = train_data["advantages"].clamp(min=clip_low) if clip_high is not None: @@ -2971,7 +2971,7 @@ def async_grpo_train( with timer.time("data_processing"): # Apply overlong filtering - mask out truncated sequences from loss computation with timer.time("overlong_filter"): - use_overlong_filtering = master_config["grpo"]["overlong_filtering"] + use_overlong_filtering = master_config.grpo.get("overlong_filtering", False) if use_overlong_filtering: loss_multiplier = repeated_batch["loss_multiplier"].clone() truncated = repeated_batch["truncated"] @@ -3127,8 +3127,8 @@ def async_grpo_train( ) # Clip advantages to prevent extreme values from small std normalization - clip_low = master_config["grpo"].get("advantage_clip_low") - clip_high = master_config["grpo"].get("advantage_clip_high") + clip_low = master_config.grpo.get("advantage_clip_low", None) + clip_high = master_config.grpo.get("advantage_clip_high", None) if clip_low is not None: train_data["advantages"] = train_data["advantages"].clamp(min=clip_low) if clip_high is not None: From b824131ddd141a757eea4bc6adddfd7041afbf13 Mon Sep 17 00:00:00 2001 From: Arnav Kundu Date: Thu, 4 Jun 2026 10:00:19 -0700 Subject: [PATCH 3/5] using the helper function to mask prompts Signed-off-by: Arnav Kundu --- nemo_rl/algorithms/grpo.py | 17 +++-------------- 1 file changed, 3 insertions(+), 14 deletions(-) diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py index b7b25316191..f0cd1caae93 100644 --- a/nemo_rl/algorithms/grpo.py +++ b/nemo_rl/algorithms/grpo.py @@ -2985,20 +2985,9 @@ def async_grpo_train( # Add loss mask to each message # Only unmask assistant messages that were actually generated (have generation_logprobs), # not assistant messages that were part of the prompt history - for i, message_log in enumerate(repeated_batch["message_log"]): - for j, message in enumerate(message_log): - if message["role"] == "assistant" and "generation_logprobs" in message: - message["token_loss_mask"] = torch.ones_like( - message["token_ids"] - ) - else: - message["token_loss_mask"] = torch.zeros_like( - message["token_ids"] - ) - if "generation_logprobs" not in message: - message["generation_logprobs"] = torch.zeros_like( - message["token_ids"], dtype=torch.float32 - ) + add_grpo_token_loss_masks_and_generation_logprobs( + repeated_batch["message_log"] + ) # Convert to flat format for training flat_messages, input_lengths = batched_message_log_to_flat_message( From 70fd5bd0a5bd0302ef62ce6b941ed6b0e72b9e05 Mon Sep 17 00:00:00 2001 From: Arnav Kundu Date: Thu, 4 Jun 2026 11:03:01 -0700 Subject: [PATCH 4/5] ruff format Signed-off-by: Arnav Kundu --- nemo_rl/algorithms/grpo.py | 20 +++++++++++++++----- 1 file changed, 15 insertions(+), 5 deletions(-) diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py index f0cd1caae93..176a3ba4da6 100644 --- a/nemo_rl/algorithms/grpo.py +++ b/nemo_rl/algorithms/grpo.py @@ -1985,9 +1985,13 @@ def grpo_train( clip_low = master_config.grpo.get("advantage_clip_low", None) clip_high = master_config.grpo.get("advantage_clip_high", None) if clip_low is not None: - train_data["advantages"] = train_data["advantages"].clamp(min=clip_low) + train_data["advantages"] = train_data["advantages"].clamp( + min=clip_low + ) if clip_high is not None: - train_data["advantages"] = train_data["advantages"].clamp(max=clip_high) + train_data["advantages"] = train_data["advantages"].clamp( + max=clip_high + ) memory_tracker.snapshot_start_of_stage("Policy train", dir()) print("▶ Preparing for training...", flush=True) @@ -2971,7 +2975,9 @@ def async_grpo_train( with timer.time("data_processing"): # Apply overlong filtering - mask out truncated sequences from loss computation with timer.time("overlong_filter"): - use_overlong_filtering = master_config.grpo.get("overlong_filtering", False) + use_overlong_filtering = master_config.grpo.get( + "overlong_filtering", False + ) if use_overlong_filtering: loss_multiplier = repeated_batch["loss_multiplier"].clone() truncated = repeated_batch["truncated"] @@ -3119,9 +3125,13 @@ def async_grpo_train( clip_low = master_config.grpo.get("advantage_clip_low", None) clip_high = master_config.grpo.get("advantage_clip_high", None) if clip_low is not None: - train_data["advantages"] = train_data["advantages"].clamp(min=clip_low) + train_data["advantages"] = train_data["advantages"].clamp( + min=clip_low + ) if clip_high is not None: - train_data["advantages"] = train_data["advantages"].clamp(max=clip_high) + train_data["advantages"] = train_data["advantages"].clamp( + max=clip_high + ) print("▶ Preparing for training...") with timer.time("training_prep"): From 9217dac448d12e59a942ce7042f70a4a21b34b31 Mon Sep 17 00:00:00 2001 From: Arnav Kundu Date: Thu, 4 Jun 2026 22:28:08 -0700 Subject: [PATCH 5/5] adding changes to grpo_sync, update yaml configs, tests and documentation Signed-off-by: Arnav Kundu --- docs/guides/grpo.md | 19 +++ examples/configs/grpo_math_1B.yaml | 2 + examples/nemo_gym/grpo_nanov3.yaml | 2 + ...rkplace_assistant_nemotron_nano_v2_9b.yaml | 2 + nemo_rl/algorithms/grpo.py | 46 +++--- nemo_rl/algorithms/grpo_sync.py | 2 + tests/unit/algorithms/test_grpo.py | 137 ++++++++++++++++++ .../unit/reference_configs/grpo_math_1B.yaml | 2 + 8 files changed, 189 insertions(+), 23 deletions(-) diff --git a/docs/guides/grpo.md b/docs/guides/grpo.md index 0cddc5c95d6..e9785730310 100755 --- a/docs/guides/grpo.md +++ b/docs/guides/grpo.md @@ -456,6 +456,25 @@ grpo: Set `overlong_filtering` to true when training on tasks where truncation at the maximum sequence length is expected, such as long-form reasoning or mathematical proofs. +#### Advantage Clipping + +After advantage normalization, per-token advantages can become very large when the per-prompt reward standard deviation is small. The optional `advantage_clip_low` and `advantage_clip_high` parameters clamp normalized advantages to a bounded range before policy training. + +When both values are `null` (the default), no clipping is applied. When set, advantages are clamped after normalization: + +$$ +A_{\text{clipped}} = \text{clamp}(A,\ \text{advantage\_clip\_low},\ \text{advantage\_clip\_high}) +$$ + +To configure: +```yaml +grpo: + advantage_clip_low: null # default: no lower bound + advantage_clip_high: null # default: no upper bound +``` + +Set explicit bounds (for example, `-5.0` and `5.0`) when small reward variance produces extreme normalized advantages that destabilize training. + #### Top-p and top-k filtering The implementation aligns with vLLM’s top-p and top-k filtering by applying an equivalent process to the logits. diff --git a/examples/configs/grpo_math_1B.yaml b/examples/configs/grpo_math_1B.yaml index d1fa48cc2ae..be98da0edb3 100644 --- a/examples/configs/grpo_math_1B.yaml +++ b/examples/configs/grpo_math_1B.yaml @@ -11,6 +11,8 @@ grpo: val_at_start: false val_at_end: false overlong_filtering: false + advantage_clip_low: null + advantage_clip_high: null max_val_samples: 256 val_batch_size: 256 seed: 42 diff --git a/examples/nemo_gym/grpo_nanov3.yaml b/examples/nemo_gym/grpo_nanov3.yaml index da25f752645..7b03f99f2ac 100644 --- a/examples/nemo_gym/grpo_nanov3.yaml +++ b/examples/nemo_gym/grpo_nanov3.yaml @@ -12,6 +12,8 @@ grpo: val_at_start: False val_at_end: False overlong_filtering: true + advantage_clip_low: null + advantage_clip_high: null max_val_samples: null val_batch_size: 256 seed: 42 diff --git a/examples/nemo_gym/grpo_workplace_assistant_nemotron_nano_v2_9b.yaml b/examples/nemo_gym/grpo_workplace_assistant_nemotron_nano_v2_9b.yaml index 7a854e64b78..e7a539903d3 100644 --- a/examples/nemo_gym/grpo_workplace_assistant_nemotron_nano_v2_9b.yaml +++ b/examples/nemo_gym/grpo_workplace_assistant_nemotron_nano_v2_9b.yaml @@ -10,6 +10,8 @@ grpo: val_at_start: true val_at_end: false overlong_filtering: false + advantage_clip_low: null + advantage_clip_high: null max_val_samples: null # inferred from size of val dataset. for multi evals, repeat val ds via `num_repeats` in `ng_prepare_data`. val_batch_size: null seed: 42 diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py index 176a3ba4da6..865cf842e6a 100644 --- a/nemo_rl/algorithms/grpo.py +++ b/nemo_rl/algorithms/grpo.py @@ -1184,6 +1184,20 @@ def _create_advantage_estimator(master_config: MasterConfig): return adv_estimator +def _clip_grpo_advantages( + advantages: torch.Tensor, + grpo_config: dict[str, Any], +) -> torch.Tensor: + """Clamp normalized advantages when clip bounds are configured.""" + clip_low = grpo_config.get("advantage_clip_low", None) + clip_high = grpo_config.get("advantage_clip_high", None) + if clip_low is not None: + advantages = advantages.clamp(min=clip_low) + if clip_high is not None: + advantages = advantages.clamp(max=clip_high) + return advantages + + def refit_policy_generation( policy: ColocatablePolicyInterface, policy_generation: GenerationInterface, @@ -1982,16 +1996,9 @@ def grpo_train( del baseline_for_log # Clip advantages to prevent extreme values from small std normalization - clip_low = master_config.grpo.get("advantage_clip_low", None) - clip_high = master_config.grpo.get("advantage_clip_high", None) - if clip_low is not None: - train_data["advantages"] = train_data["advantages"].clamp( - min=clip_low - ) - if clip_high is not None: - train_data["advantages"] = train_data["advantages"].clamp( - max=clip_high - ) + train_data["advantages"] = _clip_grpo_advantages( + train_data["advantages"], master_config.grpo + ) memory_tracker.snapshot_start_of_stage("Policy train", dir()) print("▶ Preparing for training...", flush=True) @@ -2975,9 +2982,9 @@ def async_grpo_train( with timer.time("data_processing"): # Apply overlong filtering - mask out truncated sequences from loss computation with timer.time("overlong_filter"): - use_overlong_filtering = master_config.grpo.get( - "overlong_filtering", False - ) + use_overlong_filtering = master_config.grpo[ + "overlong_filtering" + ] if use_overlong_filtering: loss_multiplier = repeated_batch["loss_multiplier"].clone() truncated = repeated_batch["truncated"] @@ -3122,16 +3129,9 @@ def async_grpo_train( ) # Clip advantages to prevent extreme values from small std normalization - clip_low = master_config.grpo.get("advantage_clip_low", None) - clip_high = master_config.grpo.get("advantage_clip_high", None) - if clip_low is not None: - train_data["advantages"] = train_data["advantages"].clamp( - min=clip_low - ) - if clip_high is not None: - train_data["advantages"] = train_data["advantages"].clamp( - max=clip_high - ) + train_data["advantages"] = _clip_grpo_advantages( + train_data["advantages"], master_config.grpo + ) print("▶ Preparing for training...") with timer.time("training_prep"): diff --git a/nemo_rl/algorithms/grpo_sync.py b/nemo_rl/algorithms/grpo_sync.py index 42ee0689134..e78b6e43906 100644 --- a/nemo_rl/algorithms/grpo_sync.py +++ b/nemo_rl/algorithms/grpo_sync.py @@ -48,6 +48,7 @@ from nemo_rl.algorithms.grpo import ( GRPOSaveState, MasterConfig, + _clip_grpo_advantages, _create_advantage_estimator, _log_mixed_rewards_and_advantages_information, _should_log_nemo_gym_responses, @@ -844,6 +845,7 @@ def grpo_train_sync( # ── Driver delta-write: advantages + (post-masking) # sample_mask under the same meta.sample_ids so workers fetch # the union via train_presharded. + advantages = _clip_grpo_advantages(advantages, master_config.grpo) policy.write_to_dataplane( meta, fields={ diff --git a/tests/unit/algorithms/test_grpo.py b/tests/unit/algorithms/test_grpo.py index 410d406b9cc..c538b5d3f67 100644 --- a/tests/unit/algorithms/test_grpo.py +++ b/tests/unit/algorithms/test_grpo.py @@ -346,6 +346,8 @@ def val_iter(self): "use_leave_one_out_baseline": False, "normalize_rewards": False, "overlong_filtering": False, + "advantage_clip_low": None, + "advantage_clip_high": None, "reward_scaling": {"enabled": False}, "reward_shaping": {"enabled": False}, "use_dynamic_sampling": False, @@ -1689,6 +1691,141 @@ def test_grpo_train_skips_reference_policy_logprobs_when_configured( ) +def _run_single_grpo_train_step(mock_grpo_components, train_func, monkeypatch): + """Run one GRPO training step with rollout/logprob infrastructure mocked.""" + mock_batch = next(iter(mock_grpo_components["train_dataloader"])) + mock_rollout_metrics = {"mean_gen_tokens_per_sample": 2.0} + policy = mock_grpo_components["policy"] + master_config = mock_grpo_components["master_config"] + master_config.grpo["max_num_steps"] = 1 + master_config.grpo["max_num_epochs"] = 1 + master_config.grpo["val_period"] = 0 + master_config.grpo["val_at_start"] = False + master_config.grpo["use_dynamic_sampling"] = False + + if train_func == async_grpo_train: + master_config.policy["generation"]["colocated"]["enabled"] = False + with ( + mock_async_grpo_infrastructure(mock_batch, mock_rollout_metrics), + _patched_logprob_phase(policy), + ): + train_func( + policy, + None, + mock_grpo_components["train_dataloader"], + mock_grpo_components["val_dataloader"], + mock_grpo_components["tokenizer"], + mock_grpo_components["loss_fn"], + mock_grpo_components["task_to_env"], + mock_grpo_components["val_task_to_env"], + mock_grpo_components["logger"], + mock_grpo_components["checkpointer"], + _default_grpo_save_state(), + master_config, + ) + else: + with ( + _patched_logprob_phase(policy), + patch( + "nemo_rl.algorithms.grpo.run_multi_turn_rollout", + return_value=(mock_batch, mock_rollout_metrics), + ), + patch( + "nemo_rl.algorithms.grpo.run_async_multi_turn_rollout", + return_value=(mock_batch, mock_rollout_metrics), + ), + patch( + "nemo_rl.algorithms.grpo.compute_and_apply_seq_logprob_error_masking", + return_value=_mock_seq_logprob_error_result(), + ), + ): + train_func( + policy, + None, + mock_grpo_components["train_dataloader"], + mock_grpo_components["val_dataloader"], + mock_grpo_components["tokenizer"], + mock_grpo_components["loss_fn"], + mock_grpo_components["task_to_env"], + mock_grpo_components["val_task_to_env"], + mock_grpo_components["logger"], + mock_grpo_components["checkpointer"], + _default_grpo_save_state(), + master_config, + ) + + +@pytest.mark.parametrize("train_func", [grpo_train, async_grpo_train]) +def test_grpo_train_clips_advantages_when_configured( + mock_grpo_components, train_func, monkeypatch +): + """Advantages passed to policy.train are clamped when clip bounds are set.""" + extreme_advantages = torch.tensor([[-10.0, 15.0]]) + mock_adv_estimator = MagicMock() + mock_adv_estimator.compute_advantage.return_value = extreme_advantages.clone() + monkeypatch.setattr( + "nemo_rl.algorithms.grpo._create_advantage_estimator", + lambda _cfg: mock_adv_estimator, + ) + + master_config = mock_grpo_components["master_config"] + master_config.grpo["advantage_clip_low"] = -2.0 + master_config.grpo["advantage_clip_high"] = 3.0 + + _run_single_grpo_train_step(mock_grpo_components, train_func, monkeypatch) + + policy = mock_grpo_components["policy"] + policy.train.assert_called_once() + clipped = policy.train.call_args[0][0]["advantages"] + assert clipped.min().item() == -2.0 + assert clipped.max().item() == 3.0 + + +@pytest.mark.parametrize("train_func", [grpo_train, async_grpo_train]) +def test_grpo_train_preserves_advantages_when_clipping_disabled( + mock_grpo_components, train_func, monkeypatch +): + """Advantages are unchanged when advantage_clip_low/high are null.""" + extreme_advantages = torch.tensor([[-10.0, 15.0]]) + mock_adv_estimator = MagicMock() + mock_adv_estimator.compute_advantage.return_value = extreme_advantages.clone() + monkeypatch.setattr( + "nemo_rl.algorithms.grpo._create_advantage_estimator", + lambda _cfg: mock_adv_estimator, + ) + + master_config = mock_grpo_components["master_config"] + master_config.grpo["advantage_clip_low"] = None + master_config.grpo["advantage_clip_high"] = None + + _run_single_grpo_train_step(mock_grpo_components, train_func, monkeypatch) + + policy = mock_grpo_components["policy"] + policy.train.assert_called_once() + advantages = policy.train.call_args[0][0]["advantages"] + assert torch.equal(advantages, extreme_advantages) + + +def test_clip_grpo_advantages_respects_config_bounds(): + """Shared clip helper clamps only when bounds are configured.""" + from nemo_rl.algorithms.grpo import _clip_grpo_advantages + + extreme_advantages = torch.tensor([[-10.0, 15.0], [0.0, 5.0]]) + + clipped = _clip_grpo_advantages( + extreme_advantages.clone(), + {"advantage_clip_low": -2.0, "advantage_clip_high": 3.0}, + ) + assert clipped.min().item() == -2.0 + assert clipped.max().item() == 3.0 + + unclipped = _clip_grpo_advantages( + extreme_advantages.clone(), + {"advantage_clip_low": None, "advantage_clip_high": None}, + ) + assert torch.equal(unclipped, extreme_advantages) + + @pytest.mark.parametrize("train_func", [grpo_train, async_grpo_train]) def test_grpo_train_skips_prev_logprobs_when_force_on_policy_ratio( mock_grpo_components, train_func diff --git a/tests/unit/reference_configs/grpo_math_1B.yaml b/tests/unit/reference_configs/grpo_math_1B.yaml index 60d3a02398a..ddd51438ad4 100644 --- a/tests/unit/reference_configs/grpo_math_1B.yaml +++ b/tests/unit/reference_configs/grpo_math_1B.yaml @@ -11,6 +11,8 @@ grpo: val_at_start: false val_at_end: false overlong_filtering: false + advantage_clip_low: null + advantage_clip_high: null max_val_samples: 256 val_batch_size: 256 seed: 42