| 1 | #!/bin/bash |
| 2 | |
| 3 | # This script aims to tune the best server parameter combinations to maximize throughput for given requirement. |
| 4 | # See details in README (benchmarks/auto_tune/README.md). |
| 5 | |
| 6 | TAG=$(date +"%Y_%m_%d_%H_%M") |
| 7 | SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd ) |
| 8 | VLLM_LOGGING_LEVEL=${VLLM_LOGGING_LEVEL:-INFO} |
| 9 | BASE=${BASE:-"$SCRIPT_DIR/../../.."} |
| 10 | MODEL=${MODEL:-"meta-llama/Llama-3.1-8B-Instruct"} |
| 11 | SYSTEM=${SYSTEM:-"TPU"} |
| 12 | TP=${TP:-1} |
| 13 | DOWNLOAD_DIR=${DOWNLOAD_DIR:-""} |
| 14 | INPUT_LEN=${INPUT_LEN:-4000} |
| 15 | OUTPUT_LEN=${OUTPUT_LEN:-16} |
| 16 | MAX_MODEL_LEN=${MAX_MODEL_LEN:-4096} |
| 17 | MIN_CACHE_HIT_PCT=${MIN_CACHE_HIT_PCT:-0} |
| 18 | MAX_LATENCY_ALLOWED_MS=${MAX_LATENCY_ALLOWED_MS:-100000000000} |
| 19 | NUM_SEQS_LIST=${NUM_SEQS_LIST:-"128 256"} |
| 20 | NUM_BATCHED_TOKENS_LIST=${NUM_BATCHED_TOKENS_LIST:-"512 1024 2048 4096"} |
| 21 | HOSTNAME=$(hostname) |
| 22 | if [[ -z "$HOSTNAME" ]]; then |
| 23 | echo "Error: Failed to determine hostname." >&2 |
| 24 | exit 1 |
| 25 | fi |
| 26 | |
| 27 | LOG_FOLDER="$BASE/auto-benchmark/$TAG" |
| 28 | RESULT="$LOG_FOLDER/result.txt" |
| 29 | PROFILE_PATH="$LOG_FOLDER/profile" |
| 30 | |
| 31 | echo "====================== AUTO TUNE PARAMETERS ====================" |
| 32 | echo "SCRIPT_DIR=$SCRIPT_DIR" |
| 33 | echo "BASE=$BASE" |
| 34 | echo "MODEL=$MODEL" |
| 35 | echo "SYSTEM=$SYSTEM" |
| 36 | echo "TP=$TP" |
| 37 | echo "DOWNLOAD_DIR=$DOWNLOAD_DIR" |
| 38 | echo "INPUT_LEN=$INPUT_LEN" |
| 39 | echo "OUTPUT_LEN=$OUTPUT_LEN" |
| 40 | echo "MAX_MODEL_LEN=$MAX_MODEL_LEN" |
| 41 | echo "MIN_CACHE_HIT_PCT=$MIN_CACHE_HIT_PCT" |
| 42 | echo "MAX_LATENCY_ALLOWED_MS=$MAX_LATENCY_ALLOWED_MS" |
| 43 | echo "NUM_SEQS_LIST=$NUM_SEQS_LIST" |
| 44 | echo "NUM_BATCHED_TOKENS_LIST=$NUM_BATCHED_TOKENS_LIST" |
| 45 | echo "VLLM_LOGGING_LEVEL=$VLLM_LOGGING_LEVEL" |
| 46 | echo "RESULT_FILE=$RESULT" |
| 47 | echo "====================== AUTO TUNEPARAMETERS ====================" |
| 48 | |
| 49 | rm -rf "$LOG_FOLDER" |
| 50 | rm -rf "$PROFILE_PATH" |
| 51 | mkdir -p "$LOG_FOLDER" |
| 52 | mkdir -p "$PROFILE_PATH" |
| 53 | |
| 54 | cd "$BASE/vllm" |
| 55 | |
| 56 | pip install -q datasets |
| 57 | |
| 58 | current_hash=$(git rev-parse HEAD) |
| 59 | echo "hash:$current_hash" >> "$RESULT" |
| 60 | echo "current_hash: $current_hash" |
| 61 | |
| 62 | TOTAL_LEN=$((INPUT_LEN + OUTPUT_LEN)) |
| 63 | RED='\033[0;31m' |
| 64 | if (( TOTAL_LEN > MAX_MODEL_LEN )); then |
| 65 | echo -e "${RED}FAILED: INPUT_LEN($INPUT_LEN) + OUTPUT_LEN($OUTPUT_LEN) = $TOTAL_LEN, which is > MAX_MODEL_LEN = $MAX_MODEL_LEN.\033[0m" >&2 |
| 66 | exit 1 |
| 67 | fi |
| 68 | |
| 69 | best_throughput=0 |
| 70 | best_max_num_seqs=0 |
| 71 | best_num_batched_tokens=0 |
| 72 | best_goodput=0 |
| 73 | best_request_rate=0 |
| 74 | |
| 75 | start_server() { |
| 76 | local gpu_memory_utilization=$1 |
| 77 | local max_num_seqs=$2 |
| 78 | local max_num_batched_tokens=$3 |
| 79 | local vllm_log=$4 |
| 80 | local profile_dir=$5 |
| 81 | |
| 82 | pkill -if "vllm serve" || true |
| 83 | |
| 84 | # Define the common arguments as a bash array. |
| 85 | # Each argument and its value are separate elements. |
| 86 | local common_args_array=( |
| 87 | "$MODEL" |
| 88 | "--port" "8004" |
| 89 | "--host" "$HOSTNAME" |
| 90 | "--gpu-memory-utilization" "$gpu_memory_utilization" |
| 91 | "--max-num-seqs" "$max_num_seqs" |
| 92 | "--max-num-batched-tokens" "$max_num_batched_tokens" |
| 93 | "--tensor-parallel-size" "$TP" |
| 94 | "--enable-prefix-caching" |
| 95 | "--load-format" "dummy" |
| 96 | "--download-dir" "$DOWNLOAD_DIR" |
| 97 | "--max-model-len" "$MAX_MODEL_LEN" |
| 98 | ) |
| 99 | |
| 100 | # Use the array expansion "${common_args_array[@]}" |
| 101 | # This correctly passes each element as a separate argument. |
| 102 | if [[ -n "$profile_dir" ]]; then |
| 103 | # Start server with profiling enabled |
| 104 | local profile_config_json="{\"profiler\": \"torch\", \"torch_profiler_dir\": \"$profile_dir\"}" |
| 105 | VLLM_SERVER_DEV_MODE=1 \ |
| 106 | vllm serve --profiler-config "$profile_config_json" "${common_args_array[@]}" > "$vllm_log" 2>&1 & |
| 107 | else |
| 108 | # Start server without profiling |
| 109 | VLLM_SERVER_DEV_MODE=1 \ |
| 110 | vllm serve "${common_args_array[@]}" > "$vllm_log" 2>&1 & |
| 111 | fi |
| 112 | local server_pid=$! |
| 113 | |
| 114 | # wait for 10 minutes... |
| 115 | server_started=0 |
| 116 | for _ in {1..60}; do |
| 117 | # This line checks whether the server is still alive or not, |
| 118 | # since that we should always have permission to send signal to the server process. |
| 119 | kill -0 $server_pid 2> /dev/null || break |
| 120 | |
| 121 | RESPONSE=$(curl -s -X GET "http://${HOSTNAME}:8004/health" -w "%{http_code}" -o /dev/stdout) |
| 122 | STATUS_CODE=$(echo "$RESPONSE" | tail -n 1) |
| 123 | if [[ "$STATUS_CODE" -eq 200 ]]; then |
| 124 | server_started=1 |
| 125 | break |
| 126 | else |
| 127 | sleep 10 |
| 128 | fi |
| 129 | done |
| 130 | |
| 131 | if (( ! server_started )); then |
| 132 | echo "server did not start within 10 minutes or crashed. Please check server log at $vllm_log". |
| 133 | return 1 |
| 134 | else |
| 135 | return 0 |
| 136 | fi |
| 137 | } |
| 138 | |
| 139 | run_benchmark() { |
| 140 | local max_num_seqs=$1 |
| 141 | local max_num_batched_tokens=$2 |
| 142 | local gpu_memory_utilization=$3 |
| 143 | echo "max_num_seq: $max_num_seqs, max_num_batched_tokens: $max_num_batched_tokens" |
| 144 | local vllm_log="$LOG_FOLDER/vllm_log_${max_num_seqs}_${max_num_batched_tokens}.txt" |
| 145 | echo "vllm_log: $vllm_log" |
| 146 | echo |
| 147 | rm -f "$vllm_log" |
| 148 | pkill -if "vllm serve" || true |
| 149 | |
| 150 | echo "starting server..." |
| 151 | # Call start_server without a profile_dir to avoid profiling overhead |
| 152 | start_server "$gpu_memory_utilization" "$max_num_seqs" "$max_num_batched_tokens" "$vllm_log" "" |
| 153 | result=$? |
| 154 | if [[ "$result" -eq 1 ]]; then |
| 155 | echo "server failed to start. gpu_memory_utilization:$gpu_memory_utilization, max_num_seqs:$max_num_seqs, max_num_batched_tokens: $max_num_batched_tokens" |
| 156 | else |
| 157 | echo "server started." |
| 158 | fi |
| 159 | echo |
| 160 | |
| 161 | echo "run benchmark test..." |
| 162 | meet_latency_requirement=0 |
| 163 | # get a basic qps by using request-rate inf |
| 164 | bm_log="$LOG_FOLDER/bm_log_${max_num_seqs}_${max_num_batched_tokens}_requestrate_inf.txt" |
| 165 | prefix_len=$(( INPUT_LEN * MIN_CACHE_HIT_PCT / 100 )) |
| 166 | adjusted_input_len=$(( INPUT_LEN - prefix_len )) |
| 167 | # --profile flag is removed from this call |
| 168 | vllm bench serve \ |
| 169 | --backend vllm \ |
| 170 | --model "$MODEL" \ |
| 171 | --dataset-name random \ |
| 172 | --random-input-len $adjusted_input_len \ |
| 173 | --random-output-len "$OUTPUT_LEN" \ |
| 174 | --ignore-eos \ |
| 175 | --disable-tqdm \ |
| 176 | --request-rate inf \ |
| 177 | --percentile-metrics ttft,tpot,itl,e2el \ |
| 178 | --goodput e2el:"$MAX_LATENCY_ALLOWED_MS" \ |
| 179 | --num-prompts 1000 \ |
| 180 | --random-prefix-len $prefix_len \ |
| 181 | --host "$HOSTNAME" \ |
| 182 | --port 8004 &> "$bm_log" |
| 183 | throughput=$(grep "Request throughput (req/s):" "$bm_log" | sed 's/[^0-9.]//g') |
| 184 | e2el=$(grep "P99 E2EL (ms):" "$bm_log" | awk '{print $NF}') |
| 185 | goodput=$(grep "Request goodput (req/s):" "$bm_log" | sed 's/[^0-9.]//g') |
| 186 | |
| 187 | if (( $(echo "$e2el <= $MAX_LATENCY_ALLOWED_MS" | bc -l) )); then |
| 188 | meet_latency_requirement=1 |
| 189 | request_rate=inf |
| 190 | fi |
| 191 | |
| 192 | if (( ! meet_latency_requirement )); then |
| 193 | # start from request-rate as int(throughput) + 1 |
| 194 | request_rate=$((${throughput%.*} + 1)) |
| 195 | while ((request_rate > 0)); do |
| 196 | # clear prefix cache |
| 197 | curl -X POST http://"${HOSTNAME}":8004/reset_prefix_cache |
| 198 | sleep 5 |
| 199 | bm_log="$LOG_FOLDER/bm_log_${max_num_seqs}_${max_num_batched_tokens}_requestrate_${request_rate}.txt" |
| 200 | vllm bench serve \ |
| 201 | --backend vllm \ |
| 202 | --model "$MODEL" \ |
| 203 | --dataset-name random \ |
| 204 | --random-input-len $adjusted_input_len \ |
| 205 | --random-output-len "$OUTPUT_LEN" \ |
| 206 | --ignore-eos \ |
| 207 | --disable-tqdm \ |
| 208 | --request-rate $request_rate \ |
| 209 | --percentile-metrics ttft,tpot,itl,e2el \ |
| 210 | --goodput e2el:"$MAX_LATENCY_ALLOWED_MS" \ |
| 211 | --num-prompts 100 \ |
| 212 | --random-prefix-len $prefix_len \ |
| 213 | --host "$HOSTNAME" \ |
| 214 | --port 8004 &> "$bm_log" |
| 215 | throughput=$(grep "Request throughput (req/s):" "$bm_log" | sed 's/[^0-9.]//g') |
| 216 | e2el=$(grep "P99 E2EL (ms):" "$bm_log" | awk '{print $NF}') |
| 217 | goodput=$(grep "Request goodput (req/s):" "$bm_log" | sed 's/[^0-9.]//g') |
| 218 | if (( $(echo "$e2el <= $MAX_LATENCY_ALLOWED_MS" | bc -l) )); then |
| 219 | meet_latency_requirement=1 |
| 220 | break |
| 221 | fi |
| 222 | request_rate=$((request_rate-1)) |
| 223 | done |
| 224 | fi |
| 225 | # write the results and update the best result. |
| 226 | if ((meet_latency_requirement)); then |
| 227 | echo "max_num_seqs: $max_num_seqs, max_num_batched_tokens: $max_num_batched_tokens, request_rate: $request_rate, e2el: $e2el, throughput: $throughput, goodput: $goodput" |
| 228 | echo "max_num_seqs: $max_num_seqs, max_num_batched_tokens: $max_num_batched_tokens, request_rate: $request_rate, e2el: $e2el, throughput: $throughput, goodput: $goodput" >> "$RESULT" |
| 229 | if (( $(echo "$throughput > $best_throughput" | bc -l) )); then |
| 230 | best_throughput=$throughput |
| 231 | best_max_num_seqs=$max_num_seqs |
| 232 | best_num_batched_tokens=$max_num_batched_tokens |
| 233 | best_goodput=$goodput |
| 234 | best_request_rate=$request_rate |
| 235 | fi |
| 236 | else |
| 237 | echo "max_num_seqs: $max_num_seqs, max_num_batched_tokens: $max_num_batched_tokens does not meet latency requirement ${MAX_LATENCY_ALLOWED_MS}" |
| 238 | echo "max_num_seqs: $max_num_seqs, max_num_batched_tokens: $max_num_batched_tokens does not meet latency requirement ${MAX_LATENCY_ALLOWED_MS}" >> "$RESULT" |
| 239 | fi |
| 240 | |
| 241 | echo "best_max_num_seqs: $best_max_num_seqs, best_num_batched_tokens: $best_num_batched_tokens, best_throughput: $best_throughput" |
| 242 | |
| 243 | pkill -if "vllm serve" || true |
| 244 | sleep 10 |
| 245 | echo "====================" |
| 246 | return 0 |
| 247 | } |
| 248 | |
| 249 | read -r -a num_seqs_list <<< "$NUM_SEQS_LIST" |
| 250 | read -r -a num_batched_tokens_list <<< "$NUM_BATCHED_TOKENS_LIST" |
| 251 | |
| 252 | # first find out the max gpu-memory-utilization without HBM OOM. |
| 253 | gpu_memory_utilization=0.98 |
| 254 | find_gpu_memory_utilization=0 |
| 255 | while (( $(echo "$gpu_memory_utilization >= 0.9" | bc -l) )); do |
| 256 | # Pass empty string for profile_dir argument |
| 257 | start_server "$gpu_memory_utilization" "${num_seqs_list[-1]}" "${num_batched_tokens_list[-1]}" "$LOG_FOLDER/vllm_log_gpu_memory_utilization_$gpu_memory_utilization.log" "" |
| 258 | result=$? |
| 259 | if [[ "$result" -eq 0 ]]; then |
| 260 | find_gpu_memory_utilization=1 |
| 261 | break |
| 262 | else |
| 263 | gpu_memory_utilization=$(echo "$gpu_memory_utilization - 0.01" | bc) |
| 264 | fi |
| 265 | done |
| 266 | |
| 267 | if [[ "$find_gpu_memory_utilization" -eq 1 ]]; then |
| 268 | echo "Using gpu_memory_utilization=$gpu_memory_utilization to serve model." |
| 269 | else |
| 270 | echo "Cannot find a proper gpu_memory_utilization over 0.9 to serve the model, please check logs in $LOG_FOLDER." |
| 271 | exit 1 |
| 272 | fi |
| 273 | |
| 274 | for num_seqs in "${num_seqs_list[@]}"; do |
| 275 | for num_batched_tokens in "${num_batched_tokens_list[@]}"; do |
| 276 | run_benchmark "$num_seqs" "$num_batched_tokens" "$gpu_memory_utilization" |
| 277 | done |
| 278 | done |
| 279 | echo "finish permutations" |
| 280 | |
| 281 | # ================================================================================= |
| 282 | # FINAL PROFILING RUN FOR THE BEST CONFIGURATION |
| 283 | # ================================================================================= |
| 284 | if (( $(echo "$best_throughput > 0" | bc -l) )); then |
| 285 | echo |
| 286 | echo "Benchmark tuning finished. Now running profiling on the best configuration found..." |
| 287 | echo "Best config: max_num_seqs: $best_max_num_seqs, max_num_batched_tokens: $best_num_batched_tokens, throughput: $best_throughput, goodput: $best_goodput" |
| 288 | echo |
| 289 | |
| 290 | vllm_log="$LOG_FOLDER/vllm_log_BEST_PROFILE.txt" |
| 291 | bm_log="$LOG_FOLDER/bm_log_BEST_PROFILE.txt" |
| 292 | |
| 293 | # Start server with the best params and profiling ENABLED |
| 294 | echo "Starting server for profiling..." |
| 295 | start_server "$gpu_memory_utilization" "$best_max_num_seqs" "$best_num_batched_tokens" "$vllm_log" "$PROFILE_PATH" |
| 296 | |
| 297 | # Run benchmark with the best params and the --profile flag |
| 298 | echo "Running benchmark with profiling..." |
| 299 | prefix_len=$(( INPUT_LEN * MIN_CACHE_HIT_PCT / 100 )) |
| 300 | adjusted_input_len=$(( INPUT_LEN - prefix_len )) |
| 301 | vllm bench serve \ |
| 302 | --backend vllm \ |
| 303 | --model "$MODEL" \ |
| 304 | --dataset-name random \ |
| 305 | --random-input-len $adjusted_input_len \ |
| 306 | --random-output-len "$OUTPUT_LEN" \ |
| 307 | --ignore-eos \ |
| 308 | --disable-tqdm \ |
| 309 | --request-rate "$best_request_rate" \ |
| 310 | --percentile-metrics ttft,tpot,itl,e2el \ |
| 311 | --goodput e2el:"$MAX_LATENCY_ALLOWED_MS" \ |
| 312 | --num-prompts 100 \ |
| 313 | --random-prefix-len $prefix_len \ |
| 314 | --host "$HOSTNAME" \ |
| 315 | --port 8004 \ |
| 316 | --profile &> "$bm_log" |
| 317 | else |
| 318 | echo "No configuration met the latency requirements. Skipping final profiling run." |
| 319 | fi |
| 320 | pkill -if "vllm serve" || true |
| 321 | echo "best_max_num_seqs: $best_max_num_seqs, best_num_batched_tokens: $best_num_batched_tokens, best_throughput: $best_throughput, profile saved in: $PROFILE_PATH" |
| 322 | echo "best_max_num_seqs: $best_max_num_seqs, best_num_batched_tokens: $best_num_batched_tokens, best_throughput: $best_throughput, profile saved in: $PROFILE_PATH" >> "$RESULT" |