[quantizer] skip quantization for RMSNorm

alibaba · Sep 11, 2024 · fc5eb63 · fc5eb63
1 parent fca16f3
commit fc5eb63
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/tinynn/graph/quantization/quantizer.py b/tinynn/graph/quantization/quantizer.py
@@ -200,6 +200,7 @@
     'log': None,
     'std': None,
     'var': None,
+    'norm': None,
     nn.LSTM: '1.13.0',
     nn.ConvTranspose2d: '1.7.0',
     nn.ConstantPad1d: '1.7.0',
@@ -261,6 +262,10 @@
     Q_MODULES_MAPPING.update({nn.SiLU: QSiLU})
     FUNCTIONAL_MODULE_MAPPING.update({'silu': nn.SiLU})
 
+if hasattr(nn, 'RMSNorm'):
+    UNSUPPORTED_PYTORCH_QUANTIZATION_OP_LIST.update({nn.RMSNorm: None})
+    FUNCTIONAL_MODULE_MAPPING.update({'rms_norm': nn.RMSNorm})
+
 # Processed QAT fuse rules
 processed_qat_rules = {}
 processed_ptq_rules = {}