A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 6, 00:45 UTC
[Model] Add AutoWeightsLoader support for jais (#38074)
Signed-off-by: grYe99 <guorongye99@gmail.com> Co-authored-by: grYe99 <guorongye99@gmail.com>
grYe99 committed
Mar 25, 2026 at 20:38 UTC
7ac48fd3577f35e235ef96e690c3fc9b847fd26a
1 file changed
+35
-33
vllm/model_executor/models/jais.py
+35
-33
@@ -53,6 +53,7 @@ from vllm.transformers_utils.configs.jais import JAISConfig
53
54
from .interfaces import SupportsPP
55
from .utils import (
56
+ AutoWeightsLoader,
57
is_pp_missing_parameter,
58
make_empty_intermediate_tensors_factory,
59
make_layers,
@@ -311,6 +312,35 @@ class JAISModel(nn.Module):
312
hidden_states = self.ln_f(hidden_states)
313
return hidden_states
314
315
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
316
+ params_dict = dict(self.named_parameters(remove_duplicate=False))
317
+ loaded_params: set[str] = set()
318
+ for name, loaded_weight in weights:
319
+ if ".attn.bias" in name or ".attn.masked_bias" in name:
320
+ # Skip attention mask.
321
+ # NOTE: "c_attn.bias" should not be skipped.
322
+ continue
323
+ if "relative_pe" in name:
324
+ continue
325
+
326
+ if is_pp_missing_parameter(name, self):
327
+ continue
328
+
329
+ param = params_dict[name]
330
+ # The HF's GPT-2 implementation uses Conv1D instead of Linear.
331
+ # Because of this, we need to transpose the weights.
332
+ # Note(zhuohan): the logic below might break quantized models.
333
+ for conv1d_weight_name in ["c_attn", "c_proj", "c_fc"]:
334
+ if conv1d_weight_name not in name:
335
+ continue
336
+ if not name.endswith(".weight"):
337
+ continue
338
+ loaded_weight = loaded_weight.t()
339
+ weight_loader = getattr(param, "weight_loader", default_weight_loader)
340
+ weight_loader(param, loaded_weight)
341
+ loaded_params.add(name)
342
+ return loaded_params
343
+
344
345
class JAISLMHeadModel(nn.Module, SupportsPP):
346
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
@@ -364,36 +394,8 @@ class JAISLMHeadModel(nn.Module, SupportsPP):
394
return logits
395
396
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
367
- params_dict = dict(self.named_parameters(remove_duplicate=False))
368
- loaded_params: set[str] = set()
369
- for name, loaded_weight in weights:
370
- if "lm_head.weight" in name:
371
- # GPT-2 ties the weights of the embedding layer and the final
372
- # linear layer.
373
- continue
374
- if ".attn.bias" in name or ".attn.masked_bias" in name:
375
- # Skip attention mask.
376
- # NOTE: "c_attn.bias" should not be skipped.
377
- continue
378
- if "relative_pe" in name:
379
- continue
380
- if not name.startswith("transformer."):
381
- name = "transformer." + name
382
-
383
- if is_pp_missing_parameter(name, self):
384
- continue
385
-
386
- param = params_dict[name]
387
- # The HF's GPT-2 implementation uses Conv1D instead of Linear.
388
- # Because of this, we need to transpose the weights.
389
- # Note(zhuohan): the logic below might break quantized models.
390
- for conv1d_weight_name in ["c_attn", "c_proj", "c_fc"]:
391
- if conv1d_weight_name not in name:
392
- continue
393
- if not name.endswith(".weight"):
394
- continue
395
- loaded_weight = loaded_weight.t()
396
- weight_loader = getattr(param, "weight_loader", default_weight_loader)
397
- weight_loader(param, loaded_weight)
398
- loaded_params.add(name)
399
- return loaded_params
397
+ loader = AutoWeightsLoader(
398
+ self,
399
+ skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
400
+ )
401
+ return loader.load_weights(weights)