Add on-policy distillation example (Qwen3-4B + Qwen3-32B vLLM teacher)#328
Add on-policy distillation example (Qwen3-4B + Qwen3-32B vLLM teacher)#328CalvinXKY wants to merge 1 commit into
Conversation
Documentation build overview
33 files changed ·
|
Signed-off-by: kaiyuan <kyxiezju@163.com>
7ef363f to
8441b80
Compare
There was a problem hiding this comment.
Code Review
This pull request introduces an on-policy distillation (OPD) example, adding a comprehensive README and two shell scripts (run-qwen3-4b-32b-opd.sh and run-qwen3-8b-opd-megatron.sh) to demonstrate both vLLM and Megatron teacher modes. The feedback highlights several shell script improvements: replacing overly aggressive pkill -9 python commands with targeted process termination to avoid disrupting unrelated workloads, adding set -o pipefail to prevent masked pipeline failures when piping to tee, and ensuring previous vLLM processes are cleaned up at startup to avoid port and GPU conflicts.
| pkill -9 ray 2>/dev/null || true | ||
| pkill -9 python 2>/dev/null || true |
There was a problem hiding this comment.
Using pkill -9 python is extremely aggressive and dangerous as it will kill all Python processes running on the system, which can disrupt other users' workloads or even kill the parent process/CI runner if it is Python-based. It is much safer to use a more targeted command like pkill -9 -f "train.py".
| pkill -9 ray 2>/dev/null || true | |
| pkill -9 python 2>/dev/null || true | |
| pkill -9 ray 2>/dev/null || true | |
| pkill -9 -f "train.py" 2>/dev/null || true |
| pkill -9 ray 2>/dev/null || true | ||
| pkill -9 python 2>/dev/null || true |
There was a problem hiding this comment.
Similarly, avoid using pkill -9 python at the end of the script to prevent killing unrelated Python processes on the system. Use pkill -9 -f "train.py" instead.
| pkill -9 ray 2>/dev/null || true | |
| pkill -9 python 2>/dev/null || true | |
| pkill -9 ray 2>/dev/null || true | |
| pkill -9 -f "train.py" 2>/dev/null || true |
| # 8×GPU colocate OPD: Qwen3-4B student (GPUs 0-3) + Qwen3-32B vLLM teacher (GPUs 4-7). | ||
| # See README.md for checkpoint conversion and data prerequisites. | ||
|
|
||
| set -ex |
There was a problem hiding this comment.
The script uses set -ex but does not set set -o pipefail. Since the training command on line 189 pipes its output to tee (2>&1 | tee ...), any failure in ray job submit or train.py will be masked because tee will exit with status 0, causing the script to incorrectly report success. It is highly recommended to add set -o pipefail to ensure pipeline failures are properly propagated.
| set -ex | |
| set -ex | |
| set -o pipefail |
| echo "=== Step 1: Clean up previous Ray / training processes ===" | ||
| ray stop --force 2>/dev/null || true |
There was a problem hiding this comment.
If a previous run of this script failed or was interrupted, the teacher vLLM server might still be running on port 13141 and holding GPU memory. It is recommended to clean up any existing vLLM processes at the start of the script to prevent port/GPU conflicts.
| echo "=== Step 1: Clean up previous Ray / training processes ===" | |
| ray stop --force 2>/dev/null || true | |
| echo "=== Step 1: Clean up previous Ray / training processes ===" | |
| pkill -9 -f "vllm" 2>/dev/null || true | |
| ray stop --force 2>/dev/null || true |
Summary
un-qwen3-4b-32b-opd.sh: 8×GPU colocate demo — Qwen3-4B student (GPUs 0–3) + external Qwen3-32B vLLM teacher (GPUs 4–7) on GSM8K.
un-qwen3-8b-opd-megatron.sh: Megatron-loaded teacher path; student rollout still uses vLLM.
Related to #327 (OPD E2E spike: Qwen3-4B student + Qwen3-32B teacher, +6.8 pp GSM8K lift).
Test plan