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_ds_zero3_offload_autotp2.json

@@ -0,0 +1,36 @@
+{
+    "bf16": {
+        "enabled": true
+    },
+    "zero_optimization": {
+        "stage": 3,
+        "offload_optimizer": {
+            "device": "cpu"
+        },
+        "offload_param": {
+            "device": "cpu"
+        },
+        "overlap_comm": true,
+        "contiguous_gradients": true,
+        "reduce_bucket_size": 5e7,
+        "stage3_prefetch_bucket_size": 5e7,
+        "stage3_param_persistence_threshold": 1e6,
+        "stage3_max_live_parameters": 1e9,
+        "stage3_max_reuse_distance": 1e9
+    },
+    "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/configs/smoke_hybrid_autotp_zero3.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": true,
+        "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_zero3_ckpt",
+        "seed": 42
+    },
+    "data": {
+        "path": "data/prompts.jsonl",
+        "prompt_field": "prompt",
+        "shuffle": true
+    },
+    "deepspeed_config": "configs/smoke_ds_zero3_offload_autotp2.json"
+}

training/opsd/main.py

@@ -106,9 +106,24 @@ 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
+    if dp_size > 1:
+        # TP groups are contiguous, so DP rank = global_rank // tp_size.
+        # Pass num_replicas/rank explicitly so the sampler shards only along
+        # the real DP dimension (not the full world, which would split data
+        # across TP ranks and deadlock the all-reduce).
+        sampler = DistributedSampler(
+            dataset,
+            shuffle=cfg.data.shuffle,
+            num_replicas=dp_size,
+            rank=deepspeed.comm.get_rank() // tp_size,
+        )
+    else:
+        sampler = 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

@@ -115,11 +115,11 @@ class TeacherWrapper:
     """Frozen teacher, always routed through DeepSpeed.
 
     The ZeRO stage is chosen to match the run: **stage 3** when there is more
-    than one rank (so the frozen params shard across GPUs instead of being
-    replicated on every one) or when ``cfg.offload_to_cpu`` is set (so those
-    shards can live on the host between forwards); otherwise **stage 0** — on
-    a single GPU with no offload, ZeRO-3's per-forward gather is pure
-    overhead. The optimizer slot is unused (no trainable params); ZeRO-3 here
+    than one data-parallel rank (i.e. ``world_size > autotp_size``, so the
+    frozen params shard across the DP group instead of being replicated) or
+    when ``cfg.offload_to_cpu`` is set (so those shards can live on the host
+    between forwards); otherwise **stage 0** — in pure-TP mode with no offload,
+    ZeRO-3's per-forward gather is pure overhead. The optimizer slot is unused (no trainable params); ZeRO-3 here
     only buys per-forward parameter gather/release.
 
     The full checkpoint is loaded on each rank before DeepSpeed partitions it;
@@ -145,11 +145,16 @@ 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
-        # 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
+        # Route through DeepSpeed. ZeRO-3 shards params across the
+        # data-parallel group, so it only pays off when there is a real DP
+        # dimension (world_size > autotp_size) or host memory to offload to.
+        # In pure-TP mode (autotp_size == world_size) with no offload there is
+        # nothing to shard and ZeRO-3's per-forward gather is pure overhead,
+        # so we drop to stage 0 there. AutoTP + ZeRO-3 is supported by recent
+        # DeepSpeed, so the two may be combined when a DP dimension exists.
+        autotp_size = getattr(cfg, 'autotp_size', 1)
+        use_autotp = autotp_size > 1
+        use_zero3 = cfg.offload_to_cpu or world_size > autotp_size
         zero_opt = {"stage": 3 if use_zero3 else 0}
         if cfg.offload_to_cpu:
             zero_opt["offload_param"] = {"device": "cpu"}
@@ -159,6 +164,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()

training/opsd/test_student_autotp_zero3.py

@@ -0,0 +1,62 @@
+"""Isolated student test: AutoTP + ZeRO-3 (+offload) forward/backward/step.
+
+Env knobs:
+  STUDENT_MODEL (default local path), OFFLOAD (1/0), AUTOTP (int)
+Launch: deepspeed --num_gpus N test_student_autotp_zero3.py
+"""
+import os
+import sys
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+
+import torch
+import deepspeed
+from deepspeed.accelerator import get_accelerator
+from transformers import AutoModelForCausalLM
+
+MODEL = os.environ.get("STUDENT_MODEL", "Qwen/Qwen2.5-0.5B-Instruct")
+
+
+def build_ds_config():
+    offload = os.environ.get("OFFLOAD", "1") == "1"
+    autotp = int(os.environ.get("AUTOTP", "2"))
+    z = {"stage": 3}
+    if offload:
+        z["offload_param"] = {"device": "cpu"}
+        z["offload_optimizer"] = {"device": "cpu"}
+    return {
+        "train_micro_batch_size_per_gpu": 1,
+        "bf16": {"enabled": True},
+        "zero_optimization": z,
+        "tensor_parallel": {"autotp_size": autotp},
+        "optimizer": {"type": "AdamW", "params": {"lr": 1e-6, "torch_adam": True}},
+    }
+
+
+def main():
+    deepspeed.init_distributed()
+    model = AutoModelForCausalLM.from_pretrained(MODEL, dtype=torch.bfloat16)
+    engine, *_ = deepspeed.initialize(
+        model=model, model_parameters=model.parameters(), config=build_ds_config()
+    )
+
+    dev = get_accelerator().current_device_name()
+    input_ids = torch.tensor([[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]], device=dev)
+    attn = torch.ones_like(input_ids)
+    out = engine(input_ids=input_ids, attention_mask=attn)
+    loss = out.logits.float().mean()
+    engine.backward(loss)
+    engine.step()
+
+    if deepspeed.comm.get_rank() == 0:
+        print(
+            f"STUDENT_OK loss={loss.item():.4f} "
+            f"mem={torch.cuda.max_memory_allocated()/1024**3:.2f}GB "
+            f"offload={os.environ.get('OFFLOAD','1')} autotp={os.environ.get('AUTOTP','2')} "
+            f"ws={os.environ.get('WORLD_SIZE','1')}",
+            flush=True,
+        )
+
+
+if __name__ == "__main__":
+    main()

training/opsd/test_teacher_autotp_zero3.py

@@ -0,0 +1,49 @@
+"""Isolated teacher test: AutoTP + ZeRO-3 (offload) forward_to_cache.
+
+Env knobs:
+  TEACHER_MODEL (default local path), OFFLOAD (1/0), AUTOTP (int)
+Launch: deepspeed --num_gpus N test_teacher_autotp_zero3.py
+"""
+import os
+import sys
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+
+import torch
+import deepspeed
+from deepspeed.accelerator import get_accelerator
+
+from teacher import TeacherWrapper
+from config import TeacherConfig
+
+MODEL = os.environ.get("TEACHER_MODEL", "Qwen/Qwen2.5-1.5B-Instruct")
+
+
+def main():
+    deepspeed.init_distributed()
+    cfg = TeacherConfig(
+        model_name_or_path=MODEL,
+        dtype="bfloat16",
+        offload_to_cpu=os.environ.get("OFFLOAD", "1") == "1",
+        autotp_size=int(os.environ.get("AUTOTP", "2")),
+    )
+    ws = int(os.environ.get("WORLD_SIZE", 1))
+    teacher = TeacherWrapper(cfg, world_size=ws)
+
+    dev = get_accelerator().current_device_name()
+    input_ids = torch.tensor([[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]], device=dev)
+    attn = torch.ones_like(input_ids)
+    cache = teacher.forward_to_cache(input_ids, attn)
+
+    if deepspeed.comm.get_rank() == 0:
+        print(
+            f"TEACHER_OK shape={cache.shape} "
+            f"mem={torch.cuda.max_memory_allocated()/1024**3:.2f}GB "
+            f"offload={cfg.offload_to_cpu} autotp={cfg.autotp_size} ws={ws}",
+            flush=True,
+        )
+    cache.free()
+
+
+if __name__ == "__main__":
+    main()