kanshan/TensorRT-LLMs
1
0
Fork 0
mirror of https://github.com/NVIDIA/TensorRT-LLM.git synced 2026-01-14 06:27:45 +08:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

[None][fix] Add TP information in weight scale loading in WeightOnlyQuantLinearMethod (#7732)

Browse Source 
Signed-off-by: Stefan Niebler <82932102+stnie@users.noreply.github.com>
This commit is contained in:
Stefan Niebler 2025-09-18 10:30:50 +02:00 committed by GitHub
parent a7ca0fff54
commit a55251bf75
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 4 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
tensorrt_llm/_torch/modules/linear.py
View File

@ -1297,7 +1297,10 @@ class WeightOnlyQuantLinearMethod(LinearMethodBase):
copy_weight(module.weight, fused_weight)
weight_scales = self.load_weight_scales(weights)
weight_scales = self.load_weight_scales(weights,
tp_size=module.tp_size,
tp_rank=module.tp_rank,
tp_mode=module.tp_mode)
# Create concatenated weight scale tensor
cat_weight_scale = torch.cat(weight_scales, dim=0)