Forward fix nightly errors from #44589 (#47151)

Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
This commit is contained in:
Harry Mellor
2026-06-30 14:02:34 +00:00
committed by GitHub
parent ab80b3dff4
commit 62c7d8009f
3 changed files with 27 additions and 17 deletions
+5 -4
View File
@@ -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__()
+16 -10
View File
@@ -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",
+6 -3
View File
@@ -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)