feat: data plane transfer queue integration

docs/design-docs/nemo-gym-integration.md

@@ -181,7 +181,7 @@ sequenceDiagram
     GRPO->>Policy: Compute loss and train
 ```
 
-> **NeMo Gym server types** (see [Core Components](https://docs.nvidia.com/nemo/gym/latest/about/concepts/core-components.html)):
+> **NeMo Gym server types** (see [Core Components](https://docs.nvidia.com/nemo/gym/about/core-components)):
 > - **Agent Server**: Orchestrates the rollout loop
 > - **Model Server**: HTTP proxy to vLLM; translates Responses API ↔ Chat Completions
 > - **Resource Server**: Provides tools and rewards

examples/configs/grpo_math_1B.yaml

@@ -391,3 +391,17 @@ logger:
 cluster:
   gpus_per_node: 1
   num_nodes: 1
+
+# TransferQueue-mediated data plane for sync GRPO.
+# Off by default — the legacy grpo_train trainer never engages this.
+# Flip enabled=true and run grpo_train_sync to use TQ-mediated bulk
+# transfer between rollout and train. See nemo_rl/data_plane/README.md.
+data_plane:
+  enabled: false
+  impl: transfer_queue
+  # backend: "simple"                 # NotRequired: TQ storage backend ('simple' or 'mooncake_cpu')
+  # storage_capacity: 1000000         # NotRequired
+  # num_storage_units: 2              # NotRequired
+  # claim_meta_poll_interval_s: 0.5   # NotRequired: blocking-claim poll cadence
+  # observability:                    # NotRequired
+  #   enabled: false

examples/run_grpo.py

@@ -99,6 +99,20 @@ def main() -> None:
         val_task_to_env,
     ) = setup_response_data(tokenizer, config["data"], config["env"])
 
+    # Pick the policy factory at the launcher level so the legacy trainer
+    # stays data-plane-agnostic (architectural invariant — see
+    # tests/data_plane/unit/test_architecture_invariants.py).
+    _dp_cfg = config.get("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
+
     (
         policy,
         policy_generation,
@@ -110,7 +124,13 @@ def main() -> None:
         checkpointer,
         grpo_state,
         master_config,
-    ) = setup(config, tokenizer, dataset, val_dataset)
+    ) = setup(
+        config,
+        tokenizer,
+        dataset,
+        val_dataset,
+        policy_factory=_policy_factory,
+    )
 
     # Check if async mode is enabled
     if "async_grpo" in config["grpo"] and config["grpo"]["async_grpo"]["enabled"]:
@@ -164,10 +184,22 @@ def main() -> None:
             max_trajectory_age_steps=async_config["max_trajectory_age_steps"],
         )
     else:
-        print("🚀 Running synchronous GRPO training")
-
-        # Run standard GRPO training
-        grpo_train(
+        # Two parallel synchronous trainers (verl-style — main_ppo.py vs
+        # main_ppo_sync.py). data_plane.enabled selects which one runs:
+        # the legacy in-memory path or the TransferQueue-mediated fork.
+        # Same model, same data, same seed → diff the wandb runs to
+        # validate parity.
+        dp_cfg = master_config.get("data_plane", {})
+        if dp_cfg.get("enabled", False):
+            from nemo_rl.algorithms.grpo_sync import grpo_train_sync
+
+            print("🚀 Running synchronous GRPO training (TransferQueue)")
+            trainer = grpo_train_sync
+        else:
+            print("🚀 Running synchronous GRPO training (legacy)")
+            trainer = grpo_train
+
+        trainer(
             policy,
             policy_generation,
             dataloader,

nemo_rl/algorithms/grpo.py

@@ -17,7 +17,7 @@
 import warnings
 from concurrent.futures import ThreadPoolExecutor
 from contextlib import nullcontext
-from typing import Any, NotRequired, Optional, TypedDict, TypeVar, cast
+from typing import Any, Callable, NotRequired, Optional, TypedDict, TypeVar, cast
 
 import numpy as np
 import ray
@@ -58,6 +58,7 @@
     get_keys_from_message_log,
 )
 from nemo_rl.data.utils import extract_necessary_env_names
+from nemo_rl.data_plane.interfaces import DataPlaneConfig
 from nemo_rl.distributed.batched_data_dict import BatchedDataDict
 from nemo_rl.distributed.ray_actor_environment_registry import get_actor_python_env
 from nemo_rl.distributed.virtual_cluster import ClusterConfig, RayVirtualCluster
@@ -206,6 +207,7 @@ class MasterConfig(TypedDict):
     logger: GRPOLoggerConfig
     cluster: ClusterConfig
     checkpointing: CheckpointingConfig
+    data_plane: NotRequired[DataPlaneConfig]
 
 
 # ===============================================================================
@@ -219,6 +221,7 @@ def setup(
     dataset: AllTaskProcessedDataset | dict[str, AllTaskProcessedDataset],
     val_dataset: Optional[AllTaskProcessedDataset],
     processor: Optional[AutoProcessor] = None,
+    policy_factory: Optional[Callable[..., ColocatablePolicyInterface]] = None,
 ) -> tuple[
     ColocatablePolicyInterface,
     Optional[GenerationInterface],
@@ -582,10 +585,15 @@ def init_train_dataloader(dataset, suffix: str = ""):
             "(reference model is not loaded)."
         )
 
+    # Caller-supplied factory lets the sync trainer swap in a TQ-mediated
+    # Policy subclass without this shared setup needing to know the data
+    # plane exists. Default is the plain Policy class — legacy behavior.
+    _make_policy = policy_factory if policy_factory is not None else Policy
+
     def init_policy():
         """Initialize policy training workers."""
         t0 = time.perf_counter()
-        p = Policy(
+        p = _make_policy(
             cluster=train_cluster,
             config=policy_config,
             tokenizer=tokenizer,
@@ -2554,7 +2562,7 @@ def async_grpo_train(
     )
 
     replay_buffer = ReplayBuffer.options(runtime_env=_replay_runtime_env).remote(
-        max_size=optimal_buffer_size
+        max_size=optimal_buffer_size,
     )
 
     _tc_py_exec = get_actor_python_env(