Skip to content

fix: skip prev_logprobs computation when force_on_policy_ratio is true #2177

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
yfw wants to merge 1 commit into from
Closed

fix: skip prev_logprobs computation when force_on_policy_ratio is true #2177

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
83 changes: 61 additions & 22 deletions nemo_rl/algorithms/grpo.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"):
Expand All @@ -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"
Expand All @@ -1761,16 +1769,28 @@ 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,
masked_correct_pct,
) = 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
Expand Down Expand Up @@ -2779,33 +2799,52 @@ 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,
masked_correct_pct,
) = 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
Expand Down
8 changes: 7 additions & 1 deletion nemo_rl/algorithms/loss/loss_functions.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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:]
Expand All @@ -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)
Expand Down
50 changes: 50 additions & 0 deletions tests/unit/algorithms/test_loss_functions.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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."""
Expand Down
Loading