Update qwen3-4B example(reward converges, colocate + non-colocate.)

[CalvinXKY]

docs(qwen3-4B): migrate example to vLLM and report e2e smoke (issue #11)

Summary

Progress toward issue #11 Qwen3-4B PPO (dense, baseline):

• Migrate the official Qwen3-4B example from SGLang to vLLM-only (strategy B: update scripts/run-qwen3-4B.sh in place).
• Update docs/zh/examples/qwen3-4B.md and docs/en/examples/qwen3-4B.md to match the vLLM runtime on main (no --rollout-backend; Ray owns GPU placement via --actor-num-gpus-per-node / --rollout-num-gpus).
• Document colocate (CUDA IPC) vs decoupled (NCCL in-process) weight sync behavior.

This PR covers documentation + e2e smoke (load ckpt → rollout → weight sync → train step). Full 3000-rollout reward convergence is tracked separately.

Code changes

File	Change
scripts/run-qwen3-4B.sh	vLLM colocate baseline: --rollout-num-gpus-per-engine 2, --vllm-gpu-memory-utilization 0.7, --train-memory-margin-bytes, PYTHONPATH includes repo root for colocate worker extension
docs/zh/examples/qwen3-4B.md	SGLANG_ARGS → VLLM_ARGS; decoupled/colocate sections updated
docs/en/examples/qwen3-4B.md	Same as zh

Test setup

Item	Value
Hardware	8× A100 (80GB), single node
Container	vime on A100-server, NFS code at /data/nfs_87/xky/RL/
Model / data	Qwen3-4B, dapo-math-17k, aime-2024
Branch for decoupled smoke	main (4c74767) — uses upstream UpdateWeightFromDistributed + NCCL in-process weight transfer
Branch for colocate smoke	test/with_ipc — IPC colocate path from #22 is not merged to main yet; colocate e2e was validated on this branch

Note: Colocate IPC implementation lives in #22 / test/with_ipc. After #22 lands, the same smoke commands should pass on main without switching branches.

E2E test scripts (manual smoke, outside this PR diff)

Scripts live in the RL workspace (run_scripts/), aligned with scripts/run-qwen3-4B.sh hyperparameters; default NUM_ROLLOUT=5 for fast validation.

1) Decoupled — 2 train + 6 rollout (main)

cd /data/nfs_87/xky/RL
git -C vime checkout main   # weight sync: NCCL in-process (update_weight_from_distributed)

CONCURRENT_RAY=0 NUM_ROLLOUT=5 bash run_scripts/qwen3-4B-vllm-verify.sh

Topology: --actor-num-gpus-per-node 2, --rollout-num-gpus 6, --rollout-num-gpus-per-engine 2 (3 vLLM engines).

Key log output (2026-05-23, train_qwen3_4b_vllm_verify_20260523_012819.log):

(MegatronTrainRayActor) update_weight_from_distributed.py:
  vLLM in-process weight transfer: addr=10.155.68.38 port=46403 world_size=7

(VLLMEngine) POST /init_weight_transfer_engine HTTP/1.1" 200 OK
(VLLMEngine) POST /start_weight_update HTTP/1.1" 200 OK
(VLLMEngine) POST /update_weights HTTP/1.1" 200 OK          [repeated ~45x across cluster]
(VLLMEngine) POST /finish_weight_update HTTP/1.1" 200 OK
(VLLMEngine) POST /resume HTTP/1.1" 200 OK

(MegatronTrainRayActor) Timer update_weights end (elapsed: ~20s)

Job 'raysubmit_8meS3UJ73QjjYBcZ' succeeded




### 2) Colocate — 8 GPUs shared (`test/with_ipc`)

```bash
cd /data/nfs_87/xky/RL
git -C vime checkout test/with_ipc

CONCURRENT_RAY=0 NUM_ROLLOUT=5 bash run_scripts/qwen3-4B-vllm-verify-colocate.sh

Topology: --colocate, --actor-num-gpus-per-node 8, 4 engines × TP=2.

Key log output (2026-05-23, verify_colocate_retry.log):

(VLLMEngine) non-default args: {
  ...
  'worker_extension_cls': 'slime.backends.megatron_utils.update_weight.update_weight_from_tensor.vLLMColocateWorkerExtension',
  'weight_transfer_config': WeightTransferConfig(backend='ipc')
}

(MegatronTrainRayActor) Timer update_weights end (elapsed: 7.0s)
(VLLMEngine) POST /update_weights HTTP/1.1" 200 OK          [repeated ~63x across cluster]
(VLLMEngine) POST /finish_weight_update HTTP/1.1" 200 OK
(VLLMEngine) POST /wake_up?tags=kv_cache HTTP/1.1" 200 OK

Job 'raysubmit_j9EcunSqa5aJDeHV' succeeded

Unit test on test/with_ipc (container):

cd /data/nfs_87/xky/RL/vime
python -m pytest tests/unit/backends/megatron_utils/update_weight/test_update_weight_from_tensor.py -q
# 6 passed in 1.69s

---

3) Colocate long-run smoke (test/with_ipc, optional)

CONCURRENT_RAY=0 START_FROM_REF=1 NUM_ROLLOUT=3000 \
  bash run_scripts/qwen3-4B-vllm-colocate-long.sh

First rollout completed after cold start from Qwen3-4B_torch_dist (avoid lr_decay_steps mismatch when resuming short-verify ckpt):

LOAD_DIR=/data/nfs_87/xky/models/Qwen3-4B_torch_dist
Timer update_weights end
Rollout generation: 0%|          | 0/256 [00:00<?, ?it/s]

Issue #11 checklist mapping

Acceptance (issue #11)	Status in this PR
Docs: docs/zh/examples/qwen3-4B.md → vLLM	✅
Script: scripts/run-qwen3-4B.sh → vLLM colocate baseline	✅
E2E smoke — non-colocate	✅ on main
E2E smoke — colocate	✅ on test/with_ipc (blocked on main until #22)
tests/test_qwen3_4B_ppo.py on vLLM + reward convergence	⏳ follow-up

Related

• Closes / contributes to: [RFC] VIME Roadmap #11 (Qwen3-4B PPO dense baseline — docs + smoke)
• Colocate IPC dependency: [feat] Support Colocated Weight Sync via CUDA IPC for vime #22 (open)
• Smoke branches: main (decoupled), test/with_ipc (colocate)

Test plan

• Read through updated zh/en qwen3-4B docs
• Decoupled smoke on A100 (main)
• Colocate smoke on A100 (test/with_ipc)
• Colocate IPC unit tests (6 passed)
• Wire tests/test_qwen3_4B_ppo.py to vLLM in CI

[gemini-code-assist]

Code Review

This pull request migrates the rollout backend from SGLang to vLLM for the Qwen3-4B example, updating the documentation and execution scripts. The changes include renaming configuration variables to VLLM_ARGS, adding memory management parameters, and providing detailed explanations for co-located and decoupled training modes. A review comment pointed out that the Chinese documentation is inconsistent with the English version, as it retains outdated SGLang references and lacks a complete example for decoupled training configuration.

[gemini-code-assist]

The Chinese documentation in this section is inconsistent with the English version and contains outdated information. It incorrectly mentions "2 ways" (2 种方式) to adjust concurrency but only provides one, and it refers to a "160" limit which was specific to SGLang's CUDA graph implementation.

To align with the English version, this section should be updated to provide an example of VLLM_ARGS for decoupled training and inference. For example:

对于训推分离模式，VLLM_ARGS 仅需配置推理后端参数，例如：

VLLM_ARGS=(
   --rollout-num-gpus-per-engine 2
   --vllm-gpu-memory-utilization 0.9
)

[CalvinXKY]

the zh decoupled section now matches en — removed the old SGLang cuda-graph notes and added the Ray resource + VLLM_ARGS example.

[CalvinXKY]

non colocate:

colocate:

[CalvinXKY]

PPO test:

CONCURRENT_RAY="${CONCURRENT_RAY:-1}"
export RAY_TMPDIR="${RAY_TMPDIR:-/tmp/ray_vime_4b_vllm_ppo_colocate}"
export RAY_PORT="${RAY_PORT:-6385}"
export RAY_DASHBOARD_PORT="${RAY_DASHBOARD_PORT:-8271}"
RAY_AGENT_PORT="${RAY_AGENT_PORT:-52381}"
RAY_ADDR="http://127.0.0.1:${RAY_DASHBOARD_PORT}"
unset RAY_ADDRESS RAY_REDIS_ADDRESS || true
mkdir -p "${RAY_TMPDIR}"

if [ "${CONCURRENT_RAY}" = "1" ]; then
  ray stop --force >/dev/null 2>&1 || true
  rm -rf "${RAY_TMPDIR}"/session_* "${RAY_TMPDIR}"/ray_current_cluster 2>/dev/null || true
  sleep 2
else
  pkill -9 sglang 2>/dev/null || true
  pkill -9 vllm 2>/dev/null || true
  sleep 3
  ray stop --force 2>/dev/null || true
  pkill -9 ray 2>/dev/null || true
  pkill -9 python 2>/dev/null || true
  sleep 3
  pkill -9 ray 2>/dev/null || true
  pkill -9 python 2>/dev/null || true
  if command -v nvidia-smi >/dev/null 2>&1; then
    nvidia-smi --query-compute-apps=pid --format=csv,noheader 2>/dev/null | xargs -r kill -9 2>/dev/null || true
  fi
  rm -rf /tmp/ray/session_* /tmp/ray/* 2>/dev/null || true
  rm -rf "${RAY_TMPDIR}"/session_* "${RAY_TMPDIR}"/ray_current_cluster 2>/dev/null || true
  sleep 2
fi

set -ex

export PYTHONUNBUFFERED=1
LOG_ROOT="${LOG_ROOT:-/data/nfs_87/xky/logs}"
TS="$(date +%Y%m%d_%H%M%S)"
export TENSORBOARD_DIR="${TENSORBOARD_DIR:-${LOG_ROOT}/tb_qwen3_4b_vllm_ppo_colocate_${TS}}"
LOG_FILE="${LOG_FILE:-${LOG_ROOT}/train_qwen3_4b_vllm_ppo_colocate_${TS}.log}"
mkdir -p "${TENSORBOARD_DIR}" "${LOG_ROOT}"

export PYTHONBUFFERED=16

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 (detected $NVLINK_COUNT NVLink references)"

if command -v nvidia-smi >/dev/null 2>&1; then
    DETECTED_GPUS=$(nvidia-smi -L 2>/dev/null | wc -l | tr -d ' ')
else
    DETECTED_GPUS=0
fi
NUM_GPUS=${NUM_GPUS:-${DETECTED_GPUS}}
if [ -z "$NUM_GPUS" ] || [ "$NUM_GPUS" -le 0 ]; then
    NUM_GPUS=8
fi
ROLLOUT_GPUS_PER_ENGINE="${ROLLOUT_GPUS_PER_ENGINE:-2}"
echo "NUM_GPUS (colocate PPO): ${NUM_GPUS}  rollout_num_gpus_per_engine=${ROLLOUT_GPUS_PER_ENGINE}"

REPO_ROOT="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")/.." &>/dev/null && pwd)"
VIME_ROOT="${VIME_ROOT:-${REPO_ROOT}/vime}"
cd "${VIME_ROOT}"

SCRIPT_DIR="${VIME_ROOT}/scripts"
source "${SCRIPT_DIR}/models/qwen3-4B.sh"

MEGATRON_CONFIG="${MEGATRON_CONFIG:-${REPO_ROOT}/run_scripts/configs/qwen3-4B-ppo-megatron.yaml}"

HF_CKPT="${HF_CKPT:-/data/nfs_87/xky/models/Qwen3-4B}"
REF_LOAD="${REF_LOAD:-/data/nfs_87/xky/models/Qwen3-4B_torch_dist}"
LOAD_DIR="${LOAD_DIR:-/data/nfs_87/xky/RL/vime_checkpoints/Qwen3-4B_vllm_ppo_colocate}"
SAVE_DIR="${SAVE_DIR:-/data/nfs_87/xky/RL/vime_checkpoints/Qwen3-4B_vllm_ppo_colocate}"
mkdir -p "${LOAD_DIR}" "${SAVE_DIR}"

CKPT_ARGS=(
   --hf-checkpoint "${HF_CKPT}"
   --ref-load "${REF_LOAD}"
   --load "${LOAD_DIR}"
   --save "${SAVE_DIR}"
   --save-interval 20
)

PROMPT_DATA="${PROMPT_DATA:-/data/nfs_87/xky/datasets/dapo-math-17k/dapo-math-17k.jsonl}"
# 默认与 test_qwen3_4B_ppo.py 一致（短跑）；可与已跑通的 GRPO verify 一样加大
NUM_ROLLOUT="${NUM_ROLLOUT:-3}"
ROLLOUT_BATCH_SIZE="${ROLLOUT_BATCH_SIZE:-8}"
N_SAMPLES_PER_PROMPT="${N_SAMPLES_PER_PROMPT:-4}"
GLOBAL_BATCH_SIZE="${GLOBAL_BATCH_SIZE:-32}"
ROLLOUT_TEMPERATURE="${ROLLOUT_TEMPERATURE:-0.8}"

ROLLOUT_ARGS=(
   --prompt-data "${PROMPT_DATA}"
   --input-key prompt
   --label-key label
   --apply-chat-template
   --rollout-shuffle
   --rm-type deepscaler
   --num-rollout "${NUM_ROLLOUT}"
   --rollout-batch-size "${ROLLOUT_BATCH_SIZE}"
   --n-samples-per-prompt "${N_SAMPLES_PER_PROMPT}"
   --rollout-max-response-len 8192
   --rollout-temperature "${ROLLOUT_TEMPERATURE}"
   --num-steps-per-rollout 1
   --global-batch-size "${GLOBAL_BATCH_SIZE}"
   --balance-data
)

ENABLE_EVAL="${ENABLE_EVAL:-1}"
EVAL_DATA="${EVAL_DATA:-/data/nfs_87/xky/datasets/aime-2024/aime-2024.jsonl}"
EVAL_INTERVAL="${EVAL_INTERVAL:-20}"
EVAL_ARGS=()
if [ "${ENABLE_EVAL}" = "1" ]; then
  EVAL_ARGS=(
     --eval-interval "${EVAL_INTERVAL}"
     --eval-prompt-data aime "${EVAL_DATA}"
     --n-samples-per-eval-prompt 1
     --eval-max-response-len 16384
     --eval-top-k 1
  )
fi

TIGHT_HOST_MEMORY="${TIGHT_HOST_MEMORY:-0}"
if [ "${TIGHT_HOST_MEMORY}" = "1" ]; then
  MAX_TOKENS_PER_GPU=2048
else
  MAX_TOKENS_PER_GPU=9216
fi

PERF_ARGS=(
   --tensor-model-parallel-size 2
   --sequence-parallel
   --pipeline-model-parallel-size 1
   --context-parallel-size 2
   --expert-model-parallel-size 1
   --expert-tensor-parallel-size 1

   --recompute-granularity full
   --recompute-method uniform
   --recompute-num-layers 1

   --use-dynamic-batch-size
   --max-tokens-per-gpu "${MAX_TOKENS_PER_GPU}"
)

PPO_ARGS=(
   --advantage-estimator ppo
   --kl-loss-coef 0.00
   --kl-loss-type k1
   --kl-coef 0.00
   --entropy-coef 0.00
   --eps-clip 4e-4
   --num-critic-only-steps 1
   --normalize-advantages
)
if [ "${TIGHT_HOST_MEMORY}" != "1" ] && [ "${USE_KL_LOSS:-1}" = "1" ]; then
  PPO_ARGS+=(--use-kl-loss)
fi

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

WANDB_ARGS=(
   # --use-wandb
   # --wandb-project slime-dev
   # --wandb-group qwen3-4B-vllm-ppo-colocate
)

VLLM_ARGS=(
   --rollout-num-gpus-per-engine "${ROLLOUT_GPUS_PER_ENGINE}"
   --rollout-backend vllm
   --vllm-gpu-memory-utilization 0.7
)
export SLIME_VLLM_SERVER_HEALTH_TIMEOUT_SEC="${SLIME_VLLM_SERVER_HEALTH_TIMEOUT_SEC:-900}"

MISC_ARGS=(
   --attention-dropout 0.0
   --hidden-dropout 0.0
   --accumulate-allreduce-grads-in-fp32
   --attention-softmax-in-fp32
   --attention-backend flash
   --train-memory-margin-bytes 2147483648
   --use-tensorboard
)

VIME_PYTHONPATH="${VIME_ROOT}:/root/Megatron-LM/"

export MASTER_ADDR="${MASTER_ADDR:-127.0.0.1}"
echo "=== Qwen3-4B vLLM PPO colocate verify ==="
echo "MEGATRON_CONFIG=${MEGATRON_CONFIG}"
echo "NUM_ROLLOUT=${NUM_ROLLOUT} ENABLE_EVAL=${ENABLE_EVAL}"
echo "Ray: port=${RAY_PORT} dashboard=${RAY_DASHBOARD_PORT}"
echo "LOG_FILE=${LOG_FILE}"

ray start --head \
  --node-ip-address "${MASTER_ADDR}" \
  --port "${RAY_PORT}" \
  --num-gpus "${NUM_GPUS}" \
  --disable-usage-stats \
  --dashboard-host=0.0.0.0 \
  --dashboard-port="${RAY_DASHBOARD_PORT}" \
  --dashboard-agent-listen-port="${RAY_AGENT_PORT}" \
  --temp-dir "${RAY_TMPDIR}"

RUNTIME_ENV_JSON="{
  \"env_vars\": {
    \"PYTHONPATH\": \"${VIME_PYTHONPATH}\",
    \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\",
    \"NCCL_NVLS_ENABLE\": \"${HAS_NVLINK}\",
    \"TENSORBOARD_DIR\": \"${TENSORBOARD_DIR}\",
    \"SLIME_VLLM_SERVER_HEALTH_TIMEOUT_SEC\": \"${SLIME_VLLM_SERVER_HEALTH_TIMEOUT_SEC}\"
  }
}"

ray job submit --address="${RAY_ADDR}" \
   --runtime-env-json="${RUNTIME_ENV_JSON}" \
   -- python3 train.py \
   --train-backend megatron \
   --megatron-config-path "${MEGATRON_CONFIG}" \
   --actor-num-nodes 1 \
   --actor-num-gpus-per-node "${NUM_GPUS}" \
   --rollout-num-gpus "${NUM_GPUS}" \
   --colocate \
   ${MODEL_ARGS[@]} \
   ${CKPT_ARGS[@]} \
   ${ROLLOUT_ARGS[@]} \
   ${OPTIMIZER_ARGS[@]} \
   ${PPO_ARGS[@]} \
   ${WANDB_ARGS[@]} \
   ${PERF_ARGS[@]} \
   ${EVAL_ARGS[@]} \
   ${VLLM_ARGS[@]} \
   ${MISC_ARGS[@]} \
   2>&1 | tee -a "${LOG_FILE}"

[aoshen02]

pkill -9 vllm likely matches nothing?

[CalvinXKY]

Good catch, removed pkill -9 vllm. vLLM workers are Ray/python processes, so rerun cleanup now uses ray stop --force + pkill ray/python, same as the original script.

[aoshen02]

We can add some similar things like --vllm-max-num-seqs, --vllm-max-num-batched-tokens, --vllm-enforce-eage

[CalvinXKY]

Done

[CalvinXKY]

The iter result GRPO:

PPO: