add quant norm

qwopqwop200 · May 8, 2023 · 3f1c582 · 3f1c582
1 parent ad22256
commit 3f1c582
Showing 1 changed file with 5 additions and 4 deletions.
diff --git a/llama_inference.py b/llama_inference.py
@@ -57,10 +57,11 @@ def noop(*args, **kwargs):
     else:
         model.load_state_dict(torch.load(checkpoint), strict=False)
 
-    quant.make_quant_attn(model)
-    if eval and fused_mlp:
-        quant.make_fused_mlp(model)
-
+    if eval:
+        quant.make_quant_attn(model)
+        quant.make_quant_norm(model)
+        if fused_mlp:
+            quant.make_fused_mlp(model)
     if warmup_autotune:
         quant.autotune_warmup_linear(model, transpose=not (eval))
         if eval and fused_mlp: