Add AutoTP=2 smoke test config and script

training/opsd/config.py

@@ -38,6 +38,7 @@ class TeacherConfig:
     dtype: str = "bfloat16"
     trust_remote_code: bool = False
     offload_to_cpu: bool = True
+    autotp_size: int = 1
 
 
 @dataclass

training/opsd/configs/smoke_ds_zero0_autotp2.json

@@ -0,0 +1,23 @@
+{
+    "bf16": {
+        "enabled": true
+    },
+    "zero_optimization": {
+        "stage": 0
+    },
+    "tensor_parallel": {
+        "autotp_size": 2
+    },
+    "optimizer": {
+        "type": "AdamW",
+        "params": {
+            "lr": 1e-6,
+            "betas": [0.9, 0.95],
+            "eps": 1e-8,
+            "weight_decay": 0.0,
+            "torch_adam": true
+        }
+    },
+    "gradient_clipping": 1.0,
+    "wall_clock_breakdown": false
+}

training/opsd/configs/smoke_hybrid_autotp.json

@@ -0,0 +1,45 @@
+{
+    "student": {
+        "model_name_or_path": "Qwen/Qwen2.5-0.5B-Instruct",
+        "dtype": "bfloat16",
+        "trust_remote_code": false
+    },
+    "teacher": {
+        "model_name_or_path": "Qwen/Qwen2.5-1.5B-Instruct",
+        "dtype": "bfloat16",
+        "trust_remote_code": false,
+        "offload_to_cpu": false,
+        "autotp_size": 2
+    },
+    "rollout": {
+        "engine": "hybrid_engine",
+        "max_prompt_length": 128,
+        "max_response_length": 32,
+        "temperature": 0,
+        "n_samples_per_prompt": 1
+    },
+    "distillation": {
+        "loss_type": "reverse_kl",
+        "temperature": 1.0,
+        "chunk_size": 128
+    },
+    "training": {
+        "micro_batch_size_per_gpu": 1,
+        "gradient_accumulation_steps": 1,
+        "learning_rate": 1e-6,
+        "weight_decay": 0.0,
+        "num_train_epochs": 1,
+        "max_steps": 3,
+        "warmup_steps": 0,
+        "save_steps": 10000,
+        "logging_steps": 1,
+        "save_dir": "./opsd_smoke_autotp_ckpt",
+        "seed": 42
+    },
+    "data": {
+        "path": "data/prompts.jsonl",
+        "prompt_field": "prompt",
+        "shuffle": true
+    },
+    "deepspeed_config": "configs/smoke_ds_zero0_autotp2.json"
+}

training/opsd/main.py

@@ -106,9 +106,12 @@ def main() -> None:
         chat_template=cfg.data.chat_template,
     )
     collator = LeftPaddedPromptCollator(tokenizer=tokenizer, max_prompt_length=cfg.rollout.max_prompt_length)
-    # Shard the dataset across data-parallel ranks. Without this, every rank
-    # iterates the full set and the run is pure redundant compute on >1 GPU.
-    sampler = DistributedSampler(dataset, shuffle=cfg.data.shuffle) if dist_world_size() > 1 else None
+    # In pure-TP mode (autotp_size == world_size), every rank must see the
+    # same data — a DistributedSampler would split prompts across TP ranks
+    # and deadlock the TP all-reduce.  Only shard when there is real DP.
+    tp_size = student_engine.autotp_size() if hasattr(student_engine, 'autotp_size') else 1
+    dp_size = dist_world_size() // tp_size
+    sampler = DistributedSampler(dataset, shuffle=cfg.data.shuffle) if dp_size > 1 else None
     loader = DataLoader(
         dataset,
         batch_size=cfg.training.micro_batch_size_per_gpu,

training/opsd/scripts/smoke_autotp.sh

@@ -0,0 +1,5 @@
+#!/bin/bash
+# Smoke test: OPSD with AutoTP=2 (requires 2 GPUs)
+# Usage: bash scripts/smoke_autotp.sh
+export PYTHONPATH=.
+deepspeed --num_gpus 2 main.py --config configs/smoke_hybrid_autotp.json

training/opsd/teacher.py

@@ -145,11 +145,13 @@ def __init__(self, cfg: TeacherConfig, world_size: int):
         for p in model.parameters():
             p.requires_grad_(False)
 
-        # Always route through DeepSpeed. ZeRO-3 only pays off when there is
+        # Route through DeepSpeed. ZeRO-3 only pays off when there is
         # another rank to shard across (world_size > 1) or host memory to
         # offload to; on a single GPU with no offload it is pure per-forward
-        # gather overhead, so we drop to stage 0 there.
-        use_zero3 = cfg.offload_to_cpu or world_size > 1
+        # gather overhead, so we drop to stage 0 there. When using AutoTP,
+        # force ZeRO-0 because AutoTP + ZeRO-3 is not yet supported.
+        use_autotp = getattr(cfg, 'autotp_size', 1) > 1
+        use_zero3 = (cfg.offload_to_cpu or world_size > 1) and not use_autotp
         zero_opt = {"stage": 3 if use_zero3 else 0}
         if cfg.offload_to_cpu:
             zero_opt["offload_param"] = {"device": "cpu"}
@@ -159,6 +161,8 @@ def __init__(self, cfg: TeacherConfig, world_size: int):
             "fp16": {"enabled": dtype is torch.float16},
             "zero_optimization": zero_opt,
         }
+        if use_autotp:
+            ds_config["tensor_parallel"] = {"autotp_size": cfg.autotp_size}
         self._callable, *_ = deepspeed.initialize(model=model, config=ds_config)
 
     @torch.no_grad()