|
| 1 | +# On-Policy Distillation (OPSD) on DeepSpeed |
| 2 | + |
| 3 | +A DeepSpeed-native port of [HJSang/OPSD_OnPolicyDistillation](https://github.com/HJSang/OPSD_OnPolicyDistillation), |
| 4 | +removing the verl dependency and building directly on DeepSpeed primitives |
| 5 | +(ZeRO-3, hybrid engine, `deepspeed.initialize`). |
| 6 | + |
| 7 | +On-policy distillation trains a small **student** model to imitate a large |
| 8 | +frozen **teacher** on the student's *own* generated rollouts. Each training |
| 9 | +step has three phases: |
| 10 | + |
| 11 | +``` |
| 12 | +┌────────────┐ prompts ┌──────────────────┐ prompt+response ┌────────────┐ |
| 13 | +│ Dataloader │ ──────────▶ │ Student rollout │ ──────────────────▶ │ Teacher │ |
| 14 | +└────────────┘ │ (hybrid / vLLM) │ │ forward │ |
| 15 | + └──────────────────┘ └─────┬──────┘ |
| 16 | + │ logits → CPU cache |
| 17 | + ▼ |
| 18 | + ┌─────────────────────┐ |
| 19 | + │ Student forward + │ |
| 20 | + │ streamed KL / JSD + │ |
| 21 | + │ backward / step │ |
| 22 | + └─────────────────────┘ |
| 23 | +``` |
| 24 | + |
| 25 | +Loss = per-token divergence (`forward_kl` | `reverse_kl` | `jsd`) between |
| 26 | +student and teacher distributions on the student's generated tokens, chunked |
| 27 | +over the sequence axis so the full `[B, T, V]` teacher tensor never |
| 28 | +co-resides with the student logits on the training device. |
| 29 | + |
| 30 | +## Layout |
| 31 | + |
| 32 | +``` |
| 33 | +examples/opsd/ |
| 34 | +├── main.py # entry point (deepspeed launcher) |
| 35 | +├── opsd/ |
| 36 | +│ ├── config.py # OPSDConfig dataclass + JSON loader |
| 37 | +│ ├── losses.py # chunked / streamed KL & JSD |
| 38 | +│ ├── teacher.py # frozen teacher + CPU logit cache |
| 39 | +│ ├── trainer.py # three-phase training loop |
| 40 | +│ ├── data.py # JSONL prompt dataset + left-pad collator |
| 41 | +│ ├── utils.py # response-mask + shift helpers |
| 42 | +│ └── rollout/ |
| 43 | +│ ├── base.py # RolloutEngine ABC, request/batch dataclasses |
| 44 | +│ ├── hybrid_engine.py # DeepSpeed hybrid-engine rollout |
| 45 | +│ └── vllm.py # vLLM rollout on disjoint GPUs |
| 46 | +├── configs/ |
| 47 | +│ ├── ds_zero3.json # base DeepSpeed ZeRO-3 + hybrid engine |
| 48 | +│ ├── opsd_hybrid_engine.json # production-ish hybrid-engine OPSD config |
| 49 | +│ ├── opsd_vllm_disjoint.json # vLLM rollout on a disjoint GPU group |
| 50 | +│ ├── smoke_hybrid.json # 5-step smoke test with Qwen2.5-0.5B / 1.5B |
| 51 | +│ ├── smoke_vllm.json # same but with vLLM rollout |
| 52 | +│ └── smoke_ds_zero3.json # ZeRO-3 config tuned for smoke runs |
| 53 | +├── scripts/ |
| 54 | +│ ├── train_opsd_hybrid.sh # launch hybrid-engine training |
| 55 | +│ └── train_opsd_vllm.sh # launch vLLM training |
| 56 | +└── tests/ # CPU-only unit tests (run with pytest) |
| 57 | +``` |
| 58 | + |
| 59 | +## Quick start |
| 60 | + |
| 61 | +### Install |
| 62 | + |
| 63 | +``` |
| 64 | +pip install deepspeed transformers datasets accelerate |
| 65 | +# Optional, only for the vLLM rollout backend: |
| 66 | +pip install 'vllm>=0.6.4' |
| 67 | +``` |
| 68 | + |
| 69 | +### Hybrid-engine training (single-node, no vLLM) |
| 70 | + |
| 71 | +``` |
| 72 | +cd examples/opsd |
| 73 | +NUM_GPUS=8 bash scripts/train_opsd_hybrid.sh configs/opsd_hybrid_engine.json |
| 74 | +``` |
| 75 | + |
| 76 | +The hybrid engine path lives entirely within DeepSpeed: the student engine |
| 77 | +both trains and generates, sharing weights without a copy step. Easiest to |
| 78 | +get running; slower generation than vLLM. |
| 79 | + |
| 80 | +### vLLM training (disjoint GPU group) |
| 81 | + |
| 82 | +``` |
| 83 | +cd examples/opsd |
| 84 | +# Train on GPUs 0..5, run vLLM on 6,7 (matches default config) |
| 85 | +NUM_TRAIN_GPUS=6 INCLUDE_GPUS=0,1,2,3,4,5 \ |
| 86 | + bash scripts/train_opsd_vllm.sh configs/opsd_vllm_disjoint.json |
| 87 | +``` |
| 88 | + |
| 89 | +vLLM gets dedicated GPUs (`rollout.gpus` in the config). Training rank 0 |
| 90 | +constructs the `LLM` handle; other training ranks receive generated token |
| 91 | +ids via NCCL broadcast. |
| 92 | + |
| 93 | +### Smoke tests (5 steps, small models) |
| 94 | + |
| 95 | +The `smoke_*.json` configs run on 2 GPUs in a few minutes with Qwen2.5-0.5B |
| 96 | +(student) and Qwen2.5-1.5B (teacher), so the full pipeline can be validated |
| 97 | +end-to-end before scaling up. |
| 98 | + |
| 99 | +``` |
| 100 | +cd examples/opsd |
| 101 | +deepspeed --num_gpus 2 main.py --config configs/smoke_hybrid.json |
| 102 | +# For vLLM (uses GPUs 0,1 for training and 2,3 for vLLM): |
| 103 | +NUM_TRAIN_GPUS=2 INCLUDE_GPUS=0,1 deepspeed --num_gpus 2 --include localhost:0,1 \ |
| 104 | + main.py --config configs/smoke_vllm.json |
| 105 | +``` |
| 106 | + |
| 107 | +## Unit tests |
| 108 | + |
| 109 | +The CPU-runnable test suite exercises the loss math, teacher caching, rollout |
| 110 | +contract, and vLLM stitch logic. Run with: |
| 111 | + |
| 112 | +``` |
| 113 | +cd examples/opsd |
| 114 | +python -m pytest tests/ -v |
| 115 | +``` |
| 116 | + |
| 117 | +## Configuration |
| 118 | + |
| 119 | +`OPSDConfig` is a plain dataclass loaded from JSON (no Hydra). The schema: |
| 120 | + |
| 121 | +```json |
| 122 | +{ |
| 123 | + "student": { "model_name_or_path": "...", "dtype": "bfloat16", "arch": "qwen2" }, |
| 124 | + "teacher": { "model_name_or_path": "...", "dtype": "bfloat16", "offload_to_cpu": true }, |
| 125 | + "rollout": { "engine": "hybrid_engine | vllm", ... }, |
| 126 | + "distillation": { "loss_type": "reverse_kl", "temperature": 1.0, "chunk_size": 512 }, |
| 127 | + "training": { "train_batch_size": 8, "learning_rate": 1e-6, ... }, |
| 128 | + "data": { "path": "data/prompts.jsonl", "prompt_field": "prompt" }, |
| 129 | + "deepspeed_config": "configs/ds_zero3.json" |
| 130 | +} |
| 131 | +``` |
| 132 | + |
| 133 | +See `configs/opsd_hybrid_engine.json` and `configs/opsd_vllm_disjoint.json` |
| 134 | +for fully-populated examples. |
| 135 | + |
| 136 | +## Adding a new model architecture |
| 137 | + |
| 138 | +No special steps are needed for new model architectures. vLLM's RLHF weight |
| 139 | +transfer API handles TP slicing internally; the caller only needs to send full |
| 140 | +tensors. |
| 141 | + |
| 142 | +## Design notes |
| 143 | + |
| 144 | +* **Why CPU-cache the teacher logits?** Holding both student and teacher |
| 145 | + `[B, T, V]` tensors on GPU at once doubles memory pressure. Staging the |
| 146 | + teacher to host between the teacher forward and the student backward halves |
| 147 | + the worst-case GPU footprint of the loss path. The streamed loss |
| 148 | + (`losses.streamed_distillation_loss`) pulls teacher chunks back to GPU |
| 149 | + one sequence slice at a time so the full tensor never re-materialises. |
| 150 | + |
| 151 | +* **Why an abstract `RolloutEngine`?** The hybrid-engine and vLLM backends |
| 152 | + have very different lifecycles (hybrid engine reads student weights live; |
| 153 | + vLLM holds its own copy and must be synced) but the trainer should not |
| 154 | + care. The ABC keeps the trainer engine-agnostic so additional backends |
| 155 | + (e.g. a future colocated-vLLM-with-`sleep_mode`) drop in without touching |
| 156 | + the loop. |
| 157 | + |
| 158 | +* **vLLM topology = disjoint, not colocated (v1).** The disjoint topology is |
| 159 | + simpler to debug — failures in vLLM don't take down training and vice |
| 160 | + versa. A colocated topology using vLLM 0.6.4+'s `sleep_mode` is planned as |
| 161 | + a follow-up. |
| 162 | + |
| 163 | +* **Weight sync uses vLLM's RLHF API.** vLLM 0.22.0+ exposes |
| 164 | + ``/update_weights`` which handles TP slicing internally. The trainer |
| 165 | + sends full tensors and vLLM distributes them. |
| 166 | + |
| 167 | +## vLLM status |
| 168 | + |
| 169 | +The vLLM rollout (`opsd/rollout/vllm.py`) is **written and unit-tested but |
| 170 | +not yet usable under the DeepSpeed launcher**. During live validation on |
| 171 | +4× H200 we hit a blocking issue: |
| 172 | + |
| 173 | +> vLLM's worker init calls `new_group(...)` on the global process group as |
| 174 | +> a collective. Under `deepspeed --num_gpus N`, the world is all `N` |
| 175 | +> training ranks but only rank 0 calls into vLLM, so the constructor hangs |
| 176 | +> waiting on the other ranks. Reproduced with vllm 0.6.6 + deepspeed 0.15.4 + |
| 177 | +> torch 2.5.1. Standalone vLLM (world size 1) works in seconds. |
| 178 | +
|
| 179 | +The fix requires running vLLM in a **separate top-level Python process** |
| 180 | +with its own world, accessed over HTTP/RPC from the trainer — the pattern |
| 181 | +used by TRL and OpenRLHF. That's a larger refactor than fits in this PR; |
| 182 | +the current `VLLMRollout` will be the basis for it once landed. |
| 183 | + |
| 184 | +What's verified for the vLLM path today: |
| 185 | +* `tests/test_vllm_stitch.py` — prompt + response stitching (CPU unit test) |
| 186 | +* `vllm.LLM` itself runs fine standalone on Qwen2.5-0.5B (validated) |
| 187 | + |
| 188 | +What's **not** verified: |
| 189 | +* End-to-end training loop with `rollout.engine = "vllm"` in `OPSDConfig` |
| 190 | +* `LLM.collective_rpc("load_weights", ...)` weight sync at training time |
| 191 | + |
| 192 | +The hybrid-engine path (`rollout.engine = "hybrid_engine"`) is validated |
| 193 | +end-to-end on the same hardware. |
| 194 | + |
| 195 | +## Other known limitations (v1) |
| 196 | + |
| 197 | +* **vLLM weight sync (when it works) goes through pickle** — |
| 198 | + `LLM.collective_rpc("load_weights", args=((name, tensor_on_cpu),))`. |
| 199 | + Expect several seconds per sync on a 7B model. A faster v2 would broadcast |
| 200 | + tensors via NCCL on a shared trainer↔vLLM process group — see verl's |
| 201 | + `bucketed_weight_transfer.py` for a reference design. |
| 202 | +* **vLLM `tensor_parallel_size > 1` is untested.** The weight bridge's |
| 203 | + slicing math is unit-tested but no live run exists. |
| 204 | +* **Reward-weighted distillation** (OPSD's `opd.reward_beta` knob) is not |
| 205 | + ported. Easy to add: scale `per_tok` by a reward weight in the loss path. |
| 206 | +* **GRPO and other on-policy RL recipes** are out of scope. The |
| 207 | + `RolloutEngine` / `WeightBridge` abstractions are reusable, but a GRPO |
| 208 | + trainer would add its own advantage / KL-to-reference logic on top. |
| 209 | +* **Qwen3-MoE** is not covered. Add `weight_bridge/qwen3_moe.py` when needed. |
| 210 | +* **Hybrid engine on Qwen-family models uses a ZeRO-3 fallback** (no |
| 211 | + hybrid-engine inference acceleration), since DeepSpeed's inference policy |
| 212 | + list only covers GPT2/GPT-NeoX/OPT/BLOOM/LLAMA/LLAMA2/InternLM as of 0.15. |
| 213 | + The fallback gathers params via `GatheredParameters` and calls the HF |
| 214 | + model's `generate` directly — correct, just ~3-5x slower than the |
| 215 | + accelerated path. |
| 216 | + |
| 217 | +## References |
| 218 | + |
| 219 | +* OPSD reference repo: <https://github.com/HJSang/OPSD_OnPolicyDistillation> |
| 220 | +* DeepSpeed hybrid engine: `deepspeed/runtime/hybrid_engine.py` |
| 221 | +* verl rollout / weight-sync design (used as a cross-check): |
| 222 | + <https://github.com/volcengine/verl/tree/main/verl/workers/rollout/vllm_rollout> |
0 commit comments