[Tests] Combine quantization and dequantization tests (#443)

Signed-off-by: Kyle Sayers <[email protected]>

Author: kylesayrs

tests/test_quantization/lifecycle/test_forward.py

@@ -19,9 +19,8 @@
 import torch
 from compressed_tensors.quantization.lifecycle.forward import (
     _process_quantization,
-    dequantize,
+    fake_quantize,
     forward_quantize,
-    quantize,
     wrap_module_forward_quantized,
 )
 from compressed_tensors.quantization.lifecycle.initialize import (
@@ -96,7 +95,7 @@ def test_forward_quantize(
 
 
 @pytest.mark.parametrize(
-    "num_bits,type,strategy,group_size,scale,zero_point,g_idx",
+    "num_bits,type,strategy,group_size,scale,zero_point,g_idx,global_scale",
     [
         (
             4,
@@ -106,6 +105,7 @@ def test_forward_quantize(
             torch.rand((1,)) * 0.01,
             torch.zeros((1,)),
             None,
+            None,
         ),
         (
             4,
@@ -115,6 +115,7 @@ def test_forward_quantize(
             torch.rand((512, 8)) * 0.01,
             torch.zeros((512, 8)),
             None,
+            None,
         ),
         (
             4,
@@ -124,6 +125,7 @@ def test_forward_quantize(
             torch.rand((512, 8)) * 0.01,
             torch.zeros((512, 8)),
             make_dummy_g_idx(1024, 128),
+            None,
         ),
         (
             8,
@@ -133,6 +135,7 @@ def test_forward_quantize(
             torch.rand((1,)) * 0.01,
             torch.zeros((1,)),
             None,
+            None,
         ),
         (
             8,
@@ -142,6 +145,7 @@ def test_forward_quantize(
             torch.rand((512, 8)) * 0.01,
             torch.zeros((512, 8)),
             None,
+            None,
         ),
         (
             8,
@@ -151,28 +155,8 @@ def test_forward_quantize(
             torch.rand((512, 8)) * 0.01,
             torch.zeros((512, 8)),
             make_dummy_g_idx(1024, 128),
+            None,
         ),
-    ],
-)
-def test_quantize(num_bits, type, strategy, group_size, scale, zero_point, g_idx):
-    args = QuantizationArgs(
-        num_bits=num_bits, type=type, strategy=strategy, group_size=group_size
-    )
-
-    x = torch.rand((512, 1024))
-    quantize(
-        x=x,
-        scale=scale,
-        zero_point=zero_point,
-        args=args,
-        dtype=args.pytorch_dtype(),
-        g_idx=g_idx,
-    )
-
-
-@pytest.mark.parametrize(
-    "num_bits,type,strategy,group_size,scale,zero_point,g_idx",
-    [
         (
             8,
             "int",
@@ -181,6 +165,7 @@ def test_quantize(num_bits, type, strategy, group_size, scale, zero_point, g_idx
             torch.rand((512, 8)) * 0.01,
             torch.zeros((512, 8)),
             None,
+            None,
         ),
         (
             8,
@@ -190,23 +175,26 @@ def test_quantize(num_bits, type, strategy, group_size, scale, zero_point, g_idx
             torch.rand((512, 8)) * 0.01,
             torch.zeros((512, 8)),
             make_dummy_g_idx(1024, 128),
+            None,
         ),
     ],
 )
-def test_dequantize(num_bits, type, strategy, group_size, scale, zero_point, g_idx):
+def test_fake_quantize_2d(
+    num_bits, type, strategy, group_size, scale, zero_point, g_idx, global_scale
+):
     args = QuantizationArgs(
         num_bits=num_bits, type=type, strategy=strategy, group_size=group_size
     )
 
-    x_q = torch.rand((512, 1024)).to(dtype=args.pytorch_dtype())
-    dequantize(
-        x_q=x_q,
+    x = torch.rand((512, 1024))
+    fake_quantize(
+        x=x,
         scale=scale,
         zero_point=zero_point,
         args=args,
-        dtype=None,
         g_idx=g_idx,
-    )
+        global_scale=global_scale,
+    )  # note that reconstruction loss is bad for uncalibrated scales
 
 
 def test_process_quantization_block_static():