feat(minf-mm-sc/tq): Support SC/TQ multimodal 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-1n4g-megatron-single-controller-async.v1.yaml

@@ -0,0 +1,51 @@
+# NeMo-RL v2 SingleController overlay for one-node, non-colocated Omni CLEVR.
+defaults: ./vlm_grpo-nemotron-omni-30ba3b-clevr-1n4g-megatron_generation.v1.yaml
+
+grpo:
+  async_grpo: null
+  val_period: 0
+  val_at_start: false
+  val_at_end: false
+
+data_plane:
+  enabled: true
+  impl: transfer_queue
+  backend: simple
+  claim_meta_poll_interval_s: 0.5
+  simple:
+    num_storage_units: 2
+
+async_rl:
+  sampler:
+    name: in_order
+    max_lookahead_versions: 1
+  recompute_kv_cache_after_weight_updates: false
+  min_groups_for_streaming_train: ${grpo.num_prompts_per_step}
+  max_inflight_prompts: ${mul:${grpo.num_prompts_per_step}, 2}
+  max_buffered_rollouts: ${mul:${grpo.num_prompts_per_step}, 2}
+  diagnostics: true
+
+policy:
+  generation:
+    backend: megatron
+    colocated:
+      enabled: false
+      resources:
+        num_nodes: 1
+        gpus_per_node: 2
+    mcore_generation_config:
+      transformer_impl: inference_optimized
+      moe_router_dtype: fp32
+      moe_pad_experts_for_cuda_graph_inference: false
+      cuda_graph_impl: local
+      inference_cuda_graph_scope: block
+      num_cuda_graphs: -1
+      use_cuda_graphs_for_non_decode_steps: false
+      enable_chunked_prefill: true
+      async_sched_mode: async
+      kv_cache_management_mode: persist
+      refit_backend: nccl
+
+cluster:
+  num_nodes: 1
+  gpus_per_node: 4

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

@@ -0,0 +1,72 @@
+defaults: ../../vlm_grpo_3B_megatron.yaml
+grpo:
+  num_prompts_per_step: 2
+  num_generations_per_prompt: 8
+  max_num_steps: 4
+  val_period: 0
+  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-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
+  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
+    max_new_tokens: 512
+    stop_strings:
+    - <image>
+    - <img>
+    - </img>
+    - <so_embedding>
+    - <so_start>
+    - <so_end>
+    mcore_generation_config:
+      buffer_size_gb: 8
+      num_cuda_graphs: -1
+      max_tokens: ${policy.max_total_sequence_length}
+      refit_backend: gloo
+      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
+    colocated:
+      enabled: false
+      resources:
+        gpus_per_node: 2
+        num_nodes: 1
+data:
+  num_workers: 0
+  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-1n4g
+cluster:
+  gpus_per_node: 4

examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron-single-controller-async.v1.yaml

@@ -0,0 +1,51 @@
+# NeMo-RL v2 SingleController overlay for eight-node, non-colocated Omni CLEVR.
+defaults: ./vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron_generation.v1.yaml
+
+grpo:
+  async_grpo: null
+  val_period: 0
+  val_at_start: false
+  val_at_end: false
+
+data_plane:
+  enabled: true
+  impl: transfer_queue
+  backend: simple
+  claim_meta_poll_interval_s: 0.5
+  simple:
+    num_storage_units: 16
+
+async_rl:
+  sampler:
+    name: in_order
+    max_lookahead_versions: 1
+  recompute_kv_cache_after_weight_updates: false
+  min_groups_for_streaming_train: ${grpo.num_prompts_per_step}
+  max_inflight_prompts: ${mul:${grpo.num_prompts_per_step}, 2}
+  max_buffered_rollouts: ${mul:${grpo.num_prompts_per_step}, 2}
+  diagnostics: false
+
+policy:
+  generation:
+    backend: megatron
+    colocated:
+      enabled: false
+      resources:
+        num_nodes: 6
+        gpus_per_node: 4
+    mcore_generation_config:
+      transformer_impl: inference_optimized
+      moe_router_dtype: fp32
+      moe_pad_experts_for_cuda_graph_inference: false
+      cuda_graph_impl: local
+      inference_cuda_graph_scope: block
+      num_cuda_graphs: -1
+      use_cuda_graphs_for_non_decode_steps: false
+      enable_chunked_prefill: true
+      async_sched_mode: async
+      kv_cache_management_mode: persist
+      refit_backend: nccl
+
+cluster:
+  num_nodes: 8
+  gpus_per_node: 4