packs and unpacks

src/aq.py

@@ -87,7 +87,7 @@ def __init__(
             self.scales_clusters = nn.Parameter(
                 scales_clusters, requires_grad=True
             )  #  [2**self.scale_nbits, num_in_groups] if scale_nbits > 0 else [2**self.scale_nbits, 1]
-            self.scale_indices = nn.Parameter(scale_indices, requires_grad=False)  #  [num_out_groups, 1]
+            self.scales_indices = nn.Parameter(scale_indices, requires_grad=False)  #  [num_out_groups, 1]
 
         # BIAS
         if bias:
@@ -128,7 +128,7 @@ def initialize(
             else:
                 scales_clusters, scales_indices, _ = fit_kmeans_1d(scales.flatten(1, -1), k=2**self.scale_nbits)
                 self.scales_clusters.data = scales_clusters
-                self.scales_indices = pack_int_data(scales_indices, self.scale_nbits)
+                self.scales_indices.data = pack_int_data(scales_indices, self.scale_nbits)
 
             weight_for_init = (weight_groupwise / scales).swapaxes(1, 2).reshape_as(reference_weight)
             del weight_groupwise
@@ -198,7 +198,9 @@ def get_scales(self) -> torch.Tensor:
                 dequantized_scales = dequantized_scales + (self.scales - self.scales.detach())
             return dequantized_scales
         else:  # train scale codebook only
-            return self.scales_clusters.gather(1, self.scales_indices)[:, :, None, None]
+            return self.scales_clusters.gather(1, unpack_int_data(self.scales_indices, self.scale_nbits))[
+                :, :, None, None
+            ]
 
     def reconstruct_weight(self, selection: Union[slice, ellipsis, torch.Tensor] = ...):
         """
@@ -239,13 +241,16 @@ def beam_search_update_codes_(
         :param kwargs: any additional keyword arguments are forwarded to beam_search_optimal_codes function
         :returns: the updated codes
         """
-        self.codes[selection] = beam_search_optimal_codes(
-            XTX=XTX,
-            reference_weight=reference_weight,
-            codebooks=self.get_codebooks(),
-            prev_codes=self.codes[selection],
-            scales=self.get_scales()[selection],
-            **kwargs,
+        self.codes[selection] = pack_int_data(
+            beam_search_optimal_codes(
+                XTX=XTX,
+                reference_weight=reference_weight,
+                codebooks=self.get_codebooks(),
+                prev_codes=unpack_int_data(self.codes[selection], self.nbits_per_codebook),
+                scales=self.get_scales()[selection],
+                **kwargs,
+            ),
+            self.nbits_per_codebook,
         )
         return self.codes[selection]
 

src/utils.py

@@ -61,11 +61,13 @@ def get_int_dtype(nbits: int) -> torch.dtype:
     raise ValueError(f"No dtype available for {nbits}-bit codebooks")
 
 
+@torch.inference_mode()
 def pack_int_data(data: torch.IntTensor, nbits: int) -> torch.IntTensor:
     data[data >= 2 ** (nbits - 1)] -= 2**nbits
     return data.to(get_int_dtype(nbits))
 
 
+@torch.inference_mode()
 def unpack_int_data(data: torch.IntTensor, nbits: int) -> torch.IntTensor:
     return data.to(torch.int64) % (2**nbits)