[https://nvbugs/5524714][fix] Fix TP sharding of fused-QKV weight scales in W4A16 AWQ (NVIDIA#8432)

Signed-off-by: Daniel Afrimi <[email protected]>
Signed-off-by: Mike Iovine <[email protected]>

Author: danielafrimi

tensorrt_llm/_torch/modules/linear.py

@@ -1393,7 +1393,7 @@ def create_weights(self, module: Linear, in_features: int,
         group_size = module.quant_config.group_size
         if in_features % group_size != 0:
             raise ValueError(
-                f"in_features ({self.in_features}) must be divisible by group_size ({group_size}) "
+                f"in_features ({in_features}) must be divisible by group_size ({group_size}) "
                 f"for INT4 per-group quantization scale dimensions.")
 
         module.weight_scale = Parameter(torch.empty(
@@ -1492,7 +1492,8 @@ def load_weights_fused_qkv_linear(self, module: Linear,
 
         copy_weight(module.weight, fused_weight)
 
-        weight_scales = self.load_weight_scales(weights)
+        weight_scales = self.load_weight_scales(weights, module.tp_size,
+                                                module.tp_rank, module.tp_mode)
 
         # Create concatenated weight scale tensor
         cat_weight_scale = torch.cat(weight_scales, dim=0).T.contiguous()