Add OPSD (On-Policy Distillation) training example

training/opsd/README.md

@@ -0,0 +1,144 @@
+# On-Policy Distillation (OPSD) on DeepSpeed
+
+A DeepSpeed-native port of [HJSang/OPSD_OnPolicyDistillation](https://github.com/HJSang/OPSD_OnPolicyDistillation),
+removing the verl dependency and building directly on DeepSpeed primitives
+(ZeRO-3, hybrid engine, `deepspeed.initialize`).
+
+On-policy distillation trains a small **student** model to imitate a large
+frozen **teacher** on the student's *own* generated rollouts. Each training
+step has three phases:
+
+```
+┌────────────┐   prompts   ┌──────────────────┐   prompt+response   ┌────────────┐
+│ Dataloader │ ──────────▶ │ Student rollout  │ ──────────────────▶ │  Teacher   │
+└────────────┘             │ (hybrid engine)  │                     │  forward   │
+                           └──────────────────┘                     └─────┬──────┘
+                                                                          │ logits → CPU cache
+                                                                          ▼
+                                                              ┌─────────────────────┐
+                                                              │ Student forward +   │
+                                                              │ streamed KL / JSD + │
+                                                              │ backward / step     │
+                                                              └─────────────────────┘
+```
+
+Loss = per-token divergence (`forward_kl` | `reverse_kl` | `jsd`) between
+student and teacher distributions on the student's generated tokens, chunked
+over the sequence axis so the full `[B, T, V]` teacher tensor never
+co-resides with the student logits on the training device.
+
+## Layout
+
+```
+training/opsd/
+├── main.py                            # entry point (deepspeed launcher)
+├── trainer.py                         # three-phase OPSD training loop
+├── config.py                          # OPSDConfig dataclass (JSON-loaded)
+├── teacher.py                         # frozen teacher + TeacherLogitCache
+├── losses.py                          # streamed distillation loss (fwd/rev KL, JSD)
+├── data.py                            # PromptDataset + left-padded collator
+├── utils.py                           # response-mask helpers
+├── configs/
+│   ├── ds_zero3.json                  # base DeepSpeed ZeRO-3 + hybrid engine
+│   ├── opsd_hybrid_engine.json        # production-ish hybrid-engine OPSD config
+│   ├── smoke_hybrid.json              # 5-step smoke test with Qwen2.5-0.5B / 1.5B
+│   ├── smoke_hybrid_gc.json           # smoke test with CUDA graph capture
+│   └── smoke_ds_zero0.json            # ZeRO-0 DeepSpeed config for smoke runs
+├── data/
+│   └── prompts.jsonl                  # sample math prompts
+├── benchmarks/                        # rollout / decode micro-benchmarks
+├── scripts/
+│   └── train_opsd_hybrid.sh           # launch hybrid-engine training
+├── tests/                             # CPU-only unit tests (run with pytest)
+└── requirements.txt
+```
+
+## Quick start
+
+### Install
+
+```
+pip install deepspeed transformers datasets accelerate
+```
+
+### Hybrid-engine training
+
+```
+cd training/opsd
+NUM_GPUS=8 bash scripts/train_opsd_hybrid.sh configs/opsd_hybrid_engine.json
+```
+
+The hybrid engine path lives entirely within DeepSpeed: the student engine
+both trains and generates, sharing weights without a copy step.
+
+### Smoke tests (5 steps, small models)
+
+The `smoke_hybrid.json` config runs on 2 GPUs in a few minutes with Qwen2.5-0.5B
+(student) and Qwen2.5-1.5B (teacher), so the full pipeline can be validated
+end-to-end before scaling up.
+
+```
+cd training/opsd
+deepspeed --num_gpus 2 main.py --config configs/smoke_hybrid.json
+```
+
+## Unit tests
+
+The CPU-runnable test suite exercises the loss math and teacher caching. Run with:
+
+```
+cd training/opsd
+python -m pytest tests/ -v
+```
+
+## Configuration
+
+`OPSDConfig` is a plain dataclass loaded from JSON (no Hydra). The schema:
+
+```json
+{
+  "student":    { "model_name_or_path": "...", "dtype": "bfloat16" },
+  "teacher":    { "model_name_or_path": "...", "dtype": "bfloat16", "offload_to_cpu": true },
+  "rollout":    { "engine": "hybrid_engine", ... },
+  "distillation": { "loss_type": "reverse_kl", "temperature": 1.0, "chunk_size": 512 },
+  "training":   { "train_batch_size": 8, "learning_rate": 1e-6, ... },
+  "data":       { "path": "data/prompts.jsonl", "prompt_field": "prompt" },
+  "deepspeed_config": "configs/ds_zero3.json"
+}
+```
+
+See `configs/opsd_hybrid_engine.json` for a fully-populated example.
+
+## Design notes
+
+* **Why CPU-cache the teacher logits?** Holding both student and teacher
+  `[B, T, V]` tensors on GPU at once doubles memory pressure. Staging the
+  teacher to host between the teacher forward and the student backward halves
+  the worst-case GPU footprint of the loss path. The streamed loss
+  (`losses.streamed_distillation_loss`) pulls teacher chunks back to GPU
+  one sequence slice at a time so the full tensor never re-materialises.
+
+* **Why an abstract `RolloutEngine`?** The ABC keeps the trainer
+  engine-agnostic so additional backends can be added without touching the
+  training loop. DeepSpeed provides the `HybridEngineRollout` implementation;
+  external frameworks may plug in their own.
+
+* **Hybrid engine on Qwen-family models uses a ZeRO-3 fallback** (no
+  hybrid-engine inference acceleration), since DeepSpeed's inference policy
+  list only covers GPT2/GPT-NeoX/OPT/BLOOM/LLAMA/LLAMA2/InternLM as of 0.15.
+  The fallback gathers params via `GatheredParameters` and calls the HF
+  model's `generate` directly — correct, just ~3-5x slower than the
+  accelerated path.
+
+## Other known limitations
+
+* **Reward-weighted distillation** (OPSD's `opd.reward_beta` knob) is not
+  ported. Easy to add: scale `per_tok` by a reward weight in the loss path.
+* **GRPO and other on-policy RL recipes** are out of scope. The
+  `RolloutEngine` abstraction is reusable, but a GRPO trainer would add its
+  own advantage / KL-to-reference logic on top.
+
+## References
+
+* OPSD reference repo: <https://github.com/HJSang/OPSD_OnPolicyDistillation>
+* DeepSpeed hybrid engine: `deepspeed/runtime/hybrid_engine.py`

training/opsd/benchmarks/bench_decode_1p1r.py

@@ -0,0 +1,184 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+# DeepSpeed Team
+"""Micro-benchmark for 1p1r HybridEngineRollout decode.
+
+Measures time breakdown of each decode step:
+  - model forward (attention + FFN)
+  - sampling (softmax + multinomial)
+  - Python overhead (mask concat, state update, etc.)
+
+Usage:
+  python examples/opsd/bench_decode_1p1r.py --model Qwen/Qwen2.5-0.5B-Instruct
+"""
+
+import argparse
+import time
+
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+from deepspeed.accelerator import get_accelerator
+
+from deepspeed.runtime.rollout.hybrid_engine_rollout import HybridEngineRollout
+from deepspeed.runtime.rollout.base import RolloutRequest, SamplingConfig
+
+
+def bench_decode_raw(model, tokenizer, device, prompt_len=64, max_new_tokens=64, num_warmup=3, num_iters=10):
+    """Raw decode loop benchmark — measures each component separately."""
+    model.eval()
+    model_dtype = next(model.parameters()).dtype
+
+    input_ids = torch.randint(10, 1000, (1, prompt_len), device=device)
+    attn_mask = torch.ones(1, prompt_len, dtype=torch.long, device=device)
+
+    results = {
+        "prompt_len": prompt_len,
+        "max_new_tokens": max_new_tokens,
+        "model_dtype": str(model_dtype),
+    }
+
+    timings = {"prefill": [], "decode_forward": [], "sampling": [], "overhead": [], "total": []}
+
+    for _ in range(num_warmup + num_iters):
+        with torch.no_grad():
+            t0 = time.perf_counter()
+            out = model(input_ids, attention_mask=attn_mask, use_cache=True)
+            past = out.past_key_values
+            logits = out.logits[:, -1:, :]
+            t_prefill = time.perf_counter()
+
+            generated = []
+            cur_token = logits.argmax(dim=-1)
+            generated.append(cur_token)
+            cur_mask = attn_mask
+
+            decode_times = []
+            sample_times = []
+            overhead_times = []
+
+            for step in range(max_new_tokens):
+                t_step = time.perf_counter()
+                cur_mask = torch.cat([cur_mask, torch.ones(1, 1, dtype=torch.long, device=device)], dim=1)
+                pos_ids = torch.tensor([[prompt_len + step]], device=device)
+
+                t_fwd = time.perf_counter()
+                out = model(cur_token,
+                            attention_mask=cur_mask,
+                            position_ids=pos_ids,
+                            past_key_values=past,
+                            use_cache=True)
+                past = out.past_key_values
+                t_fwd_end = time.perf_counter()
+
+                next_logits = out.logits[:, -1, :]
+                probs = torch.softmax(next_logits / 1.0, dim=-1)
+                cur_token = torch.multinomial(probs, 1)
+                t_sample = time.perf_counter()
+
+                generated.append(cur_token)
+                t_overhead = time.perf_counter()
+
+                decode_times.append(t_fwd_end - t_fwd)
+                sample_times.append(t_sample - t_fwd_end)
+                overhead_times.append(t_overhead - t_sample)
+
+            t_total = time.perf_counter()
+
+        timings["prefill"].append(t_prefill - t0)
+        timings["decode_forward"].append(decode_times)
+        timings["sampling"].append(sample_times)
+        timings["overhead"].append(overhead_times)
+        timings["total"].append(t_total - t0)
+
+    import numpy as np
+
+    def avg_last_n(lst, n):
+        return np.mean(lst[-n:])
+
+    def avg_of_avg(list_of_lists, n):
+        arrs = [np.array(ls[-n:]) for ls in list_of_lists]
+        return np.mean([a.mean() for a in arrs])
+
+    results["prefill_ms"] = avg_last_n(timings["prefill"], num_iters) * 1000
+    results["decode_forward_ms_per_step"] = avg_of_avg(timings["decode_forward"], num_iters) * 1000
+    results["sampling_ms_per_step"] = avg_of_avg(timings["sampling"], num_iters) * 1000
+    results["overhead_ms_per_step"] = avg_of_avg(timings["overhead"], num_iters) * 1000
+    results["total_ms"] = avg_last_n(timings["total"], num_iters) * 1000
+    results["decode_steps_total_ms"] = results["decode_forward_ms_per_step"] * max_new_tokens
+    results["sampling_total_ms"] = results["sampling_ms_per_step"] * max_new_tokens
+    results["overhead_total_ms"] = results["overhead_ms_per_step"] * max_new_tokens
+
+    return results
+
+
+def bench_hybrid_rollout(rollout, tokenizer, device, prompt_len=64, max_new_tokens=64, num_warmup=3, num_iters=10):
+    """Benchmark the full HybridEngineRollout.generate() path."""
+    input_ids = torch.randint(10, 1000, (1, prompt_len), device=device)
+    attn_mask = torch.ones(1, prompt_len, dtype=torch.long, device=device)
+    sampling = SamplingConfig(max_new_tokens=max_new_tokens, temperature=1.0, top_p=1.0)
+    request = RolloutRequest(prompt_ids=input_ids, prompt_attention_mask=attn_mask)
+
+    times = []
+    for _ in range(num_warmup + num_iters):
+        get_accelerator().synchronize()  #ignore-cuda
+        t0 = time.perf_counter()
+        with torch.no_grad():
+            rollout.generate(request, sampling)
+        get_accelerator().synchronize()  #ignore-cuda
+        times.append(time.perf_counter() - t0)
+
+    import numpy as np
+    avg = np.mean(times[-num_iters:]) * 1000
+    return {"rollout_total_ms": avg, "prompt_len": prompt_len, "max_new_tokens": max_new_tokens}
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model", default="Qwen/Qwen2.5-0.5B-Instruct")
+    parser.add_argument("--prompt-len", type=int, default=64)
+    parser.add_argument("--max-new-tokens", type=int, default=64)
+    parser.add_argument("--num-warmup", type=int, default=3)
+    parser.add_argument("--num-iters", type=int, default=10)
+    parser.add_argument("--graph-capture", action="store_true", help="Enable CUDA graph capture")
+    args = parser.parse_args()
+
+    device = get_accelerator().current_device()  #ignore-cuda
+
+    tokenizer = AutoTokenizer.from_pretrained(args.model, padding_side="left")
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    model = AutoModelForCausalLM.from_pretrained(args.model, torch_dtype=torch.bfloat16).to(device)
+
+    print(f"=== Raw decode loop benchmark (model={args.model}) ===")
+    raw = bench_decode_raw(model, tokenizer, device, args.prompt_len, args.max_new_tokens, args.num_warmup,
+                           args.num_iters)
+    print(f"  Prefill:              {raw['prefill_ms']:.2f} ms")
+    print(
+        f"  Decode forward/step:  {raw['decode_forward_ms_per_step']:.3f} ms  (total: {raw['decode_steps_total_ms']:.1f} ms)"
+    )
+    print(f"  Sampling/step:        {raw['sampling_ms_per_step']:.3f} ms  (total: {raw['sampling_total_ms']:.1f} ms)")
+    print(f"  Overhead/step:        {raw['overhead_ms_per_step']:.3f} ms  (total: {raw['overhead_total_ms']:.1f} ms)")
+    print(f"  Total:                {raw['total_ms']:.1f} ms")
+
+    print(f"\n=== HybridEngineRollout benchmark (graph_capture={args.graph_capture}) ===")
+    engine = type('Engine', (), {'module': model})()  # lightweight wrapper
+    from deepspeed.runtime.rollout.hybrid_engine_rollout import HybridEngineRolloutConfig
+    cfg = HybridEngineRolloutConfig(use_graph_capture=args.graph_capture)
+    rollout = HybridEngineRollout(engine, tokenizer, cfg=cfg)
+    rr = bench_hybrid_rollout(rollout, tokenizer, device, args.prompt_len, args.max_new_tokens, args.num_warmup,
+                              args.num_iters)
+    print(f"  Rollout generate:     {rr['rollout_total_ms']:.1f} ms")
+
+    print(f"\n=== Summary ===")
+    print(f"  Raw decode loop:      {raw['total_ms']:.1f} ms")
+    print(f"  HybridEngine rollout: {rr['rollout_total_ms']:.1f} ms")
+    print(f"  Overhead (rollout - raw): {rr['rollout_total_ms'] - raw['total_ms']:.1f} ms")
+    print(
+        f"  Bottleneck: decode forward = {raw['decode_forward_ms_per_step']:.3f} ms/step x {args.max_new_tokens} steps = {raw['decode_steps_total_ms']:.1f} ms"
+    )
+
+
+if __name__ == "__main__":
+    main()