feat(mcore-inference-multimodal): Support multimodal Megatron inference.

[cspades]

What does this PR do ?

Original PR: #3834 (All comments have at least been addressed, that's where we are at.)

• 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.

Cause I keep forgetting where everything is...

Non-Gym RL Multimodal Data Pipeline

Dataset preprocessing:

vlm_hf_data_processor → resolve_to_image / load_video → Hugging Face Omni processor → produces token_ids, pixel_values, imgs_sizes, optionally num_frames → stores media as PackedTensor

Files:

nemo_rl/data/processors.py
nemo_rl/data/multimodal_utils.py

Rollout generation:

RolloutManager.generate_and_push → AsyncRolloutImpl._generate_response → batched_message_log_to_flat_message → MegatronGeneration.generate_async → MegatronWorker._prepare_data_for_generation → _sample_vision_tensors → _build_prompt_and_multimodal_data → MCore InferenceClient.add_request

Files:

nemo_rl/experience/rollout_manager.py
nemo_rl/data/llm_message_utils.py
nemo_rl/models/generation/megatron/megatron_worker.py

Megatron inference:

resolve_multimodal_data_for_engine → tensor payload passes through without decoding → DynamicInferenceEngine._build_vlm_request → NemotronOmniInferenceWrapper._forward_vision_encoder → NemotronOmniModel._encode_images

Here media_tokens_preexpanded=True, because NeMo-RL already expanded the image placeholders.

Gym RL Multimodal Data Pipeline

Dataset-side preprocessing for training/recompute:

nemo_gym_data_processor → nemo_gym_example_to_video_datum_spec → load_video_frames_with_metadata → process_nemotron_video_frames → _expand_nemotron_video_placeholders → produces PackedTensor(pixel_values, imgs_sizes, num_frames)

Files:

nemo_rl/data/processors.py
nemo_rl/environments/nemo_gym_multimodal.py
nemo_rl/environments/nemotron_utils.py

Gym rollout generation:

AsyncNemoGymRolloutImpl → NemoGym.run_rollouts → normalize_media_in_examples → Gym RolloutCollectionHelper.run_examples → Gym agent → SimpleResponsesAPIModel.responses_dispatch → Megatron OpenAI HTTP server

Files:

nemo_rl/experience/rollout_manager.py
nemo_rl/environments/nemo_gym.py
3rdparty/Gym-workspace/Gym/nemo_gym/rollout_collection.py
3rdparty/Gym-workspace/Gym/nemo_gym/base_responses_api_model.py

HTTP-side preprocessing for generation:

chat_completions → _extract_multimodal_from_messages → raw image/video bytes → resolve_multimodal_data_for_engine → preprocess_video_bytes_list → DynamicInferenceEngine._build_vlm_request → NemotronOmniInferenceWrapper.expand_image_tokens → _forward_vision_encoder → NemotronOmniModel._encode_images

Files:

Megatron-LM/.../endpoints/chat_completions.py
Megatron-LM/.../image_preprocessing.py
Megatron-LM/.../dynamic_engine.py
Megatron-LM/.../nemotron_omni_inference_wrapper.py

Here media_tokens_preexpanded=False; Megatron-LM expands placeholders itself.

After Gym returns:

NemoGym._postprocess_nemo_gym_to_nemo_rl_result → attach generated token IDs/logprobs → attach_static_multimodal_payload → restore the original NeMo-RL PackedTensors for policy recompute and training.

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: 1092610 (PR #3834 from cye/rl_mllm_omni_multimodal)

✅ Submodules that are properly updated:

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

All submodule changes look good! ✨

[cspades]

/ok to test 1092610

[github-actions]

✅ Submodule Fast-Forward Check Results

Check based on commit: f755067 (PR #3920 from cye/rl_mllm_omni_multimodal)

✅ Submodules that are properly updated:

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

All submodule changes look good! ✨

[cspades]

/ok to test f755067

[cspades]

/ok to test cfa812e

[cspades]

/ok to test 7be9e08

[terrykong]

approved modulo passing pipeline

[cspades]

/ok to test 95b0a89