Revert "[SW-205334][SW-187731] llama70b vLLM fix graph breaks with torch.com…" (#87)

This reverts commit 01a5734.

Co-authored-by: Danny Semiat <[email protected]>

Author: dudilester

neural_compressor/torch/algorithms/fp8_quant/_quant_common/helper_modules.py

@@ -14,7 +14,6 @@
 
 import torch
 import torch.nn as nn
-import types
 
 from .quant_config import QuantMode, get_hqt_config, ScaleFormat
 from .._core.quant_dequant import QuantDequant as qdq
@@ -75,10 +74,13 @@ def set_attrs_from_orig_model(cls_instance, mod, mod_extra_config, *func_names):
     cls_instance.fake_quant = config.cfg["fake_quant"]
     cls_instance.use_qdq = config.cfg["use_qdq"]
     cls_instance.scale_format = config.cfg["scale_format"]
-    cls_instance.forward_orig = types.MethodType(mod.forward.__func__, cls_instance)
+    # store original module in order to invoke its functions during measurements.
+    # this may be omitted of torch remove the related validation from dynamo. see SW-187731.
+    cls_instance.__dict__["orig_mod"] = mod
+    cls_instance.forward_orig = mod.forward
     if func_names is not None:
         for func in func_names:
-            setattr(cls_instance, func, types.MethodType(getattr(mod, func).__func__, cls_instance))
+            setattr(cls_instance, func, getattr(mod, func))
 
 
 def get_current_repr(cls_instance, *member_names):
@@ -143,7 +145,7 @@ def forward_qdq(self, input, other):
 
     def forward_measure(self, input, other):
         measure_input((input, other), observer=self._mod_extra_config.inputs)
-        output = self.forward_orig(input, other)
+        output = self.orig_mod(input, other)
         measure_output((output,), self._mod_extra_config.outputs)
         return output
 
@@ -212,7 +214,7 @@ def forward_quant(self, input):
 
     def forward_measure(self, input):
         measure_input((input,), observer=self._mod_extra_config.inputs)
-        output = self.forward_orig(input)
+        output = self.orig_mod(input)
         measure_output((output,), self._mod_extra_config.outputs)
         return output
 
@@ -520,7 +522,7 @@ def forward_qdq(self, input):
         output = torch.matmul(qinput, qweight)
 
         if self.gather_output:
-            output = self.collective_func(output)
+            output = self.orig_mod.collective_func(output)
         return self.post_all_reduce(output)
 
     def forward_quant(self, input):
@@ -532,15 +534,15 @@ def forward_quant(self, input):
                                  scale_other_inv=self.scale_weight)
         dqoutput = self.dequant_output(output)
         if self.gather_output:
-            dqoutput = self.collective_func(dqoutput)
+            dqoutput = self.orig_mod.collective_func(dqoutput)
         return self.post_all_reduce(dqoutput)
 
     def forward_measure(self, input):
         measure_input((input,), observer=self._mod_extra_config.inputs)
         output = torch.matmul(input, self.weight.transpose(-1, -2))
         measure_output((output,), self._mod_extra_config.outputs)
         if self.gather_output:
-            output = self.collective_func(output)
+            output = self.orig_mod.collective_func(output)
         return self.post_all_reduce(output)
 
     def post_all_reduce(self, output):
@@ -689,7 +691,7 @@ def __init__(self, mod, mod_extra_config, *args, **kwargs):
         super().__init__()
         set_attrs_from_orig_model(self, mod, mod_extra_config)
         if self.quantization_mode in [QuantMode.QUANTIZE, QuantMode.LOAD]:
-            self.orig_fetch_from_cache = types.MethodType(mod.fetch_from_cache.__func__, self)
+            self.orig_fetch_from_cache = mod.fetch_from_cache
             self.quant_input = self._mod_extra_config.inputs[0]
             self.dequant_output = self._mod_extra_config.outputs[0]
             if self.use_qdq:
@@ -775,7 +777,7 @@ def forward_quant(self, input):
 
     def forward_measure(self, input):
         measure_input((input,), observer=self._mod_extra_config.inputs)
-        output = self.forward_orig(input)
+        output = self.orig_mod(input)
         measure_output((output,), self._mod_extra_config.outputs)
         return output
 
@@ -816,7 +818,7 @@ def forward_quant(self, x, dim=None, invAttnHead=None):
 
     def forward_measure(self, x, dim=None, invAttnHead=None):
         measure_input((x,), observer=self._mod_extra_config.inputs)
-        output = self.forward_orig(x, dim)
+        output = self.orig_mod(x, dim)
         measure_output((output,), self._mod_extra_config.outputs)
         return output
 
@@ -861,7 +863,7 @@ def forward_quant(self, input, scale: float = 1.0):
 
     def forward_measure(self, input, scale: float = 1.0):
         measure_input((input,), observer=self._mod_extra_config.inputs)
-        output = self.forward_orig(input, scale)
+        output = self.orig_mod(input, scale)
         measure_output((output,), self._mod_extra_config.outputs)
         return output
 
@@ -916,7 +918,7 @@ def forward_quant(self, input, scale: float = 1.0):
 
     def forward_measure(self, input, scale: float = 1.0):
         measure_input((input,), observer=self._mod_extra_config.inputs)
-        output = self.forward_orig(input, scale)
+        output = self.orig_mod(input, scale)
         measure_output((output,), self._mod_extra_config.outputs)
         return output
 

test/3x/torch/algorithms/fp8_quant/test_register_apis.py

@@ -231,7 +231,7 @@ def forward_measure(self, input):
                 )
 
                 measure_input((input,), observer=self._mod_extra_config.inputs)
-                output = self.forward_orig(input)
+                output = self.orig_mod(input)
                 measure_output((output,), self._mod_extra_config.outputs)
                 return output