support g_idx

Signed-off-by: Wang, Yi A <yi.a.wang@intel.com>
This commit is contained in:
Wang, Yi A 2025-07-16 22:46:05 -07:00
parent fc2405c549
commit d31bdfdbae

View File

@ -89,13 +89,31 @@ class QuantLinear(nn.Module):
g_idx_trivial = torch.tensor( g_idx_trivial = torch.tensor(
g_idx_trivial, dtype=torch.int32, device=self.g_idx.device g_idx_trivial, dtype=torch.int32, device=self.g_idx.device
) )
assert torch.equal( sort_zeros = not (torch.equal(self.g_idx, g_idx_trivial))
self.g_idx, g_idx_trivial
), "Non-trivial tensor g_idx is not supported"
self.qzeros = self.qzeros.cpu() self.qzeros = self.qzeros.cpu()
zeros = self.unpack_zeros_from_cuda_old_format() zeros = self.unpack_zeros_from_cuda_old_format()
new_qzeros = pack_tensor(zeros) if sort_zeros:
self.qzeros = new_qzeros.to(orig_device) zeros_group_1 = torch.zeros(
(self.infeatures, self.outfeatures),
dtype=zeros.dtype,
device=zeros.device,
)
scales = self.scales.cpu()
scale_group_1 = torch.zeros(
(self.infeatures, self.outfeatures),
dtype=scales.dtype,
device=scales.device,
)
for i in range(self.infeatures):
zeros_group_1[i] = zeros[self.g_idx[i]]
scale_group_1[i] = self.scales[self.g_idx[i]]
self.qzeros = pack_tensor(zeros_group_1).to(orig_device)
self.scales = scale_group_1.to(orig_device)
self.groupsize = 1
self.g_idx = None
else:
new_qzeros = pack_tensor(zeros)
self.qzeros = new_qzeros.to(orig_device)
@classmethod @classmethod
def new(cls, bits, groupsize, infeatures, outfeatures, bias): def new(cls, bits, groupsize, infeatures, outfeatures, bias):