[quantization] fix channelwise conversion with scalar weight scale (#4596)

This commit is contained in:
Yun Dai
2025-03-22 00:47:52 -07:00
committed by GitHub
parent 6a384d5c01
commit 8cd4250401
3 changed files with 58 additions and 0 deletions

View File

@@ -74,6 +74,11 @@ def convert_to_channelwise(
(sum(logical_widths), 1), dtype=torch.float32, device=weight_scale.device
)
# Handle scalar tensor case: broadcast same scale to all channels
if weight_scale.dim() == 0:
weight_scale_channel.fill_(weight_scale.item())
return weight_scale_channel
# Expand each scale to match the size of each logical matrix.
start = 0
for idx, logical_width in enumerate(logical_widths):