fix(vlm): dispatch async GRPO from VLM entrypoint

examples/run_vlm_grpo.py

@@ -19,7 +19,7 @@
 
 from omegaconf import OmegaConf
 
-from nemo_rl.algorithms.grpo import MasterConfig, grpo_train, setup
+from nemo_rl.algorithms.grpo import MasterConfig, async_grpo_train, grpo_train, setup
 from nemo_rl.algorithms.utils import get_tokenizer
 from nemo_rl.data.utils import setup_response_data
 from nemo_rl.distributed.virtual_cluster import init_ray
@@ -120,8 +120,8 @@ def main() -> None:
             checkpointer,
             grpo_state,
             master_config,
-            _teacher_worker_groups,
-            _alias_to_group_alias,
+            teacher_worker_groups,
+            alias_to_group_alias,
         ) = setup(config, tokenizer, dataset, val_dataset, processor=processor)
 
     rl_init_timer.record("total", time.perf_counter() - main_start)
@@ -133,21 +133,56 @@ def main() -> None:
             print(f"  {label}: {value:.1f}s")
     print("=" * 60 + "\n", flush=True)
 
-    grpo_train(
-        policy,
-        policy_generation,
-        dataloader,
-        val_dataloader,
-        tokenizer,
-        loss_fn,
-        task_to_env,
-        val_task_to_env,
-        logger,
-        checkpointer,
-        grpo_state,
-        master_config,
-        processor=processor,
-    )
+    if config.grpo.async_grpo.enabled:
+        if config.grpo.use_dynamic_sampling:
+            raise NotImplementedError(
+                "use_dynamic_sampling is not supported with async GRPO"
+            )
+        if config.grpo.reward_scaling.enabled:
+            raise NotImplementedError("reward_scaling is not supported with async GRPO")
+        if config.grpo.reward_shaping.enabled:
+            raise NotImplementedError("reward_shaping is not supported with async GRPO")
+        if config.data["use_multiple_dataloader"]:
+            raise NotImplementedError(
+                "use_multiple_dataloader is not supported with async GRPO"
+            )
+
+        print("🚀 Running async GRPO training")
+        async_grpo_train(
+            policy=policy,
+            policy_generation=policy_generation,
+            dataloader=dataloader,
+            val_dataloader=val_dataloader,
+            tokenizer=tokenizer,
+            loss_fn=loss_fn,
+            task_to_env=task_to_env,
+            val_task_to_env=val_task_to_env,
+            logger=logger,
+            checkpointer=checkpointer,
+            grpo_save_state=grpo_state,
+            master_config=master_config,
+            max_trajectory_age_steps=config.grpo.async_grpo.max_trajectory_age_steps,
+            teacher_worker_groups=teacher_worker_groups,
+            alias_to_group_alias=alias_to_group_alias,
+            processor=processor,
+        )
+    else:
+        print("🚀 Running synchronous GRPO training")
+        grpo_train(
+            policy,
+            policy_generation,
+            dataloader,
+            val_dataloader,
+            tokenizer,
+            loss_fn,
+            task_to_env,
+            val_task_to_env,
+            logger,
+            checkpointer,
+            grpo_state,
+            master_config,
+            processor=processor,
+        )
 
 
 if __name__ == "__main__":

tests/unit/test_config_validation.py

@@ -169,9 +169,18 @@ def test_multimodal_dedup_grpo_config_keys_default_off():
     assert GRPOConfig.model_fields["debug_payload_metrics"].default is False
 
 
-def test_nemo_gym_launcher_forwards_processor_to_both_trainers():
-    """Keep sync and async Gym image processing wired to the selected processor."""
-    launcher = Path(__file__).parents[2] / "examples/nemo_gym/run_grpo_nemo_gym.py"
+@pytest.mark.parametrize(
+    "launcher_relpath",
+    [
+        "examples/run_vlm_grpo.py",
+        "examples/nemo_gym/run_grpo_nemo_gym.py",
+    ],
+)
+def test_multimodal_launchers_forward_processor_to_both_trainers(
+    launcher_relpath: str,
+):
+    """Keep sync and async multimodal processing wired to the processor."""
+    launcher = Path(__file__).parents[2] / launcher_relpath
     tree = ast.parse(launcher.read_text())
     trainer_calls = {
         node.func.id: node
@@ -195,6 +204,37 @@ def test_nemo_gym_launcher_forwards_processor_to_both_trainers():
         )
 
 
+def test_vlm_launcher_dispatches_on_async_grpo_enabled():
+    """Keep the VLM async recipe from silently using the synchronous trainer."""
+    launcher = Path(__file__).parents[2] / "examples/run_vlm_grpo.py"
+    tree = ast.parse(launcher.read_text())
+
+    async_branches = [
+        node
+        for node in ast.walk(tree)
+        if isinstance(node, ast.If)
+        and ast.unparse(node.test) == "config.grpo.async_grpo.enabled"
+    ]
+    assert len(async_branches) == 1
+
+    async_branch = async_branches[0]
+    async_calls = {
+        node.func.id
+        for statement in async_branch.body
+        for node in ast.walk(statement)
+        if isinstance(node, ast.Call) and isinstance(node.func, ast.Name)
+    }
+    sync_calls = {
+        node.func.id
+        for statement in async_branch.orelse
+        for node in ast.walk(statement)
+        if isinstance(node, ast.Call) and isinstance(node.func, ast.Name)
+    }
+
+    assert "async_grpo_train" in async_calls
+    assert "grpo_train" in sync_calls
+
+
 def test_reward_penalty_config_requires_explicit_unwanted_token_ids():
     """Unwanted-token penalty requires explicit unwanted-token config.