Improve patches for transformers

Author: xadupre

.github/workflows/ci.yml

@@ -15,7 +15,7 @@ jobs:
     strategy:
       matrix:
         os: [ubuntu-latest]
-        python: ['3.12']
+        python: ['3.11', '3.12']
         transformers: ['4.48', 'main']
 
     steps:

_doc/examples/plot_export_with_dynamic_cache.py

@@ -210,16 +210,13 @@ def forward(self, cache, z):
 # The export is simple if ``transformers>=4.50``, otherwise,
 # transformers needs to be patched.
 # :func:`onnx_diagnostic.torch_export_patches.bypass_export_some_errors`
-# registers functions to serialize ``DynamicCache`` and another class
-# called ``patched_DynamicCache``. This one is modified to make
+# registers functions to serialize ``DynamicCache``. This one is modified to make
 # the shape inference implemented in :epkg:`torch` happy.
 
 if has_transformers("4.50"):
     ep = torch.export.export(model, inputs[0], dynamic_shapes=ds[0], strict=False)
 else:
-    with bypass_export_some_errors(
-        patch_transformers=True, replace_dynamic_cache=True
-    ) as modificator:
+    with bypass_export_some_errors(patch_transformers=True) as modificator:
         ep = torch.export.export(
             model, modificator(inputs[0]), dynamic_shapes=ds[0], strict=False
         )

_unittests/ut_torch_export_patches/test_onnx_export_errors.py

@@ -79,7 +79,7 @@ def forward(self, x: torch.Tensor, cache: MambaCache):
         model = Model()
         model(x, cache)
 
-        with bypass_export_some_errors(replace_dynamic_cache=True, verbose=1):
+        with bypass_export_some_errors(verbose=1):
             cache = MambaCache(_config(), max_batch_size=1, device="cpu")
             torch.export.export(Model(), (x, cache))
 

_unittests/ut_torch_export_patches/test_patch_base_class.py

@@ -0,0 +1,57 @@
+import unittest
+from onnx_diagnostic.ext_test_case import ExtTestCase
+
+
+class TestPatchBaseClass(ExtTestCase):
+    def test_check_that_trick_can_work_in_python(self):
+        class zero:
+            def ret(self, a):
+                return a - 100
+
+            def ok(self):
+                return self.ret(3)
+
+        class A(zero):
+            def ret(self, a):
+                return a + 1
+
+        class B:
+            def ret(self, a):
+                return a + 10
+
+        z = zero()
+        self.assertEqual(z.ret(4), -96)
+        self.assertEqual(z.ok(), -97)
+        a = A()
+        self.assertEqual(a.ret(4), 5)
+        self.assertEqual(a.ok(), 4)
+        b = B()
+        self.assertEqual(b.ret(4), 14)
+        self.assertFalse(hasattr(b, "ok"))
+        self.assertFalse(hasattr(B, "ok"))
+
+        self.assertEqual(A.__bases__, (zero,))
+        A.__bases__ = (zero, B)
+        self.assertEqual(a.ret(4), 5)
+        self.assertEqual(a.ok(), 4)
+        aa = A()
+        self.assertEqual(aa.ret(4), 5)
+        self.assertEqual(aa.ok(), 4)
+
+        A.__bases__ = (B, zero)
+        self.assertEqual(a.ret(4), 5)
+        self.assertEqual(a.ok(), 4)
+        aa = A()
+        self.assertEqual(aa.ret(4), 5)
+        self.assertEqual(aa.ok(), 4)
+
+        A.__bases__ = (zero,)
+        A.ret = B.ret
+        self.assertEqual(aa.ret(4), 14)
+        self.assertEqual(aa.ok(), 13)
+        self.assertEqual(a.ret(4), 14)
+        self.assertEqual(a.ok(), 13)
+
+
+if __name__ == "__main__":
+    unittest.main(verbosity=2)

_unittests/ut_torch_models/test_tiny_llms.py

@@ -29,9 +29,7 @@ def test_export_tiny_llm_2_bypassed(self):
         data = get_tiny_llm()
         model, inputs = data["model"], data["inputs"]
         self.assertEqual({"attention_mask", "past_key_values", "input_ids"}, set(inputs))
-        with bypass_export_some_errors(
-            patch_transformers=True, replace_dynamic_cache=True
-        ) as modificator:
+        with bypass_export_some_errors(patch_transformers=True) as modificator:
             inputs = modificator(inputs)
             ep = torch.export.export(
                 model, (), kwargs=inputs, dynamic_shapes=data["dynamic_shapes"]

_unittests/ut_torch_models/test_tiny_llms_onnx.py

@@ -57,9 +57,7 @@ def test_bypass_onnx_export_tiny_llm_official(self):
         data = get_tiny_llm()
         model, inputs = data["model"], data["inputs"]
         self.assertEqual({"attention_mask", "past_key_values", "input_ids"}, set(inputs))
-        with bypass_export_some_errors(
-            patch_transformers=True, replace_dynamic_cache=True, verbose=1
-        ) as modificator:
+        with bypass_export_some_errors(patch_transformers=True, verbose=1) as modificator:
             new_inputs = modificator(inputs)
             ep = torch.onnx.export(
                 model,
@@ -80,9 +78,7 @@ def test_bypass_onnx_export_tiny_llm_xdbg(self):
         data = get_tiny_llm()
         model, inputs = data["model"], data["inputs"]
         self.assertEqual({"attention_mask", "past_key_values", "input_ids"}, set(inputs))
-        with bypass_export_some_errors(
-            patch_transformers=True, replace_dynamic_cache=True, verbose=1
-        ) as modificator:
+        with bypass_export_some_errors(patch_transformers=True, verbose=1) as modificator:
             new_inputs = modificator(inputs)
             onx = to_onnx(
                 model, (), kwargs=new_inputs, dynamic_shapes=data["dynamic_shapes"], verbose=1

onnx_diagnostic/export/dynamic_shapes.py

@@ -311,7 +311,7 @@ def guess_dynamic_shape_object(self, *objs: Any, msg: Optional[Callable] = None)
                 shapes[i] = self.guess_dynamic_shape_object(*[o[i] for o in objs])
             return shapes
 
-        if obj.__class__.__name__ in ("DynamicCache", "patched_DynamicCache"):
+        if obj.__class__.__name__ == "DynamicCache":
             kc = set(len(o.key_cache) for o in objs)
             assert (
                 len(kc) == 1

onnx_diagnostic/helpers.py

@@ -408,7 +408,7 @@ def string_type(
     if type(obj).__name__ == "ValueInfoProto":
         return f"OT{obj.type.tensor_type.elem_type}"
 
-    if obj.__class__.__name__ in ("DynamicCache", "patched_DynamicCache"):
+    if obj.__class__.__name__ == "DynamicCache":
         kc = string_type(
             obj.key_cache,
             with_shape=with_shape,
@@ -1693,8 +1693,8 @@ def max_diff(
             flatten=flatten,
         )
 
-    if expected.__class__.__name__ in ("DynamicCache", "patched_DynamicCache"):
-        if got.__class__.__name__ in ("DynamicCache", "patched_DynamicCache"):
+    if expected.__class__.__name__ == "DynamicCache":
+        if got.__class__.__name__ == "DynamicCache":
             if verbose >= 6:
                 print(f"[max_diff] DynamicCache: {string_type(expected)} ? {string_type(got)}")
             return max_diff(