Skip to content

fix(vlm): propagate NeMo Gym dynamic image padding - #3738

Closed
youngeunkwon0405 wants to merge 3 commits into
mainfrom
codex/nemo-gym-dynamic-image-padding
Closed

fix(vlm): propagate NeMo Gym dynamic image padding#3738
youngeunkwon0405 wants to merge 3 commits into
mainfrom
codex/nemo-gym-dynamic-image-padding

Conversation

@youngeunkwon0405

Copy link
Copy Markdown
Contributor

What does this PR do?

Propagates the existing NeMo-Gym pad_dynamic_image_shapes option through the initial deduplicated multimodal attachment path.

When grpo.deduplicate_multimodal_data=true, GRPO attaches the initial image-model inputs before repeating prompts. That path previously called attach_image_model_inputs_to_message without forwarding the environment's dynamic-padding option. A heterogeneous multi-image example could therefore attempt to stack different image resolutions before rollout even when the NeMo-Gym environment requested padding.

Before this fix, propagation stopped at the algorithm-to-rollout-helper boundary:

master_config.env.nemo_gym.pad_dynamic_image_shapes
  -> grpo_train / validate / AsyncTrajectoryCollector._process_batch
  -> attach_initial_nemo_gym_image_payloads(batch, processor)
  -> attach_image_model_inputs_to_message(...)
       pad_dynamic_image_shapes=False  # implicit default

The synchronous training, validation, and asynchronous collection call sites invoked attach_initial_nemo_gym_image_payloads without the flag. That helper then invoked attach_image_model_inputs_to_message without the flag as well, so the target function silently used its existing False default. This early deduplicated attachment shortcut is separate from the normal NeMo-Gym environment preprocessing path, which already carried the option correctly.

After this fix, the complete chain is:

master_config.env.nemo_gym.pad_dynamic_image_shapes
  -> grpo_train / validate / AsyncTrajectoryCollector._process_batch
  -> attach_initial_nemo_gym_image_payloads(..., pad_dynamic_image_shapes=...)
  -> attach_image_model_inputs_to_message(..., pad_dynamic_image_shapes=...)

This change:

  • adds a keyword-only pad_dynamic_image_shapes argument to attach_initial_nemo_gym_image_payloads;
  • forwards env.nemo_gym.pad_dynamic_image_shapes from synchronous training, validation, and asynchronous trajectory collection; and
  • preserves the existing False default for every caller that does not opt in.

This PR is independent of the 32-GPU recipe PR and independent of #3721.

Issues

None.

Validation

  • uvx --from ruff==0.9.9 ruff check nemo_rl/algorithms/async_utils/trajectory_collector.py nemo_rl/algorithms/grpo.py nemo_rl/experience/rollouts.py tests/unit/experience/test_rollouts.py
  • uvx --from ruff==0.9.9 ruff format --check nemo_rl/algorithms/async_utils/trajectory_collector.py nemo_rl/algorithms/grpo.py nemo_rl/experience/rollouts.py tests/unit/experience/test_rollouts.py
  • PYTHONPATH=. uv run --extra nemo_gym --group test pytest -q --noconftest tests/unit/experience/test_rollouts.py -k attach_initial_nemo_gym_image_payloads — 3 passed
  • GPU Slurm validation job 6370288COMPLETED 0:0
  • Ten-step private NeMo-Gym sync and async runs completed without mixed-image padding failures: sync W&B, async W&B

Before your PR is "Ready for review"

  • Read and followed the contributor guidelines
  • Added focused unit coverage for flag propagation and the default behavior
  • Ran focused Linux GPU-container tests
  • No documentation update is needed because this fixes existing internal option propagation without adding a new public option

Signed-off-by: Youngeun Kwon <youngeunk@nvidia.com>
@copy-pr-bot

copy-pr-bot Bot commented Aug 20, 2026

Copy link
Copy Markdown

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.

@youngeunkwon0405 youngeunkwon0405 self-assigned this Aug 20, 2026

@youngeunkwon0405 youngeunkwon0405 left a comment

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This PR was reviewed by a coordinated team of 4 specialized agents (RL-codebase expert, test reviewer, independent bug scanner, and a devil's-advocate verification pass). Summary:

  • The fix is correct and needed. Pre-PR, the dedup attachment path always invoked the processor with stacking semantics, breaking heterogeneous multi-image gym prompts even when the environment requested padding; the gym-actor side already honored the flag. All three production call sites are updated and no others exist.
  • No bugs found. Config key path, master_config.env typing at all three call sites, and the downstream attach_image_model_inputs_to_message keyword signature were all verified against the code. Ruff check/format pass. The updated unit test was run locally and passes.
  • The 5 inline comments are all low-severity polish (config-read convention + shared helper, two optional verified test hardenings, a docstring nit) plus one pre-existing out-of-scope observation. Nothing blocks merge.

Validation evidence (W&B sync/async runs + Slurm job) is appropriate for a correctness fix — thanks for including it.

Generated by Claude Code

Comment thread nemo_rl/algorithms/grpo.py Outdated
Comment thread tests/unit/experience/test_rollouts.py
Comment thread nemo_rl/algorithms/async_utils/trajectory_collector.py Outdated
Comment thread nemo_rl/experience/rollouts.py
Comment thread nemo_rl/experience/rollouts.py Outdated
Signed-off-by: Youngeun Kwon <youngeunk@nvidia.com>
@youngeunkwon0405 youngeunkwon0405 added the CI:Lfast Runs a fast test suite and re-use nightly `main` container (but sync dependencies to PRs version) label Aug 21, 2026
@youngeunkwon0405

Copy link
Copy Markdown
Contributor Author

/ok to test 36ed044

@youngeunkwon0405
youngeunkwon0405 marked this pull request as ready for review August 21, 2026 02:44
@youngeunkwon0405
youngeunkwon0405 requested review from a team as code owners August 21, 2026 02:44
Signed-off-by: Youngeun Kwon <youngeunk@nvidia.com>
@youngeunkwon0405

Copy link
Copy Markdown
Contributor Author

/ok to test ac228c1

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CI:Lfast Runs a fast test suite and re-use nightly `main` container (but sync dependencies to PRs version)

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant