From 3bfb76dca9271c48940ca4af26cabfa48f50ca56 Mon Sep 17 00:00:00 2001 From: "Wang, Chang" Date: Wed, 15 Nov 2023 11:37:20 +0800 Subject: [PATCH] Support PyTorch eager mode BF16 MixedPrecision (#1321) Signed-off-by: changwangss Signed-off-by: yiliu30 Co-authored-by: yiliu30 --- neural_compressor/adaptor/torch_utils/bf16_convert.py | 4 ++++ neural_compressor/strategy/strategy.py | 1 + 2 files changed, 5 insertions(+) diff --git a/neural_compressor/adaptor/torch_utils/bf16_convert.py b/neural_compressor/adaptor/torch_utils/bf16_convert.py index 1c3b9ce789e..917976c2810 100644 --- a/neural_compressor/adaptor/torch_utils/bf16_convert.py +++ b/neural_compressor/adaptor/torch_utils/bf16_convert.py @@ -30,6 +30,10 @@ def __init__(self, module): super(BF16ModuleWrapper, self).__init__() self.add_module("module", module) self.train(module.training) + # WA for TransformerEncoder to access its Linear's weights and bias + if isinstance(module, nn.Linear): + self.weight = self.module.weight if hasattr(self.module, "weight") else None + self.bias = self.module.bias if hasattr(self.module, "bias") else None def forward(self, X): """Convert dtype.""" diff --git a/neural_compressor/strategy/strategy.py b/neural_compressor/strategy/strategy.py index 0a52e11da35..a81b3d16b17 100644 --- a/neural_compressor/strategy/strategy.py +++ b/neural_compressor/strategy/strategy.py @@ -1519,6 +1519,7 @@ def _set_framework_info(self, q_dataloader, q_func=None): elif self.config.backend == "default": framework = "pytorch_fx" if self.mixed_precision_mode: + framework = "pytorch" framework_specific_info.update({"approach": "post_training_dynamic_quant"}) framework_specific_info.update({"recipes": self.config.recipes}) framework_specific_info.update({"q_dataloader": q_dataloader})