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 2, 20:45 UTC
[Model] Use AutoWeightsLoader for InternLM2 (#38278)
Signed-off-by: Jesus De Jesus <dejesus.9297@gmail.com> Signed-off-by: javierdejesusda <javier.dejesusj9@gmail.com> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>
Javier De Jesus committed
May 26, 2026 at 12:39 UTC
ebd0692f80e6d0ce3f70144e0e271c3fc8f3ab40
1 file changed
+42
-34
vllm/model_executor/models/internlm2.py
+42
-34
@@ -41,6 +41,7 @@ from vllm.sequence import IntermediateTensors
41
from .interfaces import SupportsLoRA, SupportsPP
42
from .interfaces_base import default_pooling_type
43
from .utils import (
44
+ AutoWeightsLoader,
45
StageMissingLayer,
46
is_pp_missing_parameter,
47
make_empty_intermediate_tensors_factory,
@@ -308,6 +309,42 @@ class InternLM2Model(nn.Module):
309
hidden_states, _ = self.norm(hidden_states, residual)
310
return hidden_states
311
312
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
313
+ stacked_params_mapping = [
314
+ # (param_name, shard_name, shard_id)
315
+ ("gate_up_proj", "w1", 0),
316
+ ("gate_up_proj", "w3", 1),
317
+ ]
318
+ params_dict = dict(self.named_parameters())
319
+ loaded_params: set[str] = set()
320
+ for name, loaded_weight in weights:
321
+ if "rotary_emb.inv_freq" in name:
322
+ continue
323
+ for param_name, weight_name, shard_id in stacked_params_mapping:
324
+ if weight_name not in name:
325
+ continue
326
+ name = name.replace(weight_name, param_name)
327
+ # Skip loading extra bias for GPTQ models.
328
+ if name.endswith(".bias") and name not in params_dict:
329
+ continue
330
+ if is_pp_missing_parameter(name, self):
331
+ continue
332
+ param = params_dict[name]
333
+ weight_loader = param.weight_loader
334
+ weight_loader(param, loaded_weight, shard_id)
335
+ break
336
+ else:
337
+ # Skip loading extra bias for GPTQ models.
338
+ if name.endswith(".bias") and name not in params_dict:
339
+ continue
340
+ if is_pp_missing_parameter(name, self):
341
+ continue
342
+ param = params_dict[name]
343
+ weight_loader = getattr(param, "weight_loader", default_weight_loader)
344
+ weight_loader(param, loaded_weight)
345
+ loaded_params.add(name)
346
+ return loaded_params
347
+
348
349
class InternLM2ForCausalLM(nn.Module, SupportsPP, SupportsLoRA):
350
packed_modules_mapping = {
@@ -368,40 +405,11 @@ class InternLM2ForCausalLM(nn.Module, SupportsPP, SupportsLoRA):
405
return logits
406
407
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
371
- stacked_params_mapping = [
372
- # (param_name, shard_name, shard_id)
373
- ("gate_up_proj", "w1", 0),
374
- ("gate_up_proj", "w3", 1),
375
- ]
376
- params_dict = dict(self.named_parameters())
377
- loaded_params: set[str] = set()
378
- for name, loaded_weight in weights:
379
- if "rotary_emb.inv_freq" in name:
380
- continue
381
- for param_name, weight_name, shard_id in stacked_params_mapping:
382
- if weight_name not in name:
383
- continue
384
- name = name.replace(weight_name, param_name)
385
- # Skip loading extra bias for GPTQ models.
386
- if name.endswith(".bias") and name not in params_dict:
387
- continue
388
- if is_pp_missing_parameter(name, self):
389
- continue
390
- param = params_dict[name]
391
- weight_loader = param.weight_loader
392
- weight_loader(param, loaded_weight, shard_id)
393
- break
394
- else:
395
- # Skip loading extra bias for GPTQ models.
396
- if name.endswith(".bias") and name not in params_dict:
397
- continue
398
- if is_pp_missing_parameter(name, self):
399
- continue
400
- param = params_dict[name]
401
- weight_loader = getattr(param, "weight_loader", default_weight_loader)
402
- weight_loader(param, loaded_weight)
403
- loaded_params.add(name)
404
- return loaded_params
408
+ loader = AutoWeightsLoader(
409
+ self,
410
+ skip_prefixes=(["output."] if self.config.tie_word_embeddings else None),
411
+ )
412
+ return loader.load_weights(weights)
413
414
415
@default_pooling_type(tok_pooling_type="ALL")