onboard BGE-Base-EN-v1.5 and All-MiniLM-L6-v2

examples/inference/text_embedding/embeddings.py

@@ -0,0 +1,19 @@
+from openai import OpenAI
+
+# The url is located in the .vLLM_model-variant_url file in the corresponding model directory.
+client = OpenAI(base_url="http://gpu031:8081/v1", api_key="EMPTY")
+
+model_name = "bge-base-en-v1.5"
+
+input_texts = [
+    "The chef prepared a delicious meal.",
+]
+
+# test single embedding
+embedding_response = client.embeddings.create(
+    model=model_name,
+    input=input_texts,
+    encoding_format="float",
+)
+
+print(embedding_response)

vec_inf/cli/_cli.py

@@ -138,6 +138,13 @@ def launch(
 
     models_df = utils.load_models_df()
 
+    models_df = models_df.with_columns(
+        pl.col("model_type").replace("Reward Modeling", "Reward_Modeling")
+    )
+    models_df = models_df.with_columns(
+        pl.col("model_type").replace("Text Embedding", "Text_Embedding")
+    )
+
     if model_name in models_df["model_name"].to_list():
         default_args = utils.load_default_args(models_df, model_name)
         for arg in default_args:
@@ -148,7 +155,6 @@ def launch(
     else:
         model_args = models_df.columns
         model_args.remove("model_name")
-        model_args.remove("model_type")
         for arg in model_args:
             if locals()[arg] is not None:
                 renamed_arg = arg.replace("_", "-")

vec_inf/cli/_utils.py

@@ -139,7 +139,6 @@ def load_default_args(models_df: pl.DataFrame, model_name: str) -> dict:
     row_data = models_df.filter(models_df["model_name"] == model_name)
     default_args = row_data.to_dicts()[0]
     default_args.pop("model_name", None)
-    default_args.pop("model_type", None)
     return default_args
 
 

vec_inf/launch_server.sh

@@ -6,6 +6,7 @@ while [[ "$#" -gt 0 ]]; do
     case $1 in
         --model-family) model_family="$2"; shift ;;
         --model-variant) model_variant="$2"; shift ;;
+        --model-type) model_type="$2"; shift ;;
         --partition) partition="$2"; shift ;;
         --qos) qos="$2"; shift ;;
         --time) walltime="$2"; shift ;;
@@ -25,7 +26,7 @@ while [[ "$#" -gt 0 ]]; do
     shift
 done
 
-required_vars=(model_family model_variant partition qos walltime num_nodes num_gpus max_model_len vocab_size data_type venv log_dir model_weights_parent_dir)
+required_vars=(model_family model_variant model_type partition qos walltime num_nodes num_gpus max_model_len vocab_size data_type venv log_dir model_weights_parent_dir)
 
 for var in "$required_vars[@]"; do
     if [ -z "$!var" ]; then
@@ -36,6 +37,7 @@ done
 
 export MODEL_FAMILY=$model_family
 export MODEL_VARIANT=$model_variant
+export MODEL_TYPE=$model_type
 export JOB_PARTITION=$partition
 export QOS=$qos
 export WALLTIME=$walltime
@@ -48,6 +50,17 @@ export VENV_BASE=$venv
 export LOG_DIR=$log_dir
 export MODEL_WEIGHTS_PARENT_DIR=$model_weights_parent_dir
 
+if [[ "$model_type" == "LLM" || "$model_type" == "VLM" ]]; then
+    export VLLM_TASK="generate"
+elif [[ "$model_type" == "Reward_Modeling" ]]; then
+    export VLLM_TASK="reward"
+elif [[ "$model_type" == "Text_Embedding" ]]; then
+    export VLLM_TASK="embed"
+else
+    echo "Error: Unknown model_type: $model_type"
+    exit 1
+fi
+
 if [ -n "$max_num_seqs" ]; then
     export VLLM_MAX_NUM_SEQS=$max_num_seqs
 else
@@ -101,6 +114,8 @@ echo Num Nodes: $NUM_NODES
 echo GPUs per Node: $NUM_GPUS
 echo QOS: $QOS
 echo Walltime: $WALLTIME
+echo Model Type: $MODEL_TYPE
+echo Task: $VLLM_TASK
 echo Data Type: $VLLM_DATA_TYPE
 echo Max Model Length: $VLLM_MAX_MODEL_LEN
 echo Max Num Seqs: $VLLM_MAX_NUM_SEQS

vec_inf/models/README.md

@@ -194,6 +194,18 @@ More profiling metrics coming soon!
 |:----------:|:----------:|:----------:|:----------:|
 | [`e5-mistral-7b-instruct`](https://huggingface.co/intfloat/e5-mistral-7b-instruct) | 1x a40 | - tokens/s | - tokens/s |
 
+### [BAAI: bge](https://huggingface.co/BAAI)
+| Variant | Suggested resource allocation | Avg prompt throughput | Avg generation throughput |
+|:----------:|:----------:|:----------:|:----------:|
+| [`bge-base-en-v1.5`](https://huggingface.co/BAAI/bge-base-en-v1.5) | 1x A40 | - tokens/s | - tokens/s |
+
+### [Sentence Transformers: MiniLM](https://huggingface.co/sentence-transformers)
+| Variant | Suggested resource allocation | Avg prompt throughput | Avg generation throughput |
+|:----------:|:----------:|:----------:|:----------:|
+| [`all-MiniLM-L6-v2`](https://huggingface.co/sentence-transformers/all-MiniLM-L6-v2) | 1x A40 | - tokens/s | - tokens/s |
+
+
+
 ## Reward Modeling Models
 
 ### [Qwen: Qwen2.5-Math](https://huggingface.co/collections/Qwen/qwen25-math-66eaa240a1b7d5ee65f1da3e)

vec_inf/models/models.csv

@@ -71,3 +71,5 @@ Qwen2.5-Math-RM-72B,Qwen2.5,Math-RM-72B,Reward Modeling,4,1,152064,4096,256,true
 QwQ-32B-Preview,QwQ,32B-Preview,LLM,2,1,152064,32768,256,true,false,m2,08:00:00,a40,auto,singularity,default,/model-weights
 Pixtral-12B-2409,Pixtral,12B-2409,VLM,1,1,131072,8192,256,true,false,m2,08:00:00,a40,auto,singularity,default,/model-weights
 e5-mistral-7b-instruct,e5,mistral-7b-instruct,Text Embedding,1,1,32000,4096,256,true,false,m2,08:00:00,a40,auto,singularity,default,/model-weights
+bge-base-en-v1.5,bge,base-en-v1.5,Text Embedding,1,1,30522,512,256,true,false,m2,08:00:00,a40,auto,singularity,default,/model-weights
+all-MiniLM-L6-v2,all-MiniLM,L6-v2,Text Embedding,1,1,30522,512,256,true,false,m2,08:00:00,a40,auto,singularity,default,/model-weights

vec_inf/multinode_vllm.slurm

@@ -12,7 +12,7 @@ nvidia-smi
 source ${SRC_DIR}/find_port.sh
 
 if [ "$VENV_BASE" = "singularity" ]; then
-    export SINGULARITY_IMAGE=/projects/aieng/public/vector-inference_0.6.4.post1.sif
+    export SINGULARITY_IMAGE=/projects/aieng/public/vector-inference_latest.sif
     export VLLM_NCCL_SO_PATH=/vec-inf/nccl/libnccl.so.2.18.1
     module load singularity-ce/3.8.2
     singularity exec $SINGULARITY_IMAGE ray stop
@@ -103,6 +103,7 @@ if [ "$VENV_BASE" = "singularity" ]; then
     --max-logprobs ${VLLM_MAX_LOGPROBS} \
     --max-model-len ${VLLM_MAX_MODEL_LEN} \
     --max-num-seqs ${VLLM_MAX_NUM_SEQS} \
+    --task ${VLLM_TASK} \
     ${ENFORCE_EAGER}
 else
     source ${VENV_BASE}/bin/activate
@@ -118,5 +119,6 @@ else
     --max-logprobs ${VLLM_MAX_LOGPROBS} \
     --max-model-len ${VLLM_MAX_MODEL_LEN} \
     --max-num-seqs ${VLLM_MAX_NUM_SEQS} \
+    --task ${VLLM_TASK} \
     ${ENFORCE_EAGER}
 fi

vec_inf/vllm.slurm

@@ -23,7 +23,7 @@ fi
 
 # Activate vllm venv
 if [ "$VENV_BASE" = "singularity" ]; then
-    export SINGULARITY_IMAGE=/projects/aieng/public/vector-inference_0.6.4.post1.sif
+    export SINGULARITY_IMAGE=/projects/aieng/public/vector-inference_latest.sif
     export VLLM_NCCL_SO_PATH=/vec-inf/nccl/libnccl.so.2.18.1
     module load singularity-ce/3.8.2
     singularity exec $SINGULARITY_IMAGE ray stop
@@ -39,6 +39,7 @@ if [ "$VENV_BASE" = "singularity" ]; then
     --trust-remote-code \
     --max-model-len ${VLLM_MAX_MODEL_LEN} \
     --max-num-seqs ${VLLM_MAX_NUM_SEQS} \
+    --task ${VLLM_TASK} \
     ${ENFORCE_EAGER}
 else
     source ${VENV_BASE}/bin/activate
@@ -53,5 +54,6 @@ else
     --trust-remote-code \
     --max-model-len ${VLLM_MAX_MODEL_LEN} \
     --max-num-seqs ${VLLM_MAX_NUM_SEQS} \
+    --task ${VLLM_TASK} \
     ${ENFORCE_EAGER}
 fi