From 0eaff522c27cbd1f6b6ab0bb15c62c77ec09e4c9 Mon Sep 17 00:00:00 2001 From: Yi-Fu Wu Date: Mon, 30 Mar 2026 16:53:15 -0700 Subject: [PATCH] fix: skip prev_logprobs computation when force_on_policy_ratio is true When force_on_policy_ratio=True, the importance sampling ratio is forced to 1.0, so prev_logprobs are unnecessary. Skip the expensive prepare_for_lp_inference() and get_logprobs() calls in both sync and async GRPO paths. In the loss function, use curr_logprobs.detach() as prev_logprobs instead of loading placeholder zeros from data. Also guards against incompatible use of seq_logprob_error_threshold with force_on_policy_ratio (the threshold requires real prev_logprobs). Part of https://github.com/NVIDIA-NeMo/RL/issues/1906 Co-Authored-By: Jiaqi Zeng Co-Authored-By: Claude Opus 4.6 (1M context) Signed-off-by: Yi-Fu Wu --- nemo_rl/algorithms/grpo.py | 83 ++++++++++++++------ nemo_rl/algorithms/loss/loss_functions.py | 8 +- tests/unit/algorithms/test_loss_functions.py | 50 ++++++++++++ 3 files changed, 118 insertions(+), 23 deletions(-) diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py index e550429ce2c..35d17c2f9f6 100644 --- a/nemo_rl/algorithms/grpo.py +++ b/nemo_rl/algorithms/grpo.py @@ -1728,9 +1728,14 @@ def grpo_train( metrics_logging_data["content"] = flat_messages["content"] memory_tracker.snapshot_start_of_stage("Computing logprobs", dir()) - print("▶ Preparing for logprob inference...", flush=True) - with timer.time("logprob_inference_prep"): - policy.prepare_for_lp_inference() + # Skip prev_logprobs computation when force_on_policy_ratio=True + skip_prev_logprobs = master_config["loss_fn"].get("force_on_policy_ratio", False) + if skip_prev_logprobs: + print("▶ Skipping prev_logprobs (force_on_policy_ratio=True)...", flush=True) + else: + print("▶ Preparing for logprob inference...", flush=True) + with timer.time("logprob_inference_prep"): + policy.prepare_for_lp_inference() print("▶ Computing logprobs...", flush=True) with timer.time("policy_and_reference_logprobs"): @@ -1744,9 +1749,12 @@ def grpo_train( **extra_multimodal_data, } ) - train_data["prev_logprobs"] = policy.get_logprobs( - logprob_data, timer=timer - )["logprobs"] + if not skip_prev_logprobs: + train_data["prev_logprobs"] = policy.get_logprobs( + logprob_data, timer=timer + )["logprobs"] + else: + train_data["prev_logprobs"] = torch.zeros_like(train_data["generation_logprobs"]) if not master_config["grpo"].get( "skip_reference_policy_logprobs_calculation" @@ -1761,6 +1769,20 @@ def grpo_train( del logprob_data del extra_multimodal_data + # Seq-level logprob error metrics/masking require real prev_logprobs + seq_logprob_error_threshold = master_config["grpo"].get( + "seq_logprob_error_threshold", None + ) + if skip_prev_logprobs: + # Cannot compute seq-level metrics with placeholder prev_logprobs + assert seq_logprob_error_threshold is None, ( + "seq_logprob_error_threshold requires prev_logprobs computation; " + "cannot use with force_on_policy_ratio=True" + ) + max_seq_mult_prob_error = 0.0 + num_masked_seqs = 0 + masked_correct_pct = 0.0 + else: ( max_seq_mult_prob_error, num_masked_seqs, @@ -1768,9 +1790,7 @@ def grpo_train( ) = compute_and_apply_seq_logprob_error_masking( train_data=train_data, rewards=rewards, - seq_logprob_error_threshold=master_config["grpo"][ - "seq_logprob_error_threshold" - ], + seq_logprob_error_threshold=seq_logprob_error_threshold, ) # Compute advantages with adv_estimator using correct mask and logprobs @@ -2779,23 +2799,44 @@ def async_grpo_train( train_data.to("cpu") # Training phase (same as sync version) - print("▶ Preparing for logprob inference...") - with timer.time("logprob_inference_prep"): - policy.prepare_for_lp_inference() + # Skip prev_logprobs computation when force_on_policy_ratio=True + skip_prev_logprobs = master_config["loss_fn"].get("force_on_policy_ratio", False) + if skip_prev_logprobs: + print("▶ Skipping prev_logprobs (force_on_policy_ratio=True)...", flush=True) + fprop_logprobs = torch.zeros_like(train_data["generation_logprobs"]) + else: + print("▶ Preparing for logprob inference...") + with timer.time("logprob_inference_prep"): + policy.prepare_for_lp_inference() - print("▶ Computing logprobs...") + print("▶ Computing logprobs...", flush=True) with timer.time("policy_and_reference_logprobs"): - fprop_logprobs = policy.get_logprobs( - train_data, - timer=timer, - )["logprobs"] + if not skip_prev_logprobs: + fprop_logprobs = policy.get_logprobs( + train_data, + timer=timer, + )["logprobs"] + reference_logprobs = policy.get_reference_policy_logprobs( train_data, timer=timer, )["reference_logprobs"] - train_data["prev_logprobs"] = fprop_logprobs - train_data["reference_policy_logprobs"] = reference_logprobs + train_data["prev_logprobs"] = fprop_logprobs + train_data["reference_policy_logprobs"] = reference_logprobs + # Seq-level logprob error metrics/masking require real prev_logprobs + seq_logprob_error_threshold = master_config["grpo"].get( + "seq_logprob_error_threshold", None + ) + if skip_prev_logprobs: + assert seq_logprob_error_threshold is None, ( + "seq_logprob_error_threshold requires prev_logprobs computation; " + "cannot use with force_on_policy_ratio=True" + ) + max_seq_mult_prob_error = 0.0 + num_masked_seqs = 0 + masked_correct_pct = 0.0 + else: ( max_seq_mult_prob_error, num_masked_seqs, @@ -2803,9 +2844,7 @@ def async_grpo_train( ) = compute_and_apply_seq_logprob_error_masking( train_data=train_data, rewards=rewards, - seq_logprob_error_threshold=master_config["grpo"][ - "seq_logprob_error_threshold" - ], + seq_logprob_error_threshold=seq_logprob_error_threshold, ) # Compute advantages with adv_estimator using correct mask and logprobs diff --git a/nemo_rl/algorithms/loss/loss_functions.py b/nemo_rl/algorithms/loss/loss_functions.py index 49860b541c3..9bdba41f43a 100755 --- a/nemo_rl/algorithms/loss/loss_functions.py +++ b/nemo_rl/algorithms/loss/loss_functions.py @@ -200,7 +200,8 @@ def __call__( token_mask = data["token_mask"][:, 1:] sample_mask = data["sample_mask"] advantages = data["advantages"][:, 1:] - prev_logprobs = data["prev_logprobs"][:, 1:] + # Skip loading prev_logprobs when force_on_policy_ratio=True (will use curr_logprobs instead) + prev_logprobs = None if self.force_on_policy_ratio else data["prev_logprobs"][:, 1:] generation_logprobs = data["generation_logprobs"][:, 1:] if self.reference_policy_kl_penalty != 0: reference_policy_logprobs = data["reference_policy_logprobs"][:, 1:] @@ -210,6 +211,11 @@ def __call__( mask = token_mask * sample_mask.unsqueeze(-1) + # For truly on-policy training, use curr_logprobs as prev_logprobs + # This avoids computing prev_logprobs upstream + if self.force_on_policy_ratio: + prev_logprobs = curr_logprobs.detach() + # token_mult_prob_error # See more details and other metrics in docs/guides/grpo.md#metrics lp_error = torch.abs(generation_logprobs - prev_logprobs) # noqa: F841 (precommit ignore for now) diff --git a/tests/unit/algorithms/test_loss_functions.py b/tests/unit/algorithms/test_loss_functions.py index a49eff739d4..8b2009fa213 100644 --- a/tests/unit/algorithms/test_loss_functions.py +++ b/tests/unit/algorithms/test_loss_functions.py @@ -637,6 +637,56 @@ def test_clipped_pg_loss_force_on_policy_ratio(): assert metrics["probs_ratio_clamped_max"] == 1.0 +def test_clipped_pg_loss_force_on_policy_ratio_ignores_prev_logprobs(): + """Tests that force_on_policy_ratio ignores prev_logprobs from data and uses curr_logprobs instead. + + When force_on_policy_ratio=True, the loss function should use curr_logprobs.detach() + as prev_logprobs, so the actual prev_logprobs in data are irrelevant. This allows + skipping the expensive prev_logprobs computation upstream. + """ + if not torch.cuda.is_available(): + pytest.skip("No GPU available") + + device = "cuda" + data, batch_size, seq_len, vocab_size = _setup_clipped_pg_test_data(device=device) + + cfg = deepcopy(basic_pg_loss_test_config) + cfg["force_on_policy_ratio"] = True + loss_fn = ClippedPGLossFn(cfg) + + curr_lp = torch.tensor([[-1.0, -1.0, -1.0]], device=device) + input_ids = data["input_ids"] + dummy_logits = _create_exact_logits( + curr_lp, input_ids, batch_size, seq_len, vocab_size, device + ) + + # Run with correct prev_logprobs + data_1, _, _, _ = _setup_clipped_pg_test_data(device=device) + data_1["prev_logprobs"][0, 1:] = curr_lp + loss_input_1, data_1 = prepare_loss_input(dummy_logits.clone(), data_1, loss_fn) + loss_1, metrics_1 = loss_fn( + data=data_1, + global_valid_seqs=torch.sum(data_1["sample_mask"]), + global_valid_toks=torch.sum(data_1["sample_mask"].unsqueeze(-1) * data_1["token_mask"]), + **loss_input_1, + ) + + # Run with wildly different prev_logprobs (should be ignored) + data_2, _, _, _ = _setup_clipped_pg_test_data(device=device) + data_2["prev_logprobs"][0, 1:] = torch.tensor([-10.0, -10.0, -10.0], device=device) + loss_input_2, data_2 = prepare_loss_input(dummy_logits.clone(), data_2, loss_fn) + loss_2, metrics_2 = loss_fn( + data=data_2, + global_valid_seqs=torch.sum(data_2["sample_mask"]), + global_valid_toks=torch.sum(data_2["sample_mask"].unsqueeze(-1) * data_2["token_mask"]), + **loss_input_2, + ) + + # Both should produce identical loss and ratios since prev_logprobs is ignored + torch.testing.assert_close(loss_1, loss_2) + assert metrics_1["probs_ratio"] == metrics_2["probs_ratio"] == 1.0 + + @pytest.mark.parametrize("kl_type", ["k1", "k2", "k3"]) def test_calculate_kl(kl_type): """Tests KL calculations."""