Merge branch 'pr-2580' into mxin/qaopd-pr2442-pr2580

Signed-off-by: Meng Xin <mxin@nvidia.com>

# Conflicts:
#	examples/configs/distillation_math.yaml
#	tests/unit/data_plane/test_architecture_invariants.py

Author: mxinO

examples/configs/distillation_math.yaml

@@ -281,3 +281,18 @@ cluster:
     num_nodes: 1
     master_port_range_low: 25000
     master_port_range_high: 28000
+
+# TransferQueue-mediated data plane for on-policy distillation.
+# Off by default; set enabled=true to keep teacher top-k logits in TQ
+# instead of returning them to the driver.
+data_plane:
+    enabled: false
+    impl: transfer_queue
+    backend: "simple"
+    storage_capacity: 1000000
+    num_storage_units: 2
+    claim_meta_poll_interval_s: 0.5
+    global_segment_size: 549755813888
+    local_buffer_size: 68719476736
+    # observability:
+    #   enabled: false

examples/configs/distillation_math_tq.yaml

@@ -0,0 +1,17 @@
+# TransferQueue-enabled variant of distillation_math.yaml.
+# Uses the same model/data defaults and selects the TQ distillation trainer
+# through examples/run_distillation.py when data_plane.enabled=true.
+defaults: distillation_math.yaml
+
+checkpointing:
+    checkpoint_dir: "checkpoints/distillation-tq-${policy.model_name}"
+
+logger:
+    log_dir: "logs/distillation-tq"
+    wandb:
+        name: "distillation-tq-${data.train.dataset_name}-${teacher.model_name}-${policy.model_name}-${loss_fn.kl_type}-${distillation.topk_logits_k}"
+    tensorboard:
+        log_dir: "tb_logs-distillation-tq-${data.train.dataset_name}"
+
+data_plane:
+    enabled: true

examples/configs/recipes/llm/distillation-qwen3-32b-to-1.7b-base-1n8g-fsdp2tp1-tq.v1.yaml

@@ -0,0 +1,58 @@
+defaults: ../../distillation_math_tq.yaml
+distillation:
+  num_prompts_per_step: 64
+  max_num_steps: 20
+  val_batch_size: 32
+  val_period: 10
+  max_val_samples: 256
+loss_fn:
+  kl_type: reverse
+checkpointing:
+  checkpoint_dir: checkpoints/distillation-qwen3-32b-to-1.7b-base-tq
+policy:
+  train_global_batch_size: 32
+  generation_batch_size: 32
+  dtensor_cfg:
+    tensor_parallel_size: 1
+    context_parallel_size: 1
+  dynamic_batching:
+    enabled: false
+  make_sequence_length_divisible_by: 1
+  scheduler:
+  - name: torch.optim.lr_scheduler.LinearLR
+    kwargs:
+      start_factor: 0.1
+      end_factor: 1.0
+      total_iters: 20
+  - name: torch.optim.lr_scheduler.ConstantLR
+    kwargs:
+      factor: 1.0
+      total_iters: 10000000000
+  - milestones:
+    - 20
+teacher:
+  model_name: Qwen/Qwen3-32B
+  train_global_batch_size: 32
+  generation_batch_size: 32
+  dtensor_cfg:
+    context_parallel_size: 1
+  dynamic_batching:
+    enabled: false
+  make_sequence_length_divisible_by: 1
+  scheduler:
+  - name: torch.optim.lr_scheduler.LinearLR
+    kwargs:
+      start_factor: 0.1
+      end_factor: 1.0
+      total_iters: 20
+  - name: torch.optim.lr_scheduler.ConstantLR
+    kwargs:
+      factor: 1.0
+      total_iters: 10000000000
+  - milestones:
+    - 20
+logger:
+  log_dir: logs/distillation-qwen3-32b-to-1.7b-base-tq
+  wandb:
+    project: nemo-rl
+    name: distillation-qwen3-32b-to-1.7b-base-tq

examples/run_distillation.py

@@ -30,6 +30,18 @@
 from nemo_rl.utils.logger import get_next_experiment_dir
 
 
+def _select_trainer(master_config: MasterConfig):
+    """Pick the distillation trainer based on ``data_plane.enabled``."""
+    dp_cfg = master_config.data_plane or {}
+    if dp_cfg.get("enabled", False):
+        from nemo_rl.algorithms.distillation_sync import distillation_train_sync
+
+        print("🚀 Running on-policy distillation training (TransferQueue)")
+        return distillation_train_sync
+    print("🚀 Running on-policy distillation training (legacy)")
+    return distillation_train
+
+
 def parse_args() -> tuple[argparse.Namespace, list[str]]:
     """Parse command line arguments."""
     parser = argparse.ArgumentParser(
@@ -86,6 +98,17 @@ def main() -> None:
         val_task_to_env,
     ) = setup_response_data(tokenizer, config.data, config.env)
 
+    _dp_cfg = config.data_plane or {}
+    if _dp_cfg.get("enabled", False):
+        from nemo_rl.models.policy.tq_policy import TQPolicy
+
+        def _make_student_policy(**kwargs):
+            return TQPolicy(**kwargs, dp_cfg=_dp_cfg)
+
+        _student_policy_factory = _make_student_policy
+    else:
+        _student_policy_factory = None
+
     (
         student_policy,
         teacher_policy,
@@ -97,9 +120,16 @@ def main() -> None:
         checkpointer,
         distillation_state,
         master_config,
-    ) = setup(config, tokenizer, dataset, val_dataset)
+    ) = setup(
+        config,
+        tokenizer,
+        dataset,
+        val_dataset,
+        student_policy_factory=_student_policy_factory,
+    )
 
-    distillation_train(
+    trainer = _select_trainer(master_config)
+    trainer(
         student_policy,
         teacher_policy,
         student_generation,

nemo_rl/algorithms/distillation.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 import os
 import warnings
-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
@@ -44,6 +44,11 @@
     get_keys_from_message_log,
 )
 from nemo_rl.data.utils import load_dataloader_state
+from nemo_rl.data_plane.interfaces import DataPlaneConfig
+from nemo_rl.data_plane.transport_metrics import (
+    add_byte_metric_derivatives,
+    topk_payload_nbytes,
+)
 from nemo_rl.distributed.batched_data_dict import BatchedDataDict
 from nemo_rl.distributed.virtual_cluster import (
     ClusterConfig,
@@ -129,6 +134,7 @@ class MasterConfig(BaseModel, extra="allow"):
     logger: LoggerConfig  # Logger configuration
     cluster: ClusterConfig  # Cluster configuration
     checkpointing: CheckpointingConfig  # Checkpointing configuration
+    data_plane: Optional[DataPlaneConfig] = None
 
 
 # ===============================================================================
@@ -167,6 +173,7 @@ def setup(
     tokenizer: TokenizerType,
     train_dataset: AllTaskProcessedDataset,
     val_dataset: Optional[AllTaskProcessedDataset],
+    student_policy_factory: Optional[Callable[..., ColocatablePolicyInterface]] = None,
 ) -> tuple[
     ColocatablePolicyInterface,  # student_policy
     ColocatablePolicyInterface,  # teacher_policy
@@ -181,6 +188,14 @@ def setup(
 ]:
     """Main entry point for distillation algorithm.
 
+    Args:
+        master_config: Fully resolved distillation configuration.
+        tokenizer: Student tokenizer or processor.
+        train_dataset: Training dataset.
+        val_dataset: Optional validation dataset.
+        student_policy_factory: Optional constructor for the student
+            policy. ``None`` keeps the legacy plain ``Policy`` path.
+
     Returns:
         tuple of student_policy, teacher_policy, student_generation,
         train_dataloader, val_dataloader,
@@ -456,7 +471,10 @@ def setup(
         )
         policy_config["megatron_cfg"]["train_iters"] = total_train_iters
 
-    student_policy = Policy(
+    make_student_policy = (
+        student_policy_factory if student_policy_factory is not None else Policy
+    )
+    student_policy = make_student_policy(
         name_prefix="student",
         cluster=train_cluster,
         config=policy_config,
@@ -745,6 +763,33 @@ def distillation_train(
                     )
                     train_data["teacher_topk_logits"] = teacher_topk["topk_logits"]
                     train_data["teacher_topk_indices"] = teacher_topk["topk_indices"]
+                    teacher_topk_payload_bytes = topk_payload_nbytes(
+                        teacher_topk["topk_logits"],
+                        teacher_topk["topk_indices"],
+                    )
+                    teacher_topk_valid_payload_bytes = topk_payload_nbytes(
+                        teacher_topk["topk_logits"],
+                        teacher_topk["topk_indices"],
+                        [int(length) for length in input_lengths.tolist()],
+                    )
+                    transport_metrics: dict[str, float | int] = {
+                        "teacher_topk_payload_bytes": teacher_topk_payload_bytes,
+                        "teacher_topk_valid_payload_bytes": (
+                            teacher_topk_valid_payload_bytes
+                        ),
+                        "teacher_topk_padding_overhead_bytes": max(
+                            teacher_topk_payload_bytes
+                            - teacher_topk_valid_payload_bytes,
+                            0,
+                        ),
+                        "driver_rx_teacher_topk_bytes": teacher_topk_payload_bytes,
+                        "driver_tx_teacher_topk_bytes": teacher_topk_payload_bytes,
+                        "driver_teacher_topk_bytes": 2 * teacher_topk_payload_bytes,
+                        "driver_teacher_topk_bytes_avoided": 0,
+                        "tq_teacher_topk_write_bytes": 0,
+                        "tq_teacher_topk_write_ms_sum": 0.0,
+                        "tq_teacher_topk_write_ms_max": 0.0,
+                    }
 
                 print("▶ Preparing for training...", flush=True)
                 with timer.time("training_prep"):
@@ -964,8 +1009,13 @@ def distillation_train(
             timing_metrics["valid_tokens_per_sec_per_gpu"] = (
                 metrics["global_valid_toks"] / total_time / total_num_gpus
             )
+            add_byte_metric_derivatives(
+                transport_metrics,
+                token_count=metrics["total_num_tokens"],
+            )
             logger.log_metrics(metrics, total_steps + 1, prefix="train")
             logger.log_metrics(timing_metrics, total_steps + 1, prefix="timing/train")
+            logger.log_metrics(transport_metrics, total_steps + 1, prefix="transport")
 
             timer.reset()
             current_step += 1