From ba9a585fcc954d781b4654f9d83fbba513245cbb Mon Sep 17 00:00:00 2001 From: guyang3532 <62738430+guyang3532@users.noreply.github.com> Date: Tue, 22 Nov 2022 17:32:19 +0800 Subject: [PATCH] Fix the tensor save for backward release problem (#13679) Motivation: PythonOp is saving input for backward, it's risky since ONNX Runtime backend is not aware of this, the tensor buffer may be "released" by ORT, then potentially modified by other operators before backward function executes. Fix: This pr just clone all input of PythonOp before forward is invoked. This may be high overhead, it's just a workaround before a better fix. --- .../_custom_autograd_function_runner.py | 14 +++-- .../orttraining_test_ortmodule_autograd.py | 61 +++++++++++++++++++ 2 files changed, 71 insertions(+), 4 deletions(-) diff --git a/orttraining/orttraining/python/training/ortmodule/_custom_autograd_function_runner.py b/orttraining/orttraining/python/training/ortmodule/_custom_autograd_function_runner.py index acbfdb8b0d..63daf53266 100644 --- a/orttraining/orttraining/python/training/ortmodule/_custom_autograd_function_runner.py +++ b/orttraining/orttraining/python/training/ortmodule/_custom_autograd_function_runner.py @@ -30,7 +30,13 @@ def wrap_as_dlpack_or_not(grad_flag, tensor_flag, inplace_flag, training_mode_fl if tensor_flag: # Got a tensor. Assume it's a DLPack tensor # and convert it to Pytorch tensor. - wrapped_arg = from_dlpack(arg) + if training_mode_flag: + wrapped_arg = from_dlpack(arg).detach().clone() + # TODO: This clone is just a workround to fix the bug that + # input saved for backward may be "released" by ORT. + # we need a follow up fix to avoid the copy overhead. + else: + wrapped_arg = from_dlpack(arg) # Only requires gradient when running under training mode # and the associated tensor has grad_flag=True (i.e., @@ -38,9 +44,9 @@ def wrap_as_dlpack_or_not(grad_flag, tensor_flag, inplace_flag, training_mode_fl wrapped_arg.requires_grad = training_mode_flag and grad_flag return wrapped_arg - else: - # Use non-tensor as is. It's a PyObject*. - return arg + + # Use non-tensor as is. It's a PyObject*. + return arg def call_python_forward_function( diff --git a/orttraining/orttraining/test/python/orttraining_test_ortmodule_autograd.py b/orttraining/orttraining/test/python/orttraining_test_ortmodule_autograd.py index 6bbf6146c5..c1d27806c2 100644 --- a/orttraining/orttraining/test/python/orttraining_test_ortmodule_autograd.py +++ b/orttraining/orttraining/test/python/orttraining_test_ortmodule_autograd.py @@ -1272,3 +1272,64 @@ def test_pythonop_training_mode(): ortmodule = ORTModule(TestModel(output_size)).eval() _ = ortmodule(torch.randn(output_size, dtype=torch.float)) check_pythonop_training_mode(ortmodule, is_eval_mode=True) + + +def test_python_op_save_input_for_backward(): + class GeLUFunctionTakeActivationInput(torch.autograd.Function): + @staticmethod + def forward(ctx, x): + ctx.save_for_backward(x) + return x * 0.5 * (1.0 + torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x))) + + @staticmethod + def backward(ctx, grad_output): + (x,) = ctx.saved_tensors + tanh_out = torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)) + ff = 0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * (1 + tanh_out) + g = ff * grad_output + return g + + class TestLayer(torch.nn.Module): + def __init__(self, output_size): + super().__init__() + self.relu = GeLUFunctionTakeActivationInput.apply + self._output_size = output_size + self.bias = Parameter(torch.empty(output_size, device=torch.cuda.current_device(), dtype=torch.float)) + self.w = Parameter( + torch.empty(output_size, output_size, device=torch.cuda.current_device(), dtype=torch.float) + ) + with torch.no_grad(): + self.bias.uniform_() + self.w.uniform_() + + def forward(self, model_input): + activation0 = torch.add(model_input, 0.4) + activation1 = activation0.view(self._output_size, -1) + activation2 = torch.add(self.relu(activation1), self.bias) + activation3 = torch.mul(activation2, 0.3) + activation3 = torch.matmul(self.w, activation3) + activation4 = torch.div(activation3, 1000) + return activation4 + + class TestModule(torch.nn.Module): + def __init__(self, output_size) -> None: + super().__init__() + self.layers = torch.nn.ModuleList([TestLayer(output_size) for i in range(10)]) + + def forward(self, x): + # ModuleList can act as an iterable, or be indexed using ints + for layer in self.layers: + x = x.view(-1) + x = torch.nn.functional.relu(layer(x)) + return x + + output_size = 1024 + + def model_builder(): + return TestModule(output_size) + + def input_generator(): + return torch.randn(output_size, output_size, dtype=torch.float).requires_grad_() + + label_input = torch.ones([output_size]) + run_training_test_and_compare(model_builder, input_generator, label_input)