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_ds_zero3_offload_autotp2.json b/training/opsd/configs/smoke_ds_zero3_offload_autotp2.json new file mode 100644 index 000000000..8e2058eaa --- /dev/null +++ b/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 +} 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/configs/smoke_hybrid_autotp_zero3.json b/training/opsd/configs/smoke_hybrid_autotp_zero3.json new file mode 100644 index 000000000..349753c2b --- /dev/null +++ b/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" +} diff --git a/training/opsd/main.py b/training/opsd/main.py index 4136333d1..dde605f30 100644 --- a/training/opsd/main.py +++ b/training/opsd/main.py @@ -106,9 +106,25 @@ 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 + tp_size = tp_size or 1 # autotp_size() returns 0 when AutoTP is unconfigured + 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, 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..fb5ca9d5c 100644 --- a/training/opsd/teacher.py +++ b/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() diff --git a/training/opsd/test_student_autotp_zero3.py b/training/opsd/test_student_autotp_zero3.py new file mode 100644 index 000000000..79a65bc2c --- /dev/null +++ b/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() diff --git a/training/opsd/test_teacher_autotp_zero3.py b/training/opsd/test_teacher_autotp_zero3.py new file mode 100644 index 000000000..371aef457 --- /dev/null +++ b/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()