Fix minor bugs in onnx_plug

Author: xadupre

_unittests/ut_tasks/try_export.py

@@ -45,6 +45,9 @@ def test_imagetext2text_qwen_2_5_vl_instruct_visual(self):
         exporter = os.environ.get("EXPORTER", "custom")
 
         from transformers import AutoModel, AutoProcessor
+        from onnx_diagnostic.torch_export_patches.patches._patch_transformers_qwen2_5 import (
+            PLUGS,
+        )
 
         model_id = "Qwen/Qwen2.5-VL-7B-Instruct"
         # model_id = "Qwen/Qwen2.5-VL-3B-Instruct"
@@ -82,11 +85,17 @@ def _config_reduction(config, task):
         processor = AutoProcessor.from_pretrained(model_id, use_fast=True)
         print(f"-- processor={type(processor)}")
 
+        big_inputs = dict(
+            hidden_states=torch.rand((14308, 1176), dtype=torch_dtype).to(device),
+            grid_thw=torch.tensor([[1, 98, 146]], dtype=torch.int64).to(device),
+        )
+        print("-- save inputs")
         inputs = dict(
             hidden_states=torch.rand((1292, 1176), dtype=torch_dtype).to(device),
             grid_thw=torch.tensor([[1, 34, 38]], dtype=torch.int64).to(device),
         )
         print("-- save inputs")
+        torch.save(big_inputs, self.get_dump_file("qwen_2_5_vl_instruct_visual.inputs.big.pt"))
         torch.save(inputs, self.get_dump_file("qwen_2_5_vl_instruct_visual.inputs.pt"))
 
         print(f"-- inputs: {self.string_type(inputs, with_shape=True)}")
@@ -115,15 +124,6 @@ def _config_reduction(config, task):
             verbose=1,
             stop_if_static=2,
         ):
-            if exporter == "onnx-dynamo":
-                # The exported program in ONNXProgram cannot be restored.
-                ep2 = torch.export.export(
-                    model.visual,
-                    (),
-                    kwargs=export_inputs,
-                    dynamic_shapes=self.use_dyn_not_str(dynamic_shapes),
-                )
-                torch.export.save(ep2, f"{fileep}.backup.pt2")
             to_onnx(
                 model.visual,
                 kwargs=export_inputs,
@@ -134,6 +134,7 @@ def _config_reduction(config, task):
                 save_ep=(fileep, 2**35),
                 target_opset=22,
                 optimize=True,
+                onnx_plugs=PLUGS,
             )
 
         pt2_files = [f"{fileep}.backup.pt2", f"{fileep}.ep.pt2", f"{fileep}.pt2"]

onnx_diagnostic/export/control_flow.py

@@ -134,7 +134,11 @@ def make_custom_loop_for(
     assert body_outputs is not None, "body_outputs cannot be None"
     srank = "_".join("x".join(map(str, s.shape)) for s in body_outputs)
     sred = "x".join(map(str, reduction_dim)) if reduction_dim else ""
-    full_name = body_fn.__qualname__.replace("<locals>", "L").replace(".", "_")
+    full_name = (
+        body_fn.__qualname__.replace("<locals>", "L")
+        .replace("<lambda>", "l")
+        .replace(".", "_")
+    )
     name = f"loop_for_{full_name}_{srank}_{sred}"
     if name in _REGISTERED_SCHEMA:
         return name, _REGISTERED_SCHEMA[name][0]

onnx_diagnostic/export/onnx_plug.py

@@ -1,6 +1,6 @@
 import inspect
 from dataclasses import dataclass
-from typing import Any, Callable, Dict, List, Optional, Tuple
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 import onnx
 import torch
 from ..helpers import max_diff
@@ -49,6 +49,7 @@ class EagerDirectReplacementWithOnnx:
     :param n_outputs: same for the number of outputs,
         only tensors must be counted
     :param name: the name of the custom op, the function name if not specified
+    :param kwargs: constants
 
     Here is an example:
 
@@ -141,6 +142,7 @@ def __init__(
         n_inputs: Optional[int] = None,
         n_outputs: Optional[int] = None,
         name: Optional[str] = None,
+        kwargs: Optional[Dict[str, Union[int, float]]] = None,
     ):
         assert isinstance(
             function_proto, onnx.FunctionProto
@@ -152,7 +154,14 @@ def __init__(
         self.function_proto = function_proto
         self.n_inputs = n_inputs
         self.n_outputs = n_outputs
-        self.name = name or eager_fn.__name__
+        self.name = name or eager_fn.__qualname__.replace("<local>", "L").replace(
+            "<lambda>", "l"
+        ).replace(".", "_")
+        self.kwargs = kwargs
+        assert kwargs is None or all(isinstance(v, (int, float)) for v in kwargs.values()), (
+            f"Only int or floats are allowed for kwargs={kwargs}, one of them "
+            f"does not respect that constraint."
+        )
         sig = inspect.signature(self.eager_fn)
         params = list(sig.parameters)
         assert (
@@ -190,7 +199,7 @@ def torch_op(self) -> Callable:
     def __call__(self, *args):
         """Calls eager_fn or shape_fn if the model is being exported."""
         if is_exporting():
-            return self.shape_fn(*args)
+            return self.torch_op(*args)
         return self.eager_fn(*args)
 
     def _registers(self):
@@ -266,10 +275,16 @@ def converter(
             outputs: List[str],
             *args,
         ) -> Any:
-            if not g.has_local_function(self.name, self.domain):
+            if not g.has_local_function(
+                self.function_proto.name, domain=self.function_proto.domain
+            ):
                 g.add_function(self.function_proto)
             res = g.make_node(
-                self.name, args, outputs, domain=self.domain, name=self.target_name
+                self.function_proto.name,
+                args,
+                outputs,
+                domain=self.function_proto.domain,
+                name=self.target_name,
             )
             if not sts:
                 new_shapes = self.shape_fn(*args)
@@ -290,8 +305,8 @@ def onnx_dynamo_converter(self) -> Callable:
         """
         import onnxscript
 
-        onnx_plug_op = onnxscript.values.Opset(domain=self.domain, version=1)
-        schema = onnx_plug_op[self.name]
+        onnx_plug_op = onnxscript.values.Opset(domain=self.function_proto.domain, version=1)
+        schema = onnx_plug_op[self.function_proto.name]
         if schema is None:
             all_types = [
                 "tensor(float)",
@@ -307,8 +322,8 @@ def onnx_dynamo_converter(self) -> Callable:
             for i in range(self.n_outputs):
                 type_constraints.append((f"U{i}", all_types, ""))
             schema = onnx.defs.OpSchema(
-                self.name,
-                self.domain,
+                self.function_proto.name,
+                self.function_proto.domain,
                 1,
                 inputs=[
                     onnx.defs.OpSchema.FormalParameter(f"arg_{i}", f"T{i}")
@@ -321,7 +336,7 @@ def onnx_dynamo_converter(self) -> Callable:
                 type_constraints=type_constraints,
             )
             onnx.defs.register_schema(schema)
-        op = onnxscript.values.Op(onnx_plug_op, self.name, schema)
+        op = onnxscript.values.Op(onnx_plug_op, self.function_proto.name, schema)
 
         def converter(*cargs):
             return op(*cargs, n_outputs=self.n_outputs)