@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
sh 136 lines 4.8 KB
Raw
1 #!/bin/bash
2
3 # SPDX-License-Identifier: Apache-2.0
4 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
5
6 # vLLM Embedding Server with Enhanced Chunked Processing
7 # This script starts a vLLM server with chunked processing enabled for long text embedding.
8 # Now supports proper pooling type validation and model-specific configurations.
9
10 set -euo pipefail
11
12 # Configuration
13 MODEL_NAME=${MODEL_NAME:-"intfloat/multilingual-e5-large"}
14 MODEL_CODE=${MODEL_CODE:-"multilingual-e5-large"}
15
16 PORT=${PORT:-31090}
17 GPU_COUNT=${GPU_COUNT:-1}
18 MAX_EMBED_LEN=${MAX_EMBED_LEN:-3072000}
19 API_KEY=${API_KEY:-"your-api-key"}
20
21 # Enhanced pooling configuration with model-specific defaults
22 POOLING_TYPE=${POOLING_TYPE:-"auto"} # auto, MEAN, CLS, LAST
23 export VLLM_ENABLE_CHUNKED_PROCESSING=true
24 export CUDA_VISIBLE_DEVICES=2,3,4,5
25
26 echo "🚀 Starting vLLM Embedding Server with Enhanced Chunked Processing"
27 echo "=================================================================="
28
29 # Environment variables for optimization
30 export VLLM_WORKER_MULTIPROC_METHOD=spawn
31
32 # Function to determine optimal pooling type for known models
33 get_optimal_pooling_type() {
34 local model="$1"
35 case "$model" in
36 *"e5-"* | *"multilingual-e5"*)
37 echo "MEAN" # E5 series native pooling
38 ;;
39 *"bge-"*)
40 echo "CLS" # BGE series native pooling
41 ;;
42 *"gte-"*)
43 echo "LAST" # GTE series native pooling
44 ;;
45 *"sentence-t5"* | *"st5"*)
46 echo "MEAN" # Sentence-T5 native pooling
47 ;;
48 *"jina-embeddings"*)
49 echo "MEAN" # Jina embeddings native pooling
50 ;;
51 *"Qwen"*"Embedding"*)
52 echo "LAST" # Qwen embeddings native pooling
53 ;;
54 *)
55 echo "MEAN" # Default native pooling for unknown models
56 ;;
57 esac
58 }
59
60 # Auto-detect pooling type if not explicitly set
61 if [ "$POOLING_TYPE" = "auto" ]; then
62 POOLING_TYPE=$(get_optimal_pooling_type "$MODEL_NAME")
63 echo "🔍 Auto-detected pooling type: $POOLING_TYPE for model $MODEL_NAME"
64 fi
65
66 # Display configuration
67 echo "📋 Configuration:"
68 echo " - Model: $MODEL_NAME"
69 echo " - Port: $PORT"
70 echo " - GPU Count: $GPU_COUNT"
71 echo " - Enhanced Chunked Processing: ${VLLM_ENABLE_CHUNKED_PROCESSING}"
72 echo " - Max Embed Length: ${MAX_EMBED_LEN} tokens"
73 echo " - Native Pooling Type: $POOLING_TYPE + Normalization"
74 echo " - Cross-chunk Aggregation: MEAN (automatic)"
75 echo ""
76
77 # Validate GPU availability
78 if command -v nvidia-smi &> /dev/null; then
79 gpu_count=$(nvidia-smi --list-gpus | wc -l)
80 echo "🖥️ Available GPUs: $gpu_count"
81 if [ "$GPU_COUNT" -gt "$gpu_count" ]; then
82 echo "⚠️ Warning: Requested $GPU_COUNT GPUs but only $gpu_count available"
83 echo " Adjusting to use $gpu_count GPUs"
84 GPU_COUNT=$gpu_count
85 fi
86 else
87 echo "⚠️ Warning: nvidia-smi not found. GPU detection skipped."
88 fi
89
90 # Chunked processing uses unified MEAN aggregation
91 echo "ℹ️ Chunked Processing: Using $POOLING_TYPE pooling within chunks, MEAN aggregation across chunks"
92 echo " - All chunks processed for complete semantic coverage"
93 echo " - Weighted averaging based on chunk token counts"
94
95 echo ""
96 echo "🔧 Starting server with enhanced chunked processing configuration..."
97
98 # Build pooler config JSON
99 POOLER_CONFIG="{\"pooling_type\": \"$POOLING_TYPE\", \"use_activation\": true, \"enable_chunked_processing\": ${VLLM_ENABLE_CHUNKED_PROCESSING}, \"max_embed_len\": ${MAX_EMBED_LEN}}"
100
101 # Start vLLM server with enhanced chunked processing
102 vllm serve "$MODEL_NAME" \
103 --tensor-parallel-size "$GPU_COUNT" \
104 --enforce-eager \
105 --pooler-config "$POOLER_CONFIG" \
106 --served-model-name "${MODEL_CODE}" \
107 --api-key "$API_KEY" \
108 --trust-remote-code \
109 --port "$PORT" \
110 --host 0.0.0.0
111
112 echo ""
113 echo "✅ vLLM Embedding Server started successfully!"
114 echo ""
115 echo "📡 Server Information:"
116 echo " - Base URL: http://localhost:$PORT"
117 echo " - Model Code: ${MODEL_CODE}"
118 echo " - API Key: $API_KEY"
119 echo " - Native Pooling: $POOLING_TYPE | Cross-chunk: MEAN"
120 echo ""
121 echo "🧪 Test the server with:"
122 echo " python examples/pooling/embed/openai_embedding_long_text/client.py"
123 echo ""
124 echo "📚 Enhanced features enabled:"
125 echo " ✅ Intelligent native pooling type detection"
126 echo " ✅ Unified MEAN aggregation for chunked processing"
127 echo " ✅ Model-specific native pooling optimization"
128 echo " ✅ Enhanced max embedding length (${MAX_EMBED_LEN} tokens)"
129 echo " ✅ Complete semantic coverage for all pooling types"
130 echo " ✅ OpenAI-compatible API"
131 echo " ✅ GPU acceleration"
132 echo ""
133 echo "🔧 Advanced usage:"
134 echo " - Set POOLING_TYPE=MEAN|CLS|LAST to override auto-detection"
135 echo " - Set MAX_EMBED_LEN to adjust maximum input length"
136 echo " - All pooling types use MEAN aggregation across chunks"