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 15, 16:00 UTC
[Model] use AutoWeightsLoader in model load_weights (#15770)
Signed-off-by: rongfu.leng <rongfu.leng@daocloud.io>
rongfu.leng committed
Apr 2, 2025 at 22:47 UTC
e86c414d6ac540166781f203b1ec48826aab9f7c
4 files changed
+189
-165
docs/source/models/supported_models.md
+5
@@ -218,6 +218,11 @@ See [this page](#generative-models) for more information on how to use generativ
218
* `baichuan-inc/Baichuan2-13B-Chat`, `baichuan-inc/Baichuan-7B`, etc.
219
* ✅︎
220
* ✅︎
221
+- * `BambaForCausalLM`
222
+ * Bamba
223
+ * `ibm-ai-platform/Bamba-9B-fp8`, `ibm-ai-platform/Bamba-9B`
224
+ *
225
+ *
226
- * `BloomForCausalLM`
227
* BLOOM, BLOOMZ, BLOOMChat
228
* `bigscience/bloom`, `bigscience/bloomz`, etc.
vllm/model_executor/models/bamba.py
+55
-50
@@ -34,7 +34,7 @@ from vllm.utils import LayerBlockType
34
35
from .interfaces import (HasInnerState, IsHybrid, SupportsLoRA, SupportsPP,
36
SupportsQuant, SupportsV0Only)
37
-from .utils import (is_pp_missing_parameter,
37
+from .utils import (AutoWeightsLoader, is_pp_missing_parameter,
38
make_empty_intermediate_tensors_factory, make_layers,
39
maybe_prefix)
40
@@ -363,6 +363,58 @@ class BambaModel(nn.Module):
363
hidden_states, _ = self.final_layernorm(hidden_states, residual)
364
return hidden_states
365
366
+ def load_weights(self, weights: Iterable[Tuple[str,
367
+ torch.Tensor]]) -> Set[str]:
368
+ stacked_params_mapping = [
369
+ # (param_name, shard_name, shard_id)
370
+ ("qkv_proj", "q_proj", "q"),
371
+ ("qkv_proj", "k_proj", "k"),
372
+ ("qkv_proj", "v_proj", "v"),
373
+ ("gate_up_proj", "gate_proj", 0),
374
+ ("gate_up_proj", "up_proj", 1),
375
+ ]
376
+
377
+ params_dict = dict(self.named_parameters())
378
+ loaded_params: Set[str] = set()
379
+ for name, loaded_weight in weights:
380
+ if "rotary_emb.inv_freq" in name:
381
+ continue
382
+
383
+ if "A_log" in name:
384
+ name = name.replace("A_log", "A")
385
+
386
+ if ".self_attn." in name:
387
+ name = name.replace(".self_attn", "")
388
+
389
+ for param_name, weight_name, shard_id in stacked_params_mapping:
390
+ if weight_name not in name:
391
+ continue
392
+
393
+ name = name.replace(weight_name, param_name)
394
+ # Skip loading extra bias for GPTQ models.
395
+ if name.endswith(".bias") and name not in params_dict:
396
+ continue
397
+ # Skip layers on other devices.
398
+ if is_pp_missing_parameter(name, self):
399
+ continue
400
+ param = params_dict[name]
401
+ weight_loader = param.weight_loader
402
+ weight_loader(param, loaded_weight, shard_id)
403
+ break
404
+ else:
405
+ # Skip loading extra bias for GPTQ models.
406
+ if name.endswith(".bias") and name not in params_dict:
407
+ continue
408
+ if is_pp_missing_parameter(name, self):
409
+ continue
410
+
411
+ param = params_dict[name]
412
+ weight_loader = getattr(param, "weight_loader",
413
+ default_weight_loader)
414
+ weight_loader(param, loaded_weight)
415
+ loaded_params.add(name)
416
+ return loaded_params
417
+
418
419
class BambaForCausalLM(nn.Module, HasInnerState, SupportsLoRA, SupportsPP,
420
IsHybrid, SupportsV0Only, SupportsQuant):
@@ -502,52 +554,5 @@ class BambaForCausalLM(nn.Module, HasInnerState, SupportsLoRA, SupportsPP,
554
555
def load_weights(self, weights: Iterable[Tuple[str,
556
torch.Tensor]]) -> Set[str]:
505
- stacked_params_mapping = [
506
- # (param_name, shard_name, shard_id)
507
- ("qkv_proj", "q_proj", "q"),
508
- ("qkv_proj", "k_proj", "k"),
509
- ("qkv_proj", "v_proj", "v"),
510
- ("gate_up_proj", "gate_proj", 0),
511
- ("gate_up_proj", "up_proj", 1),
512
- ]
513
-
514
- params_dict = dict(self.named_parameters())
515
- loaded_params: Set[str] = set()
516
- for name, loaded_weight in weights:
517
- if "rotary_emb.inv_freq" in name:
518
- continue
519
-
520
- if "A_log" in name:
521
- name = name.replace("A_log", "A")
522
-
523
- if ".self_attn." in name:
524
- name = name.replace(".self_attn", "")
525
-
526
- for param_name, weight_name, shard_id in stacked_params_mapping:
527
- if weight_name not in name:
528
- continue
529
-
530
- name = name.replace(weight_name, param_name)
531
- # Skip loading extra bias for GPTQ models.
532
- if name.endswith(".bias") and name not in params_dict:
533
- continue
534
- # Skip layers on other devices.
535
- if is_pp_missing_parameter(name, self):
536
- continue
537
- param = params_dict[name]
538
- weight_loader = param.weight_loader
539
- weight_loader(param, loaded_weight, shard_id)
540
- break
541
- else:
542
- # Skip loading extra bias for GPTQ models.
543
- if name.endswith(".bias") and name not in params_dict:
544
- continue
545
- if is_pp_missing_parameter(name, self):
546
- continue
547
-
548
- param = params_dict[name]
549
- weight_loader = getattr(param, "weight_loader",
550
- default_weight_loader)
551
- weight_loader(param, loaded_weight)
552
- loaded_params.add(name)
553
- return loaded_params
557
+ loader = AutoWeightsLoader(self)
558
+ return loader.load_weights(weights)
\ No newline at end of file
vllm/model_executor/models/exaone.py
+74
-66
@@ -51,7 +51,7 @@ from vllm.sequence import IntermediateTensors
51
from vllm.transformers_utils.configs.exaone import ExaoneConfig
52
53
from .interfaces import SupportsLoRA, SupportsPP
54
-from .utils import (PPMissingLayer, is_pp_missing_parameter,
54
+from .utils import (AutoWeightsLoader, PPMissingLayer, is_pp_missing_parameter,
55
make_empty_intermediate_tensors_factory, make_layers,
56
maybe_prefix)
57
@@ -313,6 +313,7 @@ class ExaoneModel(nn.Module):
313
lora_config = vllm_config.lora_config
314
315
self.config = config
316
+ self.quant_config = quant_config
317
lora_vocab = ((lora_config.lora_extra_vocab_size *
318
(lora_config.max_loras or 1)) if lora_config else 0)
319
self.vocab_size = config.vocab_size + lora_vocab
@@ -384,6 +385,72 @@ class ExaoneModel(nn.Module):
385
hidden_states, _ = self.ln_f(hidden_states, residual)
386
return hidden_states
387
388
+ def load_weights(self, weights: Iterable[Tuple[str,
389
+ torch.Tensor]]) -> Set[str]:
390
+ stacked_params_mapping = [
391
+ # (param_name, shard_name, shard_id)
392
+ (".qkv_proj", ".q_proj", "q"),
393
+ (".qkv_proj", ".k_proj", "k"),
394
+ (".qkv_proj", ".v_proj", "v"),
395
+ (".gate_up_proj", ".c_fc_0", 0),
396
+ (".gate_up_proj", ".c_fc_1", 1),
397
+ ]
398
+ params_dict = dict(self.named_parameters())
399
+ loaded_params: Set[str] = set()
400
+ for name, loaded_weight in weights:
401
+ if "rotary_emb.inv_freq" in name:
402
+ continue
403
+ if ("rotary_emb.cos_cached" in name
404
+ or "rotary_emb.sin_cached" in name):
405
+ # Models trained using ColossalAI may include these tensors in
406
+ # the checkpoint. Skip them.
407
+ continue
408
+ if (self.quant_config is not None and
409
+ (scale_name := self.quant_config.get_cache_scale(name))):
410
+ # Loading kv cache quantization scales
411
+ param = params_dict[scale_name]
412
+ weight_loader = getattr(param, "weight_loader",
413
+ default_weight_loader)
414
+ loaded_weight = (loaded_weight if loaded_weight.dim() == 0 else
415
+ loaded_weight[0])
416
+ weight_loader(param, loaded_weight)
417
+ loaded_params.add(scale_name)
418
+ continue
419
+ for param_name, weight_name, shard_id in stacked_params_mapping:
420
+ if weight_name not in name:
421
+ continue
422
+ name = name.replace(weight_name, param_name)
423
+ # Skip loading extra bias for GPTQ models.
424
+ if name.endswith(".bias") and name not in params_dict:
425
+ continue
426
+
427
+ if is_pp_missing_parameter(name, self):
428
+ continue
429
+
430
+ param = params_dict[name]
431
+ weight_loader = param.weight_loader
432
+ weight_loader(param, loaded_weight, shard_id)
433
+
434
+ break
435
+ else:
436
+ # Skip loading extra bias for GPTQ models.
437
+ if name.endswith(".bias") and name not in params_dict:
438
+ continue
439
+ # Remapping the name of FP8 kv-scale.
440
+ name = maybe_remap_kv_scale_name(name, params_dict)
441
+ if name is None:
442
+ continue
443
+
444
+ if is_pp_missing_parameter(name, self):
445
+ continue
446
+
447
+ param = params_dict[name]
448
+ weight_loader = getattr(param, "weight_loader",
449
+ default_weight_loader)
450
+ weight_loader(param, loaded_weight)
451
+ loaded_params.add(name)
452
+ return loaded_params
453
+
454
455
class ExaoneForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
456
packed_modules_mapping = {
@@ -481,71 +548,12 @@ class ExaoneForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
548
549
def load_weights(self, weights: Iterable[Tuple[str,
550
torch.Tensor]]) -> Set[str]:
484
- stacked_params_mapping = [
485
- # (param_name, shard_name, shard_id)
486
- (".qkv_proj", ".q_proj", "q"),
487
- (".qkv_proj", ".k_proj", "k"),
488
- (".qkv_proj", ".v_proj", "v"),
489
- (".gate_up_proj", ".c_fc_0", 0),
490
- (".gate_up_proj", ".c_fc_1", 1),
491
- ]
492
- params_dict = dict(self.named_parameters())
493
- loaded_params: Set[str] = set()
494
- for name, loaded_weight in weights:
495
- if "rotary_emb.inv_freq" in name:
496
- continue
497
- if ("rotary_emb.cos_cached" in name
498
- or "rotary_emb.sin_cached" in name):
499
- # Models trained using ColossalAI may include these tensors in
500
- # the checkpoint. Skip them.
501
- continue
551
+ loader = AutoWeightsLoader(
552
+ self,
553
# With tie_word_embeddings, we can skip lm_head.weight
554
# The weight might appear unnecessarily in the files if the model is
555
# processed with quantization, LoRA, fine-tuning, etc.
505
- if self.config.tie_word_embeddings and "lm_head.weight" in name:
506
- continue
507
- if (self.quant_config is not None and
508
- (scale_name := self.quant_config.get_cache_scale(name))):
509
- # Loading kv cache quantization scales
510
- param = params_dict[scale_name]
511
- weight_loader = getattr(param, "weight_loader",
512
- default_weight_loader)
513
- loaded_weight = (loaded_weight if loaded_weight.dim() == 0 else
514
- loaded_weight[0])
515
- weight_loader(param, loaded_weight)
516
- loaded_params.add(scale_name)
517
- continue
518
- for param_name, weight_name, shard_id in stacked_params_mapping:
519
- if weight_name not in name:
520
- continue
521
- name = name.replace(weight_name, param_name)
522
- # Skip loading extra bias for GPTQ models.
523
- if name.endswith(".bias") and name not in params_dict:
524
- continue
525
-
526
- if is_pp_missing_parameter(name, self):
527
- continue
528
-
529
- param = params_dict[name]
530
- weight_loader = param.weight_loader
531
- weight_loader(param, loaded_weight, shard_id)
532
-
533
- break
534
- else:
535
- # Skip loading extra bias for GPTQ models.
536
- if name.endswith(".bias") and name not in params_dict:
537
- continue
538
- # Remapping the name of FP8 kv-scale.
539
- name = maybe_remap_kv_scale_name(name, params_dict)
540
- if name is None:
541
- continue
542
-
543
- if is_pp_missing_parameter(name, self):
544
- continue
545
-
546
- param = params_dict[name]
547
- weight_loader = getattr(param, "weight_loader",
548
- default_weight_loader)
549
- weight_loader(param, loaded_weight)
550
- loaded_params.add(name)
551
- return loaded_params
556
+ skip_prefixes=(["lm_head."]
557
+ if self.config.tie_word_embeddings else None),
558
+ )
559
+ return loader.load_weights(weights)
vllm/model_executor/models/falcon.py
+55
-49
@@ -49,7 +49,7 @@ from vllm.sequence import IntermediateTensors
49
from vllm.transformers_utils.configs import RWConfig
50
51
from .interfaces import SupportsPP
52
-from .utils import (is_pp_missing_parameter,
52
+from .utils import (AutoWeightsLoader, is_pp_missing_parameter,
53
make_empty_intermediate_tensors_factory, make_layers,
54
maybe_prefix)
55
@@ -395,6 +395,54 @@ class FalconModel(nn.Module):
395
hidden_states = self.ln_f(hidden_states)
396
return hidden_states
397
398
+ def load_weights(self, weights: Iterable[Tuple[str,
399
+ torch.Tensor]]) -> Set[str]:
400
+ total_num_heads = self.config.num_attention_heads
401
+ if self.config.new_decoder_architecture:
402
+ total_num_kv_heads = self.config.num_kv_heads
403
+ elif self.config.multi_query:
404
+ total_num_kv_heads = 1
405
+ else:
406
+ total_num_kv_heads = total_num_heads
407
+ num_query_heads_per_kv_head = total_num_heads // total_num_kv_heads
408
+ params_dict = dict(self.named_parameters(remove_duplicate=False))
409
+ loaded_params: Set[str] = set()
410
+ for name, loaded_weight in weights:
411
+ # Skip loading extra bias for GPTQ models.
412
+ if name.endswith(".bias") and name not in params_dict:
413
+ continue
414
+ if is_pp_missing_parameter(name, self):
415
+ continue
416
+ param = params_dict[name]
417
+ if "query_key_value" in name:
418
+ output_dim = getattr(param, "output_dim", None)
419
+ loaded_weight_shape = loaded_weight.shape
420
+ if output_dim is not None:
421
+ loaded_weight = loaded_weight.view(
422
+ loaded_weight_shape[:output_dim] +
423
+ (total_num_kv_heads, num_query_heads_per_kv_head + 2,
424
+ -1) + loaded_weight_shape[output_dim + 1:])
425
+ wq = loaded_weight.narrow(
426
+ output_dim + 1, 0,
427
+ num_query_heads_per_kv_head).reshape(
428
+ *loaded_weight_shape[:output_dim], -1,
429
+ *loaded_weight_shape[output_dim + 1:])
430
+ wk = loaded_weight.narrow(
431
+ output_dim + 1, num_query_heads_per_kv_head,
432
+ 1).reshape(*loaded_weight_shape[:output_dim], -1,
433
+ *loaded_weight_shape[output_dim + 1:])
434
+ wv = loaded_weight.narrow(
435
+ output_dim + 1, num_query_heads_per_kv_head + 1,
436
+ 1).reshape(*loaded_weight_shape[:output_dim], -1,
437
+ *loaded_weight_shape[output_dim + 1:])
438
+ loaded_weight = torch.cat([wq, wk, wv], dim=output_dim)
439
+
440
+ weight_loader = getattr(param, "weight_loader",
441
+ default_weight_loader)
442
+ weight_loader(param, loaded_weight)
443
+ loaded_params.add(name)
444
+ return loaded_params
445
+
446
447
class FalconForCausalLM(nn.Module, SupportsPP):
448
packed_modules_mapping = {
@@ -462,51 +510,9 @@ class FalconForCausalLM(nn.Module, SupportsPP):
510
511
def load_weights(self, weights: Iterable[Tuple[str,
512
torch.Tensor]]) -> Set[str]:
465
- total_num_heads = self.config.num_attention_heads
466
- if self.config.new_decoder_architecture:
467
- total_num_kv_heads = self.config.num_kv_heads
468
- elif self.config.multi_query:
469
- total_num_kv_heads = 1
470
- else:
471
- total_num_kv_heads = total_num_heads
472
- num_query_heads_per_kv_head = total_num_heads // total_num_kv_heads
473
- params_dict = dict(self.named_parameters(remove_duplicate=False))
474
- loaded_params: Set[str] = set()
475
- for name, loaded_weight in weights:
476
- if name == "lm_head.weight" and self.tie_word_embeddings:
477
- # Falcon uses tied embeddings except Falcon-11b.
478
- continue
479
- # Skip loading extra bias for GPTQ models.
480
- if name.endswith(".bias") and name not in params_dict:
481
- continue
482
- if is_pp_missing_parameter(name, self):
483
- continue
484
- param = params_dict[name]
485
- if "query_key_value" in name:
486
- output_dim = getattr(param, "output_dim", None)
487
- loaded_weight_shape = loaded_weight.shape
488
- if output_dim is not None:
489
- loaded_weight = loaded_weight.view(
490
- loaded_weight_shape[:output_dim] +
491
- (total_num_kv_heads, num_query_heads_per_kv_head + 2,
492
- -1) + loaded_weight_shape[output_dim + 1:])
493
- wq = loaded_weight.narrow(
494
- output_dim + 1, 0,
495
- num_query_heads_per_kv_head).reshape(
496
- *loaded_weight_shape[:output_dim], -1,
497
- *loaded_weight_shape[output_dim + 1:])
498
- wk = loaded_weight.narrow(
499
- output_dim + 1, num_query_heads_per_kv_head,
500
- 1).reshape(*loaded_weight_shape[:output_dim], -1,
501
- *loaded_weight_shape[output_dim + 1:])
502
- wv = loaded_weight.narrow(
503
- output_dim + 1, num_query_heads_per_kv_head + 1,
504
- 1).reshape(*loaded_weight_shape[:output_dim], -1,
505
- *loaded_weight_shape[output_dim + 1:])
506
- loaded_weight = torch.cat([wq, wk, wv], dim=output_dim)
507
-
508
- weight_loader = getattr(param, "weight_loader",
509
- default_weight_loader)
510
- weight_loader(param, loaded_weight)
511
- loaded_params.add(name)
512
- return loaded_params
513
+ loader = AutoWeightsLoader(
514
+ self,
515
+ skip_prefixes=(["lm_head."]
516
+ if self.config.tie_word_embeddings else None),
517
+ )
518
+ return loader.load_weights(weights)