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, 14:45 UTC
[Model] use AutoWeightsLoader for olmoe,opt,orion,persimmon,phi3_small (#16548)
Signed-off-by: rongfu.leng <rongfu.leng@daocloud.io>
rongfu.leng committed
Apr 18, 2025 at 01:48 UTC
5125d72f024aa0da545e4da03ec67109353f48e4
5 files changed
+216
-193
vllm/model_executor/models/olmoe.py
+58
-54
@@ -39,7 +39,7 @@ from vllm.model_executor.sampling_metadata import SamplingMetadata
39
from vllm.sequence import IntermediateTensors
40
41
from .interfaces import SupportsPP
42
-from .utils import (is_pp_missing_parameter,
42
+from .utils import (AutoWeightsLoader, is_pp_missing_parameter,
43
make_empty_intermediate_tensors_factory, make_layers,
44
maybe_prefix)
45
@@ -255,7 +255,7 @@ class OlmoeModel(nn.Module):
255
quant_config = vllm_config.quant_config
256
257
self.vocab_size = config.vocab_size
258
-
258
+ self.config = config
259
self.embed_tokens = VocabParallelEmbedding(
260
config.vocab_size,
261
config.hidden_size,
@@ -308,56 +308,6 @@ class OlmoeModel(nn.Module):
308
hidden_states, _ = self.norm(hidden_states, residual)
309
return hidden_states
310
311
-
312
-class OlmoeForCausalLM(nn.Module, SupportsPP):
313
-
314
- fall_back_to_pt_during_load = False
315
-
316
- def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
317
- super().__init__()
318
- config = vllm_config.model_config.hf_config
319
- quant_config = vllm_config.quant_config
320
- self.config = config
321
- self.quant_config = quant_config
322
- self.model = OlmoeModel(vllm_config=vllm_config,
323
- prefix=maybe_prefix(prefix, "model"))
324
- self.lm_head = ParallelLMHead(config.vocab_size,
325
- config.hidden_size,
326
- quant_config=quant_config)
327
- self.logits_processor = LogitsProcessor(config.vocab_size)
328
- self.sampler = get_sampler()
329
-
330
- self.make_empty_intermediate_tensors = (
331
- self.model.make_empty_intermediate_tensors)
332
-
333
- def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
334
- return self.model.get_input_embeddings(input_ids)
335
-
336
- def forward(
337
- self,
338
- input_ids: torch.Tensor,
339
- positions: torch.Tensor,
340
- intermediate_tensors: Optional[IntermediateTensors] = None,
341
- inputs_embeds: Optional[torch.Tensor] = None,
342
- ) -> Union[torch.Tensor, IntermediateTensors]:
343
- hidden_states = self.model(input_ids, positions, intermediate_tensors,
344
- inputs_embeds)
345
- return hidden_states
346
-
347
- def compute_logits(self, hidden_states: torch.Tensor,
348
- sampling_metadata: SamplingMetadata) -> torch.Tensor:
349
- logits = self.logits_processor(self.lm_head, hidden_states,
350
- sampling_metadata)
351
- return logits
352
-
353
- def sample(
354
- self,
355
- logits: Optional[torch.Tensor],
356
- sampling_metadata: SamplingMetadata,
357
- ) -> Optional[SamplerOutput]:
358
- next_tokens = self.sampler(logits, sampling_metadata)
359
- return next_tokens
360
-
311
def load_weights(self, weights: Iterable[Tuple[str,
312
torch.Tensor]]) -> Set[str]:
313
stacked_params_mapping = [
@@ -380,8 +330,6 @@ class OlmoeForCausalLM(nn.Module, SupportsPP):
330
params_dict = dict(self.named_parameters())
331
loaded_params: Set[str] = set()
332
for name, loaded_weight in weights:
383
- if "rotary_emb.inv_freq" in name:
384
- continue
333
for (param_name, weight_name, shard_id) in stacked_params_mapping:
334
# Skip non-stacked layers and experts (experts handled below).
335
if weight_name not in name:
@@ -453,3 +401,59 @@ class OlmoeForCausalLM(nn.Module, SupportsPP):
401
weight_loader(param, loaded_weight)
402
loaded_params.add(name)
403
return loaded_params
404
+
405
+
406
+class OlmoeForCausalLM(nn.Module, SupportsPP):
407
+
408
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
409
+ super().__init__()
410
+ config = vllm_config.model_config.hf_config
411
+ quant_config = vllm_config.quant_config
412
+ self.config = config
413
+ self.quant_config = quant_config
414
+ self.model = OlmoeModel(vllm_config=vllm_config,
415
+ prefix=maybe_prefix(prefix, "model"))
416
+ self.lm_head = ParallelLMHead(config.vocab_size,
417
+ config.hidden_size,
418
+ quant_config=quant_config)
419
+ self.logits_processor = LogitsProcessor(config.vocab_size)
420
+ self.sampler = get_sampler()
421
+
422
+ self.make_empty_intermediate_tensors = (
423
+ self.model.make_empty_intermediate_tensors)
424
+
425
+ def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
426
+ return self.model.get_input_embeddings(input_ids)
427
+
428
+ def forward(
429
+ self,
430
+ input_ids: torch.Tensor,
431
+ positions: torch.Tensor,
432
+ intermediate_tensors: Optional[IntermediateTensors] = None,
433
+ inputs_embeds: Optional[torch.Tensor] = None,
434
+ ) -> Union[torch.Tensor, IntermediateTensors]:
435
+ hidden_states = self.model(input_ids, positions, intermediate_tensors,
436
+ inputs_embeds)
437
+ return hidden_states
438
+
439
+ def compute_logits(self, hidden_states: torch.Tensor,
440
+ sampling_metadata: SamplingMetadata) -> torch.Tensor:
441
+ logits = self.logits_processor(self.lm_head, hidden_states,
442
+ sampling_metadata)
443
+ return logits
444
+
445
+ def sample(
446
+ self,
447
+ logits: Optional[torch.Tensor],
448
+ sampling_metadata: SamplingMetadata,
449
+ ) -> Optional[SamplerOutput]:
450
+ next_tokens = self.sampler(logits, sampling_metadata)
451
+ return next_tokens
452
+
453
+ def load_weights(self, weights: Iterable[Tuple[str,
454
+ torch.Tensor]]) -> Set[str]:
455
+ loader = AutoWeightsLoader(
456
+ self,
457
+ skip_prefixes=["rotary_emb.inv_freq"],
458
+ )
459
+ return loader.load_weights(weights)
vllm/model_executor/models/opt.py
+48
-40
@@ -43,7 +43,7 @@ from vllm.model_executor.sampling_metadata import SamplingMetadata
43
from vllm.sequence import IntermediateTensors
44
45
from .interfaces import SupportsPP
46
-from .utils import (is_pp_missing_parameter,
46
+from .utils import (AutoWeightsLoader, WeightsMapper, is_pp_missing_parameter,
47
make_empty_intermediate_tensors_factory, make_layers,
48
maybe_prefix)
49
@@ -313,6 +313,43 @@ class OPTModel(nn.Module):
313
intermediate_tensors,
314
inputs_embeds=inputs_embeds)
315
316
+ def load_weights(self, weights: Iterable[Tuple[str,
317
+ torch.Tensor]]) -> Set[str]:
318
+ stacked_params_mapping = [
319
+ # (param_name, shard_name, shard_id)
320
+ ("qkv_proj", "q_proj", "q"),
321
+ ("qkv_proj", "k_proj", "k"),
322
+ ("qkv_proj", "v_proj", "v"),
323
+ ]
324
+ params_dict = dict(self.named_parameters(remove_duplicate=False))
325
+ loaded_params: Set[str] = set()
326
+ for name, loaded_weight in weights:
327
+ for (param_name, weight_name, shard_id) in stacked_params_mapping:
328
+ if weight_name not in name:
329
+ continue
330
+ name = name.replace(weight_name, param_name)
331
+ # Skip loading extra bias for GPTQ models.
332
+ if name.endswith(".bias") and name not in params_dict:
333
+ continue
334
+ if is_pp_missing_parameter(name, self):
335
+ continue
336
+ param = params_dict[name]
337
+ weight_loader = param.weight_loader
338
+ weight_loader(param, loaded_weight, shard_id)
339
+ break
340
+ else:
341
+ # Skip loading extra bias for GPTQ models.
342
+ if name.endswith(".bias") and name not in params_dict:
343
+ continue
344
+ if is_pp_missing_parameter(name, self):
345
+ continue
346
+ param = params_dict[name]
347
+ weight_loader = getattr(param, "weight_loader",
348
+ default_weight_loader)
349
+ weight_loader(param, loaded_weight)
350
+ loaded_params.add(name)
351
+ return loaded_params
352
+
353
354
class OPTForCausalLM(nn.Module, SupportsPP):
355
packed_modules_mapping = {
@@ -320,6 +357,10 @@ class OPTForCausalLM(nn.Module, SupportsPP):
357
"gate_up_proj": ["gate_proj", "up_proj"]
358
}
359
360
+ hf_to_vllm_mapper = WeightsMapper(orig_to_new_prefix={
361
+ "decoder.": "model.decoder.",
362
+ })
363
+
364
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
365
super().__init__()
366
config = vllm_config.model_config.hf_config
@@ -371,42 +412,9 @@ class OPTForCausalLM(nn.Module, SupportsPP):
412
413
def load_weights(self, weights: Iterable[Tuple[str,
414
torch.Tensor]]) -> Set[str]:
374
- stacked_params_mapping = [
375
- # (param_name, shard_name, shard_id)
376
- ("qkv_proj", "q_proj", "q"),
377
- ("qkv_proj", "k_proj", "k"),
378
- ("qkv_proj", "v_proj", "v"),
379
- ]
380
- params_dict = dict(self.named_parameters(remove_duplicate=False))
381
- loaded_params: Set[str] = set()
382
- for name, loaded_weight in weights:
383
- if "lm_head.weight" in name and self.config.tie_word_embeddings:
384
- continue
385
- if name.startswith("decoder."):
386
- name = "model." + name
387
-
388
- for (param_name, weight_name, shard_id) in stacked_params_mapping:
389
- if weight_name not in name:
390
- continue
391
- name = name.replace(weight_name, param_name)
392
- # Skip loading extra bias for GPTQ models.
393
- if name.endswith(".bias") and name not in params_dict:
394
- continue
395
- if is_pp_missing_parameter(name, self):
396
- continue
397
- param = params_dict[name]
398
- weight_loader = param.weight_loader
399
- weight_loader(param, loaded_weight, shard_id)
400
- break
401
- else:
402
- # Skip loading extra bias for GPTQ models.
403
- if name.endswith(".bias") and name not in params_dict:
404
- continue
405
- if is_pp_missing_parameter(name, self):
406
- continue
407
- param = params_dict[name]
408
- weight_loader = getattr(param, "weight_loader",
409
- default_weight_loader)
410
- weight_loader(param, loaded_weight)
411
- loaded_params.add(name)
412
- return loaded_params
415
+ loader = AutoWeightsLoader(
416
+ self,
417
+ skip_prefixes=(["lm_head.weight"]
418
+ if self.config.tie_word_embeddings else None),
419
+ )
420
+ return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
vllm/model_executor/models/orion.py
+49
-42
@@ -30,7 +30,7 @@ from vllm.model_executor.sampling_metadata import SamplingMetadata
30
from vllm.sequence import IntermediateTensors
31
32
from .interfaces import SupportsPP
33
-from .utils import (is_pp_missing_parameter,
33
+from .utils import (AutoWeightsLoader, is_pp_missing_parameter,
34
make_empty_intermediate_tensors_factory, make_layers,
35
maybe_prefix)
36
@@ -260,6 +260,45 @@ class OrionModel(nn.Module):
260
hidden_states = self.norm(hidden_states)
261
return hidden_states
262
263
+ def load_weights(self, weights: Iterable[Tuple[str,
264
+ torch.Tensor]]) -> Set[str]:
265
+ stacked_params_mapping = [
266
+ # (param_name, shard_name, shard_id)
267
+ ("qkv_proj", "q_proj", "q"),
268
+ ("qkv_proj", "k_proj", "k"),
269
+ ("qkv_proj", "v_proj", "v"),
270
+ ("gate_up_proj", "gate_proj", 0),
271
+ ("gate_up_proj", "up_proj", 1),
272
+ ]
273
+ params_dict = dict(self.named_parameters())
274
+ loaded_params: Set[str] = set()
275
+ for name, loaded_weight in weights:
276
+ for (param_name, weight_name, shard_id) in stacked_params_mapping:
277
+ if weight_name not in name:
278
+ continue
279
+ name = name.replace(weight_name, param_name)
280
+ # Skip loading extra bias for GPTQ models.
281
+ if name.endswith(".bias") and name not in params_dict:
282
+ continue
283
+ if is_pp_missing_parameter(name, self):
284
+ continue
285
+ param = params_dict[name]
286
+ weight_loader = param.weight_loader
287
+ weight_loader(param, loaded_weight, shard_id)
288
+ break
289
+ else:
290
+ # Skip loading extra bias for GPTQ models.
291
+ if name.endswith(".bias") and name not in params_dict:
292
+ continue
293
+ if is_pp_missing_parameter(name, self):
294
+ continue
295
+ param = params_dict[name]
296
+ weight_loader = getattr(param, "weight_loader",
297
+ default_weight_loader)
298
+ weight_loader(param, loaded_weight)
299
+ loaded_params.add(name)
300
+ return loaded_params
301
+
302
303
class OrionForCausalLM(nn.Module, SupportsPP):
304
@@ -314,46 +353,14 @@ class OrionForCausalLM(nn.Module, SupportsPP):
353
354
def load_weights(self, weights: Iterable[Tuple[str,
355
torch.Tensor]]) -> Set[str]:
317
- stacked_params_mapping = [
318
- # (param_name, shard_name, shard_id)
319
- ("qkv_proj", "q_proj", "q"),
320
- ("qkv_proj", "k_proj", "k"),
321
- ("qkv_proj", "v_proj", "v"),
322
- ("gate_up_proj", "gate_proj", 0),
323
- ("gate_up_proj", "up_proj", 1),
324
- ]
325
- params_dict = dict(self.named_parameters())
326
- loaded_params: Set[str] = set()
327
- for name, loaded_weight in weights:
328
- if "rotary_emb.inv_freq" in name:
329
- continue
330
- if ("rotary_emb.cos_cached" in name
331
- or "rotary_emb.sin_cached" in name):
356
+ loader = AutoWeightsLoader(
357
+ self,
358
+ skip_prefixes=([
359
+ "rotary_emb.inv_freq",
360
# Models trained using ColossalAI may include these tensors in
361
# the checkpoint. Skip them.
334
- continue
335
- for (param_name, weight_name, shard_id) in stacked_params_mapping:
336
- if weight_name not in name:
337
- continue
338
- name = name.replace(weight_name, param_name)
339
- # Skip loading extra bias for GPTQ models.
340
- if name.endswith(".bias") and name not in params_dict:
341
- continue
342
- if is_pp_missing_parameter(name, self):
343
- continue
344
- param = params_dict[name]
345
- weight_loader = param.weight_loader
346
- weight_loader(param, loaded_weight, shard_id)
347
- break
348
- else:
349
- # Skip loading extra bias for GPTQ models.
350
- if name.endswith(".bias") and name not in params_dict:
351
- continue
352
- if is_pp_missing_parameter(name, self):
353
- continue
354
- param = params_dict[name]
355
- weight_loader = getattr(param, "weight_loader",
356
- default_weight_loader)
357
- weight_loader(param, loaded_weight)
358
- loaded_params.add(name)
359
- return loaded_params
362
+ "rotary_emb.cos_cached",
363
+ "rotary_emb.sin_cached"
364
+ ]),
365
+ )
366
+ return loader.load_weights(weights)
vllm/model_executor/models/persimmon.py
+36
-38
@@ -46,7 +46,7 @@ from vllm.model_executor.sampling_metadata import SamplingMetadata
46
from vllm.sequence import IntermediateTensors
47
48
from .interfaces import SupportsPP
49
-from .utils import (is_pp_missing_parameter,
49
+from .utils import (AutoWeightsLoader, is_pp_missing_parameter,
50
make_empty_intermediate_tensors_factory, make_layers,
51
maybe_prefix)
52
@@ -221,7 +221,7 @@ class PersimmonModel(nn.Module):
221
quant_config = vllm_config.quant_config
222
223
self.vocab_size = config.vocab_size
224
-
224
+ self.config = config
225
self.embed_tokens = VocabParallelEmbedding(config.vocab_size,
226
config.hidden_size)
227
self.start_layer, self.end_layer, self.layers = make_layers(
@@ -260,6 +260,38 @@ class PersimmonModel(nn.Module):
260
hidden_states = self.final_layernorm(hidden_states)
261
return hidden_states
262
263
+ def load_weights(self, weights: Iterable[Tuple[str,
264
+ torch.Tensor]]) -> Set[str]:
265
+ params_dict = dict(self.named_parameters(remove_duplicate=False))
266
+ loaded_params: Set[str] = set()
267
+ for name, loaded_weight in weights:
268
+ if is_pp_missing_parameter(name, self):
269
+ continue
270
+ param = params_dict[name]
271
+
272
+ if "query_key_value" in name:
273
+ # copy from vllm/model_executor/models/bloom.py
274
+ # NOTE: Persimmon's fused QKV's output_dim has the shape of
275
+ # (num_heads * 3 * head_size), while the
276
+ # required shape is (3 * num_heads * head_size).
277
+ # Thus, we need weight conversion.
278
+ output_dim = getattr(param, "output_dim", None)
279
+ num_heads = self.config.num_attention_heads
280
+ if output_dim is not None:
281
+ loaded_weight_shape = loaded_weight.shape
282
+ loaded_weight = loaded_weight.view(
283
+ loaded_weight_shape[:output_dim] + (num_heads, 3, -1) +
284
+ loaded_weight_shape[output_dim + 1:])
285
+ loaded_weight = loaded_weight.transpose(
286
+ output_dim, output_dim + 1)
287
+ loaded_weight = loaded_weight.reshape(loaded_weight_shape)
288
+
289
+ weight_loader = getattr(param, "weight_loader",
290
+ default_weight_loader)
291
+ weight_loader(param, loaded_weight)
292
+ loaded_params.add(name)
293
+ return loaded_params
294
+
295
296
class PersimmonForCausalLM(nn.Module, SupportsPP):
297
@@ -315,39 +347,5 @@ class PersimmonForCausalLM(nn.Module, SupportsPP):
347
348
def load_weights(self, weights: Iterable[Tuple[str,
349
torch.Tensor]]) -> Set[str]:
318
- params_dict = dict(self.named_parameters(remove_duplicate=False))
319
- loaded_params: Set[str] = set()
320
- for name, loaded_weight in weights:
321
- if "rotary_emb.inv_freq" in name:
322
- continue
323
- if ("rotary_emb.cos_cached" in name
324
- or "rotary_emb.sin_cached" in name):
325
- # Models trained using ColossalAI may include these tensors in
326
- # the checkpoint. Skip them.
327
- continue
328
- if is_pp_missing_parameter(name, self):
329
- continue
330
- param = params_dict[name]
331
-
332
- if "query_key_value" in name:
333
- # copy from vllm/model_executor/models/bloom.py
334
- # NOTE: Persimmon's fused QKV's output_dim has the shape of
335
- # (num_heads * 3 * head_size), while the
336
- # required shape is (3 * num_heads * head_size).
337
- # Thus, we need weight conversion.
338
- output_dim = getattr(param, "output_dim", None)
339
- num_heads = self.config.num_attention_heads
340
- if output_dim is not None:
341
- loaded_weight_shape = loaded_weight.shape
342
- loaded_weight = loaded_weight.view(
343
- loaded_weight_shape[:output_dim] + (num_heads, 3, -1) +
344
- loaded_weight_shape[output_dim + 1:])
345
- loaded_weight = loaded_weight.transpose(
346
- output_dim, output_dim + 1)
347
- loaded_weight = loaded_weight.reshape(loaded_weight_shape)
348
-
349
- weight_loader = getattr(param, "weight_loader",
350
- default_weight_loader)
351
- weight_loader(param, loaded_weight)
352
- loaded_params.add(name)
353
- return loaded_params
350
+ loader = AutoWeightsLoader(self)
351
+ return loader.load_weights(weights)
vllm/model_executor/models/phi3_small.py
+25
-19
@@ -26,7 +26,7 @@ from vllm.platforms import current_platform
26
from vllm.sequence import IntermediateTensors
27
28
from .interfaces import SupportsPP
29
-from .utils import (is_pp_missing_parameter,
29
+from .utils import (AutoWeightsLoader, WeightsMapper, is_pp_missing_parameter,
30
make_empty_intermediate_tensors_factory, make_layers,
31
maybe_prefix)
32
@@ -353,10 +353,29 @@ class Phi3SmallModel(nn.Module):
353
hidden_states = self.final_layernorm(hidden_states)
354
return hidden_states
355
356
+ def load_weights(self, weights: Iterable[Tuple[str,
357
+ torch.Tensor]]) -> Set[str]:
358
+ params_dict = dict(self.named_parameters())
359
+ loaded_params: Set[str] = set()
360
+ for name, loaded_weight in weights:
361
+ if name.endswith(".bias") and name not in params_dict:
362
+ continue
363
+ if is_pp_missing_parameter(name, self):
364
+ continue
365
+ param = params_dict[name]
366
+ weight_loader = getattr(param, "weight_loader",
367
+ default_weight_loader)
368
+ weight_loader(param, loaded_weight)
369
+ loaded_params.add(name)
370
+ return loaded_params
371
+
372
373
class Phi3SmallForCausalLM(nn.Module, SupportsPP):
374
_tied_weights_keys = ["lm_head.weight"]
375
376
+ hf_to_vllm_mapper = WeightsMapper(
377
+ orig_to_new_suffix={"rotary_emb.inv_freq": None})
378
+
379
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
380
super().__init__()
381
config = vllm_config.model_config.hf_config
@@ -448,21 +467,8 @@ class Phi3SmallForCausalLM(nn.Module, SupportsPP):
467
468
def load_weights(self, weights: Iterable[Tuple[str,
469
torch.Tensor]]) -> Set[str]:
451
-
452
- params_dict = dict(self.named_parameters())
453
- loaded_params: Set[str] = set()
454
- for name, loaded_weight in weights:
455
- if "rotary_emb.inv_freq" in name:
456
- continue
457
- if name.endswith(".bias") and name not in params_dict:
458
- continue
459
- if is_pp_missing_parameter(name, self):
460
- continue
461
- if "lm_head.weight" in name and self.config.tie_word_embeddings:
462
- continue
463
- param = params_dict[name]
464
- weight_loader = getattr(param, "weight_loader",
465
- default_weight_loader)
466
- weight_loader(param, loaded_weight)
467
- loaded_params.add(name)
468
- return loaded_params
470
+ loader = AutoWeightsLoader(
471
+ self,
472
+ skip_prefixes=(["lm_head.weight"]
473
+ if self.config.tie_word_embeddings else None))
474
+ return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)