[Inference] Support video modality in Megatron-Inference.

[cspades]

• I, the PR author, have personally reviewed every line of this PR.

What does this PR do?

Associated with: NVIDIA-NeMo/RL#3834

• Extend image modality to support video modalities in Megatron-Inference.
  • Modalities
    • Images can be either shipped into the completions endpoint as raw bytes ({"image": [image1_bytes, image2_bytes]}), or passed Megatron(Async)LLM.generate as preprocessed Tensors ({"image": {"imgs": ..., "imgs_sizes": ..., "num_tiles": ..., "num_img_embeddings_per_tile": ...}}).
    • Video can be either shipped into the completions endpoint as raw bytes ({"video": [video1_bytes, video2_bytes]}) or a byte-string containing paths to video frame images (VideoProcessingConfig.frame_manifest_magic + json.dumps({"frame_paths": [...]}).encode("utf-8")), or passed to Megatron(Async)LLM.generate as preprocessed Tensors ({"video": {"imgs": ..., "imgs_sizes": ..., "num_frames": ... }}).
    • Pre-processed Tensors can be optionally pre-expanded with multimodal placeholders by RL or the user using the config media_tokens_preexpanded: bool in the multi_modal_data dictionary input to Megatron(Async)LLM.generate.
      • In a nutshell, pre-processed + pre-expanded Tensors = NeMo-RL non-Gym handles the multimodal preprocessing, while bytes = NeMo-RL Gym sends over the raw data for Megatron-Inference to handle the multimodal preprocessing. Having direct engine access in the first (Tensor) case is generally something we want to keep maintaining, not only to support non-Gym in RL but also for users to have direct control over what Tensors go into the encoder/decoder.
    • If you ever see the term "compact" describing prompt tokens it means "prior to expansion" where we have a single token placeholder representing an image or frame or video that will be expanded into many placeholder multimodal tokens for embedding injection before generation. This form is stored and used in some situations where we want to return a non-negative prompt of tokens back to RL and also in multi-turn chat where we re-tokenize and re-expand these compact prompts for chat continuation.
    • These match existing NeMo-RL vLLM API's for multimodal. (But it's honestly a mess of possible inputs depending on what each multimodal dataset prefers...)
  • Just like images, the video embeddings are computed in _build_vlm_request before the controller takes over decoding, and the embedding installed in the request can be used for checkpointing, etc. I hope I didn't unintentionally change too much for the trunk of the decoder engine here.
  • Caching / Routing - A cache key is computed here depending on the bytes: inference_request.py/compute_media_cache_key
    • Support media affinity via cache key in coordinator routing, which mixes into the affinity score produced by prefix hits for text.
    • Support vision embedding caching to skip the encoder pass on subsequent generations. The cache is LRU-cleared and cleared on suspend.
    • Salt in a media cache key in compute_block_hashes_batched to make prefix caching aware of the multimodal data in the prompt as well.
      • Note to self: DynamicVLMInferenceRequest saves the media key into the context and this is used later (_compute_block_hashes) to salt and recompute the engine-internal KV cache hash, which uses the expanded prompt with ``
  • Inference wrappers have a few more methods that are common to multimodal: supports_<modality> for quick-checking what a wrapper can support for a model, MultimodalPromptConfig and resolve_media_token_id for looking up what image/video placeholders and token IDs should be for a model so that the inference engine can use them as embedding placeholders in expand_image_tokens called by _build_vlm_request. This is a very experimental API for inference wrappers.
    • Added a new inference wrapper for Nemotron Omni: NemotronOmniInferenceWrapper.
    • Also updated VLMInferenceWrapper for video and LLaVA.
  • dynamic_res_preprocess has been updated to support multiple ways of patching frames, and video_maintain_aspect_ratio=True is required for parity with RL + vLLM.
  • For video frames, we only sample up to num_frames so we can sample first, then convert the sampled frames to PIL / Tensor.

For both images and video, we have vLLM parity: https://wandb.ai/adlr/mllm-rl-dev?nw=nwusercye_nv

Log prob error, gradient norms, and gen KL error all look reasonable as well, with some noise.

In the RL NSight profile, multimodal embeddings are computed once per prompt group and cached for future use, and then subsequent forward passes are decoder-only, all overlapped with RL training for async non-colocated (trajectory age = 2):

There's a known issue with this gap appearing in the generation loop, will debug in a future PR.

Bugs

• Added back LLaVA audio support which seems to have been reverted for not much reason, referring to this PR: Add LLaVA audio (sound) model support #4402 -> Add VLM support to the dynamic-batching inference server #6260

Testing

• RL
  • https://github.com/cspades/RL/blob/cye/rl_mllm_omni_multimodal_scripts/scripts has 1N and 8N scripts with easy swapping between vLLM and MLLM.
• E2E Multimodal Inference
  • Run examples/inference/*_multimodal_infer.py on 2 GPUs with at least 40 GB. (For me, my two A6000's.)

uv run python -m torch.distributed.run --nproc-per-node 8 -m pytest -vvs tests/unit_tests/inference

========= 2601 passed, 108 skipped, 2872 warnings in 971.37s (0:16:11) ==========

⚠️ For major changes (either in lines of code or in its impact), please make sure to first share a design doc with the team. If you're unsure what's the best way to do so, contact @NVIDIA/mcore-oncall.

Contribution process

Pre-checks

• I have added relevant unit tests
• I have added relevant functional tests
• I have added proper typing to my code Typing guidelines
• I have added relevant documentation
• I have run the autoformatter.sh on my PR

Code review

Feel free to message or comment @NVIDIA/mcore-oncall to help accelerate your merge into main. The less complex your PR is, the faster it will be approved and merged!

All PRs start as draft. If you open a non-draft PR, it will be automatically converted to draft.

Step 1: Mark PR as "Ready for Review"

1. When your PR is ready, click Ready for Review.
2. An oncall reviewer is auto-assigned and expert reviewers are notified based on your changes.
   • Some PRs may jump straight to step 2. This is determined by .github/CODEOWNERS.

⚠️ Only mark as ready once merge-conflicts are resolved and the CI is passing.
Final Review might get declined if these requirements are not fulfilled.

Step 2: Final Review

For PRs that change megatron/core, once all expert reviewers have approved, the Final Review label is applied automatically and final reviewers are assigned.

For PRs outside megatron/core, this step is skipped.

Step 3: Approved

Once all required reviewers have approved, the Approved label is applied automatically.

Merge

Any member of mcore-engineers will be able to merge your PR.

[github-actions]

This PR has been automatically converted to draft because all PRs must start as drafts.

When you are ready for review, click Ready for Review to begin the review process. This will:

1. Add the oncall reviewer (optional reviewer)
2. Add required review teams based on your changes

See the contribution guide for more details.

[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]

/claude strict-review

[RPrenger]

LGTM

[nemo-automation-bot]

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/33105170466

[nemo-automation-bot]

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/33107050424