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 5, 13:30 UTC
Refactor Arctic loading to use AutoWeightsLoader (#38955)
Signed-off-by: Lalit Laxminarayan Bangad <lalitbangad@gmail.com> Co-authored-by: Lalit Laxminarayan Bangad <lalitbangad@meta.com>
lalit10 committed
Apr 3, 2026 at 22:01 UTC
93726b2a1c766a27f10ad7c3f709ffbacb94bcf9
1 file changed
+85
-85
vllm/model_executor/models/arctic.py
+85
-85
@@ -16,7 +16,6 @@ from vllm.distributed import (
16
get_tensor_model_parallel_world_size,
17
tensor_model_parallel_all_reduce,
18
)
19
-from vllm.logger import init_logger
19
from vllm.model_executor.layers.activation import SiluAndMul
20
from vllm.model_executor.layers.attention import Attention
21
from vllm.model_executor.layers.fused_moe import fused_experts, fused_topk
@@ -42,6 +41,7 @@ from vllm.transformers_utils.configs.arctic import ArcticConfig
41
42
from .interfaces import SupportsPP, SupportsQuant
43
from .utils import (
44
+ AutoWeightsLoader,
45
extract_layer_index,
46
is_pp_missing_parameter,
47
make_empty_intermediate_tensors_factory,
@@ -49,8 +49,6 @@ from .utils import (
49
maybe_prefix,
50
)
51
52
-logger = init_logger(__name__)
53
-
52
53
class ArcticMLP(nn.Module):
54
def __init__(
@@ -384,6 +382,7 @@ class ArcticModel(nn.Module):
382
cache_config = vllm_config.cache_config
383
quant_config = vllm_config.quant_config
384
385
+ self.config = config
386
self.vocab_size = config.vocab_size
387
self.embed_tokens = VocabParallelEmbedding(
388
self.vocab_size, config.hidden_size, org_num_embeddings=self.vocab_size
@@ -426,57 +425,6 @@ class ArcticModel(nn.Module):
425
hidden_states = self.norm(hidden_states)
426
return hidden_states
427
429
-
430
-class ArcticForCausalLM(nn.Module, SupportsPP, SupportsQuant):
431
- packed_modules_mapping = {"qkv_proj": ["q_proj", "k_proj", "v_proj"]}
432
-
433
- def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
434
- super().__init__()
435
- config = vllm_config.model_config.hf_config
436
- quant_config = vllm_config.quant_config
437
- self.config = config
438
- self.model = ArcticModel(
439
- vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
440
- )
441
- self.vocab_size = config.vocab_size
442
- self.lm_head = ParallelLMHead(
443
- self.vocab_size,
444
- config.hidden_size,
445
- quant_config=quant_config,
446
- prefix=maybe_prefix(prefix, "lm_head"),
447
- )
448
- if self.config.tie_word_embeddings:
449
- self.lm_head.weight = self.model.embed_tokens.weight
450
- self.num_experts = config.num_local_experts
451
- self.num_experts_per_tok = config.num_experts_per_tok
452
-
453
- self.logits_processor = LogitsProcessor(config.vocab_size)
454
- self.make_empty_intermediate_tensors = (
455
- self.model.make_empty_intermediate_tensors
456
- )
457
-
458
- def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
459
- return self.model.embed_input_ids(input_ids)
460
-
461
- def forward(
462
- self,
463
- input_ids: torch.Tensor | None,
464
- positions: torch.Tensor,
465
- intermediate_tensors: IntermediateTensors | None = None,
466
- inputs_embeds: torch.Tensor | None = None,
467
- ) -> torch.Tensor | IntermediateTensors:
468
- hidden_states = self.model(
469
- input_ids, positions, intermediate_tensors, inputs_embeds
470
- )
471
- return hidden_states
472
-
473
- def compute_logits(
474
- self,
475
- hidden_states: torch.Tensor,
476
- ) -> torch.Tensor | None:
477
- logits = self.logits_processor(self.lm_head, hidden_states)
478
- return logits
479
-
428
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
429
stacked_params_mapping = [
430
# (param_name, shard_name, shard_id)
@@ -487,41 +435,26 @@ class ArcticForCausalLM(nn.Module, SupportsPP, SupportsQuant):
435
436
mlp_params_mapping: list[tuple[str, str, int]] = []
437
expert_params_mapping: list[tuple[str, str, int]] = []
490
- num_layers = self.config.num_hidden_layers
491
-
492
- for layer in range(num_layers):
493
- mlp_params_mapping.append(
494
- (
495
- f"layers.{layer}.residual_mlp.w13.weight",
496
- f"layers.{layer}.residual_mlp.w1.weight",
497
- 0,
498
- )
499
- )
500
- mlp_params_mapping.append(
501
- (
502
- f"layers.{layer}.residual_mlp.w13.weight",
503
- f"layers.{layer}.residual_mlp.w3.weight",
504
- 1,
505
- )
506
- )
507
- if layer % 2 == 0:
508
- # MLP layers
438
+
439
+ for layer in range(self.config.num_hidden_layers):
440
+ is_moe_layer = (layer + 1) % self.config.moe_layer_frequency == 0
441
+ if is_moe_layer and self.config.use_residual:
442
mlp_params_mapping.append(
443
(
511
- f"layers.{layer}.block_sparse_moe.mlp.w13.weight",
512
- f"layers.{layer}.block_sparse_moe.mlp.w1.weight",
444
+ f"layers.{layer}.residual_mlp.w13.weight",
445
+ f"layers.{layer}.residual_mlp.w1.weight",
446
0,
447
)
448
)
449
mlp_params_mapping.append(
450
(
518
- f"layers.{layer}.block_sparse_moe.mlp.w13.weight",
519
- f"layers.{layer}.block_sparse_moe.mlp.w3.weight",
451
+ f"layers.{layer}.residual_mlp.w13.weight",
452
+ f"layers.{layer}.residual_mlp.w3.weight",
453
1,
454
)
455
)
523
- else:
524
- # MoE layers
456
+
457
+ if is_moe_layer:
458
for expert_id in range(self.config.num_local_experts):
459
expert_params_mapping.append(
460
("ws", f"experts.{expert_id}.w1.weight", expert_id)
@@ -532,15 +465,25 @@ class ArcticForCausalLM(nn.Module, SupportsPP, SupportsQuant):
465
expert_params_mapping.append(
466
("ws", f"experts.{expert_id}.w3.weight", expert_id)
467
)
468
+ else:
469
+ mlp_params_mapping.append(
470
+ (
471
+ f"layers.{layer}.block_sparse_moe.mlp.w13.weight",
472
+ f"layers.{layer}.block_sparse_moe.mlp.w1.weight",
473
+ 0,
474
+ )
475
+ )
476
+ mlp_params_mapping.append(
477
+ (
478
+ f"layers.{layer}.block_sparse_moe.mlp.w13.weight",
479
+ f"layers.{layer}.block_sparse_moe.mlp.w3.weight",
480
+ 1,
481
+ )
482
+ )
483
484
params_dict = dict(self.named_parameters())
485
loaded_params: set[str] = set()
486
539
- logger.info(
540
- "It will take ~10 minutes loading from the 16-bit weights. "
541
- "Alternatively, use the prequantized 8-bit weights of arctic "
542
- "and set load-format to `sharded_state` will accelerate loading."
543
- )
487
for name, loaded_weight in weights:
488
for param_name, weight_name, shard_id in stacked_params_mapping:
489
if weight_name not in name:
@@ -585,10 +528,67 @@ class ArcticForCausalLM(nn.Module, SupportsPP, SupportsQuant):
528
if is_pp_missing_parameter(name, self):
529
continue
530
param = params_dict[name]
588
-
531
weight_loader = getattr(
532
param, "weight_loader", default_weight_loader
533
)
534
weight_loader(param, loaded_weight)
535
loaded_params.add(name)
536
return loaded_params
537
+
538
+
539
+class ArcticForCausalLM(nn.Module, SupportsPP, SupportsQuant):
540
+ packed_modules_mapping = {"qkv_proj": ["q_proj", "k_proj", "v_proj"]}
541
+
542
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
543
+ super().__init__()
544
+ config = vllm_config.model_config.hf_config
545
+ quant_config = vllm_config.quant_config
546
+ self.config = config
547
+ self.model = ArcticModel(
548
+ vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
549
+ )
550
+ self.vocab_size = config.vocab_size
551
+ self.lm_head = ParallelLMHead(
552
+ self.vocab_size,
553
+ config.hidden_size,
554
+ quant_config=quant_config,
555
+ prefix=maybe_prefix(prefix, "lm_head"),
556
+ )
557
+ if self.config.tie_word_embeddings:
558
+ self.lm_head.weight = self.model.embed_tokens.weight
559
+ self.num_experts = config.num_local_experts
560
+ self.num_experts_per_tok = config.num_experts_per_tok
561
+
562
+ self.logits_processor = LogitsProcessor(config.vocab_size)
563
+ self.make_empty_intermediate_tensors = (
564
+ self.model.make_empty_intermediate_tensors
565
+ )
566
+
567
+ def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
568
+ return self.model.embed_input_ids(input_ids)
569
+
570
+ def forward(
571
+ self,
572
+ input_ids: torch.Tensor | None,
573
+ positions: torch.Tensor,
574
+ intermediate_tensors: IntermediateTensors | None = None,
575
+ inputs_embeds: torch.Tensor | None = None,
576
+ ) -> torch.Tensor | IntermediateTensors:
577
+ hidden_states = self.model(
578
+ input_ids, positions, intermediate_tensors, inputs_embeds
579
+ )
580
+ return hidden_states
581
+
582
+ def compute_logits(
583
+ self,
584
+ hidden_states: torch.Tensor,
585
+ ) -> torch.Tensor | None:
586
+ logits = self.logits_processor(self.lm_head, hidden_states)
587
+ return logits
588
+
589
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
590
+ loader = AutoWeightsLoader(
591
+ self,
592
+ skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
593
+ )
594
+ return loader.load_weights(weights)