Unified quantizer (#624)

* remove debug print statements and run linter

* use unified quantizer architecture

* use unified quantizer architecture

* use unified quantizer architecture

* typos & lint

* typos & lint

Author: mikekgfb

qops.py

@@ -12,7 +12,8 @@
     state_dict_device,
     use_et_backend,
 )
-from torch.nn.parameter import Parameter
+
+# from torch.nn.parameter import Parameter
 
 
 def linear_int8_aoti(input, weight, scales):
@@ -147,7 +148,9 @@ def __init__(
         ), "must specify both weights and scales, or neither"
         if weight is None:
             weight = torch.empty(
-                (out_features, in_features), dtype=torch.int8, device=device
+                (out_features, in_features),
+                dtype=torch.int8,
+                device=device,
             )
             if groupsize is None or (groupsize == 0):
                 scales = torch.empty(out_features, dtype=dtype, device=device)
@@ -180,56 +183,56 @@ def __init__(
         *,
         bitwidth: int,
         groupsize: Optional[int] = None,
+        weight: Optional[torch.Tensor] = None,
+        scales: Optional[torch.Tensor] = None,
     ) -> None:
         super().__init__()
         if dtype is None:
-            dtype = torch.half
-
+            dtype = get_precision()
         if groupsize is None or groupsize == 0:
             groupsize = embedding_dim
         self.groupsize = groupsize
         self.dtype = dtype
         self.bitwidth = bitwidth
 
-        if use_et_backend():
-            self.forward = self.et_forward
-        else:
-            self.forward = self.aoti_forward
+        assert (weight is None) == bool(
+            scales is None
+        ), "must specify both weights and scales, or neither"
 
-        if bitwidth == 8:
-            self.register_buffer(
-                "weight",
-                torch.empty(
-                    (num_embeddings, embedding_dim), dtype=torch.int8, device=device
-                ),
-            )
-        elif bitwidth == 4:  # packed
-            self.register_buffer(
-                "weight",
-                torch.empty(
-                    (num_embeddings, embedding_dim // 2),
-                    dtype=torch.uint8,
-                    device=device,
-                ),
-            )
-        else:
+        if bitwidth not in [4, 8]:
             raise RuntimeError(
                 f"QUantized embedding does not support bitwidth={bitwidth}"
             )
 
-        groups_per_row = (embedding_dim + groupsize - 1) // groupsize
-        if groups_per_row > 1:
-            self.register_buffer(
-                "scales",
-                torch.ones(
-                    (num_embeddings, groups_per_row), dtype=torch.float16, device=device
+        if weight is None:
+            groups_per_row = (embedding_dim + groupsize - 1) // groupsize
+            weight = torch.empty(
+                (
+                    num_embeddings,
+                    (embedding_dim * bitwidth) // 8,
                 ),
+                dtype=torch.int8,
+                device=device,
             )
+            scales = torch.empty(
+                (num_embeddings, groups_per_row),
+                dtype=dtype,
+                device=device,
+            ).squeeze(dim=-1)
+
+        self.register_buffer(
+            "weight",
+            weight,
+        )
+        self.register_buffer(
+            "scales",
+            scales,
+        )
+
+        if use_et_backend():
+            self.forward = self.et_forward
         else:
-            self.register_buffer(
-                "scales",
-                torch.ones((num_embeddings,), dtype=torch.float16, device=device),
-            )
+            self.forward = self.aoti_forward
 
     @torch.no_grad()
     def et_forward(self, indices: torch.Tensor) -> torch.Tensor:

quantize.py

@@ -438,7 +438,7 @@ def quantize(self, module):
                         ),
                     )
                 else:
-                    self.quantize(module)
+                    self.quantize(child)
 
         return module
 
@@ -450,31 +450,6 @@ def quantized_model(self) -> nn.Module:
 #####                   embedding table quantization               ######
 
 
-def replace_embedding_weight_only_grouped_int8_per_channel(
-    module, device, bitwidth: int, groupsize: Optional[int]
-):
-    for name, child in module.named_children():
-        # print(f"name: {name}")
-        if isinstance(child, nn.Embedding):
-            # print(f"{name, child}")
-            # print(f"weights size: {child.weight.size()}")
-            setattr(
-                module,
-                name,
-                QuantizedEmbedding(
-                    device=device,
-                    num_embeddings=child.weight.shape[0],
-                    embedding_dim=child.weight.shape[1],
-                    bitwidth=bitwidth,
-                    groupsize=groupsize,
-                ),
-            )
-        else:
-            replace_embedding_weight_only_grouped_int8_per_channel(
-                child, device, bitwidth, groupsize
-            )
-
-
 class EmbeddingOnlyInt8QuantHandler(QuantHandler):
     def __init__(
         self,
@@ -492,9 +467,11 @@ def __init__(
         self.bitwidth = bitwidth
 
     @torch.no_grad()
-    def create_quantized_state_dict(self) -> Dict:
-        cur_state_dict = state_dict_device(self.model_.state_dict())
-        dict_device = "cpu"  # self.device
+    def quantize(self, module):
+        # cur_state_dict = state_dict_device(self.model_.state_dict())
+        # dict_device = "cpu"  # self.device
+
+        device = self.device
 
         if self.bitwidth == 4:
             range_min = -8
@@ -505,22 +482,23 @@ def create_quantized_state_dict(self) -> Dict:
         else:
             raise ValueError(f"Unsupported bitwidth {self.bitwidth}")
 
-        for fqn, mod in self.model_.named_modules():
-            if isinstance(mod, nn.Embedding):
+        for name, child in module.named_children():
+            # print(f"name: {name}")
+            if isinstance(child, nn.Embedding):
                 # print(f"Embedding identified: {fqn, mod}")
-                # print(f"weights size: {mod.weight.size()}")
+                # print(f"weights size: {child.weight.size()}")
                 # print(f"quantize {fqn}...")
 
                 # print(
                 #     f"quantize {fqn, mod} with groupsize {self.groupsize}, bitwidth {self.bitwidth}"
                 # )
                 weight, scales, _ = dynamically_quantize_per_channel(
-                    mod.weight.float(),
+                    child.weight.float(),
                     range_min,
                     range_max,
                     torch.int8,
                     self.groupsize,
-                    scales_dtype=mod.weight.dtype,
+                    scales_dtype=child.weight.dtype,
                 )
 
                 if self.bitwidth == 4:
@@ -536,26 +514,31 @@ def create_quantized_state_dict(self) -> Dict:
                     weight_packed = weight_even + weight_odd
                     weight = weight_packed
 
-                weight = weight.to(device=dict_device)
-                scales = scales.to(device=dict_device)
-                # Update state dict
-                cur_state_dict[f"{fqn}.weight"] = weight
-                # squeeze makes groupsize=rowsize unidimensional
-                cur_state_dict[f"{fqn}.scales"] = scales.squeeze(dim=-1)
+                weight = weight
+                scales = scales.squeeze(dim=-1)
 
-        return cur_state_dict
+                # print(f"{name, child}")
+                # print(f"weights size: {child.weight.size()}")
+                setattr(
+                    module,
+                    name,
+                    QuantizedEmbedding(
+                        num_embeddings=child.weight.shape[0],
+                        embedding_dim=child.weight.shape[1],
+                        device=self.device,
+                        bitwidth=self.bitwidth,
+                        groupsize=self.groupsize,
+                        weight=weight,
+                        scales=scales,
+                    ),
+                )
+            else:
+                self.quantize(child)
 
-    def convert_for_runtime(self) -> nn.Module:
-        replace_embedding_weight_only_grouped_int8_per_channel(
-            self.model_, self.device, self.bitwidth, self.groupsize
-        )
-        return self.model_
+        return module
 
     def quantized_model(self) -> nn.Module:
-        model_updated_state_dict = self.create_quantized_state_dict()
-        self.convert_for_runtime()
-        self.model_.load_state_dict(model_updated_state_dict)
-        return self.model_
+        return self.quantize(self.model_)
 
 
 #########################################################################