From ac68cb7a75798bad884b4eecc7db3b054f1172f4 Mon Sep 17 00:00:00 2001 From: Guokai Ma Date: Mon, 20 Jul 2026 23:15:28 +0800 Subject: [PATCH] Add AutoTP=2 smoke test config and script Verified on 2x RTX 5090 with Qwen2.5-0.5B-Instruct: - Student and teacher both use AutoTP=2 - HybridEngineRollout.generate() produces identical output on both ranks - OPSD training loop runs end-to-end with temperature=0 (greedy) - No DistributedSampler in pure-TP mode (all ranks see same data) - Teacher uses ZeRO-0 when AutoTP is enabled (AutoTP+ZeRO-3 not supported) Requires DeepSpeed >= commit 53a2ac4 (AutoTP KV cache consistency fix). Signed-off-by: Guokai Ma Signed-off-by: Guokai Ma --- training/opsd/config.py | 1 + .../opsd/configs/smoke_ds_zero0_autotp2.json | 23 ++++++++++ .../opsd/configs/smoke_hybrid_autotp.json | 45 +++++++++++++++++++ training/opsd/main.py | 9 ++-- training/opsd/scripts/smoke_autotp.sh | 5 +++ training/opsd/teacher.py | 10 +++-- 6 files changed, 87 insertions(+), 6 deletions(-) create mode 100644 training/opsd/configs/smoke_ds_zero0_autotp2.json create mode 100644 training/opsd/configs/smoke_hybrid_autotp.json create mode 100755 training/opsd/scripts/smoke_autotp.sh diff --git a/training/opsd/config.py b/training/opsd/config.py index 4a1c7633d..2b75edb50 100644 --- a/training/opsd/config.py +++ b/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 diff --git a/training/opsd/configs/smoke_ds_zero0_autotp2.json b/training/opsd/configs/smoke_ds_zero0_autotp2.json new file mode 100644 index 000000000..d24aa4d64 --- /dev/null +++ b/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 +} diff --git a/training/opsd/configs/smoke_hybrid_autotp.json b/training/opsd/configs/smoke_hybrid_autotp.json new file mode 100644 index 000000000..ec9e7404b --- /dev/null +++ b/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" +} diff --git a/training/opsd/main.py b/training/opsd/main.py index 4136333d1..efa3d474a 100644 --- a/training/opsd/main.py +++ b/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, diff --git a/training/opsd/scripts/smoke_autotp.sh b/training/opsd/scripts/smoke_autotp.sh new file mode 100755 index 000000000..66c98ede9 --- /dev/null +++ b/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 diff --git a/training/opsd/teacher.py b/training/opsd/teacher.py index 934b57495..ad666bba3 100644 --- a/training/opsd/teacher.py +++ b/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()