Skip to content

fix: fix skip_reference_policy_logprobs_calculation and skip_prev_logprobs #2443

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

Merged
yuki-97 merged 13 commits into from
May 12, 2026
Merged

fix: fix skip_reference_policy_logprobs_calculation and skip_prev_logprobs #2443

Show file tree
Hide file tree
Changes from 12 commits
Commits
Show all changes
13 commits
Select commit Hold shift + click to select a range
63b6a6b
fix: skip prev_logprobs computation when force_on_policy_ratio is true
yfw Mar 30, 2026
3f55765
style: apply ruff format to PR-modified files
jinglinglingling May 8, 2026
473e7e9
refactor(grpo): move skip_reference_policy_logprobs_calculation asser…
jinglinglingling May 8, 2026
2566326
fix: skip_reference_policy_logprobs_calculation=true crashes training
Mar 10, 2026
21c6591
style: apply ruff format to PR #2174 worker files
jinglinglingling May 8, 2026
bd85091
fix: skip loading reference model when KL penalty is zero
yfw Mar 30, 2026
c31332e
feat(grpo): auto-enable skip_reference_policy_logprobs when KL=0
jinglinglingling May 8, 2026
9ee54dd
test(functional): grpo_kl_zero.sh covers reference_policy_kl_penalty=0
jinglinglingling May 8, 2026
b430912
test(grpo): add unit tests for skip-ref-logprobs and force-on-policy-…
jinglinglingling May 8, 2026
d5c1532
revert: drop worker-level reference-model guards (covered by grpo.py)
jinglinglingling May 9, 2026
92aa6ba
test(grpo): add reference_policy_kl_penalty key to sglang setup test
jinglinglingling May 9, 2026
c447a0d
test: provide real tensors for grpo_train logprob phase mocks
jinglinglingling May 9, 2026
2c57451
Apply suggestions from code review
yfw May 11, 2026
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
149 changes: 118 additions & 31 deletions nemo_rl/algorithms/grpo.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -389,6 +389,16 @@ def init_train_dataloader(dataset, suffix: str = ""):
os.environ["NRL_IGNORE_TP_ACCURACY_CHECK"] = "1"
print(" ✓ force_on_policy_ratio enabled")

# Validate skip_reference_policy_logprobs_calculation
if grpo_config.get("skip_reference_policy_logprobs_calculation"):
assert loss_config["reference_policy_kl_penalty"] == 0, (
"grpo.skip_reference_policy_logprobs_calculation=True requires "
"loss_fn.reference_policy_kl_penalty == 0"
)
print(
"Reference policy logprob calculation will be skipped since `grpo.skip_reference_policy_logprobs_calculation` is set to True and `loss_fn.reference_policy_kl_penalty` is 0."
)

# ==========================
# Cluster
# ==========================
Expand Down Expand Up @@ -554,6 +564,26 @@ def init_train_dataloader(dataset, suffix: str = ""):
policy_config["megatron_cfg"]["train_iters"] = total_train_iters

# Define initialization functions that will be used in all paths
init_reference_model = master_config["loss_fn"]["reference_policy_kl_penalty"] > 0

# Auto-enable skip_reference_policy_logprobs_calculation when the reference
# model is not loaded (per yuki-97 / terrykong review on PR #2178).
# Without this, training loops would still call
# policy.get_reference_policy_logprobs() and crash inside
# use_reference_model() because reference_state_dict was never initialized
# (issue #1968 Bug 1). Affects existing recipes that have
# reference_policy_kl_penalty=0 without setting the skip flag, e.g.
# examples/configs/recipes/llm/grpo-deepscaler-1.5b-8K.yaml.
Comment thread
yfw marked this conversation as resolved.
Outdated
if not init_reference_model and not master_config["grpo"].get(
"skip_reference_policy_logprobs_calculation"
):
master_config["grpo"]["skip_reference_policy_logprobs_calculation"] = True
print(
"Auto-enabling `grpo.skip_reference_policy_logprobs_calculation=True` "
"because `loss_fn.reference_policy_kl_penalty == 0` "
"(reference model is not loaded)."
)

def init_policy():
"""Initialize policy training workers."""
t0 = time.perf_counter()
Expand All @@ -565,6 +595,7 @@ def init_policy():
weights_path=weights_path,
optimizer_path=optimizer_path,
init_optimizer=True,
init_reference_model=init_reference_model,
)
return p, time.perf_counter() - t0

Expand Down Expand Up @@ -1336,12 +1367,6 @@ def grpo_train(
POLICY_GENERATION_STALE = True # tracks if generation needs a refit before running
assert policy_generation is not None # for mypy type check

if master_config["grpo"].get("skip_reference_policy_logprobs_calculation"):
assert master_config["loss_fn"]["reference_policy_kl_penalty"] == 0
print(
"Reference policy logprob calculation will be skipped since `grpo.skip_reference_policy_logprobs_calculation` is set to True and `loss_fn.reference_policy_kl_penalty` is 0."
)

# Check if we need to sync KV cache scales
# When fallback to policy as the policy_generation, we use getattr to check.
sync_kv_scales = getattr(policy_generation, "requires_kv_scale_sync", False)
Expand Down Expand Up @@ -1728,9 +1753,19 @@ 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 +1779,14 @@ 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 @@ -1757,20 +1797,36 @@ def grpo_train(
timer=timer,
)["reference_logprobs"]
)
else:
train_data["reference_policy_logprobs"] = torch.zeros_like(
train_data["prev_logprobs"]
)

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 @@ -2789,33 +2845,64 @@ 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"]
reference_logprobs = policy.get_reference_policy_logprobs(
train_data,
timer=timer,
)["reference_logprobs"]
if not skip_prev_logprobs:
fprop_logprobs = policy.get_logprobs(
train_data,
timer=timer,
)["logprobs"]
train_data["prev_logprobs"] = fprop_logprobs
train_data["reference_policy_logprobs"] = reference_logprobs

if not master_config["grpo"].get(
"skip_reference_policy_logprobs_calculation"
):
reference_logprobs = policy.get_reference_policy_logprobs(
train_data,
timer=timer,
)["reference_logprobs"]
train_data["reference_policy_logprobs"] = reference_logprobs
else:
train_data["reference_policy_logprobs"] = torch.zeros_like(
fprop_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
10 changes: 9 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 @@ -256,7 +256,10 @@ 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 @@ -266,6 +269,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
Loading
Loading