feat(vlm): route VLM GRPO through TQ trainer when data_plane.enabled

examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n8g-automodel-ep8.v1-tq_simple.yaml

@@ -0,0 +1,3 @@
+defaults: vlm_grpo-nemotron-omni-30ba3b-clevr-1n8g-automodel-ep8.v1.yaml
+data_plane:
+  enabled: true

examples/configs/recipes/vlm/vlm_grpo-qwen3.5-35ba3b-geo3k-2n8g-automodel-ep16-tq_simple.yaml

@@ -0,0 +1,3 @@
+defaults: vlm_grpo-qwen3.5-35ba3b-geo3k-2n8g-automodel-ep16.yaml
+data_plane:
+  enabled: true

examples/run_vlm_grpo.py

@@ -34,6 +34,23 @@
 from nemo_rl.utils.timer import Timer
 
 
+def _select_trainer(master_config: MasterConfig):
+    """Pick the synchronous trainer based on ``data_plane.enabled``.
+
+    Mirrors ``run_grpo.py`` so the VLM launcher routes through the same
+    TransferQueue-backed sibling trainer (``grpo_train_sync``) when the
+    data plane is enabled, and otherwise uses the legacy ``grpo_train``.
+    """
+    dp_cfg = master_config.data_plane or {}
+    if dp_cfg.get("enabled", False):
+        from nemo_rl.algorithms.grpo_sync import grpo_train_sync
+
+        print("🚀 Running synchronous VLM GRPO training (TransferQueue)")
+        return grpo_train_sync
+    print("🚀 Running synchronous VLM GRPO training (legacy)")
+    return grpo_train
+
+
 def parse_args() -> tuple[argparse.Namespace, list[str]]:
     """Parse command line arguments."""
     parser = argparse.ArgumentParser(description="Run GRPO training with configuration")
@@ -110,6 +127,20 @@ def main() -> None:
             processor, config.data, config.env, is_vlm=True
         )
 
+    # Pick the policy factory at the launcher level so the legacy trainer
+    # stays data-plane-agnostic (architectural invariant — see
+    # tests/unit/data_plane/test_architecture_invariants.py).
+    _dp_cfg = config.data_plane or {}
+    if _dp_cfg.get("enabled", False):
+        from nemo_rl.models.policy.tq_policy import TQPolicy
+
+        def _make_policy(**kwargs):
+            return TQPolicy(**kwargs, dp_cfg=_dp_cfg)
+
+        _policy_factory = _make_policy
+    else:
+        _policy_factory = None  # setup() defaults to plain Policy
+
     with rl_init_timer.time("setup"):
         (
             policy,
@@ -125,7 +156,14 @@ def main() -> None:
             master_config,
             teacher_worker_groups,
             alias_to_group_alias,
-        ) = setup(config, tokenizer, dataset, val_dataset, processor=processor)
+        ) = setup(
+            config,
+            tokenizer,
+            dataset,
+            val_dataset,
+            processor=processor,
+            policy_factory=_policy_factory,
+        )
 
     rl_init_timer.record("total", time.perf_counter() - main_start)
     rl_init_metrics = rl_init_timer.get_timing_metrics(reduction_op="sum")
@@ -170,8 +208,9 @@ def main() -> None:
             processor=processor,
         )
     else:
-        print("🚀 Running synchronous GRPO training")
-        grpo_train(
+        # ``_select_trainer`` prints which sync trainer it picked.
+        trainer = _select_trainer(master_config)
+        trainer(
             policy,
             policy_generation,
             dataloader,

nemo_rl/algorithms/grpo.py

@@ -462,10 +462,22 @@ def _validate_multimodal_dedup_capability(master_config: MasterConfig) -> None:
             "grpo.deduplicate_multimodal_data=true is currently qualified "
             "only with policy.generation.backend=vllm."
         )
-    if (master_config.data_plane or {}).get("enabled", False):
+    # The data plane carries deduplicated payloads: ``PackedTensor.to_wire``
+    # emits one wire row per *logical* row and walks segments under dedup, so
+    # the wire format itself is not the constraint. The one gap is NeMo-Gym:
+    # ``grpo_train_sync`` does not call
+    # ``attach_initial_nemo_gym_image_payloads``, which supplies the initial
+    # image tensors a Gym dataset omits from ``extra_env_info``. That helper is
+    # itself gated on ``should_use_nemo_gym``, so non-Gym recipes never needed
+    # it and are unaffected.
+    if (master_config.data_plane or {}).get("enabled", False) and (
+        should_use_nemo_gym(master_config)
+    ):
         raise NotImplementedError(
-            "grpo.deduplicate_multimodal_data=true is currently supported "
-            "only when data_plane.enabled=false."
+            "grpo.deduplicate_multimodal_data=true with data_plane.enabled=true "
+            "is not supported for NeMo-Gym runs: the TransferQueue trainer does "
+            "not attach the initial Gym image payloads. Non-Gym recipes are "
+            "supported."
         )
 
 

nemo_rl/algorithms/grpo_sync.py

@@ -37,6 +37,8 @@
 from typing import TYPE_CHECKING, Any, Optional
 
 if TYPE_CHECKING:
+    from transformers import AutoProcessor
+
     from nemo_rl.models.policy.tq_policy import TQPolicy
 
 import numpy as np
@@ -50,6 +52,7 @@
     MasterConfig,
     _clip_grpo_advantages,
     _create_advantage_estimator,
+    _initial_policy_generation_stale,
     _log_mixed_rewards_and_advantages_information,
     _placeholder_seq_logprob_error_metrics,
     _policy_dtype,
@@ -390,6 +393,19 @@ def grpo_train_sync(
     checkpointer: CheckpointManager,
     grpo_save_state: GRPOSaveState,
     master_config: MasterConfig,
+    # Unused here, and present only so the shared VLM launcher can pass one
+    # fixed kwarg set to whichever trainer ``_select_trainer`` returns.
+    # ``grpo_train``'s sole use of it is
+    # ``attach_initial_nemo_gym_image_payloads``, gated on
+    # ``grpo.deduplicate_multimodal_data`` *and* ``should_use_nemo_gym`` — the
+    # combination ``setup()`` rejects via
+    # ``_validate_multimodal_dedup_capability``. Non-Gym dedup runs never call
+    # that helper, so they need no processor here either.
+    #
+    # TODO: replace this parity kwarg with a ``ProcessorInterface`` both
+    # trainers consume, rather than threading ``Optional[AutoProcessor]``
+    # through every signature — ``grpo.py`` repeats it at seven sites.
+    processor: Optional["AutoProcessor"] = None,
 ) -> None:
     """Run GRPO training algorithm — TransferQueue-mediated.
 
@@ -414,7 +430,13 @@ def grpo_train_sync(
 
     kv_scales_cache = None  # Cache reused for computed kv scales
 
-    POLICY_GENERATION_STALE = True
+    # Skip a redundant iter-1 refit when setup() already synced weights
+    # (synchronizer not stale, fresh run). The redundant refit resets
+    # vLLM CUDA-graph / KV-cache state and yields a step-1
+    # token_mult_prob_error spike that converges by step 3.
+    POLICY_GENERATION_STALE = _initial_policy_generation_stale(
+        policy_generation, grpo_save_state.total_steps
+    )
     assert policy_generation is not None
 
     if master_config.grpo.skip_reference_policy_logprobs_calculation:
@@ -950,9 +972,20 @@ def grpo_train_sync(
                         # (logprobs/advantages/masks) and wire-only message
                         # log bulk fields are skipped by virtue of not being
                         # in DP_CALIB_INPUT_FIELDS.
+                        # ``DP_CALIB_INPUT_FIELDS`` names a ``multi_modal_inputs``
+                        # column that is never actually written — the rollout
+                        # writes pixel_values / image_grid_thw / … individually
+                        # — so on a VLM run this filter would otherwise yield
+                        # only the text fields and calibrate image-blind.
+                        # Local import: ``tq_policy`` is TYPE_CHECKING-only at
+                        # module scope here (see the import block at the top).
+                        from nemo_rl.models.policy.tq_policy import (
+                            _present_multimodal_fields,
+                        )
+
                         _calib_fields = [
                             f for f in (meta.fields or []) if f in DP_CALIB_INPUT_FIELDS
-                        ]
+                        ] + _present_multimodal_fields(meta)
                         calibration_data = policy.read_from_dataplane(
                             meta,
                             select_fields=_calib_fields,