Skip to content

Add on-policy distillation example (Qwen3-4B + Qwen3-32B vLLM teacher)#328

Draft
CalvinXKY wants to merge 1 commit into
vllm-project:mainfrom
CalvinXKY:examples/on-policy-distillation-qwen3-4b-32b
Draft

Add on-policy distillation example (Qwen3-4B + Qwen3-32B vLLM teacher)#328
CalvinXKY wants to merge 1 commit into
vllm-project:mainfrom
CalvinXKY:examples/on-policy-distillation-qwen3-4b-32b

Conversation

@CalvinXKY

Copy link
Copy Markdown
Collaborator

Summary

  • Add \examples/on_policy_distillation/\ with README and two launch scripts for OPD in vime.

  • 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.
  • Document checkpoint conversion notes for Qwen3-4B (tied embeddings, padded vocab) and preliminary GSM8K results from the spike.

Related to #327 (OPD E2E spike: Qwen3-4B student + Qwen3-32B teacher, +6.8 pp GSM8K lift).

Test plan

  • Pre-commit passed locally (\pre-commit run --all-files)
  • Review README / script paths and GPU layout
  • (Optional) Run \�ash examples/on_policy_distillation/run-qwen3-4b-32b-opd.sh\ on 8×A800 with prepared checkpoints

Signed-off-by: kaiyuan <kyxiezju@163.com>
@CalvinXKY CalvinXKY force-pushed the examples/on-policy-distillation-qwen3-4b-32b branch from 7ef363f to 8441b80 Compare July 8, 2026 14:11

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment on lines +25 to +26
pkill -9 ray 2>/dev/null || true
pkill -9 python 2>/dev/null || true

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

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".

Suggested change
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

Comment on lines +152 to +153
pkill -9 ray 2>/dev/null || true
pkill -9 python 2>/dev/null || true

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

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.

Suggested change
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

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

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.

Suggested change
set -ex
set -ex
set -o pipefail

Comment on lines +37 to +38
echo "=== Step 1: Clean up previous Ray / training processes ==="
ray stop --force 2>/dev/null || true

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

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.

Suggested change
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

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant