feat(minf-mm-sc/tq): Support SC/TQ multimodal inference.

[cspades]

What does this PR do ?

Add a one line overview of what this PR aims to accomplish.

• Support Megatron multimodal inference with NRL V2 (TransferQueue/SingleController).

WandB: https://wandb.ai/adlr/mllm-rl-dev?nw=nwusercye_nv

V2 ("single-controller" in-order sampling) is aligned perfectly with V1 MLLM and vLLM under the same configs.

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]

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.

[github-actions]

✅ Submodule Fast-Forward Check Results

Check based on commit: b0cec04 (PR #3921 from cye/rl_mllm_omni_mm_nocolgym_rebase_pr3727)

✅ Submodules that are properly updated:

Megatron-Bridge: ✅ PR branch is ahead of main branch (fast-forward)

All submodule changes look good! ✨

[github-actions]

✅ Submodule Fast-Forward Check Results

Check based on commit: 42526bb (PR #3921 from cye/rl_mllm_omni_mm_nocolgym_rebase_pr3727)

✅ Submodules that are properly updated:

Megatron-Bridge: ✅ PR branch is ahead of main branch (fast-forward)

All submodule changes look good! ✨

[terrykong]

Reviewed by a team of specialized agents (RL, Megatron-LM, tests, bug-finding, design, plus an adversarial pass). Thanks for this — enabling multimodal on the SingleController/TransferQueue path is a good chunk of work, and the Megatron media workaround in particular is carefully done and well-marked.

Three tests fail on this branch, which is the main thing to flag. Two are small mistakes inside the PR's own new tests; the third is the recipe/test-suite consistency check. None of it was caught because /ok to test was posted on #3920 rather than on this PR, so no GitLab lane has run here. Running one will show all three.

Separately, vlm_grpo-nemotron-omni-30ba3b-vstat-1n4g-… looks like it cannot start — it keeps its 16-node parent's parallelism while shrinking to 2 GPUs per side.

A note on what we checked and found correct, so it doesn't read as unexamined: moving only imgs in _move_retained_vlm_media is right (Megatron already forces the sibling tensors to CUDA at admission); there's no race in the RESUMED→RUNNING window; all 11 mcore_generation_config keys are valid and non-deprecated at the pinned Megatron-LM SHA; and the HACK markers correctly name the upstream PR that removes them. We also walked back several findings during review — including a claim that the field-name warmup could wedge the TransferQueue controller, which is not true at the pinned TQ revision.

Everything below is either a reproduced failure or explicitly labelled as latent.

Generated by Claude Code

[terrykong]

2 action items. Both are needed — the suggestion below is not sufficient on its own.

TL;DR — this new test fails today: _generate_response is on AsyncRolloutImpl, not RolloutManager.

I ran it on this branch:

FAILED tests/unit/experience/test_rollout_manager.py::test_generate_response_forwards_message_log_media_to_generation
E   AttributeError: 'RolloutManager' object has no attribute '_generate_response'
    tests/unit/experience/test_rollout_manager.py:127
1 failed, 17 passed

_generate_response is defined on AsyncRolloutImpl (class at line 375); RolloutManager's class body starts at line 1152 and its rollout entrypoint is run_rollout.

Good news: the production change is correct. With both fixes below the file goes to 18 passed, and the test is a real one — dropping the new generation_input_data.update(...) line makes it fail, so it does check what it claims.

tests/unit/experience/ is collected by L0_Unit_Tests_Other.sh, so this goes red as soon as a lane runs.

AI-1 — add the import first

AsyncRolloutImpl is not currently imported; the import block at line 48 brings in AsyncNemoGymRolloutImpl, RolloutManager, RolloutOutcome, RolloutRetryPolicy, RolloutStats. Add AsyncRolloutImpl to it. Not a suggestion block because line 48 is outside this diff.

Please apply this before AI-2 — committing the suggestion alone swaps the AttributeError for a NameError.

AI-2 — then swap the class on this line only

Suggested change

	manager = object.__new__(RolloutManager)
	manager = object.__new__(AsyncRolloutImpl)

Note there is a second object.__new__(RolloutManager) at line 205 in an unrelated helper — that one is correct and should stay.

[terrykong]

1 action item.

TL;DR — patch is used here but never imported, so this new test raises NameError in the Megatron_Policy shard.

The file imports only MagicMock at test_megatron_worker.py:21. We AST-checked the whole module: patch is never imported, never assigned, and used exactly once — right here. There are no star-imports, and it isn't defined in tests/unit/conftest.py (and there is no tests/conftest.py at all).

The file is pytest.mark.mcore at the module level, so it runs in L0_Unit_Tests_Megatron_Policy_* — the right lane, which means it will run and it will fail. (We could not execute it locally, no megatron-core; this is AST-verified rather than observed.)

The patch target is right — setup_data_plane imports build_data_plane_client from nemo_rl.data_plane at call time, so patching the package attribute works.

AI-1

Add patch to the import on line 21. Not a suggestion block because that line is outside this diff:

from unittest.mock import MagicMock, patch

Follow-up — no action in this PR

The test name reads as "CP slicing over TQ works", but setup_data_plane never reads model_slices_context_parallel_inputs — the guard that did was removed — so the worker.model_slices_context_parallel_inputs = True line is inert and the test only proves the guard is gone. Consider test_setup_data_plane_no_longer_rejects_cp_slicing_models. See the separate comment on coverage.

[terrykong]

2 action items.

TL;DR — this recipe cannot start: it keeps its parent's parallelism but leaves only 2 GPUs on each side, so both training and generation fail at init.

grpo_nanov3.yaml sets parallelism twice — training under policy.megatron_cfg and generation under mcore_generation_config. The 16-node parent overrides only the training half, so generation TP/EP pass straight through. This overlay then shrinks the cluster without touching either.

Resolved values (cluster 1x4 = 4 GPUs, generation colocated.resources 1x2 = 2, so training = 2):

layer	GPUs	TP	EP	result
training	2	4	4	World size (2) is insufficient... PP(1)*CP(1)*TP(4) = 4
generation	2	2	8	world_size (2) is not divisible by expert_tensor_model_pipeline_parallel size

Both raise; nothing normalizes first. We confirmed this two ways — by hand and by resolving all four recipes through the repo's own load_config.

This is an inheritance slip rather than a house style: the CLEVR overlays in this PR point at size-matched parents (TP2/EP2 for 1n4g, TP8/EP8 for 8n4g) and are fine, and vstat-8n4g survives the same parent because 8 and 24 GPUs happen to divide.

AI-1

Set the parallelism for the GPU counts this recipe actually has, or point defaults: at a size-matched VSTAT parent. One way to do it — of several, since the right numbers depend on the layout you ran:

defaults: ./vlm_grpo-nemotron-omni-30ba3b-16n8g-megatron-tp4ep4-async-gym-video.v1.yaml

policy:
  megatron_cfg:
    tensor_model_parallel_size: 2
    expert_model_parallel_size: 2

(you'll also want matching tensor_model_parallel_size / expert_model_parallel_size under this file's mcore_generation_config for the 2-GPU generation fleet.)

AI-2

While you're in here: this is the only one of the four new recipes that resolves to enable_prefix_caching: true. It inherits it from grpo_nanov3.yaml:251, which is inside mcore_generation_config — the vllm_cfg copy just below it is commented out. vstat-8n4g sets false explicitly and both CLEVR recipes get false from grpo_math_1B.yaml:415. If the divergence isn't deliberate, add enable_prefix_caching: false alongside enable_chunked_prefill: true; the two VSTAT runs aren't comparable as written.

[terrykong]

1 action item. Applies to all four new recipes, not just this one.

TL;DR — the four new recipes have no driver script, which breaks an existing unit test: assert 271 == 267.

FAILED tests/unit/test_recipes_and_test_suites.py::test_all_recipe_yamls_accounted_for_in_test_suites
E   AssertionError: Recipe YAMLs should be accounted for in the test suites
E   assert 271 == 267

It is PR-introduced — counting from git, the merge-base has 267 recipe YAMLs against 267 suite entries and passes; this branch has 271. The test is pytest.mark.run_first, so it fails early in the unit lane.

The testing skill asks for three things per recipe: the YAML, a driver .sh in tests/test_suites/vlm/ matching the base name, and an entry in a suite list. This PR has the first only, and the four appear in no list at all — not even disabled.txt.

Worth noting the practical cost, since both parent recipes are currently in disabled.txt and nobody loses coverage today: because these four are in no list, re-enabling the Omni family later won't pick them up. It's also why the CP+VLM+TQ coverage gap has no route to being closed.

AI-1

Add tests/test_suites/vlm/<same-basename>.sh for each of the four (copy the shape of vlm_grpo-nemotron-omni-30ba3b-clevr-1n4g-megatron_generation.v1.sh) and list them — in disabled.txt if the family is meant to stay off for now. Multi-file, so no suggestion block.

Context — no action

This wasn't caught because /ok to test was posted on #3920 rather than on this PR, so no GitLab lane has run here. Running it will also surface the two failing tests flagged separately.

[terrykong]

if we add these to this PR, we should have a driver script. do these tests pair with a v1 one? if so, could you add wandb with evidence that convergence matches and then we deprecate the v1 nightlies of these and replace them with v2? that way we don't increase our nightly GPU hrs too much (we're actually overdue for a major trimming)

[cspades]

I disabled the V1 nightlies since I need to find a better setup than a pretrained Omni since it has basically a converged loss curve. Can I add these nightly tests after the code freeze? For now, I verified not standalone convergence but rather parity with vLLM as my true north offline, and I need to check in an L1 functional just to test against basic errors.

[terrykong]

2 action items.

TL;DR — this warmup registers processor.model_input_names, but the writer produces a wider set, so imgs_sizes and num_frames are written to TransferQueue without ever being pre-registered.

Background, in four lines:

1. TransferQueue registers field names lazily, the first time a column is written.
2. This block exists to register them all up front instead — the comment just above says "warm every known tensor field before rollout, policy, and teacher writers become concurrent".
3. It reads the bare processor.model_input_names. The writer instead uses get_multimodal_keys_from_processor, which unions the image_processor / video_processor / feature_extractor names and subtracts the tokenizer's — and then processors.py:630-639 hand-appends imgs_sizes and num_frames, with a comment saying outright that "imgs_sizes is not always declared in model_input_names by bundled image processors".
4. So those names are never warmed. RADIO / Nemotron Omni — the model in all four new recipes — is exactly the case that comment describes.

Being straight about severity: we first thought this could wedge the TQ controller, and it cannot. At the pinned TransferQueue rev (c516143) the registration and the read are both branches of one dispatch loop on a single thread, so they can't race. This is a correctness gap — the names get registered lazily instead of up front — not a hang.

AI-1 — widen the warmup

Suggested change

	getattr(processor, "model_input_names", ())
	[
	*get_multimodal_keys_from_processor(processor),
	*UNDECLARED_MULTIMODAL_MODEL_INPUTS,
	]

with the matching import added near the other nemo_rl.data imports (outside this diff, so not a suggestion):

from nemo_rl.data.multimodal_utils import (
    UNDECLARED_MULTIMODAL_MODEL_INPUTS,
    get_multimodal_keys_from_processor,
)

Dropping the getattr(..., ()) default is deliberate — a missing attribute should fail loudly rather than silently warm nothing.

AI-2 — define the names once

multimodal_utils.py:796-801 already hand-appends four names, two more than processors.py does, so a two-name fix would still under-register. Define them once in multimodal_utils.py and have that existing loop read the constant — same names, same order, same guard, so it is an exact no-op there:

# Model inputs some remote-code processors omit from ``model_input_names``
# even though their forward requires them. Defined here so the data
# processor, the Gym extractor, and the SingleController partition warmup
# all agree on what a multimodal record can contain.
UNDECLARED_MULTIMODAL_MODEL_INPUTS = (
    "imgs_sizes",
    "num_frames",
    "pixel_values_flat",
    "image_num_patches",
)

Context — no action. We deliberately did not fold these into get_multimodal_keys_from_processor itself. Doing that would make processors.py start wrapping pixel_values_flat / image_num_patches as PackedTensors where it doesn't today — harmless for every model currently in the repo, but not the no-op it looks like, so it seemed wrong to slip into this PR. Happy to be overruled if you'd rather have the single helper.

[terrykong]

1 action item — a request for evidence, not a code change. Anchored on this recipe because that is where the evidence would come from; the code it discusses is in worker_mixin.py, and the change there is a deletion, so there is no diff line to attach to.

TL;DR — the PR removes the guard that blocked VLMs on the TransferQueue path, and nothing in the tree exercises the combination it used to block.

The removed guard, at worker_mixin.py:138-143 on the base commit (linked at the base, since it no longer exists at head):

if getattr(self, "model_slices_context_parallel_inputs", False):
    raise NotImplementedError(
        "TransferQueue/SingleController does not yet support models that "
        "insert media before context-parallel input selection. Use the "
        "synchronous NeMo-RL policy path for Nemotron Omni."
    )

The support does look genuinely implemented, so this is not a guard deleted in place of doing the work. The context-parallel slicing already lives in models/megatron/data.py and is reached through the shared entry points that both the TQ and non-TQ paths call — train, get_logprobs, get_topk_logits. The missing piece was multimodal fields on the wire, which is what this PR adds. Removing the guard is the last step, not the only one.

What is missing is anything exercising it. The nearest coverage stops short on both sides:

• test_megatron_data.py:398 and test_nemotron_omni_model.py:187 cover CP slicing with model_slices_context_parallel_inputs=True, but never through TransferQueue.
• nightly.txt:89 covers CP2 + SingleController + TQ, but text-only.
• The replacement unit test cannot close it either: setup_data_plane never reads model_slices_context_parallel_inputs now that the guard is gone, so the worker.model_slices_context_parallel_inputs = True line in it is inert.

One nuance that widens this: model_slices_context_parallel_inputs is read off the model, not derived from cp_size. Nemotron Omni sets it True even at CP=1, so all four new recipes would have tripped the old guard regardless of their context_parallel_size: 1.

AI-1

Since this is the first change to enable multimodal on SC/TQ, could you share some evidence of what it does and what it costs?

• A short reward curve from one of the four new recipes against the equivalent non-SC run, so convergence is visibly unchanged. The 1n4g CLEVR one looks cheapest, once the config issues in the other comments are sorted.
• Step time and peak HBM around a refit. _move_retained_vlm_media copies every in-flight request's media host-to-device before resume and back after — once per refit, not per step, and only imgs — but for 32-frame video that is worth a number.
• Either a cp_size: 2 variant wired into a suite list, or a line in the PR description saying CP>1 + VLM + TQ is untested.

Context — no action. These four recipes currently have no driver script, so there is no route to closing this coverage gap in CI even if you wanted to — see the comment on vlm_grpo-nemotron-omni-30ba3b-clevr-1n4g-megatron-single-controller-async.v1.yaml.

[terrykong]

2 action items.

TL;DR — two things about this signature: available_fields is typed non-optional but callers pass a nullable value, and nothing records when this helper should and shouldn't be called.

1. It raises on None. KVBatchMeta.fields is list[str] | None = None, and the three call sites in tq_policy.py pass it straight through:

TypeError: 'NoneType' object is not iterable
    nemo_rl/data_plane/schema.py:143

Every meta reaching those sites today has fields populated, so we could not build a failing case — type hygiene, not a live bug. Worth doing because KVBatchMeta.concat in this same PR handles the identical case with meta.fields or [], so the two now disagree about whether None is allowed.

2. There are three situations and the reader can't tell them apart. _isolated_meta has five call sites. Three wrap with this helper, two deliberately don't:

Site	Wraps?	Why
tq_policy.py:230 / :340 / :466	yes	policy consumes media
tq_value.py:113 / :164	no — correct	critic is text-only, is_vlm=False
setup.py:1361, tq_policy.py:168 / :187	n/a	registering a schema — no record exists yet, so they use packed_tensor_wire_fields instead

The tq_value omission is right, but a reader at tq_value.py:113 sees three wrapped siblings and two bare ones with nothing explaining the difference. The next person adding a read path has to infer the rule from five examples.

AI-1

Fix the annotation and write the rule down, in one edit:

Suggested change

	available_fields: Sequence[str],
	) -> list[str]:
	"""Include PackedTensor payload/metadata columns present in a TQ record."""
	out = list(fields)
	available = set(available_fields)
	available_fields: Sequence[str] | None,
	) -> list[str]:
	"""Include PackedTensor payload/metadata columns present in a TQ record.
	Call this when narrowing an existing record's fields for a consumer that
	can use media — the ``_isolated_meta`` sites on the policy path. Do NOT
	call it when registering a partition schema (no record exists yet; use
	``packed_tensor_wire_fields`` from the processor's inputs), or when the
	consumer is text-only by design, e.g. the PPO critic in ``tq_value``,
	which is built with ``is_vlm=False`` and cannot use media columns.
	"""
	out = list(fields)
	if not available_fields:
	return out
	available = set(available_fields)

AI-2

The docstring above helps someone who opens schema.py. It does not help the more likely reader: someone sitting at tq_value.py:113 wondering why the two calls there look different from the three in tq_policy.py. Please add a half-line at both bare sites — tq_value.py:113 and tq_value.py:164:

# Critic is text-only (built with is_vlm=False), so media columns are
# deliberately not fetched here — see fields_with_packed_tensor_payload.

Not a suggestion block: tq_value.py isn't in this PR's diff. It is two comment lines in a file you didn't otherwise touch, so entirely reasonable to push back on — but without it, the next person to add a read path has to reconstruct the rule from five examples, and guessing wrong in the media direction fails silently.

Context — no action. Two fixes we considered and rejected, so they don't get tried later:

Making KVBatchMeta.fields non-optional. None is not an unset default — it is a third state meaning this meta declares no field set, which the data plane relies on to fail loudly. From DataPlaneClient.get_data: "Field-set resolution: (1) explicit select_fields; (2) meta.fields if non-None; (3) fail loudly — never silently fetch all fields." Default it to [] and the raise at noop.py:147-151 stops firing: an undeclared meta becomes a declared-empty one and get_data returns zero columns instead of refusing.

Folding the expansion into _isolated_meta itself. It already holds the meta, so this looks like it would fix all five sites at once — but it would start handing the critic pixel columns it cannot consume. The 3-vs-2 split is intentional; only the reason for it is missing.

[terrykong]

1 action item, 1 follow-up.

TL;DR — if not record: never fires, because the real record type defines neither __len__ nor __bool__.

entry.record is a DynamicInferenceRequestRecord — a @dataclass(kw_only=True) whose only dunder is __getitem__. An object with neither __len__ nor __bool__ is always truthy, so this guard catches only None; a record with an empty .requests falls through to record[-1] and raises IndexError instead of being skipped.

Harmless today — RequestEntry is only built via from_request, which always appends one request — but the guard reads as if it does something it doesn't. The new test can't catch it either: it passes record=[first_request], a plain list, which is falsy when empty. Upstream itself writes len(entry.record.requests) rather than relying on truthiness.

AI-1

Suggested change

	record = getattr(entry, "record", None)
	if not record:
	continue
	request = record[-1]
	record = getattr(entry, "record", None)
	requests = getattr(record, "requests", None)
	if not requests:
	continue
	request = requests[-1]

Follow-up

_wake has no tests at all, including the new except BaseException leg that returns media to CPU when resume fails. That's the riskiest new code here — if it doesn't run, image batches stay pinned in HBM. Worth one test asserting media goes back to CPU and _inference_engine_asleep stays True on failure.

Context — no action. The rest of this workaround checks out, and we verified it against the pinned Megatron-LM SHA. Moving only imgs is correct: _create_vlm_request forces num_tiles / imgs_sizes / num_frames to CUDA unconditionally at dynamic_engine.py:1699-1702 and moves imgs only inside the cache-miss branch, so imgs is the one field that can still be CPU-resident. We also confirmed no request can execute between RESUMED and RUNNING — RESUMED is set after the state is already PAUSED, and PAUSED only sleeps — so the move back to CPU cannot race with inference. Your comment is accurate. And Megatron-LM#6976 is still open, so this needs to stay.