Skip to content

Commit d2481d5

Browse files
author
Aegis AI Assistant
committed
test: isolate speculative efficiency in separate parallel eval run
1 parent 5a8bf24 commit d2481d5

3 files changed

Lines changed: 296 additions & 1 deletion

File tree

‎.github/workflows/ci.yml‎

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -177,3 +177,62 @@ jobs:
177177
path: /tmp/SwiftLM-test-speculative.log
178178
retention-days: 7
179179

180+
# ── Speculative Decoding Memory Evaluation ──
181+
# Runs the 9B model with NUM_DRAFT_TOKENS=2 to check peak
182+
# memory compression/efficiency. Allowed to OOM/fail.
183+
speculative-decoding-eval:
184+
runs-on: macos-15
185+
timeout-minutes: 45
186+
needs: ci
187+
continue-on-error: true
188+
steps:
189+
- uses: actions/checkout@v4
190+
with:
191+
submodules: recursive
192+
193+
- name: Setup Swift
194+
uses: swift-actions/setup-swift@v2
195+
with:
196+
swift-version: "6.0"
197+
198+
- name: Download SwiftLM executable
199+
uses: actions/download-artifact@v4
200+
with:
201+
name: SwiftLM-macOS
202+
path: .build/release/
203+
204+
- name: Make executable
205+
run: chmod +x .build/release/SwiftLM
206+
207+
- name: Install MLX Metal library
208+
run: |
209+
python3 -m venv /tmp/mlx_venv
210+
/tmp/mlx_venv/bin/pip install --quiet mlx
211+
cp /tmp/mlx_venv/lib/python*/site-packages/mlx/lib/mlx.metallib .build/release/
212+
213+
- name: Run speculative evaluation E2E
214+
env:
215+
HF_HUB_DOWNLOAD_TIMEOUT: "900"
216+
run: |
217+
chmod +x tests/test-speculative-eval.sh
218+
for attempt in 1 2 3; do
219+
echo "Attempt $attempt of 3..."
220+
if tests/test-speculative-eval.sh .build/release/SwiftLM 15414; then
221+
exit 0
222+
fi
223+
if [ "$attempt" -lt 3 ]; then
224+
echo "Test failed, retrying in 10s..."
225+
sleep 10
226+
fi
227+
done
228+
echo "All attempts failed"
229+
exit 1
230+
231+
- name: Upload speculative eval logs on failure
232+
if: failure()
233+
uses: actions/upload-artifact@v4
234+
with:
235+
name: speculative-eval-logs
236+
path: /tmp/SwiftLM-test-speculative-eval.log
237+
retention-days: 7
238+

‎tests/test-speculative-eval.sh‎

Lines changed: 236 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,236 @@
1+
#!/bin/bash
2+
# test-speculative.sh — Speculative decoding E2E verification
3+
#
4+
# Uses a small draft model (Qwen3.5-0.8B) to accelerate a larger main model
5+
# (Qwen3.5-4B) via speculative decoding. Verifies:
6+
# 1. Dual-model loading (draft + main)
7+
# 2. Speculative decoding path activation
8+
# 3. Correct token generation
9+
# 4. Server stability under dual-model memory pressure
10+
#
11+
# Usage:
12+
# ./tests/test-speculative.sh [binary_path] [port]
13+
#
14+
# Requirements:
15+
# - ~4 GB RAM (0.8B draft ~1 GB + 4B main ~3 GB)
16+
# - macos-15 (7 GB) on GitHub Actions is sufficient
17+
# - curl, jq
18+
19+
set -euo pipefail
20+
21+
BINARY="${1:-.build/release/SwiftLM}"
22+
PORT="${2:-15414}"
23+
HOST="127.0.0.1"
24+
MAIN_MODEL="${MAIN_MODEL:-mlx-community/Qwen3.5-9B-4bit}"
25+
DRAFT_MODEL="${DRAFT_MODEL:-mlx-community/Qwen3.5-0.8B-MLX-4bit}"
26+
NUM_DRAFT_TOKENS=2
27+
URL="http://${HOST}:${PORT}"
28+
PASS=0
29+
FAIL=0
30+
TOTAL=0
31+
LOG_FILE="/tmp/SwiftLM-test-speculative-eval.log"
32+
33+
# Colors
34+
GREEN='\033[0;32m'
35+
RED='\033[0;31m'
36+
YELLOW='\033[1;33m'
37+
CYAN='\033[0;36m'
38+
NC='\033[0m'
39+
40+
log() { echo -e "${YELLOW}[spec-test]${NC} $*"; }
41+
pass() { PASS=$((PASS + 1)); TOTAL=$((TOTAL + 1)); echo -e " ${GREEN}✅ PASS${NC}: $*"; }
42+
fail() { FAIL=$((FAIL + 1)); TOTAL=$((TOTAL + 1)); echo -e " ${RED}❌ FAIL${NC}: $*"; }
43+
44+
cleanup() {
45+
if [ -n "${SERVER_PID:-}" ]; then
46+
log "Stopping server (PID $SERVER_PID)"
47+
kill -9 "$SERVER_PID" 2>/dev/null || true
48+
wait "$SERVER_PID" 2>/dev/null || true
49+
fi
50+
}
51+
trap cleanup EXIT
52+
53+
# ── Check prerequisites ─────────────────────────────────────────────
54+
if [ ! -f "$BINARY" ]; then
55+
echo "Error: Binary not found at $BINARY"
56+
echo "Run 'swift build -c release' first."
57+
exit 1
58+
fi
59+
60+
if ! command -v jq &>/dev/null; then
61+
echo "Error: jq is required. Install with: brew install jq"
62+
exit 1
63+
fi
64+
65+
# ── Memory check ────────────────────────────────────────────────────
66+
TOTAL_RAM_GB=$(sysctl -n hw.memsize 2>/dev/null | awk '{printf "%.0f", $1 / 1073741824}')
67+
log "System RAM: ${TOTAL_RAM_GB} GB"
68+
69+
if [ "$TOTAL_RAM_GB" -lt 8 ] 2>/dev/null; then
70+
log "⚠️ WARNING: ${TOTAL_RAM_GB} GB RAM detected. Dual-model test requires ~6 GB."
71+
log " Consider running on a machine with ≥8 GB RAM."
72+
fi
73+
74+
# ══════════════════════════════════════════════════════════════════════
75+
echo -e "\n${CYAN}╔══════════════════════════════════════════════════════════╗${NC}"
76+
echo -e "${CYAN}║ SwiftLM Speculative Decoding Eval Test ║${NC}"
77+
echo -e "${CYAN}║ Draft: Qwen3.5-0.8B (4-bit) → Main: Qwen3.5-9B (4-bit) ║${NC}"
78+
echo -e "${CYAN}║ Draft tokens per round: ${NUM_DRAFT_TOKENS} ║${NC}"
79+
echo -e "${CYAN}╚══════════════════════════════════════════════════════════╝${NC}\n"
80+
81+
# ── Start server with dual models ───────────────────────────────────
82+
log "Starting server with speculative decoding..."
83+
log " Main model: $MAIN_MODEL"
84+
log " Draft model: $DRAFT_MODEL"
85+
log " Draft tokens per round: $NUM_DRAFT_TOKENS"
86+
87+
"$BINARY" --model "$MAIN_MODEL" --port "$PORT" --host "$HOST" \
88+
--draft-model "$DRAFT_MODEL" \
89+
--num-draft-tokens "$NUM_DRAFT_TOKENS" \
90+
> "$LOG_FILE" 2>&1 &
91+
SERVER_PID=$!
92+
93+
# Wait for server to be ready (both models need to download + load)
94+
log "Waiting for server to load both models (this may take a while on first run)..."
95+
MAX_WAIT=900 # 15 minutes for two model downloads
96+
for i in $(seq 1 "$MAX_WAIT"); do
97+
if curl -sf "$URL/health" >/dev/null 2>&1; then
98+
log "Server ready after ${i}s"
99+
break
100+
fi
101+
if ! kill -0 "$SERVER_PID" 2>/dev/null; then
102+
echo "Error: Server process died. Last 30 lines of log:"
103+
tail -30 "$LOG_FILE"
104+
exit 1
105+
fi
106+
# Print progress every 30 seconds
107+
if [ $((i % 30)) -eq 0 ]; then
108+
log " Still waiting... (${i}s elapsed)"
109+
fi
110+
sleep 1
111+
done
112+
113+
if ! curl -sf "$URL/health" >/dev/null 2>&1; then
114+
echo "Error: Server did not become ready in ${MAX_WAIT}s"
115+
echo "Last 30 lines of log:"
116+
tail -30 "$LOG_FILE"
117+
exit 1
118+
fi
119+
120+
# ── Test 1: Verify server loaded both models ────────────────────────
121+
log "Test 1: Verify dual-model loading"
122+
123+
# Check server log for draft model loading confirmation
124+
if grep -q "Draft model loaded successfully" "$LOG_FILE"; then
125+
pass "Draft model loaded successfully"
126+
else
127+
fail "Draft model loading not confirmed in server logs"
128+
fi
129+
130+
if grep -q "speculative decoding" "$LOG_FILE"; then
131+
pass "Speculative decoding mode detected in server logs"
132+
else
133+
fail "Speculative decoding not mentioned in server logs"
134+
fi
135+
136+
# ── Test 2: Health endpoint works with dual models ──────────────────
137+
log "Test 2: Health endpoint"
138+
139+
HEALTH=$(curl -sf "$URL/health")
140+
if echo "$HEALTH" | jq -e '.status == "ok"' >/dev/null 2>&1; then
141+
pass "Health endpoint returns status=ok"
142+
else
143+
fail "Health endpoint: $HEALTH"
144+
fi
145+
146+
# ── Test 3: Streaming speculative generation ────────────────────────
147+
log "Test 3: Streaming speculative generation"
148+
149+
STREAM_OUTPUT=$(curl -sf -N --max-time 120 -X POST "$URL/v1/chat/completions" \
150+
-H "Content-Type: application/json" \
151+
-d "{\"model\":\"$MAIN_MODEL\",\"stream\":true,\"max_tokens\":30,\"messages\":[{\"role\":\"user\",\"content\":\"Name three fruits.\"}]}" \
152+
2>/dev/null || true)
153+
154+
if echo "$STREAM_OUTPUT" | grep -q "data: \[DONE\]"; then
155+
pass "Streaming speculative: received [DONE] sentinel"
156+
else
157+
fail "Streaming speculative: missing [DONE] sentinel"
158+
fi
159+
160+
CHUNK_COUNT=$(echo "$STREAM_OUTPUT" | grep -c "^data: {" || true)
161+
if [ "$CHUNK_COUNT" -gt 0 ]; then
162+
pass "Streaming speculative: received $CHUNK_COUNT data chunks"
163+
else
164+
fail "Streaming speculative: no data chunks received"
165+
fi
166+
167+
# Check server log for speculative decoding activation
168+
if grep -q "Using speculative decoding" "$LOG_FILE"; then
169+
pass "Speculative decoding path activated during generation"
170+
else
171+
fail "Speculative decoding path not activated (missing log line)"
172+
fi
173+
174+
# ── Test 5: Multiple sequential requests (stability) ────────────────
175+
log "Test 5: Sequential request stability (3 requests)"
176+
177+
SEQ_PASS=true
178+
for i in 1 2 3; do
179+
SEQ_RESP=$(curl -sf --max-time 120 -X POST "$URL/v1/chat/completions" \
180+
-H "Content-Type: application/json" \
181+
-d "{\"model\":\"$MAIN_MODEL\",\"max_tokens\":10,\"messages\":[{\"role\":\"user\",\"content\":\"Say the number $i.\"}]}" 2>/dev/null || echo "")
182+
183+
SEQ_CONTENT=$(echo "$SEQ_RESP" | jq -r '.choices[0].message.content // empty' 2>/dev/null || echo "")
184+
185+
if [ -z "$SEQ_CONTENT" ]; then
186+
SEQ_PASS=false
187+
fail "Sequential request $i: empty response"
188+
break
189+
fi
190+
done
191+
192+
if [ "$SEQ_PASS" = true ]; then
193+
pass "Sequential stability: 3/3 speculative requests completed successfully"
194+
fi
195+
196+
# ── Test 6: Memory stability check ─────────────────────────────────
197+
log "Test 6: Memory stability"
198+
199+
HEALTH_FINAL=$(curl -sf "$URL/health")
200+
MEM_ACTIVE=$(echo "$HEALTH_FINAL" | jq -r '.memory.active_mb // 0')
201+
MEM_PEAK=$(echo "$HEALTH_FINAL" | jq -r '.memory.peak_mb // 0')
202+
203+
if [ "$MEM_ACTIVE" -gt 0 ] 2>/dev/null; then
204+
pass "Memory: active=${MEM_ACTIVE} MB, peak=${MEM_PEAK} MB"
205+
else
206+
fail "Memory: could not read memory stats"
207+
fi
208+
209+
# Verify server is still responsive after all tests
210+
if curl -sf "$URL/health" >/dev/null 2>&1; then
211+
pass "Server still responsive after all speculative decoding tests"
212+
else
213+
fail "Server became unresponsive"
214+
fi
215+
216+
# ── Results ──────────────────────────────────────────────────────────
217+
echo ""
218+
log "═══════════════════════════════════════"
219+
log "Speculative Decoding Test Results"
220+
log " Draft: $DRAFT_MODEL"
221+
log " Main: $MAIN_MODEL"
222+
log " Tokens/round: $NUM_DRAFT_TOKENS"
223+
log " Results: ${PASS} passed, ${FAIL} failed, ${TOTAL} total"
224+
log "═══════════════════════════════════════"
225+
226+
if [ "$FAIL" -gt 0 ]; then
227+
echo ""
228+
log "Server log tail (last 20 lines):"
229+
tail -20 "$LOG_FILE"
230+
exit 1
231+
fi
232+
233+
echo ""
234+
log "Server log tail (last 20 lines):"
235+
tail -20 "$LOG_FILE"
236+
exit 0

‎tests/test-speculative.sh‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -74,7 +74,7 @@ fi
7474
# ══════════════════════════════════════════════════════════════════════
7575
echo -e "\n${CYAN}╔══════════════════════════════════════════════════════════╗${NC}"
7676
echo -e "${CYAN}║ SwiftLM Speculative Decoding E2E Test ║${NC}"
77-
echo -e "${CYAN}║ Draft: Qwen2.5-0.5B (4-bit) → Main: Qwen2.5-3B (4-bit) ║${NC}"
77+
echo -e "${CYAN}║ Draft: Qwen3.5-0.8B (4-bit) → Main: Qwen3.5-2B (4-bit) ║${NC}"
7878
echo -e "${CYAN}║ Draft tokens per round: ${NUM_DRAFT_TOKENS} ║${NC}"
7979
echo -e "${CYAN}╚══════════════════════════════════════════════════════════╝${NC}\n"
8080

0 commit comments

Comments
 (0)