| 1 | #!/bin/bash |
| 2 | |
| 3 | # Currently FP8 benchmark is NOT enabled. |
| 4 | |
| 5 | set -x |
| 6 | server_params=$1 |
| 7 | common_params=$2 |
| 8 | |
| 9 | json2args() { |
| 10 | # transforms the JSON string to command line args, and '_' is replaced to '-' |
| 11 | # example: |
| 12 | # input: { "model": "meta-llama/Llama-2-7b-chat-hf", "tensor_parallel_size": 1 } |
| 13 | # output: --model meta-llama/Llama-2-7b-chat-hf --tensor-parallel-size 1 |
| 14 | local json_string=$1 |
| 15 | local args=$( |
| 16 | echo "$json_string" | jq -r ' |
| 17 | to_entries | |
| 18 | map("--" + (.key | gsub("_"; "-")) + " " + (.value | tostring)) | |
| 19 | join(" ") |
| 20 | ' |
| 21 | ) |
| 22 | echo "$args" |
| 23 | } |
| 24 | |
| 25 | launch_trt_server() { |
| 26 | |
| 27 | model_path=$(echo "$common_params" | jq -r '.model') |
| 28 | model_name="${model_path#*/}" |
| 29 | model_type=$(echo "$server_params" | jq -r '.model_type') |
| 30 | model_dtype=$(echo "$server_params" | jq -r '.model_dtype') |
| 31 | model_tp_size=$(echo "$common_params" | jq -r '.tp') |
| 32 | max_batch_size=$(echo "$server_params" | jq -r '.max_batch_size') |
| 33 | max_input_len=$(echo "$server_params" | jq -r '.max_input_len') |
| 34 | max_seq_len=$(echo "$server_params" | jq -r '.max_seq_len') |
| 35 | max_num_tokens=$(echo "$server_params" | jq -r '.max_num_tokens') |
| 36 | trt_llm_version=$(echo "$server_params" | jq -r '.trt_llm_version') |
| 37 | |
| 38 | # create model caching directory |
| 39 | cd ~ |
| 40 | rm -rf models |
| 41 | mkdir -p models |
| 42 | cd models |
| 43 | models_dir=$(pwd) |
| 44 | trt_model_path=${models_dir}/${model_name}-trt-ckpt |
| 45 | trt_engine_path=${models_dir}/${model_name}-trt-engine |
| 46 | |
| 47 | # clone tensorrt backend |
| 48 | cd / |
| 49 | rm -rf tensorrtllm_backend |
| 50 | git clone https://github.com/triton-inference-server/tensorrtllm_backend.git |
| 51 | git lfs install |
| 52 | cd tensorrtllm_backend |
| 53 | git checkout "$trt_llm_version" |
| 54 | git submodule update --init --recursive |
| 55 | |
| 56 | # build trtllm engine |
| 57 | cd /tensorrtllm_backend |
| 58 | cd "./tensorrt_llm/examples/${model_type}" |
| 59 | python3 convert_checkpoint.py \ |
| 60 | --model_dir "${model_path}" \ |
| 61 | --dtype "${model_dtype}" \ |
| 62 | --tp_size "${model_tp_size}" \ |
| 63 | --output_dir "${trt_model_path}" |
| 64 | trtllm-build \ |
| 65 | --checkpoint_dir "${trt_model_path}" \ |
| 66 | --use_fused_mlp \ |
| 67 | --reduce_fusion disable \ |
| 68 | --workers 8 \ |
| 69 | --gpt_attention_plugin "${model_dtype}" \ |
| 70 | --gemm_plugin "${model_dtype}" \ |
| 71 | --tp_size "${model_tp_size}" \ |
| 72 | --max_batch_size "${max_batch_size}" \ |
| 73 | --max_input_len "${max_input_len}" \ |
| 74 | --max_seq_len "${max_seq_len}" \ |
| 75 | --max_num_tokens "${max_num_tokens}" \ |
| 76 | --output_dir "${trt_engine_path}" |
| 77 | |
| 78 | # handle triton protobuf files and launch triton server |
| 79 | cd /tensorrtllm_backend |
| 80 | mkdir triton_model_repo |
| 81 | cp -r all_models/inflight_batcher_llm/* triton_model_repo/ |
| 82 | cd triton_model_repo |
| 83 | rm -rf ./tensorrt_llm/1/* |
| 84 | cp -r "${trt_engine_path}"/* ./tensorrt_llm/1 |
| 85 | python3 ../tools/fill_template.py -i tensorrt_llm/config.pbtxt triton_backend:tensorrtllm,engine_dir:/tensorrtllm_backend/triton_model_repo/tensorrt_llm/1,decoupled_mode:true,batching_strategy:inflight_fused_batching,batch_scheduler_policy:guaranteed_no_evict,exclude_input_in_output:true,triton_max_batch_size:2048,max_queue_delay_microseconds:0,max_beam_width:1,max_queue_size:2048,enable_kv_cache_reuse:false |
| 86 | python3 ../tools/fill_template.py -i preprocessing/config.pbtxt "triton_max_batch_size:2048,tokenizer_dir:$model_path,preprocessing_instance_count:5" |
| 87 | python3 ../tools/fill_template.py -i postprocessing/config.pbtxt "triton_max_batch_size:2048,tokenizer_dir:$model_path,postprocessing_instance_count:5,skip_special_tokens:false" |
| 88 | python3 ../tools/fill_template.py -i ensemble/config.pbtxt triton_max_batch_size:"$max_batch_size" |
| 89 | python3 ../tools/fill_template.py -i tensorrt_llm_bls/config.pbtxt "triton_max_batch_size:$max_batch_size,decoupled_mode:true,accumulate_tokens:False,bls_instance_count:1" |
| 90 | cd /tensorrtllm_backend |
| 91 | python3 scripts/launch_triton_server.py \ |
| 92 | --world_size="${model_tp_size}" \ |
| 93 | --model_repo=/tensorrtllm_backend/triton_model_repo & |
| 94 | |
| 95 | } |
| 96 | |
| 97 | launch_tgi_server() { |
| 98 | model=$(echo "$common_params" | jq -r '.model') |
| 99 | tp=$(echo "$common_params" | jq -r '.tp') |
| 100 | port=$(echo "$common_params" | jq -r '.port') |
| 101 | server_args=$(json2args "$server_params") |
| 102 | |
| 103 | if echo "$common_params" | jq -e 'has("fp8")' >/dev/null; then |
| 104 | echo "Key 'fp8' exists in common params." |
| 105 | server_command="/tgi-entrypoint.sh \ |
| 106 | --model-id $model \ |
| 107 | --num-shard $tp \ |
| 108 | --port $port \ |
| 109 | --quantize fp8 \ |
| 110 | $server_args" |
| 111 | else |
| 112 | echo "Key 'fp8' does not exist in common params." |
| 113 | server_command="/tgi-entrypoint.sh \ |
| 114 | --model-id $model \ |
| 115 | --num-shard $tp \ |
| 116 | --port $port \ |
| 117 | $server_args" |
| 118 | fi |
| 119 | |
| 120 | echo "Server command: $server_command" |
| 121 | eval "$server_command" & |
| 122 | |
| 123 | } |
| 124 | |
| 125 | launch_lmdeploy_server() { |
| 126 | model=$(echo "$common_params" | jq -r '.model') |
| 127 | tp=$(echo "$common_params" | jq -r '.tp') |
| 128 | port=$(echo "$common_params" | jq -r '.port') |
| 129 | server_args=$(json2args "$server_params") |
| 130 | |
| 131 | server_command="lmdeploy serve api_server $model \ |
| 132 | --tp $tp \ |
| 133 | --server-port $port \ |
| 134 | $server_args" |
| 135 | |
| 136 | # run the server |
| 137 | echo "Server command: $server_command" |
| 138 | bash -c "$server_command" & |
| 139 | } |
| 140 | |
| 141 | launch_sglang_server() { |
| 142 | |
| 143 | model=$(echo "$common_params" | jq -r '.model') |
| 144 | tp=$(echo "$common_params" | jq -r '.tp') |
| 145 | port=$(echo "$common_params" | jq -r '.port') |
| 146 | server_args=$(json2args "$server_params") |
| 147 | |
| 148 | if echo "$common_params" | jq -e 'has("fp8")' >/dev/null; then |
| 149 | echo "Key 'fp8' exists in common params. Use neuralmagic fp8 model for convenience." |
| 150 | model=$(echo "$common_params" | jq -r '.neuralmagic_quantized_model') |
| 151 | server_command="python3 \ |
| 152 | -m sglang.launch_server \ |
| 153 | --tp $tp \ |
| 154 | --model-path $model \ |
| 155 | --port $port \ |
| 156 | $server_args" |
| 157 | else |
| 158 | echo "Key 'fp8' does not exist in common params." |
| 159 | server_command="python3 \ |
| 160 | -m sglang.launch_server \ |
| 161 | --tp $tp \ |
| 162 | --model-path $model \ |
| 163 | --port $port \ |
| 164 | $server_args" |
| 165 | fi |
| 166 | |
| 167 | # run the server |
| 168 | echo "Server command: $server_command" |
| 169 | eval "$server_command" & |
| 170 | } |
| 171 | |
| 172 | launch_vllm_server() { |
| 173 | |
| 174 | export VLLM_HOST_IP=$(hostname -I | awk '{print $1}') |
| 175 | |
| 176 | model=$(echo "$common_params" | jq -r '.model') |
| 177 | tp=$(echo "$common_params" | jq -r '.tp') |
| 178 | port=$(echo "$common_params" | jq -r '.port') |
| 179 | server_args=$(json2args "$server_params") |
| 180 | |
| 181 | if echo "$common_params" | jq -e 'has("fp8")' >/dev/null; then |
| 182 | echo "Key 'fp8' exists in common params. Use neuralmagic fp8 model for convenience." |
| 183 | model=$(echo "$common_params" | jq -r '.neuralmagic_quantized_model') |
| 184 | server_command="vllm serve $model \ |
| 185 | -tp $tp \ |
| 186 | --port $port \ |
| 187 | $server_args" |
| 188 | else |
| 189 | echo "Key 'fp8' does not exist in common params." |
| 190 | server_command="vllm serve $model \ |
| 191 | -tp $tp \ |
| 192 | --port $port \ |
| 193 | $server_args" |
| 194 | fi |
| 195 | |
| 196 | # run the server |
| 197 | echo "Server command: $server_command" |
| 198 | eval "$server_command" & |
| 199 | } |
| 200 | |
| 201 | main() { |
| 202 | |
| 203 | if [[ "$CURRENT_LLM_SERVING_ENGINE" == "trt" ]]; then |
| 204 | launch_trt_server |
| 205 | fi |
| 206 | |
| 207 | if [[ "$CURRENT_LLM_SERVING_ENGINE" == "tgi" ]]; then |
| 208 | launch_tgi_server |
| 209 | fi |
| 210 | |
| 211 | if [[ "$CURRENT_LLM_SERVING_ENGINE" == "lmdeploy" ]]; then |
| 212 | launch_lmdeploy_server |
| 213 | fi |
| 214 | |
| 215 | if [[ "$CURRENT_LLM_SERVING_ENGINE" == "sglang" ]]; then |
| 216 | launch_sglang_server |
| 217 | fi |
| 218 | |
| 219 | if [[ "$CURRENT_LLM_SERVING_ENGINE" == *"vllm"* ]]; then |
| 220 | launch_vllm_server |
| 221 | fi |
| 222 | } |
| 223 | |
| 224 | main |