fix: Minor changes to quantizer, qmax, sawb

Signed-off-by: Brandon Groth <[email protected]>

Author: BrandonGroth

fms_mo/quant_refactor/base_quant.py

@@ -75,6 +75,9 @@ def __init__(
             RuntimeError: perCh or perGrp was selected without specifying Nch, Ngrp.
             RuntimeError: qscheme is not allowed, or could be a typo.
         """
+        # Init Nch/Ngrp to none incase they won't be set
+        self.Nch = None
+        self.Ngrp = None
         if isinstance(unit, torch.qscheme):
             if "per_channel" in str(unit):
                 self.q_unit = "perCh"
@@ -124,8 +127,8 @@ def __repr__(self):
         """
         q_uint_str = f"qunit={self.q_unit}"
         symmetric_str = f", symmetric={self.symmetric}"
-        Nch_str = f", Nch={self.Nch}",
-        Ngrp_str = f", Nch={self.Ngrp}",
+        Nch_str = f", Nch={self.Nch}" if self.Nch is not None else "",
+        Ngrp_str = f", Ngrp={self.Ngrp}" if self.Ngrp is not None else "",
         single_sided_str = f", single_sided={self.single_sided}"
         qlevel_lowering_str = f", qlevel_lowering={self.qlevel_lowering}"
         return (

fms_mo/quant_refactor/qmax_new.py

@@ -89,9 +89,10 @@ def __init__(
             use_PT_native_Qfunc=kwargs.get("use_PT_native_Qfunc", False),
         )
 
-        with torch.no_grad():
-            self.clip_valn.data *= init_clip_valn
-            self.clip_val.data *= init_clip_val
+        if not self.training:
+            with torch.no_grad():
+                self.clip_valn.data *= init_clip_valn
+                self.clip_val.data *= init_clip_val
 
         self.align_zero = align_zero
         self.clipSTE = clipSTE

fms_mo/quant_refactor/sawb_new.py

@@ -92,9 +92,10 @@ def __init__(
             use_PT_native_Qfunc=kwargs.get("use_PT_native_Qfunc", False),
         )
 
-        with torch.no_grad():
-            self.clip_valn.data *= init_clip_valn
-            self.clip_val.data *= init_clip_val
+        if not self.training:
+            with torch.no_grad():
+                self.clip_valn.data *= init_clip_valn
+                self.clip_val.data *= init_clip_val
 
         self.clipSTE = clipSTE
         self.align_zero = align_zero