Support PyTorch eager mode BF16 MixedPrecision (#1321)

Signed-off-by: changwangss <[email protected]> Signed-off-by: yiliu30 <[email protected]> Co-authored-by: yiliu30 <[email protected]>
intel · Nov 15, 2023 · 3bfb76d · 3bfb76d
1 parent dc9328c
commit 3bfb76d
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 0 deletions.
diff --git 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
@@ -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})