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

docs/design-docs/generation.md

@@ -143,6 +143,43 @@ The `mcore_generation_config` section controls Megatron Core inference engine be
 - **num_cuda_graphs**: Number of CUDA graphs to pre-allocate for different batch sizes. More graphs can improve performance by avoiding runtime graph capture, but consume more memory.
 - **max_tokens**: Maximum total number of tokens (across all requests) that can be processed simultaneously. This limits the maximum batch size and sequence length combinations. Increasing this might throw OOM depending on vocab size and buffer size allocated. 
 
+### Multimodal Megatron Generation
+
+Megatron inference supports image and video inputs in NeMo-RL. Enable multimodal processing with `policy.is_vlm: true`, use the `megatron` generation backend, and provide a `megatron_inference_wrapper`. The wrapper must subclass `megatron.core.inference.model_inference_wrappers.abstract_model_inference_wrapper.AbstractModelInferenceWrapper` in Megatron-Core and declare `supports_<modality> = True` for each supported modality.
+
+```yaml
+policy:
+  is_vlm: true
+  generation:
+    backend: megatron
+    mcore_generation_config:
+      megatron_inference_wrapper: megatron.core.inference.model_inference_wrappers.multimodal.nemotron_omni_inference_wrapper.NemotronOmniInferenceWrapper
+      image_dynamic_resolution: true
+      video_num_frames: 16
+      video_temporal_patch_size: 2
+      video_target_num_patches: 2048
+      video_maintain_aspect_ratio: true
+      vision_embedding_cache_max_bytes: 0
+      allow_stale_multimodal_embeddings: false
+data:
+  default:
+    num_frames: 16
+    video_temporal_patch_size: 2
+    video_target_num_patches: 2048
+    video_maintain_aspect_ratio: true
+```
+
+- `image_dynamic_resolution` preserves variable image shapes instead of forcing one fixed resolution; for example, a wide image uses a wider patch grid than a square image.
+- `vision_model_type` optionally selects the MCore vision encoder type used by image and video preprocessing. Set it to the encoder expected by the inference wrapper; when omitted, MCore uses its default (`radio`).
+- `num_frames` controls uniform video-frame sampling. Use `video_num_frames` for the corresponding MCore key.
+- `video_temporal_patch_size` groups sampled frames into temporal tubelets; for example, size `2` turns 16 frames into 8 temporal groups.
+- `video_target_num_patches` sets `num_patches_per_frame = patch_height * patch_width <= video_target_num_patches`, which produces `num_patches_per_frame * num_frames / video_temporal_patch_size` total video patches prior to spatial merging (i.e. further grouped / concatenated into MxM patch blocks) that are provided to the vision encoder.
+- `video_maintain_aspect_ratio=true` keeps `patch_width / patch_height ~= source_width / source_height`; `false` uses `patch_width = patch_height ~= sqrt(video_target_num_patches)` (for example, `sqrt(256) = 16`).
+- `vision_embedding_cache_max_bytes` limits GPU memory used to reuse vision embeddings for repeated media; `0` disables the cache, while `1073741824` permits up to 1 GiB.
+- `allow_stale_multimodal_embeddings` controls whether cached embeddings survive model-weight changes. Keep it `false` for RL refits; use `true` only when weights remain fixed.
+- `expose_http_server` should be `true` for NeMo Gym.
+
+Keep the video preprocessing values identical in `data.default` and `mcore_generation_config` to avoid disparity between the training policy and inference generation.
 
 ## Usage Examples
 

examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-16n8g-megatron-tp4ep4-async-gym-video.v1.yaml

@@ -48,6 +48,10 @@ policy:
   make_sequence_length_divisible_by: 32
   generation:
     bad_words: []
+    mcore_generation_config:
+      image_dynamic_resolution: true
+      logprobs_mode: raw_logprobs
+      megatron_inference_wrapper: megatron.core.inference.model_inference_wrappers.multimodal.nemotron_omni_inference_wrapper.NemotronOmniInferenceWrapper
     vllm_cfg:
       reset_encoder_cache_after_weight_update: false
       video:

examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-circle-click-2n8g-megatron-tp2ep8.v1.yaml

@@ -90,6 +90,9 @@ policy:
       max_tokens: ${policy.max_total_sequence_length}
       expose_http_server: true
       enable_prefix_caching: true
+      image_dynamic_resolution: true
+      logprobs_mode: raw_logprobs
+      megatron_inference_wrapper: megatron.core.inference.model_inference_wrappers.multimodal.nemotron_omni_inference_wrapper.NemotronOmniInferenceWrapper
       parsers:
       - deepseek-r1-reasoning
       - qwen3-coder-tool

examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-circle-count-1n4g-megatron_generation.v1.yaml

@@ -0,0 +1,96 @@
+defaults: ../../vlm_grpo_3B_megatron.yaml
+grpo:
+  num_prompts_per_step: 2
+  num_generations_per_prompt: 8
+  max_num_steps: 4
+  val_period: 0
+  max_val_samples: null
+  val_batch_size: null
+  async_grpo:
+    enabled: true
+    max_trajectory_age_steps: 2
+    in_flight_weight_updates: true
+loss_fn:
+  reference_policy_kl_penalty: 0.0
+  use_importance_sampling_correction: true
+checkpointing:
+  enabled: false
+  checkpoint_dir: results/nemo-rl-omni/nemotron-omni-circle-count-1n4g
+policy:
+  model_name: nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16
+  is_vlm: true
+  train_global_batch_size: 16
+  logprob_batch_size: 1
+  max_total_sequence_length: 8192
+  sequence_packing:
+    enabled: true
+  megatron_cfg:
+    env_vars:
+      TORCH_CUDA_ARCH_LIST: '10.0'
+    tensor_model_parallel_size: 2
+    expert_model_parallel_size: 2
+    sequence_parallel: true
+    bias_activation_fusion: false
+    activation_checkpointing: true
+  generation:
+    backend: megatron
+    bad_words: null
+    mcore_generation_config:
+      expose_http_server: true
+      buffer_size_gb: 8
+      num_cuda_graphs: -1
+      max_tokens: ${policy.max_total_sequence_length}
+      transformer_impl: transformer_engine
+      activation_checkpointing: false
+      tensor_model_parallel_size: 2
+      expert_model_parallel_size: 2
+      expert_tensor_parallel_size: 1
+      sequence_parallel: true
+      moe_pad_experts_for_cuda_graph_inference: true
+      image_dynamic_resolution: true
+      logprobs_mode: raw_logprobs
+      megatron_inference_wrapper: megatron.core.inference.model_inference_wrappers.multimodal.nemotron_omni_inference_wrapper.NemotronOmniInferenceWrapper
+    vllm_cfg:
+      async_engine: true
+      expose_http_server: true
+    colocated:
+      enabled: false
+      resources:
+        gpus_per_node: 2
+        num_nodes: 1
+data:
+  _override_: true
+  max_input_seq_length: null
+  shuffle: false
+  num_workers: 0
+  train:
+    data_path: 3rdparty/Gym-workspace/Gym/resources_servers/circle_count/data/example.jsonl
+  validation:
+    data_path: 3rdparty/Gym-workspace/Gym/resources_servers/circle_count/data/example.jsonl
+  default:
+    dataset_name: NemoGymDataset
+    env_name: nemo_gym
+    prompt_file: null
+    processor: nemo_gym_data_processor
+env:
+  _override_: true
+  should_use_nemo_gym: true
+  should_log_nemo_gym_responses: true
+  nemo_gym:
+    is_trajectory_collection: false
+    port_range_low: 5000
+    port_range_high: 5999
+    config_paths:
+    - responses_api_models/vllm_model/configs/vllm_model_for_training.yaml
+    - resources_servers/circle_count/configs/circle_count.yaml
+    circle_count_simple_agent:
+      responses_api_agents:
+        simple_agent:
+          max_steps: 1
+logger:
+  tensorboard_enabled: false
+  wandb:
+    project: nemo-rl-omni
+    name: nemotron-omni-circle-count-1n4g
+cluster:
+  gpus_per_node: 4

examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron_generation.v1.yaml

@@ -0,0 +1,77 @@
+defaults: ../../vlm_grpo_3B_megatron.yaml
+grpo:
+  num_prompts_per_step: 12
+  num_generations_per_prompt: 8
+  val_at_start: true
+  val_at_end: true
+  val_batch_size: 64
+  async_grpo:
+    enabled: true
+    max_trajectory_age_steps: 2
+    in_flight_weight_updates: true
+loss_fn:
+  reference_policy_kl_penalty: 0.0
+  use_importance_sampling_correction: true
+checkpointing:
+  enabled: false
+  checkpoint_dir: results/nemo-rl-omni/nemotron-omni-clevr-megatron-8n4g
+policy:
+  model_name: nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16
+  is_vlm: true
+  train_global_batch_size: ${mul:${grpo.num_prompts_per_step}, ${grpo.num_generations_per_prompt}}
+  logprob_batch_size: 1
+  max_total_sequence_length: 4096
+  sequence_packing:
+    enabled: true
+  megatron_cfg:
+    env_vars:
+      TORCH_CUDA_ARCH_LIST: '10.0'
+    tensor_model_parallel_size: 8
+    expert_model_parallel_size: 8
+    sequence_parallel: true
+    bias_activation_fusion: false
+    activation_checkpointing: true
+    optimizer:
+      exp_avg_dtype: bfloat16
+      exp_avg_sq_dtype: bfloat16
+      store_param_remainders: true
+  generation:
+    backend: megatron
+    max_new_tokens: 2048
+    bad_words: null
+    mcore_generation_config:
+      buffer_size_gb: 8
+      async_sched_mode: async
+      num_cuda_graphs: -1
+      use_cuda_graphs_for_non_decode_steps: false
+      max_tokens: ${policy.max_total_sequence_length}
+      transformer_impl: inference_optimized
+      activation_checkpointing: false
+      tensor_model_parallel_size: 8
+      expert_model_parallel_size: 8
+      expert_tensor_parallel_size: 1
+      context_parallel_size: 1
+      sequence_parallel: true
+      moe_router_dtype: fp32
+      moe_pad_experts_for_cuda_graph_inference: false
+      mamba_inference_ssm_states_dtype: float32
+      mamba_inference_conv_states_dtype: float32
+      image_dynamic_resolution: true
+      logprobs_mode: raw_logprobs
+      megatron_inference_wrapper: megatron.core.inference.model_inference_wrappers.multimodal.nemotron_omni_inference_wrapper.NemotronOmniInferenceWrapper
+    colocated:
+      enabled: false
+      resources:
+        gpus_per_node: 4
+        num_nodes: 6
+data:
+  default:
+    prompt_file: examples/prompts/clevr_cogent_cot_nemotron_omni.txt
+logger:
+  tensorboard_enabled: false
+  wandb:
+    project: nemo-rl-omni
+    name: nemotron-omni-clevr-megatron-8n4g
+cluster:
+  gpus_per_node: 4
+  num_nodes: 8

examples/nemo_gym/prepare_video_dataset.py

@@ -256,7 +256,7 @@ def convert(args: argparse.Namespace) -> None:
         row["responses_create_params"] = {
             "input": input_messages,
             "metadata": {
-                "chat_template_kwargs": {"enable_thinking": True},
+                "chat_template_kwargs": json.dumps({"enable_thinking": True}),
             },
         }
         raw_answer = source_row.get("answer")

nemo_rl/algorithms/grpo.py

@@ -77,7 +77,6 @@
 from nemo_rl.data.utils import extract_necessary_env_names, load_dataloader_state
 from nemo_rl.data_plane.interfaces import DataPlaneConfig
 from nemo_rl.distributed.batched_data_dict import BatchedDataDict
-from nemo_rl.distributed.ray_actor_environment_registry import get_actor_python_env
 from nemo_rl.distributed.virtual_cluster import (
     TOPO_RANK_UNKNOWN,
     ClusterConfig,
@@ -158,7 +157,7 @@
 )
 from nemo_rl.utils.nsys import maybe_gpu_profile_step
 from nemo_rl.utils.timer import TimeoutChecker, Timer
-from nemo_rl.utils.venvs import create_local_venv_on_each_node
+from nemo_rl.utils.venvs import make_actor_runtime_env
 from nemo_rl.weight_sync.checkpoint_engine_config import (
     checkpoint_engine_refit_config,
 )
@@ -4553,28 +4552,9 @@ def async_grpo_train(
     print(f"   - train_global_batch_size: {train_gbs}")
     print(f"   - min_trajectories_needed: {min_trajectories_needed} (async mode)")
 
-    _replay_py_exec = get_actor_python_env(
+    _replay_runtime_env = make_actor_runtime_env(
         "nemo_rl.algorithms.async_utils.ReplayBuffer"
     )
-    if _replay_py_exec.startswith("uv"):
-        # Lazily build a dedicated venv across all Ray nodes on-demand.
-        _replay_py_exec = create_local_venv_on_each_node(
-            _replay_py_exec,
-            "nemo_rl.algorithms.async_utils.ReplayBuffer",
-        )
-
-    _replay_py_venv = os.path.dirname(
-        os.path.dirname(_replay_py_exec)
-    )  # to remove the "bin/python" suffix
-
-    _replay_runtime_env = {
-        "py_executable": _replay_py_exec,
-        "env_vars": {
-            **os.environ,
-            "VIRTUAL_ENV": _replay_py_venv,
-            "UV_PROJECT_ENVIRONMENT": _replay_py_venv,
-        },
-    }
 
     # Calculate optimal buffer size based on generation limits to prevent length bias
     # Each weight version generates exactly num_prompts_per_step trajectories
@@ -4681,29 +4661,13 @@ def async_grpo_train(
         set(trained_task_indices) if frontier_restore else set()
     )
 
-    _tc_py_exec = get_actor_python_env(
-        "nemo_rl.algorithms.async_utils.AsyncTrajectoryCollector"
-    )
-    if _tc_py_exec.startswith("uv"):
-        _tc_py_exec = create_local_venv_on_each_node(
-            _tc_py_exec,
-            "nemo_rl.algorithms.async_utils.AsyncTrajectoryCollector",
-        )
-
-    _tc_py_venv = os.path.dirname(
-        os.path.dirname(_tc_py_exec)
-    )  # to remove the "bin/python" suffix
-
-    _tc_runtime_env = {
-        "py_executable": _tc_py_exec,
-        "env_vars": {
-            **os.environ,
-            "VIRTUAL_ENV": _tc_py_venv,
-            "UV_PROJECT_ENVIRONMENT": _tc_py_venv,
+    _tc_runtime_env = make_actor_runtime_env(
+        "nemo_rl.algorithms.async_utils.AsyncTrajectoryCollector",
+        extra_env_vars={
             # Names this actor's spans the way RayWorkerGroup names its groups'.
             "NRL_WORKER_GROUP": "trajectory_collector",
         },
-    }
+    )
 
     # Captured inside rl.grpo.job, so the collector's spans join this run's
     # trace instead of starting their own roots. Empty unless the job group is

nemo_rl/data/collate_fn.py

@@ -45,11 +45,10 @@ def rl_collate_fn(data_batch: list[DatumSpec]) -> BatchedDataDict[Any]:
     # Extract stop_strings if present
     stop_strings = [datum.get("stop_strings", None) for datum in data_batch]
 
-    # check if any of the data batch has vllm content and images
+    # Presence of the key selects vLLM's native-media path. Placeholder-style
+    # processors intentionally set the content to None so vLLM uses input_ids.
     extra_args = {}
-    if any(
-        [datum_spec.get("vllm_content", None) is not None for datum_spec in data_batch]
-    ):
+    if any("vllm_content" in datum_spec for datum_spec in data_batch):
         vllm_content = [
             datum_spec.get("vllm_content", None) for datum_spec in data_batch
         ]
@@ -119,11 +118,10 @@ def eval_collate_fn(data_batch: list[DatumSpec]) -> BatchedDataDict[Any]:
     idx = [datum_spec["idx"] for datum_spec in data_batch]
     task_names = [datum_spec.get("task_name", None) for datum_spec in data_batch]
 
-    # Check if any of the data batch has vllm content (multimodal data)
+    # Preserve native media when placeholder-style processors intentionally
+    # set vllm_content to None in favor of their expanded input_ids.
     extra_args = {}
-    if any(
-        datum_spec.get("vllm_content", None) is not None for datum_spec in data_batch
-    ):
+    if any("vllm_content" in datum_spec for datum_spec in data_batch):
         extra_args["vllm_content"] = [
             datum_spec.get("vllm_content", None) for datum_spec in data_batch
         ]