Fix small performance regression

0f39807b · Casper Hansen · ed618bb0 · 0f39807b
Commit 0f39807b authored Sep 20, 2023 by Casper Hansen
Hide whitespace changes
Inline Side-by-side

Showing with 4 additions and 1 deletion

awq/quantize/quantizer.py awq/quantize/quantizer.py +4 -1

No files found.
--- a/awq/quantize/quantizer.py
+++ b/awq/quantize/quantizer.py
@@ -75,8 +75,11 @@ class AwqQuantizer:

            # [STEP 4]: Quantize weights
            for name, linear_layer in named_linears.items():
+                # NOTE: small regression in perplexity if linear layer you use .cpu().float()
+                linear_layer = linear_layer.cuda().half()
+
                linear_layer.weight.data, scales, zeros = self.pseudo_quantize_tensor(
-                    linear_layer.weight.data.float(), 
+                    linear_layer.weight.data, 
                    get_scale_zp=True
                )