[ORTModule] Fix Custom Op Registry for Torch 1.13+ (#13250)

This PR has two fixes:
- https://github.com/pytorch/pytorch/pull/85636 change the behavior of
register_custom_op_symbolic to only register the symbolic function at a
single version. For ORTModule we need to pass the op_set version when
calling it.
- Since torch_1.13 the signature of einsum is changed to have a new
argument, need to change our custom op symbolic registry code
accordingly.

Without the fixes, ORTModule will not work with the nightly torch, and
the new torch version will be released.

orttraining/orttraining/python/training/ortmodule/_custom_op_symbolic_registry.py

@@ -9,6 +9,8 @@ from packaging.version import Version
 from torch.onnx import register_custom_op_symbolic
 from torch.onnx.symbolic_helper import _get_tensor_dim_size, _get_tensor_sizes, parse_args
 
+from onnxruntime.training import ortmodule
+
 
 class CustomOpSymbolicRegistry:
     _SYMBOLICS = {}
@@ -21,14 +23,23 @@ class CustomOpSymbolicRegistry:
     def register_all(cls):
         for name, fn in cls._SYMBOLICS.items():
             # Symbolic name is in format: domain::name
-            # Exporter will fail to register symbolic with non-empty domain when torch version is < 1.11.0.
-            if Version(torch.__version__) >= Version("1.11.0") or name.startswith("::"):
-                register_custom_op_symbolic(name, fn, 1)
+            register_custom_op_symbolic(
+                name,
+                fn,
+                ortmodule._defined_from_envvar("ORTMODULE_ONNX_OPSET_VERSION", ortmodule.ONNX_OPSET_VERSION),
+            )
 
 
-def register_symbolic(name, domain=""):
+def register_symbolic(name, domain="", torch_version_start=None, torch_version_end=None):
     def symbolic_wrapper(fn):
-        CustomOpSymbolicRegistry.register(name, domain, fn)
+        need_register = True
+        if torch_version_start is not None and Version(torch.__version__) < Version(torch_version_start):
+            need_register = False
+        # torch_version_end is exclusive.
+        if torch_version_end is not None and Version(torch.__version__) >= Version(torch_version_end):
+            need_register = False
+        if need_register:
+            CustomOpSymbolicRegistry.register(name, domain, fn)
         return fn
 
     return symbolic_wrapper
@@ -206,7 +217,8 @@ def squeeze(g, self, dim=None):
 # Exporter's prim::ConstantChunk uses multiple Slice nodes, which is fine for inference.
 # For training, the gradient graph will be multiple SliceGrad and one Sum, which is inefficient compared to
 # exporting to Split with SplitGrad as gradient graph.
-@register_symbolic("ConstantChunk", "prim")
+# Exporter will fail to register symbolic with non-empty domain when torch version is < 1.11.0.
+@register_symbolic("ConstantChunk", "prim", torch_version_start="1.11.0")
 def prim_ConstantChunk(g, self, chunks, dim):
     if chunks == 1:
         return self
@@ -422,9 +434,19 @@ def permute_and_reshape_tensor(
     return new_tensor, shape_tensor
 
 
-@register_symbolic("einsum")
+@register_symbolic("einsum", torch_version_end="1.13.0")
 @parse_args("s", "v")
-def einsum(g, equation, tensor_list):
+def einsum_pre_troch_113(g, equation, tensor_list):
+    return einsum_internal(g, equation, tensor_list)
+
+
+@register_symbolic("einsum", torch_version_start="1.13.0")
+@parse_args("s", "v", "is")
+def einsum_torch_113(g, equation, tensor_list, path=None):
+    return einsum_internal(g, equation, tensor_list)
+
+
+def einsum_internal(g, equation, tensor_list):
     tensors = sym_help._unpack_list(tensor_list)
     num_ops = len(tensors)
     assert num_ops > 0

orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py

@@ -1383,12 +1383,12 @@ def test_gradient_correctness_chunk(dim, chunks):
         os.remove(os.path.join(os.getcwd(), "chunk_model_execution_model_training.onnx"))
 
 
-# In PyTorch 1.11.0, there is issue during reduce node shape handling for exporter, so any sub-graph that
+# In PyTorch 1.11 to 1.12, there is issue during reduce node shape handling for exporter, so any sub-graph that
 # contains ReduceProd will fail to run, for example, "sec,sm->ecm", "sec,ecm->sm".
-# Currently skip these cases and test_gradient_correctness_einsum_2,
-# will enable these tests again once the issue in PyTorch is fixed.
-skip_torch_1_11 = pytest.mark.skipif(
-    Version(torch.__version__) >= Version("1.11.0"), reason="PyTorch 1.11 incompatible"
+# Skip these cases and test_gradient_correctness_einsum_2 for these versions.
+skip_einsum_test_if = pytest.mark.skipif(
+    Version(torch.__version__) >= Version("1.11.0") and Version(torch.__version__) < Version("1.13.0"),
+    reason="PyTorch 1.11 and 1.12 incompatible",
 )
 
 
@@ -1401,8 +1401,8 @@ skip_torch_1_11 = pytest.mark.skipif(
         "ks,ksm->sm",
         "kes,ems->mek",
         "kes,ksm->ms",
-        pytest.param("sec,sm->ecm", marks=[skip_torch_1_11]),
-        pytest.param("sec,ecm->sm", marks=[skip_torch_1_11]),
+        pytest.param("sec,sm->ecm", marks=[skip_einsum_test_if]),
+        pytest.param("sec,ecm->sm", marks=[skip_einsum_test_if]),
     ],
 )
 def test_gradient_correctness_einsum(equation):
@@ -1453,7 +1453,7 @@ def test_gradient_correctness_einsum(equation):
         _test_helpers.assert_gradients_match_and_reset_gradient(ort_model, pt_model, atol=1e-3, rtol=1e-3)
 
 
-@skip_torch_1_11
+@skip_einsum_test_if
 def test_gradient_correctness_einsum_2():
     class NeuralNetEinsum(torch.nn.Module):
         def __init__(self, bias_size):