feat(mllm-v2): Support single-controller for Megatron-Inference.

docs/guides/single-controller.md

@@ -234,6 +234,9 @@ Do not carry `max_num_epochs: -1` across either. [ppo.md](./ppo.md#asynchronous-
 
 The SC path is still under active development. Feature gaps are tracked in [issue #2625](https://github.com/NVIDIA-NeMo/RL/issues/2625). Notable items:
 
+- Multimodal/VLM GRPO is supported with Megatron generation. Set
+  `policy.is_vlm: true`; see the
+  [CLEVR Single-Controller recipe](../../examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-8n4g-megatron-single-controller-async.v1.yaml).
 - Multi-Teacher On-Policy Distillation (MOPD) is supported for text-only NeMo
   Gym rollouts; multimodal/VLM MOPD is not yet supported. See
   [Multi-Teacher On-Policy Distillation](../about/algorithms/mopd.md#running-mopd).

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

@@ -49,6 +49,14 @@ policy:
   generation:
     bad_words: []
     mcore_generation_config:
+      buffer_size_gb: 8
+      use_cuda_graphs_for_non_decode_steps: false
+      moe_pad_experts_for_cuda_graph_inference: false
+      moe_router_dtype: fp32
+      vision_embedding_cache_max_bytes: 536870912
+      video_num_frames: ${data.default.num_frames}
+      video_temporal_patch_size: ${data.default.video_temporal_patch_size}
+      video_target_num_patches: ${data.default.video_target_num_patches}
       image_dynamic_resolution: true
       logprobs_mode: raw_logprobs
       megatron_inference_wrapper: megatron.core.inference.model_inference_wrappers.multimodal.nemotron_omni_inference_wrapper.NemotronOmniInferenceWrapper
@@ -91,6 +99,9 @@ data:
   max_input_seq_length: ${policy.max_total_sequence_length}
   num_workers: 0
   default:
+    num_frames: 32
+    video_sampling_style: nemotron_vl
+    video_temporal_patch_size: 2
     video_target_num_patches: 1024
     video_maintain_aspect_ratio: true
 env:

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

@@ -0,0 +1,18 @@
+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
+  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}

examples/run_grpo_single_controller.py

@@ -126,7 +126,12 @@ def main() -> None:
     maybe_configure_data_plane_env(config.data_plane)
     init_ray()
 
-    tokenizer = get_tokenizer(config.policy["tokenizer"])
+    processor = None
+    if config.policy.get("is_vlm"):
+        processor = get_tokenizer(config.policy["tokenizer"], get_processor=True)
+        tokenizer = processor.tokenizer
+    else:
+        tokenizer = get_tokenizer(config.policy["tokenizer"])
     assert config.policy["generation"] is not None, (
         "A generation config is required for SC-driven async GRPO"
     )
@@ -144,7 +149,9 @@ def main() -> None:
     if bool(config.env.get("should_use_nemo_gym")):
         setup_nemo_gym_config(config, tokenizer)
 
-    actor_args, setup_timing_metrics = setup_single_controller(config, tokenizer)
+    actor_args, setup_timing_metrics = setup_single_controller(
+        config, tokenizer, processor=processor
+    )
 
     print("🚀 Launching SingleControllerActor")
     sc = SingleControllerActor.remote(

nemo_rl/algorithms/single_controller_utils/setup.py

@@ -70,6 +70,7 @@
 )
 from nemo_rl.algorithms.utils import set_seed
 from nemo_rl.data.collate_fn import rl_collate_fn
+from nemo_rl.data.multimodal_utils import WIRE_MULTIMODAL_FIELDS
 from nemo_rl.data.utils import load_dataloader_state, setup_response_data
 from nemo_rl.data_plane import (
     DATA_PLANE_CHECKPOINT_SCHEMA_VERSION,
@@ -997,9 +998,10 @@ def setup_single_controller(
         policy_config["pretrained_checkpoint"] = checkpointing_pretrained
 
     # Token capture: validate the supported combination loudly at setup
-    # (NeMo-Gym rollout path, vLLM backend, async_engine=true) and give
-    # capture-enabled vLLM workers a venv that carries nemo_gym (the
-    # worker hosts Gym's capture core + adapter in-process).
+    # (NeMo-Gym rollout path, vLLM backend, async_engine=true). The vLLM
+    # worker venv always carries nemo_gym (see VLLM_EXECUTABLE in
+    # ray_actor_environment_registry.py), so nothing here needs to change the
+    # worker's environment.
     token_capture_cfg = master_config.token_capture
     if token_capture_cfg.enabled:
         if not should_use_nemo_gym(master_config):
@@ -1020,14 +1022,6 @@ def setup_single_controller(
                 "policy.generation.vllm_cfg.async_engine=true (the capture "
                 "host is the worker's in-process HTTP server)"
             )
-        from nemo_rl.distributed.ray_actor_environment_registry import (
-            ACTOR_ENVIRONMENT_REGISTRY,
-        )
-        from nemo_rl.distributed.virtual_cluster import PY_EXECUTABLES
-
-        ACTOR_ENVIRONMENT_REGISTRY[
-            "nemo_rl.models.generation.vllm.vllm_worker_async.VllmAsyncGenerationWorker"
-        ] = PY_EXECUTABLES.VLLM_GYM
 
         # Fill the derived ledger-hosting fields (see TokenCaptureConfig): a
         # per-run control-plane bearer token and the process-shared capture
@@ -1074,6 +1068,8 @@ def setup_single_controller(
     # ==========================
     # TODO: add validate dataset wiring.
     use_nemo_gym = should_use_nemo_gym(master_config)
+    data_tokenizer = processor if processor is not None else tokenizer
+    is_vlm = processor is not None
     if use_nemo_gym and generation_config["backend"] not in ("vllm", "megatron"):
         raise NotImplementedError(
             "SC NeMo-Gym integration currently supports the vllm and megatron backends only; got "
@@ -1084,13 +1080,18 @@ def setup_single_controller(
     if use_nemo_gym:
         # NeMo-Gym creates the env actor outside setup_response_data; we wire
         # it in after generation is up (it needs the OpenAI server URLs).
-        response_data = setup_response_data(tokenizer, data_config, env_configs=None)
+        response_data = setup_response_data(
+            data_tokenizer, data_config, env_configs=None, is_vlm=is_vlm
+        )
         assert len(response_data) == 2
         dataset, _val_dataset = response_data
         env_handles: dict[str, EnvironmentInterface] = {}
     else:
         response_data = setup_response_data(
-            tokenizer, data_config, env_configs=master_config.env
+            data_tokenizer,
+            data_config,
+            env_configs=master_config.env,
+            is_vlm=is_vlm,
         )
         assert len(response_data) == 4
         dataset, _val_dataset, env_handles, _val_env_handles = response_data
@@ -1154,6 +1155,7 @@ def setup_single_controller(
     megatron_reserved_url = None
     megatron_port_holder = None
     reserved_http_server_port = None
+    weight_synchronizer: Optional[WeightSynchronizer] = None
     if megatron_backend:
         generation_config["model_name"] = master_config.policy["model_name"]
 
@@ -1313,7 +1315,6 @@ def _build_generation_then_trainer(
         build_tasks["trainer"] = _build_trainer_and_value
 
     # Submit build tasks and get results
-    weight_synchronizer: Optional[WeightSynchronizer] = None
     try:
         with ThreadPoolExecutor(max_workers=len(build_tasks)) as executor:
             submitted = {k: executor.submit(fn) for k, fn in build_tasks.items()}
@@ -1340,7 +1341,9 @@ def _build_generation_then_trainer(
                     train_cluster=train_cluster,
                     inference_cluster=inference_cluster,
                     refit_buffer_size_gb=policy_config.get("refit_buffer_size_gb"),
+                    refit_timeout_s=master_config.async_rl.generation_fleet_health.refit_timeout_s,
                 )
+                generation.weight_synchronizer = weight_synchronizer
                 weight_synchronizer.init_communicator()
                 setup_timing_metrics.collective_init_time_s = time.perf_counter() - t0
                 t0 = time.perf_counter()
@@ -1427,12 +1430,19 @@ def _build_generation_then_trainer(
         # SingleController reuses one partition for the run. Warm every known
         # tensor field before rollout, policy, and teacher writers become
         # concurrent; TransferQueue otherwise registers field names lazily.
+        partition_fields = fields_with_optional_routed_experts(
+            SC_ROLLOUT_SCHEMA_FIELDS,
+            enabled=router_replay_enabled(policy_config),
+        )
+        if processor is not None:
+            partition_fields.extend(
+                field
+                for field in sorted(WIRE_MULTIMODAL_FIELDS)
+                if field not in partition_fields
+            )
         dp_client.register_partition(
             partition_id=partition_id,
-            fields=fields_with_optional_routed_experts(
-                SC_ROLLOUT_SCHEMA_FIELDS,
-                enabled=router_replay_enabled(policy_config),
-            ),
+            fields=partition_fields,
             num_samples=(
                 master_config.async_rl.max_buffered_rollouts
                 * algo_cfg.num_generations_per_prompt
@@ -1457,13 +1467,19 @@ def _build_generation_then_trainer(
             )
         group_size = algo_cfg.num_generations_per_prompt
         num_rollout_samples = master_config.async_rl.max_buffered_rollouts * group_size
+        partition_fields = fields_with_optional_routed_experts(
+            DP_TRAIN_FIELDS,
+            enabled=r3_enabled and not token_capture_cfg.defer_routed_experts_to_policy,
+        )
+        if processor is not None:
+            partition_fields.extend(
+                field
+                for field in sorted(WIRE_MULTIMODAL_FIELDS)
+                if field not in partition_fields
+            )
         dp_client.register_partition(
             partition_id=partition_id,
-            fields=fields_with_optional_routed_experts(
-                DP_TRAIN_FIELDS,
-                enabled=r3_enabled
-                and not token_capture_cfg.defer_routed_experts_to_policy,
-            ),
+            fields=partition_fields,
             num_samples=num_rollout_samples,
             consumer_tasks=["prev_lp", "ref_lp", "train"],
             grpo_group_size=group_size,
@@ -1478,23 +1494,7 @@ def _build_generation_then_trainer(
         # Host Gym's capture core in every vLLM DP leader (in-worker DP
         # client + TQTokenSink + the single install_capture call), and give
         # workers the initial weight version to stamp on captured calls.
-        try:
-            generation.setup_token_capture(
-                dp_config, token_capture_cfg.staging_partition
-            )
-        except Exception as error:
-            if "No module named 'nemo_gym'" in str(error):
-                # Worker venvs are cached by actor class name
-                # (nemo_rl/utils/venvs.py), so a venv prebuilt before token
-                # capture predates the nemo_gym extra and is reused as-is.
-                raise RuntimeError(
-                    "token_capture.enabled requires nemo_gym inside the vLLM "
-                    "worker venv, but the cached worker venv predates it. "
-                    "Rebuild worker venvs (NRL_FORCE_REBUILD_VENVS=true) or "
-                    "delete $NEMO_RL_VENV_DIR/nemo_rl.models.generation.vllm."
-                    "vllm_worker_async.VllmAsyncGenerationWorker and rerun."
-                ) from error
-            raise
+        generation.setup_token_capture(dp_config, token_capture_cfg.staging_partition)
         generation.set_rollout_weight_version(0)
 
     if weight_synchronizer is None:
@@ -1509,6 +1509,7 @@ def _build_generation_then_trainer(
             refit_buffer_size_gb=policy_config.get("refit_buffer_size_gb"),
             refit_timeout_s=master_config.async_rl.generation_fleet_health.refit_timeout_s,
         )
+        generation.weight_synchronizer = weight_synchronizer
         weight_synchronizer.init_communicator()
         setup_timing_metrics.collective_init_time_s = time.perf_counter() - t0
 

nemo_rl/data/multimodal_utils.py

@@ -1093,6 +1093,18 @@ def encode_multimodal_for_wire(
         )
 
 
+# Model inputs some remote-code processors omit from ``model_input_names`` even
+# though their forward requires them. Consumed by
+# ``extract_multimodal_model_inputs``; membership here does NOT imply the field
+# is wire-registered (see ``PACKED_/PER_TOKEN_MULTIMODAL_FIELDS``).
+UNDECLARED_MULTIMODAL_MODEL_INPUTS = (
+    "imgs_sizes",
+    "num_frames",
+    "pixel_values_flat",
+    "image_num_patches",
+)
+
+
 def get_multimodal_keys_from_processor(processor) -> list[str]:
     """Get keys of the multimodal data that can be used as model inputs.
 
@@ -1216,12 +1228,7 @@ def extract_multimodal_model_inputs(
     # TODO(rohitrango): Let ProcessorInterface declare model-specific media inputs.
     # Some remote-code processors omit these inputs from model_input_names even
     # though their model forward requires them.
-    for key in (
-        "imgs_sizes",
-        "num_frames",
-        "pixel_values_flat",
-        "image_num_patches",
-    ):
+    for key in UNDECLARED_MULTIMODAL_MODEL_INPUTS:
         if key in processed and key not in multimodal_keys:
             multimodal_keys.append(key)
     for key in multimodal_keys:
@@ -1233,7 +1240,7 @@ def extract_multimodal_model_inputs(
                 f"Processor model input {key!r} must be a torch.Tensor, got "
                 f"{type(value).__name__}."
             )
-        if key == "imgs_sizes":
+        if key in ("imgs_sizes", "num_frames"):
             value = value.to(dtype=torch.int32)
         extracted[key] = PackedTensor(
             value,

nemo_rl/data/processors.py

@@ -460,9 +460,8 @@ def vlm_hf_data_processor(
     from nemo_rl.data.datasets.response_datasets.refcoco import format_refcoco_dataset
     from nemo_rl.data.multimodal_utils import (
         PackedTensor,
-        get_dim_to_pack_along,
+        extract_multimodal_model_inputs,
         get_multimodal_default_settings_from_processor,
-        get_multimodal_keys_from_processor,
         resolve_to_image,
         uses_image_placeholder,
     )
@@ -608,50 +607,13 @@ def vlm_hf_data_processor(
 
     # add this for backward compatibility
     user_message["token_ids"] = message["input_ids"][0]
-    # add all keys and values to the user message, and the list of keys
-    multimodal_keys = list(get_multimodal_keys_from_processor(processor))
-    # Current Nemotron Omni processors emit imgs_sizes. Historical MMPR
-    # checkpoints instead emit a batch of fixed-size image tiles and only
-    # declare pixel_values. Treat each tile as one dynamic-resolution image so
-    # the Nemotron Omni path can patchify it and preserve the processor's exact
-    # placeholder count.
-    if (
-        uses_placeholder
-        and "pixel_values" in message
-        and "imgs_sizes" not in message
-        and message["pixel_values"].ndim == 4
-    ):
-        pixel_values = message["pixel_values"]
-        num_tiles, _, height, width = pixel_values.shape
-        message["imgs_sizes"] = torch.tensor(
-            [[height, width]] * num_tiles, dtype=torch.long
-        )
-
-    # imgs_sizes is not always declared in model_input_names by bundled image
-    # processors, so append it explicitly when present. RADIO uses temporal
-    # patching even for still images and requires one num_frames=1 entry per
-    # image/tile.
-    if "imgs_sizes" in message and "imgs_sizes" not in multimodal_keys:
-        multimodal_keys.append("imgs_sizes")
-    if "imgs_sizes" in message and "num_frames" not in message:
-        message["num_frames"] = torch.ones(len(message["imgs_sizes"]), dtype=torch.long)
-    if "num_frames" in message and "num_frames" not in multimodal_keys:
-        multimodal_keys.append("num_frames")
-    for key in multimodal_keys:
-        if key in message:
-            user_message[key] = PackedTensor(
-                message[key],
-                dim_to_pack=get_dim_to_pack_along(processor, key),
-                pad_to_max_shape=uses_placeholder and key == "pixel_values",
-            )
-
-    # specifically for gemma, we need to add token_type_ids to the user message as a sequence-type value
-    if "token_type_ids" in message:
-        user_message["token_type_ids"] = message["token_type_ids"][0]
-
-    # for qwen2.5-vl (transformers>=5.3), mm_token_type_ids tells the model which tokens are text/image/video for 3D RoPE
-    if "mm_token_type_ids" in message:
-        user_message["mm_token_type_ids"] = message["mm_token_type_ids"][0]
+    # Single source of truth for media extraction: the MMPR imgs_sizes
+    # fallback, RADIO num_frames synthesis, PackedTensor wrapping (incl. the
+    # imgs_sizes int32 cast) and the gemma / qwen2.5-vl sequence-type maps all
+    # live in ``extract_multimodal_model_inputs``, which the NeMo-Gym path also
+    # uses. One implementation is what stops the two paths from handing the
+    # same model differently-typed inputs.
+    user_message.update(extract_multimodal_model_inputs(processor, message))
 
     ### append to user message
     message_log.append(user_message)