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, 09:45 UTC
[Model] use AutoWeightsLoader for solar (#18113)
rongfu.leng committed
May 17, 2025 at 15:24 UTC
9214e60631a79506e7669650de87806a123e0b0b
5 files changed
+299
-260
vllm/model_executor/models/mixtral_quant.py
+53
-47
@@ -51,7 +51,7 @@ from vllm.model_executor.sampling_metadata import SamplingMetadata
51
from vllm.sequence import IntermediateTensors
52
53
from .interfaces import SupportsPP
54
-from .utils import (is_pp_missing_parameter,
54
+from .utils import (AutoWeightsLoader, is_pp_missing_parameter,
55
make_empty_intermediate_tensors_factory, make_layers,
56
maybe_prefix)
57
@@ -354,50 +354,6 @@ class MixtralModel(nn.Module):
354
hidden_states, _ = self.norm(hidden_states, residual)
355
return hidden_states
356
357
-
358
-class MixtralForCausalLM(nn.Module, SupportsPP):
359
- fall_back_to_pt_during_load = False
360
-
361
- def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
362
- super().__init__()
363
- config = vllm_config.model_config.hf_config
364
- quant_config = vllm_config.quant_config
365
- self.config = config
366
- self.quant_config = quant_config
367
- self.model = MixtralModel(vllm_config=vllm_config,
368
- prefix=maybe_prefix(prefix, "model"))
369
- self.lm_head = ParallelLMHead(config.vocab_size,
370
- config.hidden_size,
371
- quant_config=quant_config)
372
- if self.config.tie_word_embeddings:
373
- self.lm_head.weight = self.model.embed_tokens.weight
374
- self.logits_processor = LogitsProcessor(config.vocab_size)
375
- self.make_empty_intermediate_tensors = (
376
- self.model.make_empty_intermediate_tensors)
377
-
378
- def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
379
- return self.model.get_input_embeddings(input_ids)
380
-
381
- def forward(
382
- self,
383
- input_ids: torch.Tensor,
384
- positions: torch.Tensor,
385
- intermediate_tensors: Optional[IntermediateTensors] = None,
386
- inputs_embeds: Optional[torch.Tensor] = None,
387
- ) -> Union[torch.Tensor, IntermediateTensors]:
388
- hidden_states = self.model(input_ids, positions, intermediate_tensors,
389
- inputs_embeds)
390
- return hidden_states
391
-
392
- def compute_logits(
393
- self,
394
- hidden_states: torch.Tensor,
395
- sampling_metadata: SamplingMetadata,
396
- ) -> Optional[torch.Tensor]:
397
- logits = self.logits_processor(self.lm_head, hidden_states,
398
- sampling_metadata)
399
- return logits
400
-
357
def load_weights(self, weights: Iterable[tuple[str,
358
torch.Tensor]]) -> set[str]:
359
stacked_params_mapping = [
@@ -410,8 +366,6 @@ class MixtralForCausalLM(nn.Module, SupportsPP):
366
params_dict = dict(self.named_parameters())
367
loaded_params: set[str] = set()
368
for name, loaded_weight in weights:
413
- if "rotary_emb.inv_freq" in name:
414
- continue
369
if name.endswith("scale"):
370
# Remapping the name of FP8 kv-scale.
371
name = maybe_remap_kv_scale_name(name, params_dict)
@@ -446,3 +400,55 @@ class MixtralForCausalLM(nn.Module, SupportsPP):
400
weight_loader(param, loaded_weight)
401
loaded_params.add(name)
402
return loaded_params
403
+
404
+
405
+class MixtralForCausalLM(nn.Module, SupportsPP):
406
+ fall_back_to_pt_during_load = False
407
+
408
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
409
+ super().__init__()
410
+ config = vllm_config.model_config.hf_config
411
+ quant_config = vllm_config.quant_config
412
+ self.config = config
413
+ self.quant_config = quant_config
414
+ self.model = MixtralModel(vllm_config=vllm_config,
415
+ prefix=maybe_prefix(prefix, "model"))
416
+ self.lm_head = ParallelLMHead(config.vocab_size,
417
+ config.hidden_size,
418
+ quant_config=quant_config)
419
+ if self.config.tie_word_embeddings:
420
+ self.lm_head.weight = self.model.embed_tokens.weight
421
+ self.logits_processor = LogitsProcessor(config.vocab_size)
422
+ self.make_empty_intermediate_tensors = (
423
+ self.model.make_empty_intermediate_tensors)
424
+
425
+ def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
426
+ return self.model.get_input_embeddings(input_ids)
427
+
428
+ def forward(
429
+ self,
430
+ input_ids: torch.Tensor,
431
+ positions: torch.Tensor,
432
+ intermediate_tensors: Optional[IntermediateTensors] = None,
433
+ inputs_embeds: Optional[torch.Tensor] = None,
434
+ ) -> Union[torch.Tensor, IntermediateTensors]:
435
+ hidden_states = self.model(input_ids, positions, intermediate_tensors,
436
+ inputs_embeds)
437
+ return hidden_states
438
+
439
+ def compute_logits(
440
+ self,
441
+ hidden_states: torch.Tensor,
442
+ sampling_metadata: SamplingMetadata,
443
+ ) -> Optional[torch.Tensor]:
444
+ logits = self.logits_processor(self.lm_head, hidden_states,
445
+ sampling_metadata)
446
+ return logits
447
+
448
+ def load_weights(self, weights: Iterable[tuple[str,
449
+ torch.Tensor]]) -> set[str]:
450
+ loader = AutoWeightsLoader(
451
+ self,
452
+ skip_prefixes=(["rotary_emb.inv_freq"]),
453
+ )
454
+ return loader.load_weights(weights)
vllm/model_executor/models/nemotron.py
+68
-60
@@ -48,7 +48,7 @@ from vllm.sequence import IntermediateTensors
48
from vllm.transformers_utils.configs import NemotronConfig
49
50
from .interfaces import SupportsLoRA, SupportsPP
51
-from .utils import (PPMissingLayer, is_pp_missing_parameter,
51
+from .utils import (AutoWeightsLoader, PPMissingLayer, is_pp_missing_parameter,
52
make_empty_intermediate_tensors_factory, make_layers,
53
maybe_prefix)
54
@@ -300,6 +300,7 @@ class NemotronModel(nn.Module):
300
lora_config = vllm_config.lora_config
301
302
self.config = config
303
+ self.quant_config = quant_config
304
lora_vocab = (lora_config.lora_extra_vocab_size *
305
(lora_config.max_loras or 1)) if lora_config else 0
306
self.vocab_size = config.vocab_size + lora_vocab
@@ -362,6 +363,63 @@ class NemotronModel(nn.Module):
363
hidden_states, _ = self.norm(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
+ ]
374
+ params_dict = dict(self.named_parameters())
375
+ loaded_params: set[str] = set()
376
+ for name, loaded_weight in weights:
377
+ if (self.quant_config is not None and
378
+ (scale_name := self.quant_config.get_cache_scale(name))):
379
+ # Loading kv cache quantization scales
380
+ param = params_dict[scale_name]
381
+ weight_loader = getattr(param, "weight_loader",
382
+ default_weight_loader)
383
+ loaded_weight = (loaded_weight if loaded_weight.dim() == 0 else
384
+ loaded_weight[0])
385
+ weight_loader(param, loaded_weight)
386
+ loaded_params.add(scale_name)
387
+ continue
388
+ for (param_name, weight_name, shard_id) in stacked_params_mapping:
389
+ if weight_name not in name:
390
+ continue
391
+ name = name.replace(weight_name, param_name)
392
+ # Skip loading extra bias for GPTQ models.
393
+ if name.endswith(".bias") and name not in params_dict:
394
+ continue
395
+
396
+ if is_pp_missing_parameter(name, self):
397
+ continue
398
+
399
+ param = params_dict[name]
400
+ weight_loader = param.weight_loader
401
+ weight_loader(param, loaded_weight, shard_id)
402
+
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
+ # Remapping the name of FP8 kv-scale.
409
+ name = maybe_remap_kv_scale_name(name, params_dict)
410
+ if name is None:
411
+ continue
412
+
413
+ if is_pp_missing_parameter(name, self):
414
+ continue
415
+
416
+ param = params_dict[name]
417
+ weight_loader = getattr(param, "weight_loader",
418
+ default_weight_loader)
419
+ weight_loader(param, loaded_weight)
420
+ loaded_params.add(name)
421
+ return loaded_params
422
+
423
424
class NemotronForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
425
packed_modules_mapping = {
@@ -444,64 +502,14 @@ class NemotronForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
502
503
def load_weights(self, weights: Iterable[tuple[str,
504
torch.Tensor]]) -> set[str]:
447
- stacked_params_mapping = [
448
- # (param_name, shard_name, shard_id)
449
- (".qkv_proj", ".q_proj", "q"),
450
- (".qkv_proj", ".k_proj", "k"),
451
- (".qkv_proj", ".v_proj", "v"),
452
- ]
453
- params_dict = dict(self.named_parameters())
454
- loaded_params: set[str] = set()
455
- for name, loaded_weight in weights:
456
- if "rotary_emb.inv_freq" in name:
457
- continue
458
- if ("rotary_emb.cos_cached" in name
459
- or "rotary_emb.sin_cached" in name):
505
+ loader = AutoWeightsLoader(
506
+ self,
507
+ skip_prefixes=([
508
+ "rotary_emb.inv_freq",
509
# Models trained using ColossalAI may include these tensors in
510
# the checkpoint. Skip them.
462
- continue
463
- if (self.quant_config is not None and
464
- (scale_name := self.quant_config.get_cache_scale(name))):
465
- # Loading kv cache quantization scales
466
- param = params_dict[scale_name]
467
- weight_loader = getattr(param, "weight_loader",
468
- default_weight_loader)
469
- loaded_weight = (loaded_weight if loaded_weight.dim() == 0 else
470
- loaded_weight[0])
471
- weight_loader(param, loaded_weight)
472
- loaded_params.add(scale_name)
473
- continue
474
- for (param_name, weight_name, shard_id) in stacked_params_mapping:
475
- if weight_name not in name:
476
- continue
477
- name = name.replace(weight_name, param_name)
478
- # Skip loading extra bias for GPTQ models.
479
- if name.endswith(".bias") and name not in params_dict:
480
- continue
481
-
482
- if is_pp_missing_parameter(name, self):
483
- continue
484
-
485
- param = params_dict[name]
486
- weight_loader = param.weight_loader
487
- weight_loader(param, loaded_weight, shard_id)
488
-
489
- break
490
- else:
491
- # Skip loading extra bias for GPTQ models.
492
- if name.endswith(".bias") and name not in params_dict:
493
- continue
494
- # Remapping the name of FP8 kv-scale.
495
- name = maybe_remap_kv_scale_name(name, params_dict)
496
- if name is None:
497
- continue
498
-
499
- if is_pp_missing_parameter(name, self):
500
- continue
501
-
502
- param = params_dict[name]
503
- weight_loader = getattr(param, "weight_loader",
504
- default_weight_loader)
505
- weight_loader(param, loaded_weight)
506
- loaded_params.add(name)
507
- return loaded_params
511
+ "rotary_emb.cos_cached",
512
+ "rotary_emb.sin_cached"
513
+ ]),
514
+ )
515
+ return loader.load_weights(weights)
vllm/model_executor/models/olmo.py
+56
-47
@@ -47,7 +47,7 @@ from vllm.model_executor.sampling_metadata import SamplingMetadata
47
from vllm.sequence import IntermediateTensors
48
49
from .interfaces import SupportsPP
50
-from .utils import (is_pp_missing_parameter,
50
+from .utils import (AutoWeightsLoader, is_pp_missing_parameter,
51
make_empty_intermediate_tensors_factory, make_layers,
52
maybe_prefix)
53
@@ -285,6 +285,45 @@ class OlmoModel(nn.Module):
285
hidden_states = self.norm(hidden_states)
286
return hidden_states
287
288
+ def load_weights(self, weights: Iterable[tuple[str,
289
+ torch.Tensor]]) -> set[str]:
290
+ stacked_params_mapping = [
291
+ # (param_name, shard_name, shard_id)
292
+ ("qkv_proj", "q_proj", "q"),
293
+ ("qkv_proj", "k_proj", "k"),
294
+ ("qkv_proj", "v_proj", "v"),
295
+ ("gate_up_proj", "gate_proj", 0),
296
+ ("gate_up_proj", "up_proj", 1),
297
+ ]
298
+ params_dict = dict(self.named_parameters(remove_duplicate=False))
299
+ loaded_params: set[str] = set()
300
+ for name, loaded_weight in weights:
301
+ for (param_name, weight_name, shard_id) in stacked_params_mapping:
302
+ if weight_name not in name:
303
+ continue
304
+ name = name.replace(weight_name, param_name)
305
+ # Skip loading extra bias for GPTQ models.
306
+ if name.endswith(".bias") and name not in params_dict:
307
+ continue
308
+ if is_pp_missing_parameter(name, self):
309
+ continue
310
+ param = params_dict[name]
311
+ weight_loader = param.weight_loader
312
+ weight_loader(param, loaded_weight, shard_id)
313
+ break
314
+ else:
315
+ # Skip loading extra bias for GPTQ models.
316
+ if name.endswith(".bias") and name not in params_dict:
317
+ continue
318
+ if is_pp_missing_parameter(name, self):
319
+ continue
320
+ param = params_dict[name]
321
+ weight_loader = getattr(param, "weight_loader",
322
+ default_weight_loader)
323
+ weight_loader(param, loaded_weight)
324
+ loaded_params.add(name)
325
+ return loaded_params
326
+
327
328
class OlmoForCausalLM(nn.Module, SupportsPP):
329
"""
@@ -341,51 +380,21 @@ class OlmoForCausalLM(nn.Module, SupportsPP):
380
381
def load_weights(self, weights: Iterable[tuple[str,
382
torch.Tensor]]) -> set[str]:
344
- stacked_params_mapping = [
345
- # (param_name, shard_name, shard_id)
346
- ("qkv_proj", "q_proj", "q"),
347
- ("qkv_proj", "k_proj", "k"),
348
- ("qkv_proj", "v_proj", "v"),
349
- ("gate_up_proj", "gate_proj", 0),
350
- ("gate_up_proj", "up_proj", 1),
351
- ]
352
- params_dict = dict(self.named_parameters(remove_duplicate=False))
353
- loaded_params: set[str] = set()
354
- for name, loaded_weight in weights:
355
- if "rotary_emb.inv_freq" in name:
356
- continue
357
- if ("rotary_emb.cos_cached" in name
358
- or "rotary_emb.sin_cached" in name):
383
+ loader = AutoWeightsLoader(
384
+ self,
385
+ skip_prefixes=([
386
+ "rotary_emb.inv_freq",
387
# Models trained using ColossalAI may include these tensors in
388
# the checkpoint. Skip them.
361
- continue
362
- # With tie_word_embeddings, we can skip lm_head.weight
363
- # The weight might appear unnecessarily in the files if the model is
364
- # processed with quantization, LoRA, fine-tuning, etc.
365
- if self.config.tie_word_embeddings and "lm_head.weight" in name:
366
- continue
367
- for (param_name, weight_name, shard_id) in stacked_params_mapping:
368
- if weight_name not in name:
369
- continue
370
- name = name.replace(weight_name, param_name)
371
- # Skip loading extra bias for GPTQ models.
372
- if name.endswith(".bias") and name not in params_dict:
373
- continue
374
- if is_pp_missing_parameter(name, self):
375
- continue
376
- param = params_dict[name]
377
- weight_loader = param.weight_loader
378
- weight_loader(param, loaded_weight, shard_id)
379
- break
380
- else:
381
- # Skip loading extra bias for GPTQ models.
382
- if name.endswith(".bias") and name not in params_dict:
383
- continue
384
- if is_pp_missing_parameter(name, self):
385
- continue
386
- param = params_dict[name]
387
- weight_loader = getattr(param, "weight_loader",
388
- default_weight_loader)
389
- weight_loader(param, loaded_weight)
390
- loaded_params.add(name)
391
- return loaded_params
389
+ "rotary_emb.cos_cached",
390
+ "rotary_emb.sin_cached",
391
+ "lm_head.weight"
392
+ ] if self.config.tie_word_embeddings else [
393
+ "rotary_emb.inv_freq",
394
+ # Models trained using ColossalAI may include these tensors in
395
+ # the checkpoint. Skip them.
396
+ "rotary_emb.cos_cached",
397
+ "rotary_emb.sin_cached"
398
+ ]),
399
+ )
400
+ return loader.load_weights(weights)
vllm/model_executor/models/olmo2.py
+52
-44
@@ -49,8 +49,8 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
49
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
50
from vllm.model_executor.models.interfaces import SupportsPP
51
from vllm.model_executor.models.utils import (
52
- is_pp_missing_parameter, make_empty_intermediate_tensors_factory,
53
- make_layers, maybe_prefix)
52
+ AutoWeightsLoader, is_pp_missing_parameter,
53
+ make_empty_intermediate_tensors_factory, make_layers, maybe_prefix)
54
from vllm.model_executor.sampling_metadata import SamplingMetadata
55
from vllm.sequence import IntermediateTensors
56
@@ -314,6 +314,40 @@ class Olmo2Model(nn.Module):
314
hidden_states = self.norm(hidden_states)
315
return hidden_states
316
317
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
318
+ stacked_params_mapping = [
319
+ # (param_name, shard_name, shard_id)
320
+ ("qkv_proj", "q_proj", "q"),
321
+ ("qkv_proj", "k_proj", "k"),
322
+ ("qkv_proj", "v_proj", "v"),
323
+ ("gate_up_proj", "gate_proj", 0),
324
+ ("gate_up_proj", "up_proj", 1),
325
+ ]
326
+
327
+ params_dict = dict(self.named_parameters(remove_duplicate=False))
328
+ for name, loaded_weight in weights:
329
+ if is_pp_missing_parameter(name, self):
330
+ continue
331
+ for param_name, weight_name, shard_id in stacked_params_mapping:
332
+ if weight_name not in name:
333
+ continue
334
+ name = name.replace(weight_name, param_name)
335
+ # Skip loading extra bias for GPTQ models.
336
+ if name.endswith(".bias") and name not in params_dict:
337
+ continue
338
+ param = params_dict[name]
339
+ weight_loader = param.weight_loader # type: ignore
340
+ weight_loader(param, loaded_weight, shard_id)
341
+ break
342
+ else:
343
+ # Skip loading extra bias for GPTQ models.
344
+ if name.endswith(".bias") and name not in params_dict:
345
+ continue
346
+ param = params_dict[name]
347
+ weight_loader = getattr(param, "weight_loader",
348
+ default_weight_loader)
349
+ weight_loader(param, loaded_weight)
350
+
351
352
class Olmo2ForCausalLM(nn.Module, SupportsPP):
353
"""
@@ -367,47 +401,21 @@ class Olmo2ForCausalLM(nn.Module, SupportsPP):
401
return logits
402
403
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
370
- stacked_params_mapping = [
371
- # (param_name, shard_name, shard_id)
372
- ("qkv_proj", "q_proj", "q"),
373
- ("qkv_proj", "k_proj", "k"),
374
- ("qkv_proj", "v_proj", "v"),
375
- ("gate_up_proj", "gate_proj", 0),
376
- ("gate_up_proj", "up_proj", 1),
377
- ]
378
-
379
- params_dict = dict(self.named_parameters(remove_duplicate=False))
380
- for name, loaded_weight in weights:
381
- if "rotary_emb.inv_freq" in name:
382
- continue
383
- if ("rotary_emb.cos_cached" in name
384
- or "rotary_emb.sin_cached" in name):
404
+ loader = AutoWeightsLoader(
405
+ self,
406
+ skip_prefixes=([
407
+ "rotary_emb.inv_freq",
408
# Models trained using ColossalAI may include these tensors in
409
# the checkpoint. Skip them.
387
- continue
388
- if is_pp_missing_parameter(name, self):
389
- continue
390
- # With tie_word_embeddings, we can skip lm_head.weight
391
- # The weight might appear unnecessarily in the files if the model is
392
- # processed with quantization, LoRA, fine-tuning, etc.
393
- if self.config.tie_word_embeddings and "lm_head.weight" in name:
394
- continue
395
- for param_name, weight_name, shard_id in stacked_params_mapping:
396
- if weight_name not in name:
397
- continue
398
- name = name.replace(weight_name, param_name)
399
- # Skip loading extra bias for GPTQ models.
400
- if name.endswith(".bias") and name not in params_dict:
401
- continue
402
- param = params_dict[name]
403
- weight_loader = param.weight_loader # type: ignore
404
- weight_loader(param, loaded_weight, shard_id)
405
- break
406
- else:
407
- # Skip loading extra bias for GPTQ models.
408
- if name.endswith(".bias") and name not in params_dict:
409
- continue
410
- param = params_dict[name]
411
- weight_loader = getattr(param, "weight_loader",
412
- default_weight_loader)
413
- weight_loader(param, loaded_weight)
410
+ "rotary_emb.cos_cached",
411
+ "rotary_emb.sin_cached",
412
+ "lm_head.weight"
413
+ ] if self.config.tie_word_embeddings else [
414
+ "rotary_emb.inv_freq",
415
+ # Models trained using ColossalAI may include these tensors in
416
+ # the checkpoint. Skip them.
417
+ "rotary_emb.cos_cached",
418
+ "rotary_emb.sin_cached"
419
+ ]),
420
+ )
421
+ return loader.load_weights(weights)
vllm/model_executor/models/solar.py
+70
-62
@@ -50,7 +50,7 @@ from vllm.model_executor.sampling_metadata import SamplingMetadata
50
from vllm.sequence import IntermediateTensors
51
52
from .interfaces import SupportsLoRA, SupportsPP
53
-from .utils import (PPMissingLayer, is_pp_missing_parameter,
53
+from .utils import (AutoWeightsLoader, PPMissingLayer, is_pp_missing_parameter,
54
make_empty_intermediate_tensors_factory, make_layers,
55
maybe_prefix)
56
@@ -269,6 +269,7 @@ class SolarModel(nn.Module):
269
lora_config = vllm_config.lora_config
270
271
self.config = config
272
+ self.quant_config = quant_config
273
lora_vocab = ((lora_config.lora_extra_vocab_size *
274
(lora_config.max_loras or 1)) if lora_config else 0)
275
self.vocab_size = config.vocab_size + lora_vocab
@@ -360,6 +361,65 @@ class SolarModel(nn.Module):
361
hidden_states, _ = self.norm(hidden_states, residual)
362
return hidden_states
363
364
+ def load_weights(self, weights: Iterable[tuple[str,
365
+ torch.Tensor]]) -> set[str]:
366
+ stacked_params_mapping = [
367
+ # (param_name, shard_name, shard_id)
368
+ (".qkv_proj", ".q_proj", "q"),
369
+ (".qkv_proj", ".k_proj", "k"),
370
+ (".qkv_proj", ".v_proj", "v"),
371
+ (".gate_up_proj", ".gate_proj", 0),
372
+ (".gate_up_proj", ".up_proj", 1),
373
+ ]
374
+ params_dict = dict(self.named_parameters())
375
+ loaded_params: set[str] = set()
376
+ for name, loaded_weight in weights:
377
+ if (self.quant_config is not None and
378
+ (scale_name := self.quant_config.get_cache_scale(name))):
379
+ # Loading kv cache quantization scales
380
+ param = params_dict[scale_name]
381
+ weight_loader = getattr(param, "weight_loader",
382
+ default_weight_loader)
383
+ loaded_weight = (loaded_weight if loaded_weight.dim() == 0 else
384
+ loaded_weight[0])
385
+ weight_loader(param, loaded_weight)
386
+ loaded_params.add(scale_name)
387
+ continue
388
+ for param_name, weight_name, shard_id in stacked_params_mapping:
389
+ if weight_name not in name:
390
+ continue
391
+ name = name.replace(weight_name, param_name)
392
+ # Skip loading extra bias for GPTQ models.
393
+ if name.endswith(".bias") and name not in params_dict:
394
+ continue
395
+
396
+ if is_pp_missing_parameter(name, self):
397
+ continue
398
+
399
+ param = params_dict[name]
400
+ weight_loader = param.weight_loader
401
+ weight_loader(param, loaded_weight, shard_id)
402
+
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
+ # Remapping the name of FP8 kv-scale.
409
+ name = maybe_remap_kv_scale_name(name, params_dict)
410
+ if name is None:
411
+ continue
412
+
413
+ if is_pp_missing_parameter(name, self):
414
+ continue
415
+
416
+ param = params_dict[name]
417
+ weight_loader = getattr(param, "weight_loader",
418
+ default_weight_loader)
419
+ weight_loader(param, loaded_weight)
420
+ loaded_params.add(name)
421
+ return loaded_params
422
+
423
424
class SolarForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
425
packed_modules_mapping = {
@@ -440,66 +500,14 @@ class SolarForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
500
501
def load_weights(self, weights: Iterable[tuple[str,
502
torch.Tensor]]) -> set[str]:
443
- stacked_params_mapping = [
444
- # (param_name, shard_name, shard_id)
445
- (".qkv_proj", ".q_proj", "q"),
446
- (".qkv_proj", ".k_proj", "k"),
447
- (".qkv_proj", ".v_proj", "v"),
448
- (".gate_up_proj", ".gate_proj", 0),
449
- (".gate_up_proj", ".up_proj", 1),
450
- ]
451
- params_dict = dict(self.named_parameters())
452
- loaded_params: set[str] = set()
453
- for name, loaded_weight in weights:
454
- if "rotary_emb.inv_freq" in name:
455
- continue
456
- if ("rotary_emb.cos_cached" in name
457
- or "rotary_emb.sin_cached" in name):
503
+ loader = AutoWeightsLoader(
504
+ self,
505
+ skip_prefixes=([
506
+ "rotary_emb.inv_freq",
507
# Models trained using ColossalAI may include these tensors in
508
# the checkpoint. Skip them.
460
- continue
461
- if (self.quant_config is not None and
462
- (scale_name := self.quant_config.get_cache_scale(name))):
463
- # Loading kv cache quantization scales
464
- param = params_dict[scale_name]
465
- weight_loader = getattr(param, "weight_loader",
466
- default_weight_loader)
467
- loaded_weight = (loaded_weight if loaded_weight.dim() == 0 else
468
- loaded_weight[0])
469
- weight_loader(param, loaded_weight)
470
- loaded_params.add(scale_name)
471
- continue
472
- for param_name, weight_name, shard_id in stacked_params_mapping:
473
- if weight_name not in name:
474
- continue
475
- name = name.replace(weight_name, param_name)
476
- # Skip loading extra bias for GPTQ models.
477
- if name.endswith(".bias") and name not in params_dict:
478
- continue
479
-
480
- if is_pp_missing_parameter(name, self):
481
- continue
482
-
483
- param = params_dict[name]
484
- weight_loader = param.weight_loader
485
- weight_loader(param, loaded_weight, shard_id)
486
-
487
- break
488
- else:
489
- # Skip loading extra bias for GPTQ models.
490
- if name.endswith(".bias") and name not in params_dict:
491
- continue
492
- # Remapping the name of FP8 kv-scale.
493
- name = maybe_remap_kv_scale_name(name, params_dict)
494
- if name is None:
495
- continue
496
-
497
- if is_pp_missing_parameter(name, self):
498
- continue
499
-
500
- param = params_dict[name]
501
- weight_loader = getattr(param, "weight_loader",
502
- default_weight_loader)
503
- weight_loader(param, loaded_weight)
504
- loaded_params.add(name)
505
- return loaded_params
509
+ "rotary_emb.cos_cached",
510
+ "rotary_emb.sin_cached"
511
+ ]),
512
+ )
513
+ return loader.load_weights(weights)