Add support for Matryoshka embeddings (#126) (#11142)

Co-authored-by: Satyam Kumar <satyamk@linkedin.com>

python/sglang/test/runners.py

@@ -12,10 +12,11 @@
 # limitations under the License.
 # ==============================================================================
 
+import json
 import multiprocessing as mp
 import os
 from dataclasses import dataclass
-from typing import List, Optional, Tuple, Union
+from typing import Any, List, Optional, Tuple, Union
 
 import torch
 import torch.nn.functional as F
@@ -89,7 +90,9 @@ def get_token_ids_logprobs(logits, token_ids):
     return logprobs
 
 
-def _get_sentence_transformer_embedding_model(model_path, torch_dtype):
+def _get_sentence_transformer_embedding_model(
+    model_path, torch_dtype, matryoshka_dim: Optional[int] = None
+):
     from sentence_transformers import SentenceTransformer
     from sentence_transformers.util import is_sentence_transformer_model
 
@@ -97,6 +100,7 @@ def _get_sentence_transformer_embedding_model(model_path, torch_dtype):
         model = SentenceTransformer(
             model_path,
             model_kwargs={"torch_dtype": torch_dtype},
+            truncate_dim=matryoshka_dim,
         )
     else:  # if no pre-trained sentence-transformers model
         from sentence_transformers import models
@@ -106,7 +110,9 @@ def _get_sentence_transformer_embedding_model(model_path, torch_dtype):
             word_embedding_model.get_word_embedding_dimension(),
             pooling_mode="lasttoken",
         )
-        model = SentenceTransformer(modules=[word_embedding_model, pooling_model])
+        model = SentenceTransformer(
+            modules=[word_embedding_model, pooling_model], truncate_dim=matryoshka_dim
+        )
 
     return model.cuda()
 
@@ -135,6 +141,7 @@ class HFRunner:
         output_str_only: bool = False,
         trust_remote_code: bool = False,
         patch_model_do_sample_false: bool = False,
+        matryoshka_dim: Optional[int] = None,
     ):
         self.model_type = model_type
         self.output_str_only = output_str_only
@@ -151,6 +158,7 @@ class HFRunner:
                 self.out_queue,
                 model_path,
                 torch_dtype,
+                matryoshka_dim,
             ),
         )
         self.model_proc.start()
@@ -225,7 +233,14 @@ class HFRunner:
         embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=1)
         return embeddings.contiguous()
 
-    def start_model_process(self, in_queue, out_queue, model_path, torch_dtype):
+    def start_model_process(
+        self,
+        in_queue,
+        out_queue,
+        model_path,
+        torch_dtype,
+        matryoshka_dim: Optional[int] = None,
+    ):
         # Apply model-specific patches
         monkey_patch_gemma2_sdpa()
 
@@ -259,7 +274,7 @@ class HFRunner:
                 self.processor = AutoProcessor.from_pretrained(model_path)
             else:
                 self.model = _get_sentence_transformer_embedding_model(
-                    model_path, torch_dtype
+                    model_path, torch_dtype, matryoshka_dim=matryoshka_dim
                 )
         elif self.model_type == "reward" or self.model_type == "cross_encoder":
             from transformers import AutoModelForSequenceClassification
@@ -519,6 +534,7 @@ class SRTRunner:
         lora_target_modules: Optional[List[str]] = None,
         enable_lora: Optional[bool] = None,
         max_loaded_loras: Optional[int] = None,
+        json_model_override_args: Optional[dict[str, Any]] = None,
         lora_eviction_policy: str = "lru",
     ):
         self.model_type = model_type
@@ -566,6 +582,11 @@ class SRTRunner:
             lora_target_modules=lora_target_modules,
             enable_lora=enable_lora,
             max_loaded_loras=max_loaded_loras,
+            json_model_override_args=(
+                json.dumps(json_model_override_args)
+                if json_model_override_args
+                else "{}"
+            ),
             lora_eviction_policy=lora_eviction_policy,
             **spec_kwargs,
         )
@@ -594,6 +615,7 @@ class SRTRunner:
         logprob_start_len: int = 0,
         top_k: Optional[int] = None,
         token_ids_logprob: Optional[List[int]] = None,
+        dimensions: Optional[int] = None,
     ):
         if self.is_generation:
             return self.forward_generation_raw(
@@ -607,7 +629,9 @@ class SRTRunner:
             )
         else:
             if self.model_type == "embedding":
-                response = self.engine.encode(prompt=prompts, image_data=image_data)
+                response = self.engine.encode(
+                    prompt=prompts, image_data=image_data, dimensions=dimensions
+                )
                 if isinstance(response, list):
                     logits = [x["embedding"] for x in response]
                 else: