From fc5eb6337f62a281f3c0a0f5630b92c8a1248c0a Mon Sep 17 00:00:00 2001
From: zk1998 <zhoukun_1998@qq.com>
Date: Wed, 11 Sep 2024 17:09:29 +0800
Subject: [PATCH] [quantizer] skip quantization for RMSNorm

---
 tinynn/graph/quantization/quantizer.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/tinynn/graph/quantization/quantizer.py b/tinynn/graph/quantization/quantizer.py
index d8cdfdd..a19975a 100644
--- 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 = {}