fix(grpo): grpo_sync dict unpack of seq_logprob error (#2559 follow-up) by qiaochuz-nv · Pull Request #2607 · NVIDIA-NeMo/RL

qiaochuz-nv · 2026-05-28T17:40:50Z

Summary

PR #2559 changed compute_and_apply_seq_logprob_error_masking in
nemo_rl/algorithms/grpo.py:1294 from a 3-tuple return to a dict
return, and migrated the two existing call sites in grpo.py
(lines 1935 and 3064) to consume the dict via key access. However,
the same day, PR #2439 (data plane transfer queue integration) merged
a NEW call site at nemo_rl/algorithms/grpo_sync.py:772-781. PR #2559
never touched grpo_sync.py, so this new call site retained its
3-tuple unpack and now raises ValueError: too many values to unpack (expected 3) at GRPO sync step=1 the moment data_plane.enabled=true
is set on origin/main.

This PR replaces the 3-tuple unpack in grpo_sync.py with the same
dict-key access pattern used in grpo.py, preserving the existing local
variable names (max_seq_mult_prob_error, num_masked_seqs,
masked_correct_pct) so the downstream metrics emission at
grpo_sync.py:1014-1016 is unchanged.

-                    (
-                        max_seq_mult_prob_error,
-                        num_masked_seqs,
-                        masked_correct_pct,
-                    ) = compute_and_apply_seq_logprob_error_masking(
+                    seq_error_result = compute_and_apply_seq_logprob_error_masking(
                         train_data=masking_data,
                         rewards=rewards,
                         seq_logprob_error_threshold=master_config.grpo[
                             "seq_logprob_error_threshold"
                         ],
                     )
+                    max_seq_mult_prob_error = seq_error_result[
+                        "max_seq_mult_prob_error"
+                    ]
+                    num_masked_seqs = seq_error_result["num_masked_seqs"]
+                    masked_correct_pct = seq_error_result["masked_correct_pct"]

Root cause

Two PRs landed on origin/main within hours of each other on
2026-05-28:

64be007eb "feat(grpo): add sequence-level logprob error metrics" (feat(grpo): add sequence-level logprob error metrics #2559)
— refactored compute_and_apply_seq_logprob_error_masking to return a
dict with 8 keys (added mean_*, min_*, *_after_mask variants on
top of the original 3 fields). Updated the 2 call sites it knew
about: grpo.py:1935 and grpo.py:3064.
4454fa64e "feat: data plane transfer queue integration" (feat: data plane transfer queue integration #2439)
— added a NEW call site at grpo_sync.py:772-781 that destructures
the legacy 3-tuple. This file was never touched by feat(grpo): add sequence-level logprob error metrics #2559.

Because the two diffs do not overlap textually (different files,
different functions), neither PR's CI run saw the mismatch — each PR
passed in isolation. The bug only surfaces post-merge when both diffs
are present and data_plane.enabled=true exercises grpo_sync.py.

Why the existing test did not catch this

The repo's test_grpo_data_plane_transfer_queue_daily_pr.sh testcase
(introduced by #2439) does exercise this exact code path. It started
failing immediately on origin/main once #2559 merged. This PR is
the fix surfaced by that testcase.

Test command (functional / e2e)

Container: /lustre/fsw/coreai_dlalgo_ci/qiaochuz/containers/nemo-rl-nightly-20260528.sqsh
Cluster: EOS, coreai_dlalgo_qa, partition=batch, 1 node, 2× H100
Command (drives both BEFORE and AFTER runs through the daily-PR
testcase plumbing — _runtime_cherry_pick.sh ensures /opt/nemo-rl
matches origin/main HEAD 64be007eb post-#2559):

uv run /opt/nemo-rl/examples/run_grpo.py \
    policy.model_name=Qwen/Qwen3-0.6B \
    grpo.num_prompts_per_step=2 \
    grpo.num_generations_per_prompt=6 \
    grpo.adv_estimator.name=reinforce_plus_plus \
    policy.train_global_batch_size=6 \
    policy.train_micro_batch_size=2 \
    cluster.gpus_per_node=2 \
    grpo.max_num_steps=2 \
    logger.tensorboard_enabled=true \
    logger.wandb_enabled=false \
    checkpointing.enabled=false \
    data_plane.enabled=true \
    data_plane.impl=transfer_queue \
    data_plane.backend=simple

Before fix — run output

SLURM job 5342797 (EOS, coreai_dlalgo_qa-qiaochuz, 2× H100,
nemo-rl-nightly-20260528.sqsh, /opt/nemo-rl HEAD = 5494d14d5
post-#2559):

==== buggy tuple-unpack at grpo_sync.py:772-781 (BEFORE patch) ====
                    (
                        max_seq_mult_prob_error,
                        num_masked_seqs,
                        masked_correct_pct,
                    ) = compute_and_apply_seq_logprob_error_masking(
                        train_data=masking_data,
                        rewards=rewards,
                        seq_logprob_error_threshold=master_config.grpo[
                            "seq_logprob_error_threshold"
                        ],
                    )
                    sample_mask = masking_data["sample_mask"]
==== end buggy snippet ====

# Run reaches GRPO step=1, generates rollouts, computes rewards, then crashes:

Traceback (most recent call last):
  File "/opt/nemo-rl/examples/run_grpo.py", line 224, in <module>
    main()
  File "/opt/nemo-rl/examples/run_grpo.py", line 207, in main
    trainer(
        policy,
        ...
        master_config,
    )
  File "/opt/nemo-rl/nemo_rl/algorithms/grpo_sync.py", line 772, in grpo_train_sync
    (
    ...
    ) = compute_and_apply_seq_logprob_error_masking(
ValueError: too many values to unpack (expected 3)

# Run finished: status=failed

After fix — run output

SLURM job 5342817 (same cluster + container; patched grpo_sync.py
file-overlay'd into /opt/nemo-rl/nemo_rl/algorithms/grpo_sync.py
before launching run_grpo.py):

==== patched grpo_sync.py snippet (AFTER patch) ====
                    seq_error_result = compute_and_apply_seq_logprob_error_masking(
                        train_data=masking_data,
                        rewards=rewards,
                        seq_logprob_error_threshold=master_config.grpo[
                            "seq_logprob_error_threshold"
                        ],
                    )
                    max_seq_mult_prob_error = seq_error_result[
                        "max_seq_mult_prob_error"
                    ]
                    num_masked_seqs = seq_error_result["num_masked_seqs"]
                    masked_correct_pct = seq_error_result["masked_correct_pct"]
                    sample_mask = masking_data["sample_mask"]
==== end patched snippet ====

# Both training steps complete with valid metrics:

🔹 train/loss - 2 steps
  Step 1     Step 2
  -3.06e-05  -2.59e-05

🔹 train/sampling_importance_ratio - 2 steps
Min: 0.999982   Max: 1.00099   Avg: 1.00049

🔹 train/token_mult_prob_error - 2 steps
Min: 1.01462   Max: 1.01717   Avg: 1.01589

🔹 train/total_num_tokens - 2 steps
  Step 1   Step 2
  6144     6144

# Metric-check verdict:

                                 Metric Checks
┏━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━┓
┃ Status ┃ Check                             ┃ Value                 ┃ Message ┃
┡━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━┩
│ PASS   │ len(data["train/loss"]) >= 1      │ 2                     │         │
│ PASS   │ max(data["train/gen_kl_error"]) < │ 0.0007329469081014395 │         │
│        │ 0.02                              │                       │         │
└────────┴───────────────────────────────────┴───────────────────────┴─────────┘
PASS: AFTER-fix GRPO data-plane TransferQueue path completed

Detected by

NeMo daily-PR impact pipeline gap testcase
test_grpo_data_plane_transfer_queue_daily_pr.sh
Testcase path:
nemo_llm/test_suite/rl/testcases/algorithms/test_grpo_data_plane_transfer_queue_daily_pr.sh

Test plan

Functional / e2e test passes on the patched checkout (logs above)
Local variable names preserved → metrics emission at
grpo_sync.py:1014-1016 unchanged
Behavior identical to grpo.py:1935-1962 migration done by feat(grpo): add sequence-level logprob error metrics #2559
Codecov/patch ≥ 80% — a follow-up commit will add a unit test
covering the 3 new dict-access lines

…ng (#2559 follow-up) PR #2559 changed compute_and_apply_seq_logprob_error_masking to return a dict, and updated the two grpo.py call sites at lines 1935 and 3064. However, PR #2439 (data plane transfer queue integration) merged the same day added a NEW call site at grpo_sync.py:772-781 that #2559 did not see and therefore did not migrate. HEAD origin/main now raises ValueError: too many values to unpack (expected 3) at GRPO sync step=1 the moment data_plane.enabled=true is set. Fix replaces the 3-tuple unpack with the same dict-key access pattern used in grpo.py, preserving the existing local variable names (max_seq_mult_prob_error, num_masked_seqs, masked_correct_pct) so the downstream metrics emission at lines 1014-1016 is unchanged. Verified end-to-end on EOS H100 with nemo-rl-nightly-20260528.sqsh: - BEFORE: ValueError: too many values to unpack (expected 3) at step=1 - AFTER: 2/2 train steps PASS, train/loss valid, gen_kl_error < 0.02 Signed-off-by: Qiaochu Zhu <qiaochuz@nvidia.com>

copy-pr-bot · 2026-05-28T17:40:55Z

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

Adds CPU-runnable unit coverage for the 3 dict-access lines at grpo_sync.py:773-781 introduced by the #2559 follow-up fix. Signed-off-by: Qiaochu Zhu <qiaochuz@nvidia.com>

copy-pr-bot · 2026-05-28T17:47:32Z

/ok to test 336c80a

@qiaochuz-nv, there was an error processing your request: E2

See the following link for more information: https://docs.gha-runners.nvidia.com/cpr/e/2/

qiaochuz-nv · 2026-05-28T17:48:23Z

/ok to test 336c80a

copy-pr-bot · 2026-05-28T17:48:26Z

/ok to test 336c80a

@qiaochuz-nv, there was an error processing your request: E2

See the following link for more information: https://docs.gha-runners.nvidia.com/cpr/e/2/

qiaochuz-nv · 2026-05-28T17:49:13Z

/ok to test 799db61

ZhiyuLi-Nvidia · 2026-05-28T17:51:24Z

Thank you @qiaochuz-nv, looks like there's existing PR to fix
#2600

qiaochuz-nv · 2026-05-28T17:53:37Z

close directly thanks @ZhiyuLi-Nvidia

qiaochuz-nv requested a review from a team as a code owner May 28, 2026 17:40

qiaochuz-nv added the CI:Lfast Runs a fast test suite and re-use nightly `main` container (but sync dependencies to PRs version) label May 28, 2026

test(grpo_sync): cover dict access for seq_logprob_error_masking

799db61

Adds CPU-runnable unit coverage for the 3 dict-access lines at grpo_sync.py:773-781 introduced by the #2559 follow-up fix. Signed-off-by: Qiaochu Zhu <qiaochuz@nvidia.com>

qiaochuz-nv requested a review from a team as a code owner May 28, 2026 17:45

copy-pr-bot Bot temporarily deployed to public May 28, 2026 17:49 Inactive

qiaochuz-nv requested a review from ZhiyuLi-Nvidia May 28, 2026 17:49

copy-pr-bot Bot temporarily deployed to nemo-ci May 28, 2026 17:49 Inactive

copy-pr-bot Bot temporarily deployed to public May 28, 2026 17:49 Inactive

qiaochuz-nv requested a review from macandro96 May 28, 2026 17:50

copy-pr-bot Bot temporarily deployed to public May 28, 2026 17:50 Inactive

qiaochuz-nv closed this May 28, 2026

qiaochuz-nv deleted the qiaochuz/fix_grpo_sync_seq_err_unpack branch May 28, 2026 17:53

copy-pr-bot Bot temporarily deployed to public May 28, 2026 17:54 Inactive

Navigation Menu

Uh oh!

fix(grpo): grpo_sync dict unpack of seq_logprob error (#2559 follow-up) - #2607

fix(grpo): grpo_sync dict unpack of seq_logprob error (#2559 follow-up)#2607
qiaochuz-nv wants to merge 2 commits into
mainfrom
qiaochuz/fix_grpo_sync_seq_err_unpack

qiaochuz-nv commented May 28, 2026

Uh oh!

copy-pr-bot Bot commented May 28, 2026

Uh oh!

copy-pr-bot Bot commented May 28, 2026

Uh oh!

qiaochuz-nv commented May 28, 2026

Uh oh!

copy-pr-bot Bot commented May 28, 2026

Uh oh!

qiaochuz-nv commented May 28, 2026

Uh oh!

ZhiyuLi-Nvidia commented May 28, 2026 •

edited

Loading

Uh oh!

qiaochuz-nv commented May 28, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

qiaochuz-nv commented May 28, 2026

Summary

Root cause

Why the existing test did not catch this

Test command (functional / e2e)

Before fix — run output

After fix — run output

Detected by

Test plan

Uh oh!

copy-pr-bot Bot commented May 28, 2026

Uh oh!

copy-pr-bot Bot commented May 28, 2026

Uh oh!

qiaochuz-nv commented May 28, 2026

Uh oh!

copy-pr-bot Bot commented May 28, 2026

Uh oh!

qiaochuz-nv commented May 28, 2026

Uh oh!

ZhiyuLi-Nvidia commented May 28, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

qiaochuz-nv commented May 28, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

ZhiyuLi-Nvidia commented May 28, 2026 •

edited

Loading