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 4, 10:00 UTC
Refactor Step3Text loading to use AutoWeightsLoader (#41492)
Signed-off-by: Matthew Santiago <carag.matthew@gmail.com>
Matthew Santiago committed
May 2, 2026 at 01:22 UTC
8586369f617a964235d0d9d32d6ebb1076a4581d
1 file changed
+54
-49
vllm/model_executor/models/step3_text.py
+54
-49
@@ -41,6 +41,7 @@ from vllm.transformers_utils.configs.step3_vl import Step3TextConfig
41
42
from .interfaces import SupportsPP
43
from .utils import (
44
+ AutoWeightsLoader,
45
PPMissingLayer,
46
is_pp_missing_parameter,
47
make_empty_intermediate_tensors_factory,
@@ -382,55 +383,6 @@ class Step3TextModel(nn.Module):
383
hidden_states, _ = self.norm(hidden_states, residual)
384
return hidden_states
385
385
-
386
-class Step3TextForCausalLM(nn.Module, SupportsPP):
387
- def __init__(
388
- self,
389
- *,
390
- vllm_config: VllmConfig,
391
- prefix: str = "",
392
- ):
393
- super().__init__()
394
- config = vllm_config.model_config.hf_config
395
-
396
- self.config = config
397
- self.vllm_config = vllm_config
398
-
399
- self.model = Step3TextModel(vllm_config=vllm_config, prefix=prefix)
400
-
401
- if get_pp_group().is_last_rank:
402
- self.lm_head = ParallelLMHead(
403
- config.vocab_size,
404
- config.hidden_size,
405
- prefix=maybe_prefix(prefix, "lm_head"),
406
- )
407
- self.logits_processor = LogitsProcessor(config.vocab_size)
408
- else:
409
- self.lm_head = PPMissingLayer()
410
-
411
- self.make_empty_intermediate_tensors = (
412
- self.model.make_empty_intermediate_tensors
413
- )
414
-
415
- def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
416
- return self.model.embed_input_ids(input_ids)
417
-
418
- def forward(
419
- self,
420
- input_ids: torch.Tensor | None,
421
- positions: torch.Tensor,
422
- intermediate_tensors: IntermediateTensors | None = None,
423
- inputs_embeds: torch.Tensor | None = None,
424
- ):
425
- hidden_states = self.model(
426
- input_ids, positions, intermediate_tensors, inputs_embeds
427
- )
428
- return hidden_states
429
-
430
- def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor:
431
- logits = self.logits_processor(self.lm_head, hidden_states)
432
- return logits
433
-
386
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
387
qkv_params_mapping = [
388
# (param_name, shard_name, relative_start_idx, relative_end_idx)
@@ -553,3 +505,56 @@ class Step3TextForCausalLM(nn.Module, SupportsPP):
505
weight_loader(param, loaded_weight)
506
loaded_params.add(name)
507
return loaded_params
508
+
509
+
510
+class Step3TextForCausalLM(nn.Module, SupportsPP):
511
+ def __init__(
512
+ self,
513
+ *,
514
+ vllm_config: VllmConfig,
515
+ prefix: str = "",
516
+ ):
517
+ super().__init__()
518
+ config = vllm_config.model_config.hf_config
519
+
520
+ self.config = config
521
+ self.vllm_config = vllm_config
522
+
523
+ self.model = Step3TextModel(vllm_config=vllm_config, prefix=prefix)
524
+
525
+ if get_pp_group().is_last_rank:
526
+ self.lm_head = ParallelLMHead(
527
+ config.vocab_size,
528
+ config.hidden_size,
529
+ prefix=maybe_prefix(prefix, "lm_head"),
530
+ )
531
+ self.logits_processor = LogitsProcessor(config.vocab_size)
532
+ else:
533
+ self.lm_head = PPMissingLayer()
534
+
535
+ self.make_empty_intermediate_tensors = (
536
+ self.model.make_empty_intermediate_tensors
537
+ )
538
+
539
+ def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
540
+ return self.model.embed_input_ids(input_ids)
541
+
542
+ def forward(
543
+ self,
544
+ input_ids: torch.Tensor | None,
545
+ positions: torch.Tensor,
546
+ intermediate_tensors: IntermediateTensors | None = None,
547
+ inputs_embeds: torch.Tensor | None = None,
548
+ ):
549
+ hidden_states = self.model(
550
+ input_ids, positions, intermediate_tensors, inputs_embeds
551
+ )
552
+ return hidden_states
553
+
554
+ def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor:
555
+ logits = self.logits_processor(self.lm_head, hidden_states)
556
+ return logits
557
+
558
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
559
+ loader = AutoWeightsLoader(self)
560
+ return loader.load_weights(weights)