feat(multimodal): Support multimodal Megatron generation.

[cspades]

What does this PR do ?

Depends on NVIDIA/Megatron-LM#6809 and NVIDIA-NeMo/Megatron-Bridge#5799

• Integrates multimodal Megatron Inference into NeMo-RL using Nemotron Omni as a start, and having functional / image (CLEVR) and video (VSTAT) accuracy parity with vLLM: https://wandb.ai/adlr/mllm-rl-dev?nw=nwusercye_nv
• Refactors some multimodal utils after an earlier video GRPO PR: feat(video): add Gym support for sync and async GRPO - v2 #3648

Usage

• Many YAML configs like RL/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-16n8g-megatron-tp4ep4-async-gym-video.v1.yaml have been updated to support multimodal generation (if provided a multimodal dataset). Some example MLLM and vLLM parity scripts I used to test with are in a side-branch: https://github.com/cspades/RL/tree/cye/rl_mllm_omni_multimodal_scripts/scripts
  • To run these scripts: GENERATION_BACKEND={vllm,megatron} scripts/submit_nemotron_omni_vstat_megatron_8n4g.sh on the Slurm login node or run_nemotron_omni_clevr_megatron_1n4g.sh in an interactive node.

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.

[cspades]

/ok to test 63f916f

[cspades]

Lockfile test is not gonna be fixed until we bump MBridge and MCore. 😩

[copy-pr-bot]

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

Contributors can view more details about this message here.

[cspades]

Team review by 5 agents (rl-expert, bug-finder, test-agent, design-reviewer, devil-advocate). 8 findings staged as inline comments: 3 critical (a submodule/pyproject version mismatch that breaks uv sync --extra mcore today, 3 new tests never collected by any CI shard, and a bounded ~600s setup stall on Megatron weight-sync failure with NeMo Gym), 5 suggestions (a duplicate/hidden-default config field, a medium-confidence eval.py regression for mixed multimodal batches, a low-severity KeyError gap, a non-blocking design note, and a verified-passing test-coverage addition). Local pre-commit could not run in this environment (Python 3.13.14 interpreter unavailable) — not attributable to this PR's diff.

Generated by Claude Code

[cspades]

pyproject.toml:440

1 action item.

This bumps emerging-optimizers to @v0.3.0, but the nested Megatron-LM submodule pin at 3rdparty/Megatron-Bridge-workspace/Megatron-Bridge/3rdparty/Megatron-LM/pyproject.toml:244 still pins v0.2.0. Reproduced directly: uv lock --check fails with

Failed to resolve dependencies for `megatron-core`: Requirements contain conflicting URLs for package `emerging-optimizers`: v0.2.0 vs v0.3.0

The comment directly above this line documents the exact invariant this violates: "Keep these URLs/revs in sync with Megatron-LM's [tool.uv.sources] (uv errors loudly on mismatch after a submodule bump)." Every contributor's uv sync --extra mcore is currently broken by this.

AI-1

Revert this to v0.2.0 until the companion Megatron-LM/Megatron-Bridge submodule bump lands together, or land the submodule bump in the same PR.

[cspades]

tests/unit/models/generation/test_megatron_generation.py:43

1 action item.

This test (and the two below it, test_direct_megatron_video_request_marks_preexpanded_prompt at L71 and test_direct_megatron_multimodal_generate_round_trip at L107) is missing @pytest.mark.mcore — every other test/class in this file (from L441 onward) carries it.

This matters because none of the other L0 shards ever collect this file's content: Models_*.sh explicitly --ignores this directory, and Vllm_*.sh/Sglang.sh filter by their own backend markers. Only L0_Unit_Tests_Megatron.sh (--mcore-only) would collect it, and marker-filtering means these 3 unmarked tests are excluded there too. These 3 tests are not executed by any currently-configured CI shard.

AI-1

Add @pytest.mark.mcore above each of the 3 new tests.

Suggested change

	def test_direct_megatron_media_request_preserves_preexpanded_prompt():
	@pytest.mark.mcore
	def test_direct_megatron_media_request_preserves_preexpanded_prompt():

[cspades]

nemo_rl/algorithms/grpo.py:1402-1423

1 action item.

Background: for non-colocated (skip-load) Megatron generation, the inference engine's HTTP server does not start at construction — only during sync_weights() inside init_megatron_weight_synchronizer (defined at L1275). _spinup_nemo_gym (submitted as nemo_gym_future in the same ThreadPoolExecutor) polls that endpoint and blocks until it responds.

What breaks: init_megatron_weight_synchronizer(policy, policy_generation) is called inline (not as its own future) at L1420, between generation_future.result() and nemo_gym_future.result(), still inside the with ThreadPoolExecutor(...) block. If init_communicator() raises before sync_weights() runs, the exception unwinds through ThreadPoolExecutor.__exit__ → shutdown(wait=True), which now waits on nemo_gym_future — still polling an endpoint that will never come up. Traced through nemo_gym's own readiness wait (RunHelper.start() → _wait_for_model_endpoints): it's bounded by model_endpoint_readiness_timeout_seconds, defaulting to 600s and unset by every recipe in this repo. So the real error (and the outer finally: ray.kill(port_holder) cleanup at L1423) is delayed by up to 10 minutes behind a confusing secondary ConfigError about unreachable model endpoints — a control-plane failure now looks like a hang instead of failing fast. (If a recipe ever sets that timeout to 0, this becomes a true infinite hang.)

AI-1

Don't let a raise from init_megatron_weight_synchronizer transit ThreadPoolExecutor.__exit__ while nemo_gym_future can still block on it. Either submit it as its own executor task and use concurrent.futures.wait(..., return_when=FIRST_EXCEPTION) + cancel_futures=True, or wrap the inline call in try/except and explicitly cancel/bound nemo_gym_future (or at least call ray.kill(port_holder)) before re-raising.

[cspades]

Disputing my own finding here: @tdene reviewed this exact restructured code on the sibling PR #3864 and confirmed the 600s wait_for_model_endpoints timeout is intentional, working-as-designed behavior, not a leak: #3864 (comment) ("there is already a 600s timeout on wait_for_model_endpoints. Nothing leaks, things time out as designed."). Withdrawing the suggested code change — no action needed here.

[cspades]

nemo_rl/models/generation/megatron/config.py:56

1 action item.

logprobs_mode is declared twice in MCoreGenerationSpecificArgs: NotRequired[...] here, and required (Literal["processed_logprobs", "raw_logprobs"], pre-existing) at L85. Since TypedDict class bodies are executed like a normal namespace, the later (required) declaration wins — this new NotRequired entry is dead.

Consistent with that, megatron_worker.py:391-393 was changed from a plain mcore_generation_config["logprobs_mode"] subscript to .get("logprobs_mode", "processed_logprobs") — a hidden default at the call site that contradicts both the still-required TypedDict declaration and the config-conventions rule that defaults belong in exemplar YAML, not .get(k, default) in code. Every exemplar/recipe YAML checked already sets this key explicitly.

AI-1

Pick one: delete the stale required declaration at L85 to make the field genuinely optional (keep the .get() fallback), or delete the dead NotRequired entry here and revert megatron_worker.py:391-393 back to the plain subscript.

[cspades]

nemo_rl/evals/eval.py:352-374

1 action item (medium confidence — code path confirmed, haven't confirmed a specific recipe hits it).

For a native-content (non-placeholder) VLM processor, vllm_content[i] is populated with the chat-template-formatted dialog regardless of whether that particular row has any media (vlm_hf_data_processor sets it unconditionally when not uses_placeholder). The old gate used vllm_content whenever it was non-None; this PR's gate is if multi_modal_data: (L352) — so a text-only row inside an otherwise-multimodal batch now falls into the "Text-only fallback" branch (L374-383), which naively joins raw message["content"] strings instead of using the already-formatted vllm_content[i].

AI-1

Confirm whether mixed (some-rows-have-media, some-don't) multimodal eval batches are a real path for any current recipe. If so, use vllm_content[i] whenever it's non-None, independent of whether multi_modal_data is empty for that row.

[cspades]

nemo_rl/models/generation/megatron/utils.py:109

1 action item, low severity — no exemplar/recipe currently triggers this (all set both fields together).

build_video_preprocessing_config gates only on video_temporal_patch_size being set, then unconditionally does generation_config["video_num_frames"]. Both fields are independently NotRequired in MCoreGenerationConfig, so a config setting only video_temporal_patch_size gets a raw KeyError instead of a clear error.

AI-1

Suggested change

	num_frames=int(generation_config["video_num_frames"]),
	video_num_frames = generation_config.get("video_num_frames")
	if video_num_frames is None:
	raise ValueError(
	"video_num_frames must be set when video_temporal_patch_size is set."
	)
	return VideoProcessingConfig(
	image_config=image_config,
	num_frames=int(video_num_frames),

[cspades]

nemo_rl/algorithms/grpo.py:1362-1433

No action needed — non-blocking design suggestion, not a blocker for this PR.

Megatron's nemo_gym-parallel-init path here diverges from the shared init_tasks dict/loop shape vLLM's branch still uses (L1506-1577) into 3 hand-named futures (policy_future, generation_future, nemo_gym_future) with the weight-sync ordering baked in procedurally, to satisfy the same "sync must precede gym resolving" constraint the bug above is about. A named plausible extension: sglang/trtllm adding a similar deferred-load + parallel-gym-init mode would have no shared primitive to extend and would likely copy this hand-unrolled block a third time.

Follow-up

Consider a shared run_parallel_init_tasks(tasks, after={name: hook}) primitive — an optional after hook keyed by task name that fires once that task's future resolves, still inside the same executor block — so Megatron could rejoin the same dict/loop shape as vLLM instead of a bespoke block.

[cspades]

nemo_rl/utils/venvs.py:233-235

1 action item — new test, verified passing locally.

make_actor_runtime_env's "uv"-prefixed executable branch (the local per-node venv creation path, previously inlined in grpo.py and now centralized here) has no direct unit test — tests/unit/algorithms/test_grpo.py mocks make_actor_runtime_env wholesale instead of exercising it.

AI-1

Suggested change

	add_hf_modules_cache_to_pythonpath,
	def test_make_actor_runtime_env_builds_local_venv_for_uv_python_executable():
	"""Mirrors the inline venv-creation logic that used to live in grpo.py."""
	with (
	patch(
	"nemo_rl.distributed.ray_actor_environment_registry.get_actor_python_env",
	return_value="uv run --group vllm",
	) as mock_get_env,
	patch(
	"nemo_rl.utils.venvs.create_local_venv_on_each_node",
	return_value="/fake/venv/bin/python",
	) as mock_create_venv,
	):
	runtime_env = make_actor_runtime_env("some.module.SomeActor")
	mock_get_env.assert_called_once_with("some.module.SomeActor")
	mock_create_venv.assert_called_once_with(
	"uv run --group vllm", "some.module.SomeActor"
	)
	assert runtime_env["py_executable"] == "/fake/venv/bin/python"
	assert runtime_env["env_vars"]["VIRTUAL_ENV"] == "/fake/venv"
	assert runtime_env["env_vars"]["UV_PROJECT_ENVIRONMENT"] == "/fake/venv"

(needs make_actor_runtime_env added to this file's imports)