@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
sh 78 lines 2.51 KB
Raw
1 #!/usr/bin/env bash
2 set -euxo pipefail
3 # Nightly e2e test for prefetch offloading with a MoE model.
4 # Runs DeepSeek-V2-Lite with prefetch offloading of MoE expert weights
5 # and validates GSM8K accuracy matches baseline (no offloading).
6 #
7 # args: [THRESHOLD] [NUM_QUESTIONS] [START_PORT]
8 #
9 # Environment variables:
10 # ATTENTION_BACKEND - attention backend to use (e.g., FLASH_ATTN,
11 # ROCM_ATTN, FLASHINFER). If unset, uses vllm default.
12 THRESHOLD=${1:-0.25}
13 NUM_Q=${2:-1319}
14 PORT=${3:-8030}
15 OUT_DIR=${OUT_DIR:-/tmp/vllm-scheduled}
16 mkdir -p "${OUT_DIR}"
17
18 wait_for_server() {
19 local port=$1
20 timeout 600 bash -c '
21 until curl -sf "http://127.0.0.1:'"$port"'/health" > /dev/null; do
22 sleep 1
23 done'
24 }
25
26 MODEL="deepseek-ai/DeepSeek-V2-Lite"
27
28 # ── Build optional vllm serve flags ─────────────────────────────────────
29
30 EXTRA_ARGS=()
31 if [[ -n "${ATTENTION_BACKEND:-}" ]]; then
32 echo "Using attention backend: ${ATTENTION_BACKEND}"
33 EXTRA_ARGS+=(--attention-backend "${ATTENTION_BACKEND}")
34 fi
35
36 # ROCm: run eager to avoid intermittent HIP-graph decode corruption.
37 # See https://github.com/ROCm/clr/issues/279
38 # TODO(aarushjain29): Revert after TheRock 7.14
39 if command -v rocm-smi &> /dev/null || command -v amd-smi &> /dev/null || [[ -d /opt/rocm ]] || [[ -n "${ROCM_PATH:-}" ]]; then
40 echo "ROCm platform detected: adding --enforce-eager to avoid HIP-graph decode corruption"
41 EXTRA_ARGS+=(--enforce-eager)
42 fi
43
44 cleanup() {
45 if [[ -n "${SERVER_PID:-}" ]] && kill -0 "${SERVER_PID}" 2>/dev/null; then
46 kill "${SERVER_PID}" 2>/dev/null || true
47 for _ in {1..20}; do
48 kill -0 "${SERVER_PID}" 2>/dev/null || break
49 sleep 0.5
50 done
51 kill -9 "${SERVER_PID}" 2>/dev/null || true
52 fi
53 }
54 trap cleanup EXIT
55
56 vllm serve "$MODEL" \
57 --max-model-len 2048 \
58 --offload-group-size 8 \
59 --offload-num-in-group 2 \
60 --offload-prefetch-step 1 \
61 --offload-params w13_weight w2_weight \
62 --generation-config vllm \
63 --port "$PORT" \
64 ${EXTRA_ARGS+"${EXTRA_ARGS[@]}"} &
65 SERVER_PID=$!
66 wait_for_server "$PORT"
67
68 TAG=$(echo "$MODEL" | tr '/: \\n' '_____')
69 OUT="${OUT_DIR}/${TAG}_prefetch_offload.json"
70 python3 tests/evals/gsm8k/gsm8k_eval.py --host http://127.0.0.1 --port "$PORT" --num-questions "${NUM_Q}" --save-results "${OUT}"
71 python3 - <<PY
72 import json; acc=json.load(open('${OUT}'))['accuracy']
73 print(f"${MODEL} prefetch_offload: accuracy {acc:.3f}")
74 assert acc >= ${THRESHOLD}, f"${MODEL} prefetch_offload accuracy {acc}"
75 PY
76
77 cleanup
78 SERVER_PID=