[Spike] OPD E2E — Qwen3-4B Student + Qwen3-32B Teacher on GSM8K (vLLM teacher)

[CalvinXKY]

Summary

End-to-end spike for On-Policy Distillation (OPD) in vime:

• Student: Qwen3-4B (Megatron train + vLLM rollout, colocate)
• Teacher: Qwen3-32B (external vLLM service, --opd-type vllm)
• Task: GSM8K math reasoning
• Training: 500 rollout steps, GRPO + OPD (--use-opd --opd-kl-coef 1.0)

Goal: validate vime's OPD path (custom-rm-path → post_process_rewards → apply_opd_kl_to_advantages) and measure GSM8K accuracy lift after distillation.

---

Hardware & GPU Layout

Item	Value
GPUs	8× A800 80GB (single node)
Image	inferactinc/public:vime-vllm-latest

GPUs	Role
0–3	Student train (Megatron TP=2) + Student rollout (vLLM TP=2)
4–7	Teacher vLLM (Qwen3-32B, TP=4, port 13141)

---

Models & Data

Model	Path	Notes
Teacher	/data/nfs_87/model/Qwen3-32B	tie_word_embeddings=False
Student HF	/data/nfs_87/model/Qwen3-4B	tie_word_embeddings=True
Student torch_dist	/data/nfs_87/model/Qwen3-4B_torch_dist_padded	padded_vocab_size=152064, converted with TP=1
Train data	/data/nfs_87/data/gsm8k/train.parquet
Eval data	/data/nfs_87/data/gsm8k/test.parquet	1319 samples

Student checkpoint conversion (before training)

Qwen3-4B uses tied embeddings — do not pass --untie-embeddings-and-output-weights. Use TP=1 for conversion; training uses TP=2 at runtime.

python3 tools/convert_hf_to_torch_dist.py \
  --hf-checkpoint /data/nfs_87/model/Qwen3-4B \
  --save /data/nfs_87/model/Qwen3-4B_torch_dist_padded \
  --padded-vocab-size 152064

---

Launch Script (run-opd-qwen3-4b-32b.sh)

#!/bin/bash
set -x

export PYTHONUNBUFFERED=1

NVLINK_COUNT=$(nvidia-smi topo -m 2>/dev/null | grep -o 'NV[0-9][0-9]*' | wc -l)
if [ "$NVLINK_COUNT" -gt 0 ]; then
    HAS_NVLINK=1
else
    HAS_NVLINK=0
fi
echo "HAS_NVLINK: $HAS_NVLINK"

NUM_GPUS=8
TEACHER_GPUS=4
TEACHER_TP=4
TRAIN_GPUS=4
TEACHER_HOST=127.0.0.1
TEACHER_PORT=13141

VIME_ROOT=/root/vime
source ${VIME_ROOT}/scripts/models/qwen3-4B.sh

TEACHER_MODEL_PATH=/data/nfs_87/model/Qwen3-32B
STUDENT_HF=/data/nfs_87/model/Qwen3-4B
STUDENT_TORCH_DIST=/data/nfs_87/model/Qwen3-4B_torch_dist_padded

DATA_DIR=/data/nfs_87/data
LOG_DIR=/root/opd_logs
mkdir -p ${LOG_DIR}

echo "=== Step 1: Launch vLLM teacher server (Qwen3-32B, TP=${TEACHER_TP}) ==="
export CUDA_VISIBLE_DEVICES=4,5,6,7
python3 -m vllm.entrypoints.openai.api_server \
    --model ${TEACHER_MODEL_PATH} \
    --host 0.0.0.0 \
    --port ${TEACHER_PORT} \
    --tensor-parallel-size ${TEACHER_TP} \
    --gpu-memory-utilization 0.85 \
    --trust-remote-code \
    --dtype bfloat16 \
    --max-model-len 8192 \
    > ${LOG_DIR}/teacher_vllm.log 2>&1 &
TEACHER_PID=$!
echo "Teacher vLLM server PID: ${TEACHER_PID}"

echo "Waiting for teacher server to be ready..."
for i in $(seq 1 120); do
    if ! kill -0 ${TEACHER_PID} 2>/dev/null; then
        echo "ERROR: Teacher server process died. Check ${LOG_DIR}/teacher_vllm.log"
        exit 1
    fi
    HTTP_CODE=$(curl -s -o /dev/null -w "%{http_code}" http://${TEACHER_HOST}:${TEACHER_PORT}/health 2>/dev/null || true)
    if [ "${HTTP_CODE}" = "200" ]; then
        echo "Teacher vLLM server is ready!"
        break
    fi
    if [ $i -eq 120 ]; then
        echo "ERROR: Teacher server failed to start within 10 minutes"
        kill ${TEACHER_PID} 2>/dev/null || true
        exit 1
    fi
    sleep 5
done

echo "=== Step 2: Run OPD training (Qwen3-4B student) ==="
export CUDA_VISIBLE_DEVICES=0,1,2,3

ray stop --force 2>/dev/null || true
pkill -9 ray 2>/dev/null || true
sleep 3

ray start --head --node-ip-address 127.0.0.1 --num-gpus ${TRAIN_GPUS} --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265

echo "Waiting for Ray to be ready..."
sleep 15

for i in $(seq 1 30); do
    if curl -s http://127.0.0.1:8265/api/ray/version 2>/dev/null | grep -q .; then
        echo "Ray is ready!"
        break
    fi
    echo "Waiting for Ray... $i/30"
    sleep 2
done

CKPT_ARGS=(
    --hf-checkpoint ${STUDENT_HF}
    --ref-load ${STUDENT_TORCH_DIST}
    --load ${STUDENT_TORCH_DIST}
    --save /root/opd_checkpoints/qwen3-4b-opd
    --save-interval 50
    --megatron-to-hf-mode bridge
)

ROLLOUT_ARGS=(
    --prompt-data ${DATA_DIR}/gsm8k/train.parquet
    --input-key messages
    --label-key label
    --apply-chat-template
    --rollout-shuffle
    --rm-type math
    --num-rollout 500
    --rollout-batch-size 32
    --n-samples-per-prompt 4
    --rollout-max-response-len 4096
    --rollout-temperature 0.8
    --global-batch-size 64
)

EVAL_ARGS=(
    --eval-interval 50
    --eval-prompt-data gsm8k ${DATA_DIR}/gsm8k/test.parquet
    --n-samples-per-eval-prompt 1
    --eval-max-response-len 4096
    --eval-top-k 1
)

PERF_ARGS=(
    --tensor-model-parallel-size 2
    --sequence-parallel
    --pipeline-model-parallel-size 1
    --context-parallel-size 1
    --expert-model-parallel-size 1
    --expert-tensor-parallel-size 1
    --use-dynamic-batch-size
    --max-tokens-per-gpu 9216
)

GRPO_ARGS=(
    --advantage-estimator grpo
    --use-opd
    --opd-type vllm
    --opd-kl-coef 1.0
    --use-kl-loss
    --kl-loss-coef 0.00
    --kl-loss-type low_var_kl
    --entropy-coef 0.00
    --eps-clip 0.2
    --eps-clip-high 0.28
)

OPTIMIZER_ARGS=(
    --optimizer adam
    --lr 1e-6
    --lr-decay-style constant
    --weight-decay 0.1
    --adam-beta1 0.9
    --adam-beta2 0.98
)

VLLM_ARGS=(
    --rollout-num-gpus-per-engine 2
    --vllm-gpu-memory-utilization 0.7
    --vllm-max-num-seqs 32
    --vllm-max-cudagraph-capture-size 16
)

RM_ARGS=(
    --custom-rm-path vime.rollout.on_policy_distillation.reward_func
    --custom-reward-post-process-path vime.rollout.on_policy_distillation.post_process_rewards
    --rm-url http://${TEACHER_HOST}:${TEACHER_PORT}/inference/v1/generate
)

MISC_ARGS=(
    --attention-dropout 0.0
    --hidden-dropout 0.0
    --accumulate-allreduce-grads-in-fp32
    --attention-softmax-in-fp32
    --attention-backend flash
    --actor-num-nodes 1
    --actor-num-gpus-per-node ${TRAIN_GPUS}
    --colocate
    --make-vocab-size-divisible-by 128
)

export PYTHONPATH=${VIME_ROOT}:/root/Megatron-LM/
export CUDA_DEVICE_MAX_CONNECTIONS=1
export NCCL_NVLS_ENABLE=${HAS_NVLINK}

python3 train.py \
    --train-backend megatron \
    ${MODEL_ARGS[@]} \
    ${CKPT_ARGS[@]} \
    ${ROLLOUT_ARGS[@]} \
    ${OPTIMIZER_ARGS[@]} \
    ${GRPO_ARGS[@]} \
    ${PERF_ARGS[@]} \
    ${EVAL_ARGS[@]} \
    ${VLLM_ARGS[@]} \
    ${RM_ARGS[@]} \
    ${MISC_ARGS[@]} \
    2>&1 | tee ${LOG_DIR}/opd_training.log

echo "=== Training complete, stopping teacher server ==="
kill ${TEACHER_PID} 2>/dev/null || true
ray stop --force 2>/dev/null || true
echo "=== Done ==="

Script uses direct python3 train.py (not ray job submit). MODEL_ARGS comes from source scripts/models/qwen3-4B.sh.

---

OPD Data Flow

Student rollout (vLLM, GPU 0-3)
  → token sequence + student logprobs
Teacher vLLM (HTTP POST, GPU 4-7)
  → teacher_log_probs per token (prompt_logprobs)
post_process_rewards
  → store teacher_log_probs; scalar_rewards=[0.0]
apply_opd_kl_to_advantages
  → advantages -= opd_kl_coef * (student_logp - teacher_logp)
GRPO policy update

---

Test Procedure

1. Convert Qwen3-4B HF → torch_dist (--padded-vocab-size 152064)
2. Run launch script: start Teacher vLLM → Ray → OPD training (500 steps)
3. Export HF checkpoint from Megatron save dir
4. Evaluate on full GSM8K test (1319 samples, greedy temperature=0.0)

Eval script: eval_gsm8k.py — extract \boxed{} answer, normalize to float compare.

---

Results

Training metrics

Metric	Step 1	Step 499	Change
rollout/opd_reverse_kl	0.216	0.110	−49%
Wall time / step	~2 min	~2 min	stable

GSM8K accuracy (full test set, n=1319, greedy)

Model	Correct	Accuracy
Teacher (Qwen3-32B)	1169/1319	88.6%
Student pre-OPD (Qwen3-4B)	1039/1319	78.8%
Student post-OPD (500 steps)	1129/1319	85.6%

Metric	Value
OPD lift	+6.8 pp (78.8% → 85.6%)
Teacher–Student gap (pre)	9.8 pp
Teacher–Student gap (post)	3.0 pp
Gap closed	~70%

100-sample eval showed ±3% variance; full-set greedy eval used for reliable comparison.

---

Conclusions

1. vime OPD (--opd-type vllm) works E2E with external Teacher vLLM + colocate Student train/rollout on 8×A800.
2. 500-step OPD yields +6.8 pp GSM8K gain, closing ~70% of the 4B↔32B gap.
3. opd_reverse_kl (0.216 → 0.110) is a useful training health signal.

---

Proposed Follow-ups

• Add examples/on_policy_distillation/run-qwen3-4b-32b-opd.sh + README
• Document Qwen3-4B checkpoint conversion (tie/untie, padded vocab, TP=1 convert)
• Optional: sweep --opd-kl-coef (0.5 / 1.0 / 2.0), longer training (>500 steps)
• Optional: compare --opd-type megatron vs vllm on same student/teacher pair

[CalvinXKY]