Implement BGE-M3 Sparse Embeddings in SGLang (#10869)
Co-authored-by: Christian Bahls <christian.bahls@planet-ai.de> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
This commit is contained in:
@@ -77,6 +77,7 @@ from sglang.srt.model_loader.utils import (
|
||||
DEFAULT_GPU_MEMORY_FRACTION_FOR_CALIBRATION = (
|
||||
0.8 # Reserve 20% GPU memory headroom for ModelOpt calibration
|
||||
)
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.srt.model_loader.weight_utils import (
|
||||
download_safetensors_index_file_from_hf,
|
||||
download_weights_from_hf,
|
||||
@@ -244,10 +245,19 @@ def _initialize_model(
|
||||
quant_config = _get_quantization_config(
|
||||
model_config, load_config, packed_modules_mapping
|
||||
)
|
||||
return model_class(
|
||||
config=model_config.hf_config,
|
||||
quant_config=quant_config,
|
||||
)
|
||||
|
||||
# Build kwargs conditionally
|
||||
kwargs = {
|
||||
"config": model_config.hf_config,
|
||||
"quant_config": quant_config,
|
||||
}
|
||||
|
||||
# Only add sparse head kwargs if envs.SGLANG_EMBEDDINGS_SPARSE_HEAD.is_set()
|
||||
if envs.SGLANG_EMBEDDINGS_SPARSE_HEAD.is_set():
|
||||
kwargs["sparse_head"] = envs.SGLANG_EMBEDDINGS_SPARSE_HEAD.value
|
||||
kwargs["model_path"] = model_config.model_path
|
||||
|
||||
return model_class(**kwargs)
|
||||
|
||||
|
||||
class BaseModelLoader(ABC):
|
||||
|
||||
Reference in New Issue
Block a user