feat: Multimodal nemo gym compatible grpo pipeline

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

@@ -0,0 +1,160 @@
+defaults: ../../vlm_grpo_3B_megatron.yaml
+grpo:
+  num_prompts_per_step: 1
+  num_val_generations_per_prompt: 1
+  max_num_steps: 100
+  val_period: 500
+  overlong_filtering: true
+  max_val_samples: null
+  val_batch_size: null
+  seq_logprob_error_threshold: 2
+loss_fn:
+  reference_policy_kl_penalty: 0
+  kl_input_clamp_value: null
+  kl_output_clamp_value: null
+  ratio_clip_max: 0.28
+  use_on_policy_kl_approximation: true
+  use_importance_sampling_correction: true
+checkpointing:
+  enabled: false
+  checkpoint_dir: results/grpo-nemotron-omni-30ba3b-gymv-circle-click
+  metric_name: val:total_reward/mean
+  keep_top_k: 1000000
+  checkpoint_must_save_by: 00:03:40:00
+policy:
+  model_name: nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16
+  is_vlm: true
+  tokenizer:
+    chat_template_kwargs:
+      enable_thinking: true
+      truncate_history_thinking: false
+  train_global_batch_size: ${mul:${grpo.num_prompts_per_step}, ${grpo.num_generations_per_prompt}}
+  generation_batch_size: 64
+  logprob_batch_size: 1
+  max_total_sequence_length: 8192
+  logprob_chunk_size: 2048
+  megatron_cfg:
+    activation_checkpointing: true
+    bias_activation_fusion: false
+    tensor_model_parallel_size: 2
+    expert_model_parallel_size: 8
+    context_parallel_size: 2
+    sequence_parallel: true
+    moe_router_dtype: fp32
+    moe_router_bias_update_rate: 0.001
+    moe_aux_loss_coeff: 0.0
+    moe_router_enable_expert_bias: true
+    defer_fp32_logits: true
+    track_moe_metrics: true
+    moe_per_layer_logging: true
+    radio_force_cpe_eval_mode: true
+    clear_memory_caches_before_refit: true
+    optimizer:
+      lr: 3.0e-06
+      min_lr: 3.0e-06
+      weight_decay: 0.0
+      optimizer_cpu_offload: true
+      optimizer_offload_fraction: 1.0
+    scheduler:
+      lr_decay_iters: null
+      lr_warmup_iters: 0
+      lr_warmup_init: 3.0e-07
+    distributed_data_parallel_config:
+      overlap_param_gather: false
+      average_in_collective: false
+  sequence_packing:
+    enabled: true
+  make_sequence_length_divisible_by: 32
+  optimizer: null
+  scheduler: null
+  generation:
+    max_new_tokens: ${policy.max_total_sequence_length}
+    bad_words: []
+    mcore_generation_config:
+      transformer_impl: inference_optimized
+      activation_checkpointing: false
+      mamba_inference_ssm_states_dtype: float32
+      inference_moe_token_dispatcher_type: nccl
+      inference_grouped_gemm_backend: vllm
+      moe_router_num_groups: null
+      moe_router_group_topk: null
+      pipeline_model_parallel_size: 1
+      expert_tensor_parallel_size: 1
+      expert_model_parallel_size: 8
+      sequence_parallel: true
+      context_parallel_size: 1
+      tensor_model_parallel_size: 2
+      buffer_size_gb: 20
+      num_cuda_graphs: -1
+      max_tokens: ${policy.max_total_sequence_length}
+      async_engine: true
+      expose_http_server: true
+      enable_prefix_caching: true
+      parsers:
+      - deepseek-r1-reasoning
+      - qwen3-coder-tool
+    vllm_cfg:
+      async_engine: true
+      tensor_parallel_size: 8
+      enforce_eager: true
+      enable_prefix_caching: false
+      expose_http_server: true
+      reasoning_parser_plugin: nemo_rl/models/generation/vllm/reasoning_parsers/nano_v3_reasoning_parser.py
+      http_server_serving_chat_kwargs:
+        enable_auto_tools: true
+        tool_parser: qwen3_coder
+        reasoning_parser: nano_v3
+        chat_template_content_format: string
+        default_chat_template_kwargs:
+          enable_thinking: true
+          truncate_history_thinking: false
+    vllm_kwargs:
+      limit_mm_per_prompt:
+        image: 1
+      max_num_batched_tokens: 16384
+      mamba_ssm_cache_dtype: float32
+      compilation_config:
+        backend: eager
+    colocated:
+      enabled: false
+      resources:
+        gpus_per_node: 8
+        num_nodes: 1
+data:
+  max_input_seq_length: null
+  shuffle: false
+  train:
+    data_path: /path/to/train_dataset.jsonl
+  validation:
+    data_path: /path/to/eval_dataset.jsonl
+  default:
+    dataset_name: NemoGymDataset
+    env_name: nemo_gym
+    prompt_file: null
+    processor: nemo_gym_data_processor
+env:
+  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_click/configs/circle_click.yaml
+    circle_click_simple_agent:
+      responses_api_agents:
+        simple_agent:
+          max_steps: 1
+logger:
+  log_dir: logs/grpo-nemotron-omni-30ba3b-gymv-circle-click
+  wandb_enabled: true
+  wandb:
+    project: grpo-nemotron-omni-gymv
+    name: grpo-nemotron-omni-30ba3b-gymv-circle-click
+  mlflow:
+    experiment_name: grpo-nemotron-omni-gymv
+    run_name: grpo-nemotron-omni-30ba3b-gymv-circle-click
+cluster:
+  gpus_per_node: 8
+  num_nodes: 2

examples/nemo_gym/grpo_nanov3.yaml

@@ -240,7 +240,7 @@ policy:
       block_size_tokens: 256  # Size of each KV cache block in tokens (affects memory granularity)
       use_cuda_graphs_for_non_decode_steps: true  # Enable CUDA graphs for prefill/context processing
       enable_chunked_prefill: true
-      max_tokens: 16384 # Maximum number of tokens to use in a single step. Analogous to vllm's max_num_batched_tokens
+      max_tokens: ${policy.max_total_sequence_length} # Maximum number of tokens to use in a single step. Analogous to vllm's max_num_batched_tokens
       kv_cache_management_mode: "persist"  # KV cache lifecycle across suspend/resume. Options: "persist", "offload". To select "recompute", set grpo.async_grpo.recompute_kv_cache_after_weight_updates=true.
       materialize_only_last_token_logits: true
       num_speculative_tokens: 0

examples/nemo_gym/run_grpo_nemo_gym.py

@@ -152,8 +152,14 @@ def main() -> None:
         )
 
     with rl_init_timer.time("tokenizer"):
-        # setup tokenizer
-        tokenizer = get_tokenizer(config.policy["tokenizer"])
+        is_vlm = bool(config.policy.get("is_vlm"))
+        if is_vlm:
+            processor = get_tokenizer(config.policy["tokenizer"], get_processor=True)
+            tokenizer = processor.tokenizer
+        else:
+            processor = None
+            tokenizer = get_tokenizer(config.policy["tokenizer"])
+
         assert config.policy["generation"] is not None, (
             "A generation config is required for GRPO"
         )
@@ -171,6 +177,11 @@ def main() -> None:
             has_refit_draft_weights=has_refit_draft_weights,
             trains_mtp=trains_mtp,
         )
+        if is_vlm and "vllm_cfg" in config.policy["generation"]:
+            assert not config.policy["generation"]["vllm_cfg"]["skip_tokenizer_init"], (
+                "VLMs require tokenizer to be initialized before generation, "
+                "so skip_tokenizer_init must be set to False."
+            )
 
         # NeMo-Gym specific config setup.
         setup_nemo_gym_config(config, tokenizer)
@@ -181,8 +192,9 @@ def main() -> None:
     # NeMo-Gym environment needs to get dp_openai_server_base_urls from policy_generation, so we don't setup env here.
     with rl_init_timer.time("data"):
         print("\n▶ Setting up data...")
+        data_tokenizer = processor if processor is not None else tokenizer
         train_dataset, val_dataset = setup_response_data(
-            tokenizer, config.data, env_configs=None
+            data_tokenizer, config.data, env_configs=None
         )
 
     # Validation dataset config setup.
@@ -231,7 +243,13 @@ def main() -> None:
             master_config,
             teacher_worker_groups,
             alias_to_group_alias,
-        ) = setup(config, tokenizer, train_dataset, val_dataset)
+        ) = setup(
+            config,
+            tokenizer,
+            train_dataset,
+            val_dataset,
+            processor=processor,
+        )
 
     rl_init_timer.record("total", time.perf_counter() - main_start)
     rl_init_metrics = rl_init_timer.get_timing_metrics(reduction_op="sum")

nemo_rl/algorithms/grpo.py

@@ -2042,6 +2042,9 @@ def _build_async_grpo_train_data(
         }
     )
     _preserve_router_replay_routed_experts(train_data, flat_messages, policy_config)
+    # update multimodal data unconditionally
+    extra_multimodal_data = flat_messages.get_multimodal_dict(as_tensors=False)
+    train_data.update(extra_multimodal_data)
     return train_data
 
 
@@ -3988,6 +3991,7 @@ def async_grpo_train(
     assert master_config.loss_fn.use_importance_sampling_correction, (
         "Importance sampling correction must be enabled for async GRPO for good convergence due to off-policy samples!"
     )
+
     if router_replay_enabled(master_config.policy) and (
         master_config.data_plane or {}
     ).get("enabled", False):

nemo_rl/data/llm_message_utils.py

@@ -363,9 +363,26 @@ def batched_message_log_to_flat_message(
     result = BatchedDataDict()
     for key in all_keys:
         values = [seq.get(key) for seq in sequenced_lists]
-        # if the values are PackedTensors, create a new PackedTensor from the list of values
-        if values and isinstance(values[0], PackedTensor):
-            result[key] = PackedTensor.flattened_concat(values)
+        packed_template = next(
+            (value for value in values if isinstance(value, PackedTensor)), None
+        )
+        if packed_template is not None:
+            if any(
+                value is not None and not isinstance(value, PackedTensor)
+                for value in values
+            ):
+                raise TypeError(
+                    f"Expected PackedTensor or None for {key=}, "
+                    f"got {[type(value).__name__ for value in values]}"
+                )
+            filled_packed_values = cast(
+                list[PackedTensor],
+                [
+                    PackedTensor.empty_like(packed_template) if value is None else value
+                    for value in values
+                ],
+            )
+            result[key] = PackedTensor.flattened_concat(filled_packed_values)
             continue
         if not values or not isinstance(values[0], Tensor):
             result[key] = values

nemo_rl/data/multimodal_utils.py

@@ -45,6 +45,13 @@
     "audio": ["wav", "flac", "mp3"],
 }
 
+_PLACEHOLDER_STYLE_PROCESSOR_NAMES = frozenset(
+    {
+        "NemotronNanoVLV2Processor",
+        "NemotronH_Nano_Omni_Reasoning_V3Processor",
+    }
+)
+
 
 # different media namings maybe used in the raw dataset,
 # in which case, they need to be mapped to the allowed ones
@@ -67,6 +74,19 @@
 logger = logging.getLogger(__name__)
 
 
+def uses_image_placeholder(processor: Any) -> bool:
+    """Return whether a processor requires explicit image placeholders.
+
+    Args:
+        processor: Multimodal processor to classify.
+
+    Returns:
+        Whether the processor expands image placeholders through ``__call__``
+        rather than tokenized ``apply_chat_template``.
+    """
+    return type(processor).__name__ in _PLACEHOLDER_STYLE_PROCESSOR_NAMES
+
+
 class PackedTensor:
     """Wrapper around a list of torch tensors and a dimension along which to pack the tensors.
 
@@ -376,11 +396,79 @@ def resolve_to_image(image_path_or_image: str | Image.Image) -> Image.Image:
         header, encoded = image_path_or_image.split(",", 1)
         image_data = base64.b64decode(encoded)
         return Image.open(BytesIO(image_data)).convert("RGB")
+    elif image_path_or_image.startswith("file://"):
+        return Image.open(image_path_or_image.removeprefix("file://")).convert("RGB")
     else:
         # Handle local file path
         return Image.open(image_path_or_image).convert("RGB")
 
 
+def image_to_data_url(image: Image.Image, fmt: str = "PNG") -> str:
+    """Encode a PIL Image as a base64 ``data:`` URL.
+
+    Args:
+        image: PIL image to encode.
+        fmt: PIL image format used for serialization (e.g. ``"PNG"``, ``"JPEG"``).
+            The value is also lowercased and embedded in the MIME type of the
+            returned URL.
+
+    Returns:
+        A ``data:image/<fmt>;base64,<payload>`` URL suitable for embedding in
+        an OpenAI Responses ``input_image`` content part.
+    """
+    buf = BytesIO()
+    image.save(buf, format=fmt)
+    encoded = base64.b64encode(buf.getvalue()).decode("utf-8")
+    return f"data:image/{fmt.lower()};base64,{encoded}"
+
+
+def encode_images_in_examples(nemo_gym_examples: list[dict]) -> list[dict]:
+    """Replace local image paths in NeMo Gym examples with base64 data URLs.
+
+    Walks each example's ``responses_create_params.input[].content[]`` items
+    and rewrites any ``input_image`` part whose ``image_url`` is a local path
+    (or ``file://`` URL) into a base64 ``data:`` URL via
+    :func:`image_to_data_url`. Parts whose URL already starts with ``http://``,
+    ``https://``, or ``data:`` are left untouched. Malformed items (non-dict
+    entries, missing/empty URLs, non-list ``input``/``content``) are skipped
+    without raising.
+
+    The examples are mutated in place; the same list is also returned for
+    convenience so callers can chain the call.
+
+    Args:
+        nemo_gym_examples: List of NeMo Gym example dicts. Each example is
+            expected to contain a ``responses_create_params`` mapping with an
+            ``input`` list of Responses API messages.
+
+    Returns:
+        The same ``nemo_gym_examples`` list, with local image references
+        rewritten to base64 data URLs in place.
+    """
+    for example in nemo_gym_examples:
+        input_items = example.get("responses_create_params", {}).get("input", [])
+        if not isinstance(input_items, list):
+            continue
+        for item in input_items:
+            if not isinstance(item, dict):
+                continue
+            content = item.get("content", [])
+            if not isinstance(content, list):
+                continue
+            for part in content:
+                if not isinstance(part, dict) or part.get("type") != "input_image":
+                    continue
+                url = part.get("image_url", "")
+                if isinstance(url, dict):
+                    url = url.get("url", "")
+                if not isinstance(url, str) or not url:
+                    continue
+                if url.startswith(("http://", "https://", "data:")):
+                    continue
+                part["image_url"] = image_to_data_url(resolve_to_image(url))
+    return nemo_gym_examples
+
+
 def get_media_from_message(message: dict[str, Any]) -> dict[str, list[Any]]:
     """Get all media from a message log item."""
     # Handle None or missing content (e.g., assistant messages with only tool_calls)