From b9e23bd086af3607b650f7e0ca55c9d9ce723d87 Mon Sep 17 00:00:00 2001 From: Vincent Wang Date: Tue, 11 Oct 2022 15:20:51 +0800 Subject: [PATCH] [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. --- .../ortmodule/_custom_op_symbolic_registry.py | 38 +++++++++++++++---- .../python/orttraining_test_ortmodule_api.py | 16 ++++---- 2 files changed, 38 insertions(+), 16 deletions(-) diff --git a/orttraining/orttraining/python/training/ortmodule/_custom_op_symbolic_registry.py b/orttraining/orttraining/python/training/ortmodule/_custom_op_symbolic_registry.py index 634922379f..838da4ea1a 100644 --- a/orttraining/orttraining/python/training/ortmodule/_custom_op_symbolic_registry.py +++ b/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 diff --git a/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py b/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py index f6508f7621..ba29656e0b 100644 --- a/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py +++ b/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):