feat(mllm-v2): Support single-controller for Megatron-Inference. by cspades · Pull Request #4009 · NVIDIA-NeMo/RL

cspades · 2026-09-05T00:00:40Z

What does this PR do ?

Minor pipecleaning to support single-controller / TransferQueue / V2 for Megatron generation.
- https://wandb.ai/adlr/mllm-rl-dev?nw=nwusercye_nv - in_order-parity with vLLM and MLLM using RL V1.

Issues

List issues that this PR closes (syntax):

Usage

You can potentially add a usage example below

# Add a code snippet demonstrating how to use this

Before your PR is "Ready for review"

Pre checks:

Make sure you read and followed Contributor guidelines
Did you write any new necessary tests?
Did you run the unit tests and functional tests locally? Visit our Testing Guide for how to run tests
Did you add or update any necessary documentation? Visit our Document Development Guide for how to write, build and test the docs.

Additional Information

...

copy-pr-bot · 2026-09-05T00:00:44Z

Auto-sync is disabled for ready for review pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

cspades

Summary

This PR pipecleans multimodal (VLM) rollouts through the SingleController + TransferQueue (V2) path so Megatron-Inference can serve them. Before it, V2/SC was text-only: SC_ROLLOUT_SCHEMA_FIELDS declared no media columns and pack_payload had no multimodal encoding, so a VLM run on SC would either race TransferQueue's lazy field registration or simply drop pixels.

The change closes that in three places:

Entrypoint — run_grpo_single_controller.py builds an AutoProcessor when policy.is_vlm, and threads it into setup_single_controller(..., processor=...) / setup_response_data(..., is_vlm=True).
Declare — the SC partition warmup pre-registers the multimodal wire fields before rollout / policy / teacher writers go concurrent.
Write + train — pack_payload encodes media via encode_multimodal_for_wire with per-row geometry in KVBatchMeta.tags, _generate_response flattens through batched_message_log_to_flat_message so pixels accompany the expanded placeholders, and train_microbatch is unblocked for multimodal by attaching the media-token validity mask.

I traced all eight hops end-to-end (entry → processing → warmup → rollout → Megatron generate → TQ write → fetch → trainer forward) and the path is complete for the shipped CLEVR SC recipe. Sequence-packing safety is independently enforced at three points (truncate_tensors, get_and_validate_seqlen, and materialize all skip packed media).

Extraction unified

While reviewing, we found the Gym and non-Gym paths had two independent implementations of media extraction that had drifted: vlm_hf_data_processor shipped imgs_sizes as int64 while extract_multimodal_model_inputs cast it to int32 — the same nemotron-omni model receiving a different dtype depending on which data path built the batch. That is now fixed in two steps: the dtype was aligned to int32 (matching the two other construction sites), and then vlm_hf_data_processor's 49-line block was collapsed onto the shared helper, so there is exactly one extraction implementation and the two paths cannot drift again. This also fixed a latent bug for free — the old block did message["token_type_ids"][0] unconditionally, which for a 1-D input silently took a scalar.

About this review

Findings below are the issues not fixed in this PR's scope — mostly pre-existing or SC-wide rather than introduced here. Each is a separate comment. Several are explicitly not asks for this PR; the actionability line at the top of each says which.

No code was executed for this review (static analysis only); conclusions are from reading source, git archaeology, and git merge-base --is-ancestor.

Generated by Claude Code

cspades · 2026-09-05T02:20:53Z

/ok to test b986f14

cspades · 2026-09-06T06:32:08Z

/ok to test a27c364

terrykong

Thanks for this — the SC multimodal path is a lot of moving pieces and the wiring reads cleanly end to end. Reviewed at a27c3648, after the rebase.

What this does. Multimodal/VLM rollouts now run through SingleController + TransferQueue so Megatron-Inference can serve them: the entrypoint builds an AutoProcessor when policy.is_vlm is set, the SC partition warmup pre-registers the multimodal wire fields, pack_payload encodes media with per-row geometry in KVBatchMeta.tags, _generate_response flattens through batched_message_log_to_flat_message, and train_microbatch gains the media-token validity mask. Nearly all of it is new call sites into machinery that already shipped, which is why the diff is small for what it turns on.

This PR un-breaks main. main still carries the guard if self.media_placeholder_token_id is not None or (self.model_slices_context_parallel_inputs):, and the pre-existing parametrized case test_forwards_model_owned_packing_flags[model-owned-cp-slicing] sets model_slices_context_parallel_inputs = True and then calls train_microbatch — so it trips that guard every time it runs. Deleting the guard is the right fix rather than a convenient one: _train_microbatch_body already forwarded three of the four capabilities the guard's own message named, and this PR adds the fourth.

Also worth keeping: collapsing vlm_hf_data_processor onto extract_multimodal_model_inputs deletes a second implementation of media extraction that had already drifted between the Gym and non-Gym paths.

Status. mergeStateStatus is BLOCKED on approvals — CI restarted with the rebase and had no failures at the time of writing.

Related work. #3921 is this same work at full size — 87 files to this PR's 28, overlapping on 19, including single_controller_utils/setup.py, payload.py, multimodal_utils.py and four of the same test files. The two have taken different routes for getting packed-tensor shapes across the wire: this PR registers the static field set and puts per-row shapes in KVBatchMeta.tags as <field>__row_shapes, while #3921 adds a packed_tensor_wire_fields(...) helper that registers a companion meta column per packed field (PACKED_TENSOR_META_PREFIX + field) and builds the field list from the processor instead. Those two names are in plain text rather than linked because they exist only in #3921 — neither is on main or in this branch. Both PRs edit the same registration call in setup.py#L1450, so whichever merges second gets a conflict there plus a decision to make: which of the two should survive?

Generated by Claude Code

cspades · 2026-09-07T01:35:48Z

/ok to test b6e348f

cspades · 2026-09-07T02:09:35Z

/ok to test 5237afe

Signed-off-by: Cory Ye <cye@nvidia.com>

This reverts commit ff519e1. Signed-off-by: Cory Ye <cye@nvidia.com>

Signed-off-by: Cory Ye <cye@nvidia.com>

cspades · 2026-09-07T03:11:07Z

/ok to test 81c99e8

Signed-off-by: Cory Ye <cye@nvidia.com>

cspades · 2026-09-07T06:59:57Z

/ok to test 2b064d2

Signed-off-by: Cory Ye <cye@nvidia.com>

cspades · 2026-09-07T15:00:41Z

/ok to test 32fbaef

…wapping the env at runtime Token capture swapped the VllmAsyncGenerationWorker registry entry to VLLM_GYM at setup time. Worker venvs are cached by actor class name, so a venv prebuilt with plain --extra vllm (which is what the Dockerfile bakes) was reused as-is and the nemo_gym import failed. This is why the SingleController L1 job fails even on a freshly built container. Make VLLM_EXECUTABLE use VLLM_GYM so the baked venv already has nemo_gym, and drop the runtime registry override and the error wrapper that described the old behavior. Signed-off-by: Terry Kong <terryk@nvidia.com>

copy-pr-bot · 2026-09-07T17:17:22Z

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.

Signed-off-by: Cory Ye <cye@nvidia.com>

cspades · 2026-09-07T17:33:22Z

/ok to test 82a0dc5

#4009 pointed VLLM_EXECUTABLE at PY_EXECUTABLES.VLLM_GYM so token capture (token_capture.enabled) can import nemo_gym inside the worker. Worker venvs are cached by actor class name, so a venv prebuilt with plain --extra vllm is reused as-is and the import fails; that is the L1_Functional_Tests_SingleController failure on this branch. This branch builds the registry from ACTOR_ENVIRONMENTS rather than the literal dict on main, so the change did not carry over on rebase. Add a guard test: a PY_EXECUTABLES constant that names extras but is wired to no actor is the signature of exactly this miss. Signed-off-by: Terry Kong <terryk@nvidia.com>

The example still showed ["vllm"], but the vLLM workers carry ["vllm", "nemo_gym"] since #4009 moved them onto PY_EXECUTABLES.VLLM_GYM. Use the real value and say why, since a two-extra entry is otherwise puzzling. Signed-off-by: Terry Kong <terryk@nvidia.com>

Catches the telemetry branch up to main, which had moved 22 commits and tripped pr-branch-up-to-date-check (max 10 behind). Eleven files conflicted. Resolutions, in the order a reviewer would want them: - single_controller.py: main split _dispatch_one_prompt into a token-capture branch (generate_for_finalization + the finalizer actor pool, #3837) and moved the legacy generate_and_push loop into an else. Git interleaved our per-prompt span into the new branch, so the function was restored to main's exactly and the per_prompt_scope() + rl.sc.generate_and_push umbrella re-applied to the legacy loop only. The token-capture branch is left uninstrumented: it commits through the finalizer pool rather than here, so its attempt boundary is a different shape than this span describes. Recorded in the coverage-gap table in docs/observability/span-groups.md. - run_grpo_single_controller.py: keeps our startup_span / setup_span phases and adopts main's VLM processor (#4009), with the processor-aware tokenizer construction moved inside setup_span ("tokenizer") and processor= threaded to setup_single_controller. - run_grpo.py: adopts main's make_policy_factory() helper in place of our inline factory selection, keeping setup_span("workers"). - nemo_gym.py: keeps @accepts_trace_context and the run_rollouts / _stream_rollouts span split, alongside main's ledger control plane; both signatures widen to main's 4-tuple yield. - rollout_manager.py: keeps dispatch_with_trace_context and unpacks main's added resolved_agent_ref. - factory.py: adopts main's LocalDataPlaneConfig-aware observability lookup, keeping the telemetry_enabled_in_env() arm that installs the wrapper for its spans. - vllm_worker.py: keeps umbrella_trace_fn(U_MODEL_INIT) on _load_model beside main's _refit_with_reload_api_enabled. - config.py, worker_mixin.py, virtual_cluster.py, vllm_generation.py: both sides added adjacent imports or fields; all kept. The textually merged uv.lock was corrupt (a missing source field on opentelemetry-instrumentation-aiohttp-client, which matched more than one package), so it was regenerated from main's with uv 0.11.28 to match CI's revision 3. The result is 52 insertions: the nemo-lens rev bump and the three otel aiohttp packages, nothing else. Both uv lock --check runs pass. Verified: all 54 changed Python files compile, ruff check and format are clean, and the seven test_source_drift.py guards pass -- including the carrier guard, which confirms run_rollouts is still dispatched with trace context after the signature change. Signed-off-by: Raj Singh <rajsin@nvidia.com>

cspades requested review from a team as code owners September 5, 2026 00:00

cspades self-assigned this Sep 5, 2026

cspades changed the title ~~[mllm-v2] Support single-controller for Megatron-Inference.~~ feat(mllm-v2): Support single-controller for Megatron-Inference. Sep 5, 2026

cspades force-pushed the cye/rl_mllm_omni_mm_nocolgym_rebase_pr2957 branch 2 times, most recently from b321715 to 4295411 Compare September 5, 2026 02:02

cspades commented Sep 5, 2026

View reviewed changes

cspades added the CI:L1 Run doctests, unit tests, and functional tests label Sep 5, 2026

terrykong mentioned this pull request Sep 5, 2026

refactor(docker): read the actor venv list from a Python leaf module #4002

Merged

cspades force-pushed the cye/rl_mllm_omni_mm_nocolgym_rebase_pr2957 branch from b986f14 to a27c364 Compare September 6, 2026 06:28

terrykong reviewed Sep 6, 2026

View reviewed changes

cspades requested a review from a team as a code owner September 6, 2026 18:53

github-actions Bot added the Documentation Improvements or additions to documentation label Sep 6, 2026

cspades force-pushed the cye/rl_mllm_omni_mm_nocolgym_rebase_pr2957 branch 3 times, most recently from 7e0410e to b6e348f Compare September 7, 2026 01:35

cspades added 2 commits September 6, 2026 19:35

Pass pre-expanded media tokens to MLLM.

5bb259b

Signed-off-by: Cory Ye <cye@nvidia.com>

Add tests and lint.

62da428

Signed-off-by: Cory Ye <cye@nvidia.com>

cspades added 6 commits September 6, 2026 19:38

Restore Gym revision required by multimodal branch

11f2eb7

Signed-off-by: Cory Ye <cye@nvidia.com>

Revert "Restore Gym revision required by multimodal branch"

91286e1

This reverts commit ff519e1. Signed-off-by: Cory Ye <cye@nvidia.com>

Address more review comments.

319aaa1

Signed-off-by: Cory Ye <cye@nvidia.com>

Strengthen reserved port handoff test

f87affb

Signed-off-by: Cory Ye <cye@nvidia.com>

Lint.

56acacf

Signed-off-by: Cory Ye <cye@nvidia.com>

Fix rebase error.

81c99e8

Signed-off-by: Cory Ye <cye@nvidia.com>

cspades force-pushed the cye/rl_mllm_omni_mm_nocolgym_rebase_pr2957 branch from 5237afe to 81c99e8 Compare September 7, 2026 03:10

terrykong previously approved these changes Sep 7, 2026

View reviewed changes

terrykong enabled auto-merge (squash) September 7, 2026 03:56

cspades dismissed terrykong’s stale review via ae7ee1f September 7, 2026 06:54

Fix flaky unit tests and test coverage.

2b064d2

Signed-off-by: Cory Ye <cye@nvidia.com>

cspades force-pushed the cye/rl_mllm_omni_mm_nocolgym_rebase_pr2957 branch from ae7ee1f to 2b064d2 Compare September 7, 2026 06:59

Add a digest to the RL venv cache to automatically refresh stale venv's.

32fbaef

Signed-off-by: Cory Ye <cye@nvidia.com>

terrykong mentioned this pull request Sep 7, 2026

fix(sc): give vLLM workers the nemo_gym extra by default instead of swapping the env at runtime #4043

Merged

Revert the cache digest and refresh.

82a0dc5

Signed-off-by: Cory Ye <cye@nvidia.com>

terrykong approved these changes Sep 7, 2026

View reviewed changes

terrykong merged commit fb85ee3 into main Sep 7, 2026
107 checks passed

terrykong deleted the cye/rl_mllm_omni_mm_nocolgym_rebase_pr2957 branch September 7, 2026 21:24

terrykong mentioned this pull request Sep 7, 2026

fix: honor NEMO_RL_PY_EXECUTABLES_SYSTEM for all actor environments #4020

Merged

4 tasks

Uh oh!

Conversation

cspades commented Sep 5, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do ?

Issues

Usage

Before your PR is "Ready for review"

Additional Information

Uh oh!

copy-pr-bot Bot commented Sep 5, 2026

Uh oh!

cspades left a comment

Choose a reason for hiding this comment

Summary

Extraction unified

About this review

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

cspades commented Sep 5, 2026

Uh oh!

cspades commented Sep 6, 2026

Uh oh!

terrykong left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

cspades commented Sep 7, 2026

Uh oh!

cspades commented Sep 7, 2026

Uh oh!

cspades commented Sep 7, 2026

Uh oh!

cspades commented Sep 7, 2026

Uh oh!

cspades commented Sep 7, 2026

Uh oh!

copy-pr-bot Bot commented Sep 7, 2026

Uh oh!

cspades commented Sep 7, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

cspades commented Sep 5, 2026 •

edited

Loading