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, 07:15 UTC
[Model] Use AutoWeightsLoader for mamba2 (#18918)
Signed-off-by: iLeGend <824040212@qq.com>
iLeGend committed
May 30, 2025 at 12:50 UTC
3987e2ae963963b9edb132935deabd16dd5a7468
1 file changed
+24
-19
vllm/model_executor/models/mamba2.py
+24
-19
@@ -32,7 +32,7 @@ from vllm.model_executor.sampling_metadata import SamplingMetadata
32
from vllm.sequence import IntermediateTensors
33
from vllm.utils import LayerBlockType
34
35
-from .utils import (is_pp_missing_parameter,
35
+from .utils import (AutoWeightsLoader, is_pp_missing_parameter,
36
make_empty_intermediate_tensors_factory, make_layers,
37
maybe_prefix)
38
@@ -167,6 +167,27 @@ class Mamba2Model(nn.Module):
167
168
return hidden_states
169
170
+ def load_weights(self, weights: Iterable[tuple[str,
171
+ torch.Tensor]]) -> set[str]:
172
+ params_dict = dict(self.named_parameters())
173
+ loaded_params: set[str] = set()
174
+ for name, loaded_weight in weights:
175
+ if "A_log" in name:
176
+ name = name.replace("A_log", "A")
177
+
178
+ # Skip loading extra bias for GPTQ models.
179
+ if name.endswith(".bias") and name not in params_dict:
180
+ continue
181
+ if is_pp_missing_parameter(name, self):
182
+ continue
183
+
184
+ param = params_dict[name]
185
+ weight_loader = getattr(param, "weight_loader",
186
+ default_weight_loader)
187
+ weight_loader(param, loaded_weight)
188
+ loaded_params.add(name)
189
+ return loaded_params
190
+
191
192
class Mamba2ForCausalLM(nn.Module, HasInnerState, IsAttentionFree,
193
SupportsV0Only):
@@ -282,21 +303,5 @@ class Mamba2ForCausalLM(nn.Module, HasInnerState, IsAttentionFree,
303
304
def load_weights(self, weights: Iterable[tuple[str,
305
torch.Tensor]]) -> set[str]:
285
- params_dict = dict(self.named_parameters())
286
- loaded_params: set[str] = set()
287
- for name, loaded_weight in weights:
288
- if "A_log" in name:
289
- name = name.replace("A_log", "A")
290
-
291
- # Skip loading extra bias for GPTQ models.
292
- if name.endswith(".bias") and name not in params_dict:
293
- continue
294
- if is_pp_missing_parameter(name, self):
295
- continue
296
-
297
- param = params_dict[name]
298
- weight_loader = getattr(param, "weight_loader",
299
- default_weight_loader)
300
- weight_loader(param, loaded_weight)
301
- loaded_params.add(name)
302
- return loaded_params
306
+ loader = AutoWeightsLoader(self)
307
+ return loader.load_weights(weights)