try fix test

add debug print to debug

remove print, add bias to asym quant tests

lint

Author: lsy323

jetstream_pt/layers.py

@@ -246,7 +246,6 @@ def quantize_weight_from_nn_linear(self, weight):
         weight, (1,), self.n_bit, self.is_symmetric_weight, self.block_size
     )
     w_dq = dequantize_tensor(w_q, scale, zp)
-    print("check qweight cosine dist: ", _calc_cosine_dist(weight, w_dq))
     self._load_quantized_weights(w_q, scale, zp)
 
   def forward(self, inputs):

tests/test_quantization.py

@@ -143,13 +143,12 @@ def quantize_dequantize_weight(w, n_bit):
       w_q_asym, s_asym, zp_asym = quantize_tensor(
           w, (1,), n_bit=n_bit, symmetric=False
       )
-      # print(f"w_q_asym {w_q_asym}, s_asym {s_asym}, zp_asym {zp_asym}")
       w_dq_asym = dequantize_tensor(w_q_asym, s_asym, zp_asym)
-      # print(f"w_dq_asym {w_dq_asym}")
-      # self._print_diff(w, w_dq)
-      # self._print_diff(w, w_dq_asym)
       # Asymmetric is more accurate than symmetric.
-      self.assertLess((w - w_dq_asym).norm(), (w - w_dq).norm())
+      self.assertLess(
+          (w - w_dq_asym).norm(),
+          (w - w_dq).norm(),
+      )
       # Blockwise quant.
       w_block_q, s_block, _ = quantize_tensor(
           w, (1,), n_bit=n_bit, symmetric=True, block_size=2
@@ -169,7 +168,9 @@ def quantize_dequantize_weight(w, n_bit):
       # Blockwise asymmetric is more accurate than blockwise symmetric.
       self.assertLess((w - w_block_asym_dq).norm(), (w - w_block_dq).norm())
 
-    w = torch.randn(2, 8)
+    w = (
+        torch.randn(2, 8) + 2
+    )  # Add a bias to normal dist to test asymmetric quant.
     for bit in [4, 8]:
       with self.subTest(bit=bit):
         quantize_dequantize_weight(w, bit)