From 2347a0aca85f81ad8b30b2eb9f41f9b01bc479c2 Mon Sep 17 00:00:00 2001 From: pengwa Date: Wed, 7 Jul 2021 08:58:01 +0800 Subject: [PATCH] Autograd Function Fallback bug fix - moe support (#8105) * Support forward inputs orders like "Non_tensor/Tensor/Non_tensor". Correspondingly, support "None/Tensor_Grad/None" fpr backward outputs. * Report RuntimeError when PythonOp detected but _enable_custom_autograd_function is enabled. * Fix "PoliCheck ] - Defect : Term "hang", Component : orttraining\orttraining\python\training\ortmodule\__init__.py (1 issue)" * rename call_convention->input_convention, input_tensor_requires_grads->input_requires_grads * fix minor comment * revert polycheck fix in case of conflict * Update orttraining/orttraining/core/graph/training_op_defs.cc Co-authored-by: Tim Harris * Apply suggestions from code review Refine the schema description Co-authored-by: Tim Harris * Resolve review comments Co-authored-by: Tim Harris --- .../language_interop_ops/torch/torch_proxy.cc | 18 ++- .../language_interop_ops/torch/torch_proxy.h | 1 - .../core/graph/gradient_builder.cc | 64 ++++++---- .../core/graph/training_op_defs.cc | 71 ++++++----- .../_custom_autograd_function_exporter.py | 29 ++++- .../_custom_autograd_function_runner.py | 7 +- .../ortmodule/_graph_execution_manager.py | 4 +- .../orttraining_test_ortmodule_autograd.py | 115 ++++++++++++++++++ .../torch_custom_function_kernel_base.cc | 44 ++++--- .../torch/torch_custom_function_kernel_base.h | 9 +- 10 files changed, 273 insertions(+), 89 deletions(-) diff --git a/onnxruntime/core/language_interop_ops/torch/torch_proxy.cc b/onnxruntime/core/language_interop_ops/torch/torch_proxy.cc index b252a08706..f5c59e71b1 100644 --- a/onnxruntime/core/language_interop_ops/torch/torch_proxy.cc +++ b/onnxruntime/core/language_interop_ops/torch/torch_proxy.cc @@ -171,10 +171,15 @@ void InvokeRunner( // first element. for (; i < static_cast(PyTuple_Size(result_ptr.get())); ++i) { PyObject* dl_tensor_pointer = PyTuple_GetItem(result_ptr.get(), i); - ORT_ENFORCE(Py_REFCNT(dl_tensor_pointer) == 1, "Ref count of dl_tensor_pointer should be 1."); - // Todo: be noted we did not pass whether tensor is bool or not. - // Currently we assume we don't pass boolean data. - returned_ortvalues.push_back(training::framework::torch::FromDlpack(dl_tensor_pointer, false)); + if (dl_tensor_pointer == Py_None) { + OrtValue empty_ort_value; + returned_ortvalues.push_back(empty_ort_value); + } else { + ORT_ENFORCE(Py_REFCNT(dl_tensor_pointer) == 1, "Ref count of dl_tensor_pointer should be 1."); + // Todo (pengwa): be noted we did not pass whether tensor is bool or not. + // Currently we assume we don't pass boolean data. + returned_ortvalues.push_back(training::framework::torch::FromDlpack(dl_tensor_pointer, false)); + } } } @@ -294,7 +299,6 @@ void TorchProxy::Forward( void TorchProxy::Backward( void* callback, - const std::vector& requires_grads, const std::vector& tensor_args, const std::vector& tensor_indices, const std::vector& obj_args, @@ -308,6 +312,10 @@ void TorchProxy::Backward( // Python-related calls should happen only if guard is alive. GilGuard guard; auto runner = OrtTorchFunctionPool::GetInstance().GetBackwardRunner(); + + // Pass all zero since backward inputs don't require gradients. + const auto all_input_count = tensor_args.size() + obj_args.size(); + const std::vector requires_grads(all_input_count, 0); Invoke( runner, reinterpret_cast(callback), diff --git a/onnxruntime/core/language_interop_ops/torch/torch_proxy.h b/onnxruntime/core/language_interop_ops/torch/torch_proxy.h index 89f7d642e1..29ef4ce9fc 100644 --- a/onnxruntime/core/language_interop_ops/torch/torch_proxy.h +++ b/onnxruntime/core/language_interop_ops/torch/torch_proxy.h @@ -49,7 +49,6 @@ class TorchProxy { void Backward( void* callback, - const std::vector& requires_grads, const std::vector& tensor_args, const std::vector& tensor_indices, const std::vector& obj_args, diff --git a/orttraining/orttraining/core/graph/gradient_builder.cc b/orttraining/orttraining/core/graph/gradient_builder.cc index 1624e66e3c..6d45469351 100755 --- a/orttraining/orttraining/core/graph/gradient_builder.cc +++ b/orttraining/orttraining/core/graph/gradient_builder.cc @@ -824,24 +824,24 @@ IMPLEMENT_GRADIENT_BUILDER(GetAddSubGradient) { std::vector output; if (a.name.compare(b.name) == 0) { if (IsGradientRequiredForSrcNodeInput(0)) { - output.push_back( - NodeDef("Identity", - {GO(0)}, - {GI(0)})); + output.push_back( + NodeDef("Identity", + {GO(0)}, + {GI(0)})); } if (IsGradientRequiredForSrcNodeInput(1)) { - if (is_sub) { - output.push_back( - NodeDef("Neg", - {GO(0)}, - {GI(1)})); - } else /*is_add*/ { - output.push_back( - NodeDef("Identity", - {GO(0)}, - {GI(1)})); - } + if (is_sub) { + output.push_back( + NodeDef("Neg", + {GO(0)}, + {GI(1)})); + } else /*is_add*/ { + output.push_back( + NodeDef("Identity", + {GO(0)}, + {GI(1)})); + } } return output; } @@ -918,14 +918,14 @@ IMPLEMENT_GRADIENT_BUILDER(GetMulGradient) { std::vector output; if (a.name.compare(b.name) == 0) { - if(IsGradientRequiredForSrcNodeInput(0)) { + if (IsGradientRequiredForSrcNodeInput(0)) { output.push_back( NodeDef("Mul", {GO(0), I(1)}, {GI(0)})); } - if(IsGradientRequiredForSrcNodeInput(1)) { + if (IsGradientRequiredForSrcNodeInput(1)) { output.push_back( NodeDef("Mul", {GO(0), I(0)}, @@ -1727,6 +1727,7 @@ IMPLEMENT_GRADIENT_BUILDER(GetPythonOpGradient) { std::vector attrs; ORT_ENFORCE(utils::HasString(src_attrs.at("name"))); attrs.push_back(MakeAttribute("name", src_attrs.at("name").s())); + attrs.push_back(MakeAttribute("output_convention", src_attrs.at("input_convention").s())); attrs.push_back(MakeAttribute("inplace", src_attrs.at("inplace").i())); // input_tensor_types[i] store the type of autograd.Function.apply's ith output. @@ -1775,33 +1776,46 @@ IMPLEMENT_GRADIENT_BUILDER(GetPythonOpGradient) { } attrs.push_back(MakeAttribute("output_tensor_ranks", output_tensor_ranks)); - std::vector output_tensor_requires_grads; - for (const auto requires_grad : src_attrs["input_tensor_requires_grads"].ints()) { - output_tensor_requires_grads.push_back(requires_grad); - } - attrs.push_back(MakeAttribute("output_tensor_requires_grads", output_tensor_requires_grads)); - std::vector input_args; // Put Python context generated by PythonOp. input_args.push_back(O(0)); // Put other outputs. for (int i = 1; i < GetSrcNodeOutputSize(); ++i) { - if (input_tensor_requires_grads.at(i)) { + if (src_attrs["output_tensor_requires_grads"].ints().Get(i - 1)) { // Only add FW outputs which // 1. are tensors, // 2. needs gradients (requires_grad=True in Pytorch). input_args.push_back(GO(i)); + } else { + input_args.push_back(ArgDef()); } } // Also connect forward outputs to PythonOpGrad for random segement fault issues. + // Todo (pengwa): we should investigate whether we could avoid those outputs that are not used + // in backward computation. for (int i = 1; i < GetSrcNodeOutputSize(); ++i) { input_args.push_back(O(i)); } + // src_attrs["input_requires_grads"] stores all inputs's requires_grad attributes, + // including both tensor inputs and non-tensor inputs (e.g. constants), here we filter out + // those non-tensor inputs when constructing PythonOpGrad's outputs. + const std::string& input_convention = src_attrs.at("input_convention").s(); + const auto& fw_input_requires_grads = src_attrs["input_requires_grads"].ints(); + std::vector bw_tensor_output_requires_grads; + for (auto i = 0; i < fw_input_requires_grads.size(); ++i) { + if (input_convention[i] == 'd') { // only handle gradients for tensor type inputs. + bw_tensor_output_requires_grads.push_back(fw_input_requires_grads.Get(i)); + } + } + ORT_ENFORCE(static_cast(GetSrcNodeInputSize()) == bw_tensor_output_requires_grads.size(), + "PythonOpGrad requiring gradient output count mismatch."); + attrs.push_back(MakeAttribute("output_tensor_requires_grads", bw_tensor_output_requires_grads)); + std::vector output_args; for (int i = 0; i < GetSrcNodeInputSize(); ++i) { - if (output_tensor_requires_grads[i]) { + if (bw_tensor_output_requires_grads[i]) { output_args.push_back(GI(i)); } else { output_args.push_back(ArgDef()); diff --git a/orttraining/orttraining/core/graph/training_op_defs.cc b/orttraining/orttraining/core/graph/training_op_defs.cc index 347b546b0b..7145c496f8 100644 --- a/orttraining/orttraining/core/graph/training_op_defs.cc +++ b/orttraining/orttraining/core/graph/training_op_defs.cc @@ -155,7 +155,7 @@ TensorProto ToDimensionOneTensor(int32_t value) { bool BuildContextDependentFunctionBodyNllLossInternal( const FunctionBodyBuildContext& ctx, const OpSchema& schema, - FunctionProto& functionProto) { + FunctionProto& functionProto) { if (ctx.getInputType(0) == nullptr) { // we cannot create a correct function body without knowing the input type return false; @@ -286,10 +286,10 @@ bool BuildContextDependentFunctionBodyNllLossInternal( {MakeAttribute("value", ToDimensionOneFloatTensor(0.0f))}}); if (!float_input) { body.push_back( - {{"const_zero_casted"}, - "Cast", - {"const_zero_float"}, - {MakeAttribute("to", static_cast(input_type))}}); + {{"const_zero_casted"}, + "Cast", + {"const_zero_float"}, + {MakeAttribute("to", static_cast(input_type))}}); } body.push_back( {{"input_gather_element_transform"}, @@ -315,16 +315,16 @@ bool BuildContextDependentFunctionBodyNllLossInternal( {MakeAttribute("value", ToDimensionOneFloatTensor(1.0f))}}); if (!float_input) { body.push_back( - {{"const_one_casted"}, - "Cast", - {"const_one_float"}, - {MakeAttribute("to", static_cast(input_type))}}); + {{"const_one_casted"}, + "Cast", + {"const_one_float"}, + {MakeAttribute("to", static_cast(input_type))}}); } body.push_back( {{"weight_gather"}, "Where", - {"squeeze_mask", float_input ? "const_zero_float" : "const_zero_casted", - float_input ? "const_one_float" :"const_one_casted"}}); + {"squeeze_mask", float_input ? "const_zero_float" : "const_zero_casted", + float_input ? "const_one_float" : "const_one_casted"}}); } else { body.push_back( @@ -2887,9 +2887,14 @@ Return true if all elements are true and false otherwise. "Name of custom class.", AttributeProto::STRING) .Attr( - "call_convention", - "call_convention[i]==c means a non-tensor argument. call_convention[i]==d means a tensor.", + "input_convention", + "input_convention[i]==c means a non-tensor argument. input_convention[i]==d means a tensor.", AttributeProto::STRING) + .Attr( + "input_requires_grads", + "Flags to indicate whether the torch.autograd.apply's inputs require gradients (including flags for both tensor" + " and non-tensor inputs)", + AttributeProto::INTS) // Input Pytorch tensors. .Attr( "input_tensor_types", @@ -2899,10 +2904,6 @@ Return true if all elements are true and false otherwise. "input_tensor_ranks", "Input tensors' ranks of autograd.Function.apply.", AttributeProto::INTS) - .Attr( - "input_tensor_requires_grads", - "Flags to indicate which inputs has gradient", - AttributeProto::INTS) // Input int scalars. .Attr( "input_int_scalars", @@ -3006,7 +3007,7 @@ Return true if all elements are true and false otherwise. ORT_ENFORCE(input_tensor_types_proto, "PythonOp's must have \"input_tensor_types\" attribute."); // Check if the inferred input types match those described in the // "input_tensor_types" attributes. - int64_t input_tensor_types_count = input_tensor_types_proto->ints_size(); + const int64_t input_tensor_types_count = input_tensor_types_proto->ints_size(); ORT_ENFORCE(static_cast(input_tensor_types_count) == ctx.getNumInputs(), "PythonOp's input list and \"input_tensor_types\" attribute should have the same length."); for (auto i = 0; i < input_tensor_types_count; ++i) { @@ -3031,7 +3032,7 @@ Return true if all elements are true and false otherwise. // This is a required field. ORT_ENFORCE(output_tensor_types_proto, "PythonOp's must have \"output_tensor_types\" attribute."); - static size_t rank_count = 0; + size_t rank_count = 0; // Set inferred output types. for (auto i = 1; i < static_cast(ctx.getNumOutputs()); ++i) { updateOutputElemType(ctx, i, static_cast(output_tensor_types_proto->ints().at(i - 1))); @@ -3063,8 +3064,10 @@ Return true if all elements are true and false otherwise. .Input( 1, "inputs", - "Inputs of autograd.Function.backward. There are 2*N inputs: \ - N gradient inputs + N forward run activations of PythonOp (not including the context).", + "There are 2*N inputs: " + " N gradient inputs (as inputs of autograd.Function.backward) + " + " N forward run activations of autograd.Function.apply." + "The N forward run inputs are used as control dependency between PythonOpGrad and PythonOp", "T", OpSchema::Variadic, /*is_homogeneous*/ false, @@ -3083,37 +3086,45 @@ Return true if all elements are true and false otherwise. AttributeProto::STRING) .Attr( "inplace", - "Indicate if the output should reuse input memory.", + "Indicate if the output should reuse input memory. Todo(pengwa): do we really need it?", AttributeProto::INT, static_cast(0)) .Attr( "input_tensor_types", - "Input types of autograd.Function.backward.", + "Input types of autograd.Function.backward (including only tensor inputs)." + "This attribute is mostly used for input checks for better robustnes.", AttributeProto::INTS, false) .Attr( "input_tensor_ranks", - "Input ranks of autograd.Function.backward.", + "Input ranks of autograd.Function.backward (including only tensor inputs)." + "This attribute is mostly used for input checks for better robustness.", AttributeProto::INTS, false) .Attr( "input_tensor_requires_grads", - "Flags to indicate which inputs has gradient", + "Flags to indicate which inputs have gradients (including only tensor inputs)." + "This attribute is mostly used for input checks for better robustness.", AttributeProto::INTS) .Attr( "output_tensor_types", - "Output types of autograd.Function.backward.", + "Output types of autograd.Function.backward outputs (including only tensor outputs).", AttributeProto::INTS, false) .Attr( "output_tensor_ranks", - "Output ranks of autograd.Function.backward.", + "Output ranks of autograd.Function.backward outputs (including only tensor outputs).", AttributeProto::INTS, false) .Attr( "output_tensor_requires_grads", - "Flags to indicate which outputs has gradient", + "Flags to indicate which outputs have gradients (including only tensor outputs).", AttributeProto::INTS) + .Attr( + "output_convention", + "A string inidicating autograd.Function.backward outputs's type." + "value 'c' - non-tensor output; value 'd' - tensor output.", + AttributeProto::STRING) .TypeConstraint( "T", OpSchema::all_tensor_types(), @@ -3154,9 +3165,9 @@ Return true if all elements are true and false otherwise. ORT_ENFORCE(static_cast(output_tensor_types_proto->ints_size()) == ctx.getNumOutputs(), "PythonOpGrad's output list and \"output_tensor_types\" attribute should have the same length."); // This is a required field. - ORT_ENFORCE(output_tensor_types_proto, "PythonOp's must have \"output_tensor_types\" attribute."); + ORT_ENFORCE(output_tensor_types_proto, "PythonOpGrad's must have \"output_tensor_types\" attribute."); // Set inferred output types. - static size_t rank_count = 0; + size_t rank_count = 0; for (auto i = 0; i < static_cast(ctx.getNumOutputs()); ++i) { updateOutputElemType(ctx, i, static_cast(output_tensor_types_proto->ints().at(i))); const auto output_tensor_ranks = ctx.getAttribute("output_tensor_ranks"); diff --git a/orttraining/orttraining/python/training/ortmodule/_custom_autograd_function_exporter.py b/orttraining/orttraining/python/training/ortmodule/_custom_autograd_function_exporter.py index ee09ea59f4..5016feba8e 100644 --- a/orttraining/orttraining/python/training/ortmodule/_custom_autograd_function_exporter.py +++ b/orttraining/orttraining/python/training/ortmodule/_custom_autograd_function_exporter.py @@ -20,7 +20,7 @@ def _export(g, n, *args, **kwargs): training_mode = symbolic_helper._training_mode cconv = n.cconv() input_tensor_types = [] - input_tensor_requires_grads = [] + input_requires_grads = [] input_tensor_ranks = [] input_int_scalars = [] @@ -48,7 +48,7 @@ def _export(g, n, *args, **kwargs): tensor_args.append(arg) requires_grad = 1 if arg.requires_grad() else 0 - input_tensor_requires_grads.append(requires_grad) + input_requires_grads.append(requires_grad) scalar_type = int(symbolic_helper.cast_pytorch_to_onnx[arg.type( ).scalarType()]) @@ -57,7 +57,7 @@ def _export(g, n, *args, **kwargs): elif call_type == 'c': # Got a non-tensor variable. # Non-tensor can't have gradient. - input_tensor_requires_grads.append(0) + input_requires_grads.append(0) if isinstance(arg, float): # A float. input_float_scalar_positions.append(i) @@ -108,11 +108,11 @@ def _export(g, n, *args, **kwargs): attrs = { 'name_s': name, 'inplace_i': inplace, - 'call_convention_s': cconv, + 'input_convention_s': cconv, 'outputs': n.outputsSize(), 'input_tensor_types_i': input_tensor_types, 'input_tensor_ranks_i': input_tensor_ranks, - 'input_tensor_requires_grads_i': input_tensor_requires_grads, + 'input_requires_grads_i': input_requires_grads, 'output_tensor_types_i': output_tensor_types, 'output_tensor_ranks_i': output_tensor_ranks, 'output_tensor_requires_grads_i': output_tensor_requires_grads, @@ -145,7 +145,24 @@ def _export(g, n, *args, **kwargs): sys.stderr.flush() raise -def _post_process_after_export(exported_model): +def _post_process_after_export(exported_model, enable_custom_autograd_function): + if enable_custom_autograd_function: + return _post_process_enabling_autograd_fallback(exported_model) + + is_fallback_needed = False + for node in exported_model.graph.node: + if node.domain == 'com.microsoft' and node.op_type in ["PythonOp"]: + is_fallback_needed = True + break + + if is_fallback_needed: + raise RuntimeError('Detected autograd functions usage in current model, the run will fail \ + without enabling \'_enable_custom_autograd_function\'. Please enable it with: \ + \'module._execution_manager(is_training_mode)._enable_custom_autograd_function = True\'') + + return exported_model + +def _post_process_enabling_autograd_fallback(exported_model): index = 0 for node in exported_model.graph.node: if node.domain == 'com.microsoft' and node.op_type in ["PythonOp"]: 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 99f673e7bf..200e43f258 100644 --- a/orttraining/orttraining/python/training/ortmodule/_custom_autograd_function_runner.py +++ b/orttraining/orttraining/python/training/ortmodule/_custom_autograd_function_runner.py @@ -101,7 +101,7 @@ def call_python_forward_function( elif isinstance(result, tuple) or isinstance(result, list): ctx = extract_context(result) wrapped = [ctx] - wrapped.extend(list(to_dlpack(value) for value in result)) + wrapped.extend(list(to_dlpack(value) if value is not None else None for value in result)) # Inside the returned list, first element is context and the rest # are DLPack tensors. return wrapped @@ -157,11 +157,14 @@ def call_python_backward_function( if isinstance(result, torch.Tensor): return [to_dlpack(result)] elif isinstance(result, tuple) or isinstance(result, list): - return [to_dlpack(value) for value in result if value is not None] + return [to_dlpack(value) if value is not None else None for value in result] else: raise Exception('Unsupported returned type: ', type(result)) try: + # Backward inputs should not require gradients. + assert all(grad_flag == 0 for grad_flag in requires_grad_flags) + # Prepare inputs for calling Python function. wrapped_args = list(wrap_as_dlpack_or_not(grad_flag, tensor_flag, inplace, is_training_mode, arg) for grad_flag, tensor_flag, arg in zip(requires_grad_flags, tensor_type_flags, args)) diff --git a/orttraining/orttraining/python/training/ortmodule/_graph_execution_manager.py b/orttraining/orttraining/python/training/ortmodule/_graph_execution_manager.py index cec03077b2..27896815f5 100644 --- a/orttraining/orttraining/python/training/ortmodule/_graph_execution_manager.py +++ b/orttraining/orttraining/python/training/ortmodule/_graph_execution_manager.py @@ -283,8 +283,8 @@ class GraphExecutionManager(GraphExecutionInterface): except RuntimeError as e: raise RuntimeError('There was an error while exporting the PyTorch model to ONNX: {}'.format(e)) exported_model = onnx.load_model_from_string(f.getvalue()) - if self._enable_custom_autograd_function: - exported_model = _post_process_after_export(exported_model) + + exported_model = _post_process_after_export(exported_model, self._enable_custom_autograd_function) return exported_model diff --git a/orttraining/orttraining/test/python/orttraining_test_ortmodule_autograd.py b/orttraining/orttraining/test/python/orttraining_test_ortmodule_autograd.py index 7720ac1e67..607ebb6246 100644 --- a/orttraining/orttraining/test/python/orttraining_test_ortmodule_autograd.py +++ b/orttraining/orttraining/test/python/orttraining_test_ortmodule_autograd.py @@ -161,6 +161,53 @@ def test_ScalarAndTuple(): run_training_test_and_compare(model_builder, input_generator, label_input) +def test_ScalarAndTupleReordered(): + class ScalarAndTupleReorderedFunction(torch.autograd.Function): + @staticmethod + def forward(ctx, alpha, beta, input, gamma): + ctx.save_for_backward(input) + ctx.alpha = alpha + ctx.beta = beta + ctx.gamma = gamma + return alpha * beta[0] * beta[1] * gamma * input.clamp(min=0) + + @staticmethod + def backward(ctx, grad_output): + input, = ctx.saved_tensors + alpha = ctx.alpha + beta = ctx.beta + gamma = ctx.gamma + grad_input = grad_output.clone() + grad_input[input < 0] = 0 + return None, None, alpha * beta[0] * beta[1] * gamma * grad_input, None + + class ScalarAndTupleReorderedModel(torch.nn.Module): + def __init__(self, output_size): + super(ScalarAndTupleReorderedModel, self).__init__() + self.activation = ScalarAndTupleReorderedFunction.apply + self.linear_a = torch.nn.Linear(output_size, output_size) + self.linear_b = torch.nn.Linear(output_size, output_size) + + def forward(self, x): + h = self.linear_a(x) + h = self.activation(5.0, (-1.0, 2.0), h, -1.0) + h = self.linear_b(h) + return h + + output_size = 2 + + def model_builder(): + return ScalarAndTupleReorderedModel(output_size) + + def input_generator(): + return torch.randn(output_size, dtype=torch.float) + + # generate a label that have same shape as forward output. + label_input = torch.ones([output_size]) + + run_training_test_and_compare(model_builder, input_generator, label_input) + + @pytest.mark.skip(reason="This test is not correct. All tensors modified by in-place operattions should be mark_dirty(...).") def test_InplaceUpdateInputAsOutputNotRequireGrad(): class InplaceUpdateInputAsOutputNotRequireGradFunction(torch.autograd.Function): @@ -703,3 +750,71 @@ def test_Share_Input(): run_training_test_and_compare(model_builder, input_generator, label_input) run_training_test_and_compare(model_builder, input_generator_with_requires_grad, label_input) + + +def test_GeLU_When_Autograd_Func_Fallback_Not_Enabled(): + @torch.jit.script + def bias_gelu(bias, y): + x = bias + y + return x * 0.5 * (1.0 + torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x))) + + @torch.jit.script + def bias_gelu_backward(g, bias, y): + x = bias + y + 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) + return ff*g + + class GeLUFunction(torch.autograd.Function): + @staticmethod + def forward(ctx, input, bias): + ctx.save_for_backward(input, bias) + return bias_gelu(bias, input) + + @staticmethod + def backward(ctx, grad_output): + input, bias = ctx.saved_tensors + tmp = bias_gelu_backward(grad_output, bias, input) + return tmp, tmp + + class GeLUModel(torch.nn.Module): + def __init__(self, output_size): + super(GeLUModel, self).__init__() + self.relu = GeLUFunction.apply + self.bias = Parameter(torch.empty( + output_size, + device=torch.cuda.current_device(), + dtype=torch.float)) + + with torch.no_grad(): + self.bias.uniform_() + + def forward(self, model_input): + out = self.relu(model_input, self.bias) + return out + + output_size = 1024 + + def model_builder(): + return GeLUModel(output_size) + + def input_generator(): + return torch.randn(output_size, dtype=torch.float) + + # generate a label that have same shape as forward output. + label_input = torch.ones([output_size]) + + m_ort = model_builder() + x_ort = input_generator() + + try: + device = torch.device("cpu") + m_ort.to(device) + model = ORTModule(m_ort) + model.train() + + inputs_on_device = [x_ort.to(device)] + output = model(*inputs_on_device) + except RuntimeError as e: + assert "Detected autograd functions usage in current model, the run will fail" in str(e) \ No newline at end of file diff --git a/orttraining/orttraining/training_ops/cpu/torch/torch_custom_function_kernel_base.cc b/orttraining/orttraining/training_ops/cpu/torch/torch_custom_function_kernel_base.cc index 4f5405589d..cf38d1fcb5 100644 --- a/orttraining/orttraining/training_ops/cpu/torch/torch_custom_function_kernel_base.cc +++ b/orttraining/orttraining/training_ops/cpu/torch/torch_custom_function_kernel_base.cc @@ -28,13 +28,16 @@ std::vector CreateOrtValueArgs(OpKernelContext* context, void PythonOpBase::Init(const OpKernelInfo& info) { ORT_THROW_IF_ERROR(info.GetAttr("name", &name_)); - inplace_ = info.GetAttrOrDefault("inplace", static_cast(0)); + ORT_THROW_IF_ERROR(info.GetAttr("inplace", &inplace_)); + is_training_mode_ = static_cast(info.GetAttrOrDefault("training_mode", static_cast(0))); - ORT_THROW_IF_ERROR(info.GetAttr("call_convention", &call_convention_)); + ORT_THROW_IF_ERROR(info.GetAttr("input_convention", &input_convention_)); + + ORT_THROW_IF_ERROR(info.GetAttrs("input_requires_grads", input_requires_grads_)); + ORT_ENFORCE(input_requires_grads_.size() == input_convention_.size()); // Input tensors. - input_tensor_types_ = info.GetAttrsOrDefault("input_tensor_types", std::vector()); - input_tensor_requires_grads_ = info.GetAttrsOrDefault("input_tensor_requires_grads", std::vector()); + ORT_THROW_IF_ERROR(info.GetAttrs("input_tensor_types", input_tensor_types_)); ORT_ENFORCE(input_tensor_types_.size() == info.node().InputDefs().size()); @@ -68,10 +71,15 @@ void PythonOpBase::Init(const OpKernelInfo& info) { input_pointer_scalar_positions_ = info.GetAttrsOrDefault("input_pointer_scalar_positions", std::vector()); ORT_ENFORCE(input_pointer_scalars_.size() == input_pointer_scalar_positions_.size()); + auto non_tensor_input_count = input_int_scalars_.size() + input_float_scalars_.size() + + input_int_tuple_positions_.size() + input_float_tuple_positions_.size() + + input_pointer_scalars_.size(); + ORT_ENFORCE(non_tensor_input_count + input_tensor_types_.size() == input_convention_.size(), + "Total input (tensor + non-tensor) count did not match."); // Output tensors. - output_tensor_types_ = info.GetAttrsOrDefault("output_tensor_types", std::vector()); - output_tensor_requires_grads_ = info.GetAttrsOrDefault("output_tensor_requires_grads", std::vector()); + ORT_THROW_IF_ERROR(info.GetAttrs("output_tensor_types", output_tensor_types_)); + ORT_THROW_IF_ERROR(info.GetAttrs("output_tensor_requires_grads", output_tensor_requires_grads_)); CreateConstArgs(); CreateArgPositions(); @@ -94,7 +102,7 @@ void PythonOpBase::RunForward(OpKernelContext* context, std::string err; TorchProxy::GetInstance().Forward( OrtTorchFunctionPool::GetInstance().GetForwardCore(name_), - input_tensor_requires_grads_, + input_requires_grads_, args, arg_positions_, const_args_, @@ -225,14 +233,18 @@ void PythonOpGradBase::Init(const OpKernelInfo& info) { ORT_THROW_IF_ERROR(info.GetAttr("name", &name_)); ORT_THROW_IF_ERROR(info.GetAttr("inplace", &inplace_)); ORT_THROW_IF_ERROR(info.GetAttrs("input_tensor_types", input_tensor_types_)); + ORT_THROW_IF_ERROR(info.GetAttr("output_convention", &output_convention_)); ORT_THROW_IF_ERROR(info.GetAttrs("output_tensor_types", output_tensor_types_)); - input_tensor_requires_grads_ = info.GetAttrsOrDefault("input_tensor_requires_grads", std::vector()); output_tensor_requires_grads_ = info.GetAttrsOrDefault("output_tensor_requires_grads", std::vector()); + ORT_ENFORCE(output_tensor_types_.size() == output_tensor_requires_grads_.size(), "backward tensor output count mismatch"); + SetPositions(); } void PythonOpGradBase::RunBackward(OpKernelContext* context, std::vector& returned_ortvalues) const { + // Todo (pengwa): this is fragile once we added more inputs, re-visist this + // for more robustness. auto args = CreateOrtValueArgs(context, 1, (context->InputCount() - 1) / 2); // This is called "const" because that's how Pytorch calls all non-tensor inputs. const Tensor* context_id_tensor = context->Input(0); @@ -245,7 +257,6 @@ void PythonOpGradBase::RunBackward(OpKernelContext* context, TorchProxy::GetInstance().Backward( OrtTorchFunctionPool::GetInstance() .GetBackwardCore(name_), - input_tensor_requires_grads_, args, arg_positions_, const_args, @@ -258,13 +269,18 @@ void PythonOpGradBase::RunBackward(OpKernelContext* context, void PythonOpGradBase::SetOutputs(OpKernelContext* context, std::vector& returned_ortvalues) const { auto* ctx_internal = reinterpret_cast(context); - const auto outputs_count = static_cast(ctx_internal->OutputCount()); - for (size_t i = 0; i < outputs_count; ++i) { - if (!output_tensor_requires_grads_[i]) { - continue; + ORT_ENFORCE(output_convention_.size() == returned_ortvalues.size(), "backward output count mismatch."); + int tensor_output_index = 0; + for (size_t i = 0; i < returned_ortvalues.size(); ++i) { + if (output_convention_[i] == 'd') { + if (output_tensor_requires_grads_[tensor_output_index]) { + ORT_THROW_IF_ERROR(ctx_internal->SetOutputMLValue(tensor_output_index, returned_ortvalues.at(i))); + } + ++tensor_output_index; } - ORT_THROW_IF_ERROR(ctx_internal->SetOutputMLValue(static_cast(i), returned_ortvalues.at(i))); } + + ORT_ENFORCE(tensor_output_index == ctx_internal->OutputCount(), "backward tensor output count mismatch."); } void PythonOpGradBase::SetPositions() { diff --git a/orttraining/orttraining/training_ops/cpu/torch/torch_custom_function_kernel_base.h b/orttraining/orttraining/training_ops/cpu/torch/torch_custom_function_kernel_base.h index 6bea6ae05d..b9e2d18deb 100644 --- a/orttraining/orttraining/training_ops/cpu/torch/torch_custom_function_kernel_base.h +++ b/orttraining/orttraining/training_ops/cpu/torch/torch_custom_function_kernel_base.h @@ -52,14 +52,14 @@ class PythonOpBase { // Name of containing class. For example, MyReLU. std::string name_; int64_t inplace_; - std::string call_convention_; + std::string input_convention_; bool is_training_mode_; + // input_requires_grads_[i] indicates if the i-th inputs of apply() should have gradient. + std::vector input_requires_grads_; // Attributes of input tensors for calling MyReLU.apply(...). // Types. input_tensor_types_[i] is the element type of the i-th tensor. std::vector input_tensor_types_; - // input_tensor_types_[i] indicates if the i-th tensor should have gradient. - std::vector input_tensor_requires_grads_; // Concatenation of all floats from apply(...) 's inputs. std::vector input_int_scalars_; @@ -117,9 +117,10 @@ class PythonOpGradBase { int64_t inplace_; // Input types of MyReLU.backward(...). std::vector input_tensor_types_; + + std::string output_convention_; // Output types of MyReLU.apply(...). std::vector output_tensor_types_; - std::vector input_tensor_requires_grads_; std::vector output_tensor_requires_grads_; std::vector arg_positions_; std::vector const_arg_positions_;