lint

Author: metascroy

exir/passes/_quant_patterns_and_replacements.py

@@ -25,16 +25,16 @@
 
 from torch import Tensor
 from torch.library import custom_op
+
+
 @custom_op("quant_fusion::_pack_embedding_weight", mutates_args=())
 def _pack_embedding_weight(weight: Tensor, bitwidth: int) -> Tensor:
     num_embeddings, embedding_dim = weight.shape
 
     if bitwidth == 2:
         assert embedding_dim % 4 == 0, "embedding_dim must be divisible by 4"
         weight_range_shifted = weight.add(2).view(torch.uint8)
-        weight_view = weight_range_shifted.view(
-            num_embeddings, embedding_dim // 4, 4
-        )
+        weight_view = weight_range_shifted.view(num_embeddings, embedding_dim // 4, 4)
         weight_0 = weight_view[:, :, 0]
         weight_1 = weight_view[:, :, 1] << 2
         weight_2 = weight_view[:, :, 2] << 4
@@ -53,7 +53,7 @@ def _pack_embedding_weight(weight: Tensor, bitwidth: int) -> Tensor:
         return packed_weight
     elif bitwidth == 8:
         return weight
-    
+
     raise RuntimeError(f"Unsupported bitwidth {bitwidth}")
 
 
@@ -64,7 +64,12 @@ def _(weight, bit_width):
     num_embeddings, embedding_dim = weight.shape
     values_per_byte = 8 // bit_width
     assert embedding_dim % values_per_byte == 0
-    return torch.empty(num_embeddings, embedding_dim // values_per_byte, dtype=torch.uint8, device=weight.device)
+    return torch.empty(
+        num_embeddings,
+        embedding_dim // values_per_byte,
+        dtype=torch.uint8,
+        device=weight.device,
+    )
 
 
 # TODO: extending an existing library that is defined in OSS might be a bit
@@ -114,9 +119,10 @@ def embedding_weight_checks(weight, weight_scales, weight_zero_points):
     assert (
         weight_zero_points is None or weight_zero_points.dtype == weight_scales.dtype
     ), "Expecting weight_zero_points to be None or have same dtype as weight_scales"
-    assert (
-        weight_zero_points is None or weight_zero_points.dim() in [1, 2]
-    ), f"Expecting weight_zero_points tensor to be None or have dim()==1, but found {weight_zero_points.dim()}"
+    assert weight_zero_points is None or weight_zero_points.dim() in [
+        1,
+        2,
+    ], f"Expecting weight_zero_points tensor to be None or have dim()==1, but found {weight_zero_points.dim()}"
     assert weight_zero_points is None or weight_zero_points.size(0) == weight.size(
         0
     ), f"Expecting weight_zero_points tensor to be None or have same number of rows as weights, but found {weight.size()} and {weight_zero_points.size()}"
@@ -278,6 +284,7 @@ def embedding_2bit(
     )
     return torch.ops.aten.embedding.default(weight, indices)
 
+
 @register_fake("quantized_decomposed::embedding_2bit")
 def _(
     weight: torch.Tensor,
@@ -286,12 +293,13 @@ def _(
     weight_quant_min: int,
     weight_quant_max: int,
     indices: torch.Tensor,
-    ):
+):
     num_embeddings, packed_embedding_dim = weight.shape
     embedding_dim = packed_embedding_dim * 4
     embedding = torch.nn.Embedding(num_embeddings, embedding_dim, device=weight.device)
     return embedding(indices)
 
+
 @register_fake("quantized_decomposed::embedding_2bit.out")
 def embedding_2bit_out_meta(
     weight: torch.Tensor,
@@ -311,6 +319,7 @@ def embedding_2bit_out_meta(
         indices,
     )
 
+
 @impl(quantized_decomposed_lib, "embedding_2bit.dtype", "CompositeExplicitAutograd")
 def embedding_2bit_dtype(
     weight: torch.Tensor,
@@ -352,6 +361,7 @@ def embedding_2bit_dtype(
     )
     return torch.ops.aten.embedding.default(weight, indices)
 
+
 @register_fake("quantized_decomposed::embedding_2bit.dtype")
 def _(
     weight: torch.Tensor,
@@ -361,12 +371,13 @@ def _(
     weight_quant_max: int,
     indices: torch.Tensor,
     dtype: Optional[torch.dtype],
-    ) -> torch.Tensor:
+) -> torch.Tensor:
     num_embeddings, packed_embedding_dim = weight.shape
     embedding_dim = packed_embedding_dim * 4
     embedding = torch.nn.Embedding(num_embeddings, embedding_dim, device=weight.device)
     return embedding(indices).to(dtype)
 
+
 @register_fake("quantized_decomposed::embedding_2bit.dtype_out")
 def embedding_2bit_dtype_out_meta(
     weight: torch.Tensor,
@@ -448,6 +459,7 @@ def embedding_4bit(
     )
     return torch.ops.aten.embedding.default(weight, indices)
 
+
 @register_fake("quantized_decomposed::embedding_4bit")
 def _(
     weight: torch.Tensor,
@@ -456,12 +468,13 @@ def _(
     weight_quant_min: int,
     weight_quant_max: int,
     indices: torch.Tensor,
-    ):
+):
     num_embeddings, packed_embedding_dim = weight.shape
     embedding_dim = packed_embedding_dim * 2
     embedding = torch.nn.Embedding(num_embeddings, embedding_dim, device=weight.device)
     return embedding(indices)
 
+
 @register_fake("quantized_decomposed::embedding_4bit.out")
 def embedding_4bit_out_meta(
     weight: torch.Tensor,
@@ -521,6 +534,7 @@ def embedding_4bit_dtype(
     )
     return torch.ops.aten.embedding.default(weight, indices)
 
+
 @register_fake("quantized_decomposed::embedding_4bit.dtype")
 def _(
     weight: torch.Tensor,
@@ -530,12 +544,13 @@ def _(
     weight_quant_max: int,
     indices: torch.Tensor,
     dtype: Optional[torch.dtype],
-    ) -> torch.Tensor:
+) -> torch.Tensor:
     num_embeddings, packed_embedding_dim = weight.shape
     embedding_dim = packed_embedding_dim * 2
     embedding = torch.nn.Embedding(num_embeddings, embedding_dim, device=weight.device)
     return embedding(indices).to(dtype)
 
+
 @register_fake("quantized_decomposed::embedding_4bit.dtype_out")
 def embedding_4bit_dtype_out_meta(
     weight: torch.Tensor,
@@ -970,10 +985,16 @@ def replacement(x, dim, start, end, x_scale, x_zero_point, x_qmin, x_qmax):
         )
     ]
 
-def _get_embedding_ops_patterns_and_replacements_torchao() ->  List[Tuple[Callable, Callable, List[Callable]]]:
+
+def _get_embedding_ops_patterns_and_replacements_torchao() -> (  # noqa C901
+    List[Tuple[Callable, Callable, List[Callable]]]
+):
     def embedding_byte_pattern(indices, int_data, group_size, scale, zero_point):
-        dq = torch.ops.torchao.dequantize_affine.default(int_data, [1, group_size], scale, zero_point, torch.int8, -128, 127)
+        dq = torch.ops.torchao.dequantize_affine.default(
+            int_data, [1, group_size], scale, zero_point, torch.int8, -128, 127
+        )
         return torch.ops.aten.embedding.default(dq, indices)
+
     def embedding_byte_replacement(indices, int_data, group_size, scale, zero_point):
         zero_point_dtype_cast = torch.ops.aten.to.dtype(zero_point, scale.dtype)
         return torch.ops.quantized_decomposed.embedding_byte.default(
@@ -984,10 +1005,26 @@ def embedding_byte_replacement(indices, int_data, group_size, scale, zero_point)
             127,
             indices,
         )
-    def embedding_byte_dtype_pattern(indices, int_data, group_size, scale, zero_point, output_dtype):
-        dq = torch.ops.torchao.dequantize_affine.default(int_data, [1, group_size], scale, zero_point, torch.int8, -128, 127, 'INT', output_dtype)
+
+    def embedding_byte_dtype_pattern(
+        indices, int_data, group_size, scale, zero_point, output_dtype
+    ):
+        dq = torch.ops.torchao.dequantize_affine.default(
+            int_data,
+            [1, group_size],
+            scale,
+            zero_point,
+            torch.int8,
+            -128,
+            127,
+            "INT",
+            output_dtype,
+        )
         return torch.ops.aten.embedding.default(dq, indices)
-    def embedding_byte_dtype_replacement(indices, int_data, group_size, scale, zero_point, output_dtype):
+
+    def embedding_byte_dtype_replacement(
+        indices, int_data, group_size, scale, zero_point, output_dtype
+    ):
         zero_point_dtype_cast = torch.ops.aten.to.dtype(zero_point, scale.dtype)
         return torch.ops.quantized_decomposed.embedding_byte.dtype(
             int_data,
@@ -996,48 +1033,136 @@ def embedding_byte_dtype_replacement(indices, int_data, group_size, scale, zero_
             -128,
             127,
             indices,
-            dtype=output_dtype
+            dtype=output_dtype,
         )
-    
+
     def embedding_2bit_pattern(indices, int_data, group_size, scale, zero_point):
-        dq = torch.ops.torchao.dequantize_affine.default(int_data, [1, group_size], scale, zero_point, torch.int8, -2, 1)
+        dq = torch.ops.torchao.dequantize_affine.default(
+            int_data, [1, group_size], scale, zero_point, torch.int8, -2, 1
+        )
         return torch.ops.aten.embedding.default(dq, indices)
+
     def embedding_2bit_replacement(indices, int_data, group_size, scale, zero_point):
-        packed_int_data = torch.ops.quant_fusion._pack_embedding_weight.default(int_data, 2)
+        packed_int_data = torch.ops.quant_fusion._pack_embedding_weight.default(
+            int_data, 2
+        )
         zero_point_dtype_cast = torch.ops.aten.to.dtype(zero_point, scale.dtype)
-        return torch.ops.quantized_decomposed.embedding_2bit.default(packed_int_data, scale, zero_point_dtype_cast, -2, 1, indices)
+        return torch.ops.quantized_decomposed.embedding_2bit.default(
+            packed_int_data, scale, zero_point_dtype_cast, -2, 1, indices
+        )
 
-    def embedding_2bit_dtype_pattern(indices, int_data, group_size, scale, zero_point, output_dtype):
-        dq = torch.ops.torchao.dequantize_affine.default(int_data, [1, group_size], scale, zero_point, torch.int8, -2, 1, 'INT', output_dtype)
+    def embedding_2bit_dtype_pattern(
+        indices, int_data, group_size, scale, zero_point, output_dtype
+    ):
+        dq = torch.ops.torchao.dequantize_affine.default(
+            int_data,
+            [1, group_size],
+            scale,
+            zero_point,
+            torch.int8,
+            -2,
+            1,
+            "INT",
+            output_dtype,
+        )
         return torch.ops.aten.embedding.default(dq, indices)
-    def embedding_2bit_dtype_replacement(indices, int_data, group_size, scale, zero_point, output_dtype):
-        packed_int_data = torch.ops.quant_fusion._pack_embedding_weight.default(int_data, 2)
+
+    def embedding_2bit_dtype_replacement(
+        indices, int_data, group_size, scale, zero_point, output_dtype
+    ):
+        packed_int_data = torch.ops.quant_fusion._pack_embedding_weight.default(
+            int_data, 2
+        )
         zero_point_dtype_cast = torch.ops.aten.to.dtype(zero_point, scale.dtype)
-        return torch.ops.quantized_decomposed.embedding_2bit.dtype(packed_int_data, scale, zero_point_dtype_cast, -2, 1, indices, dtype=output_dtype)
-    
+        return torch.ops.quantized_decomposed.embedding_2bit.dtype(
+            packed_int_data,
+            scale,
+            zero_point_dtype_cast,
+            -2,
+            1,
+            indices,
+            dtype=output_dtype,
+        )
+
     def embedding_4bit_pattern(indices, int_data, group_size, scale, zero_point):
-        dq = torch.ops.torchao.dequantize_affine.default(int_data, [1, group_size], scale, zero_point, torch.int8, -8, 7)
+        dq = torch.ops.torchao.dequantize_affine.default(
+            int_data, [1, group_size], scale, zero_point, torch.int8, -8, 7
+        )
         return torch.ops.aten.embedding.default(dq, indices)
+
     def embedding_4bit_replacement(indices, int_data, group_size, scale, zero_point):
-        packed_int_data = torch.ops.quant_fusion._pack_embedding_weight.default(int_data, 4)
+        packed_int_data = torch.ops.quant_fusion._pack_embedding_weight.default(
+            int_data, 4
+        )
         zero_point_dtype_cast = torch.ops.aten.to.dtype(zero_point, scale.dtype)
-        return torch.ops.quantized_decomposed.embedding_4bit.default(packed_int_data, scale, zero_point_dtype_cast, -8, 7, indices)
-    
-    def embedding_4bit_dtype_pattern(indices, int_data, group_size, scale, zero_point, output_dtype):
-        dq = torch.ops.torchao.dequantize_affine.default(int_data, [1, group_size], scale, zero_point, torch.int8, -8, 7, 'INT', output_dtype)
+        return torch.ops.quantized_decomposed.embedding_4bit.default(
+            packed_int_data, scale, zero_point_dtype_cast, -8, 7, indices
+        )
+
+    def embedding_4bit_dtype_pattern(
+        indices, int_data, group_size, scale, zero_point, output_dtype
+    ):
+        dq = torch.ops.torchao.dequantize_affine.default(
+            int_data,
+            [1, group_size],
+            scale,
+            zero_point,
+            torch.int8,
+            -8,
+            7,
+            "INT",
+            output_dtype,
+        )
         return torch.ops.aten.embedding.default(dq, indices)
-    def embedding_4bit_dtype_replacement(indices, int_data, group_size, scale, zero_point, output_dtype):
-        packed_int_data = torch.ops.quant_fusion._pack_embedding_weight.default(int_data, 4)
+
+    def embedding_4bit_dtype_replacement(
+        indices, int_data, group_size, scale, zero_point, output_dtype
+    ):
+        packed_int_data = torch.ops.quant_fusion._pack_embedding_weight.default(
+            int_data, 4
+        )
         zero_point_dtype_cast = torch.ops.aten.to.dtype(zero_point, scale.dtype)
-        return torch.ops.quantized_decomposed.embedding_4bit.dtype(packed_int_data, scale, zero_point_dtype_cast, -8, 7, indices, dtype=output_dtype)
+        return torch.ops.quantized_decomposed.embedding_4bit.dtype(
+            packed_int_data,
+            scale,
+            zero_point_dtype_cast,
+            -8,
+            7,
+            indices,
+            dtype=output_dtype,
+        )
 
     return [
-        (_trace_and_lower_to_edge_ops(embedding_byte_pattern), _trace_and_lower_to_edge_ops(embedding_byte_replacement), []),
-        (_trace_and_lower_to_edge_ops(embedding_byte_dtype_pattern), _trace_and_lower_to_edge_ops(embedding_byte_dtype_replacement), []),
-        (_trace_and_lower_to_edge_ops(embedding_2bit_pattern), _trace_and_lower_to_edge_ops(embedding_2bit_replacement), []),
-        (_trace_and_lower_to_edge_ops(embedding_2bit_dtype_pattern), _trace_and_lower_to_edge_ops(embedding_2bit_dtype_replacement), []),
-        (_trace_and_lower_to_edge_ops(embedding_4bit_pattern), _trace_and_lower_to_edge_ops(embedding_4bit_replacement), []),
-        (_trace_and_lower_to_edge_ops(embedding_4bit_dtype_pattern), _trace_and_lower_to_edge_ops(embedding_4bit_dtype_replacement), []),
+        (
+            _trace_and_lower_to_edge_ops(embedding_byte_pattern),
+            _trace_and_lower_to_edge_ops(embedding_byte_replacement),
+            [],
+        ),
+        (
+            _trace_and_lower_to_edge_ops(embedding_byte_dtype_pattern),
+            _trace_and_lower_to_edge_ops(embedding_byte_dtype_replacement),
+            [],
+        ),
+        (
+            _trace_and_lower_to_edge_ops(embedding_2bit_pattern),
+            _trace_and_lower_to_edge_ops(embedding_2bit_replacement),
+            [],
+        ),
+        (
+            _trace_and_lower_to_edge_ops(embedding_2bit_dtype_pattern),
+            _trace_and_lower_to_edge_ops(embedding_2bit_dtype_replacement),
+            [],
+        ),
+        (
+            _trace_and_lower_to_edge_ops(embedding_4bit_pattern),
+            _trace_and_lower_to_edge_ops(embedding_4bit_replacement),
+            [],
+        ),
+        (
+            _trace_and_lower_to_edge_ops(embedding_4bit_dtype_pattern),
+            _trace_and_lower_to_edge_ops(embedding_4bit_dtype_replacement),
+            [],
+        ),
     ]
 
 

exir/passes/quant_fusion_pass.py

@@ -89,6 +89,7 @@ def _get_qparams(node):
         qnode.replace_all_uses_with(maybe_cat)
         model.graph.erase_node(qnode)
 
+
 def _remove_dtype_getattr_nodes(model: GraphModule) -> None:
     for n in model.graph.nodes:
         if n.op == "call_function" and n.target == getattr:
@@ -99,7 +100,8 @@ def _remove_dtype_getattr_nodes(model: GraphModule) -> None:
     model.graph.eliminate_dead_code()
     model.graph.lint()
     model.recompile()
-                
+
+
 class QuantFusionPass(ExportPass):
     def __init__(self, _fix_node_meta_val=False):
         super().__init__()