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 7, 09:00 UTC
[Core] Move save_tensorized_model logic to Worker (#35825)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Nick Hill committed
Mar 3, 2026 at 15:31 UTC
d15c3b90fc70ba8d787ee2b172caf5b978909fe9
2 files changed
+6
-17
vllm/v1/worker/gpu_model_runner.py
+1
-12
@@ -58,7 +58,7 @@ from vllm.model_executor.layers.rotary_embedding import (
58
MRotaryEmbedding,
59
XDRotaryEmbedding,
60
)
61
-from vllm.model_executor.model_loader import TensorizerLoader, get_model_loader
61
+from vllm.model_executor.model_loader import get_model_loader
62
from vllm.model_executor.model_loader.reload import (
63
finalize_layerwise_reload,
64
initialize_layerwise_reload,
@@ -194,7 +194,6 @@ from .utils import (
194
)
195
196
if TYPE_CHECKING:
197
- from vllm.model_executor.model_loader.tensorizer import TensorizerConfig
197
from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
198
from vllm.v1.spec_decode.ngram_proposer import NgramProposer
199
@@ -4510,16 +4509,6 @@ class GPUModelRunner(
4509
weights_not_loaded,
4510
)
4511
4513
- def save_tensorized_model(
4514
- self,
4515
- tensorizer_config: "TensorizerConfig",
4516
- ) -> None:
4517
- TensorizerLoader.save_model(
4518
- self.get_model(),
4519
- tensorizer_config=tensorizer_config,
4520
- model_config=self.model_config,
4521
- )
4522
-
4512
def _get_prompt_logprobs_dict(
4513
self,
4514
hidden_states: torch.Tensor,
vllm/v1/worker/gpu_worker.py
+5
-5
@@ -57,6 +57,7 @@ from vllm.v1.worker.utils import is_residual_scattered_for_sp
57
from vllm.v1.worker.worker_base import WorkerBase
58
from vllm.v1.worker.workspace import init_workspace_manager
59
60
+from ...model_executor.model_loader import TensorizerLoader
61
from .gpu.warmup import warmup_kernels
62
from .utils import request_memory
63
@@ -836,12 +837,11 @@ class Worker(WorkerBase):
837
max_size=max_size,
838
)
839
839
- def save_tensorized_model(
840
- self,
841
- tensorizer_config: "TensorizerConfig",
842
- ) -> None:
843
- self.model_runner.save_tensorized_model(
840
+ def save_tensorized_model(self, tensorizer_config: "TensorizerConfig") -> None:
841
+ TensorizerLoader.save_model(
842
+ self.get_model(),
843
tensorizer_config=tensorizer_config,
844
+ model_config=self.model_config,
845
)
846
847
def init_weight_transfer_engine(self, init_info: dict) -> None: