Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
This commit is contained in:
@@ -348,7 +348,11 @@ class CohereForCausalLM(nn.Module, SupportsLoRA, SupportsPP, SupportsQuant):
|
||||
".v_proj": (".qkv_proj", "v"),
|
||||
".gate_proj": (".gate_up_proj", 0),
|
||||
".up_proj": (".gate_up_proj", 1),
|
||||
}
|
||||
},
|
||||
# ModelOpt NVFP4 checkpoints carry raw quantizer-module state
|
||||
# (e.g. "*.weight_quantizer._double_scale"); drop them before loading.
|
||||
# See #41925.
|
||||
orig_to_new_substr={"_quantizer.": None},
|
||||
)
|
||||
packed_modules_mapping = {
|
||||
"qkv_proj": ["q_proj", "k_proj", "v_proj"],
|
||||
@@ -356,9 +360,6 @@ class CohereForCausalLM(nn.Module, SupportsLoRA, SupportsPP, SupportsQuant):
|
||||
}
|
||||
# LoRA specific attributes
|
||||
embedding_modules = {"embed_tokens": "input_embeddings"}
|
||||
# ModelOpt NVFP4 checkpoints carry raw quantizer-module state
|
||||
# (e.g. "*.weight_quantizer._double_scale"); drop them before loading. See #41925.
|
||||
hf_to_vllm_mapper = WeightsMapper(orig_to_new_substr={"_quantizer.": None})
|
||||
|
||||
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
|
||||
super().__init__()
|
||||
|
||||
@@ -291,6 +291,17 @@ class Gemma3DecoderLayer(nn.Module):
|
||||
|
||||
@support_torch_compile
|
||||
class Gemma3Model(nn.Module):
|
||||
hf_to_vllm_mapper = WeightsMapper(
|
||||
orig_to_new_stacked={
|
||||
# weight_name: (param_name, shard_id)
|
||||
".q_proj": (".qkv_proj", "q"),
|
||||
".k_proj": (".qkv_proj", "k"),
|
||||
".v_proj": (".qkv_proj", "v"),
|
||||
".gate_proj": (".gate_up_proj", 0),
|
||||
".up_proj": (".gate_up_proj", 1),
|
||||
}
|
||||
)
|
||||
|
||||
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
|
||||
super().__init__()
|
||||
config = vllm_config.model_config.hf_config
|
||||
@@ -361,18 +372,13 @@ class Gemma3Model(nn.Module):
|
||||
hidden_states, _ = self.norm(hidden_states, residual)
|
||||
return hidden_states
|
||||
|
||||
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
|
||||
loader = AutoWeightsLoader(self)
|
||||
return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
|
||||
|
||||
|
||||
class Gemma3ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
|
||||
hf_to_vllm_mapper = WeightsMapper(
|
||||
orig_to_new_stacked={
|
||||
# weight_name: (param_name, shard_id)
|
||||
".q_proj": (".qkv_proj", "q"),
|
||||
".k_proj": (".qkv_proj", "k"),
|
||||
".v_proj": (".qkv_proj", "v"),
|
||||
".gate_proj": (".gate_up_proj", 0),
|
||||
".up_proj": (".gate_up_proj", 1),
|
||||
}
|
||||
)
|
||||
hf_to_vllm_mapper = Gemma3Model.hf_to_vllm_mapper
|
||||
packed_modules_mapping = {
|
||||
"qkv_proj": [
|
||||
"q_proj",
|
||||
|
||||
@@ -25,7 +25,7 @@ from ..layers.pooler.tokwise import (
|
||||
from .interfaces import SupportsLateInteraction
|
||||
from .interfaces_base import VllmModelForPooling
|
||||
from .qwen3 import Qwen3ForCausalLM, Qwen3Model
|
||||
from .utils import AutoWeightsLoader, maybe_prefix
|
||||
from .utils import AutoWeightsLoader, WeightsMapper, maybe_prefix
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -193,6 +193,9 @@ class JinaEmbeddingsV5Model(Qwen3ForCausalLM, VllmModelForPooling):
|
||||
"""
|
||||
|
||||
is_pooling_model = True
|
||||
hf_to_vllm_mapper = Qwen3ForCausalLM.hf_to_vllm_mapper | WeightsMapper(
|
||||
orig_to_new_prefix={"": "model."}
|
||||
)
|
||||
|
||||
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
|
||||
super().__init__(vllm_config=vllm_config, prefix=prefix)
|
||||
@@ -254,6 +257,6 @@ class JinaEmbeddingsV5Model(Qwen3ForCausalLM, VllmModelForPooling):
|
||||
tensor = tensor + (lora_B @ lora_A) * scaling
|
||||
yield name, tensor
|
||||
|
||||
loader = AutoWeightsLoader(self.model, ignore_unexpected_prefixes=["lm_head."])
|
||||
loader = AutoWeightsLoader(self, ignore_unexpected_prefixes=["lm_head."])
|
||||
weights = _merge_weights(weights)
|
||||
return loader.load_weights(weights, mapper=self.model.hf_to_vllm_mapper)
|
||||
return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
|
||||
|
||||
Reference in New Issue
Block a user