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, 08:30 UTC
[Bugfix][Model] Fix baichuan model loader for tp (#18597)
Signed-off-by: Mengqing Cao <cmq0113@163.com>
Mengqing Cao committed
May 23, 2025 at 17:39 UTC
4ce64e2df48649c4873f828b8bf71790aa1e56ee
1 file changed
+7
-4
vllm/model_executor/models/baichuan.py
+7
-4
@@ -42,7 +42,8 @@ from vllm.model_executor.layers.quantization import QuantizationConfig
42
from vllm.model_executor.layers.rotary_embedding import get_rope
43
from vllm.model_executor.layers.vocab_parallel_embedding import (
44
ParallelLMHead, VocabParallelEmbedding)
45
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
45
+from vllm.model_executor.model_loader.weight_utils import (
46
+ default_weight_loader, row_parallel_weight_loader)
47
from vllm.model_executor.sampling_metadata import SamplingMetadata
48
from vllm.sequence import IntermediateTensors
49
@@ -384,7 +385,7 @@ class BaiChuanBaseForCausalLM(nn.Module, SupportsLoRA, SupportsPP,
385
lora_config = vllm_config.lora_config
386
self.config = config
387
self.lora_config = lora_config
387
-
388
+ self.tp_size = get_tensor_model_parallel_world_size()
389
self.quant_config = quant_config
390
self.model = BaiChuanModel(vllm_config=vllm_config,
391
prefix=prefix,
@@ -438,8 +439,10 @@ class BaiChuanBaseForCausalLM(nn.Module, SupportsLoRA, SupportsPP,
439
is_baichuan2 = self.config.vocab_size == 125696
440
if is_baichuan2:
441
loaded_weight = torch.nn.functional.normalize(loaded_weight)
441
-
442
- default_weight_loader(param, loaded_weight)
442
+ if self.tp_size > 1:
443
+ row_parallel_weight_loader(param, loaded_weight)
444
+ else:
445
+ default_weight_loader(param, loaded_weight)
446
447
448
class BaichuanForCausalLM(BaiChuanBaseForCausalLM):