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 = {}