@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 153 lines 5.83 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 # ruff: noqa: SIM117
4 import copy
5 from collections.abc import Generator
6
7 import torch
8 from torch import nn
9
10 from vllm.config import ModelConfig, ParallelConfig, VllmConfig
11 from vllm.config.load import LoadConfig
12 from vllm.logger import init_logger
13 from vllm.model_executor.model_loader.base_loader import BaseModelLoader
14 from vllm.model_executor.model_loader.tensorizer import (
15 TensorizerConfig,
16 deserialize_tensorizer_model,
17 init_tensorizer_model,
18 is_vllm_tensorized,
19 serialize_vllm_model,
20 tensorizer_weights_iterator,
21 )
22 from vllm.model_executor.model_loader.utils import (
23 get_model_architecture,
24 initialize_model,
25 )
26 from vllm.utils.torch_utils import set_default_torch_dtype
27
28 logger = init_logger(__name__)
29
30 BLACKLISTED_TENSORIZER_ARGS = {
31 "device", # vLLM decides this
32 "dtype", # vLLM decides this
33 "mode", # Not meant to be configurable by the user
34 }
35
36
37 def validate_config(config: dict):
38 for k, v in config.items():
39 if v is not None and k in BLACKLISTED_TENSORIZER_ARGS:
40 raise ValueError(f"{k} is not an allowed Tensorizer argument.")
41
42
43 class TensorizerLoader(BaseModelLoader):
44 """Model loader using CoreWeave's tensorizer library."""
45
46 def __init__(self, load_config: LoadConfig):
47 super().__init__(load_config)
48 if isinstance(load_config.model_loader_extra_config, TensorizerConfig):
49 self.tensorizer_config = load_config.model_loader_extra_config
50 else:
51 validate_config(load_config.model_loader_extra_config)
52 self.tensorizer_config = TensorizerConfig(
53 **load_config.model_loader_extra_config["tensorizer_config"]
54 )
55
56 def _verify_config(
57 self, model_config: ModelConfig, parallel_config: ParallelConfig
58 ):
59 self.tensorizer_config.verify_with_model_config(model_config)
60 self.tensorizer_config.verify_with_parallel_config(parallel_config)
61
62 def _get_weights_iterator(
63 self,
64 ) -> Generator[tuple[str, torch.Tensor], None, None]:
65 tensorizer_args = self.tensorizer_config._construct_tensorizer_args()
66 return tensorizer_weights_iterator(tensorizer_args)
67
68 def _load_model_serialized_cpu(
69 self,
70 vllm_config: VllmConfig,
71 prefix: str = "",
72 ) -> nn.Module:
73 """Load a serialized model with tensorizer to the CPU.
74
75 This is only necessary when the model isn't vLLM-tensorized (see
76 examples/features/tensorize_vllm_model.py) This should still
77 be faster than default HuggingFace loading, but will be slower than
78 loading a vLLM-tensorized model.
79 """
80 device_config = vllm_config.device_config
81 model_config = vllm_config.model_config
82 with set_default_torch_dtype(model_config.dtype):
83 with torch.device(device_config.device):
84 model = initialize_model(vllm_config=vllm_config, prefix=prefix)
85
86 model.load_weights(self._get_weights_iterator())
87 return model.eval()
88
89 def download_model(self, model_config: ModelConfig) -> None:
90 self.tensorizer_config.verify_with_model_config(model_config)
91
92 with self.tensorizer_config.open_stream():
93 pass
94
95 def _patch_tensorizer_config(self, model_config: ModelConfig) -> TensorizerConfig:
96 model_class = get_model_architecture(model_config)[0]
97 tensorizer_config = copy.copy(self.tensorizer_config)
98 tensorizer_config.model_class = model_class
99 tensorizer_config.hf_config = model_config.hf_config
100 tensorizer_config.dtype = model_config.dtype
101 return tensorizer_config
102
103 def load_weights(self, model: nn.Module, model_config: ModelConfig) -> None:
104 """Load serialized model weights with tensorizer.
105
106 Expects a vLLM-tensorized model. See the
107 examples/features/tensorize_vllm_model.py example script
108 for serializing vLLM models."""
109 if is_vllm_tensorized(self.tensorizer_config):
110 tensorizer_config = self._patch_tensorizer_config(model_config)
111 deserialize_tensorizer_model(model, tensorizer_config)
112 else:
113 model.load_weights(self._get_weights_iterator())
114
115 def load_model(
116 self, vllm_config: VllmConfig, model_config: ModelConfig, prefix: str = ""
117 ) -> nn.Module:
118 parallel_config = vllm_config.parallel_config
119 self._verify_config(model_config, parallel_config)
120
121 if parallel_config.tensor_parallel_size > 1:
122 from vllm.distributed import get_tensor_model_parallel_rank
123
124 assert self.tensorizer_config.tensorizer_uri is not None
125 self.tensorizer_config.tensorizer_uri = (
126 self.tensorizer_config.tensorizer_uri % get_tensor_model_parallel_rank()
127 )
128
129 if is_vllm_tensorized(self.tensorizer_config):
130 tensorizer_config = self._patch_tensorizer_config(model_config)
131 device_config = vllm_config.device_config
132 with set_default_torch_dtype(model_config.dtype):
133 with torch.device(device_config.device):
134 model = init_tensorizer_model(
135 tensorizer_config=tensorizer_config, vllm_config=vllm_config
136 )
137 self.load_weights(model, model_config)
138 return model
139 return self._load_model_serialized_cpu(vllm_config=vllm_config, prefix=prefix)
140
141 @staticmethod
142 def save_model(
143 model: torch.nn.Module,
144 tensorizer_config: TensorizerConfig | dict,
145 model_config: ModelConfig,
146 ) -> None:
147 if isinstance(tensorizer_config, dict):
148 tensorizer_config = TensorizerConfig(**tensorizer_config)
149 serialize_vllm_model(
150 model=model,
151 tensorizer_config=tensorizer_config,
152 model_config=model_config,
153 )