mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-26 19:52:38 +00:00
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 <tiharr@microsoft.com> * Apply suggestions from code review Refine the schema description Co-authored-by: Tim Harris <tiharr@microsoft.com> * Resolve review comments Co-authored-by: Tim Harris <tiharr@microsoft.com>
This commit is contained in:
pengwa
GitHub
parent
40e5279f8f
commit
2347a0aca8
10 changed files with 273 additions and 89 deletions
|
|
@ -171,10 +171,15 @@ void InvokeRunner(
|
|||
// first element.
|
||||
for (; i < static_cast<size_t>(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<int64_t>& requires_grads,
|
||||
const std::vector<OrtValue>& tensor_args,
|
||||
const std::vector<int64_t>& tensor_indices,
|
||||
const std::vector<void*>& 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<int64_t> requires_grads(all_input_count, 0);
|
||||
Invoke(
|
||||
runner,
|
||||
reinterpret_cast<PyObject*>(callback),
|
||||
|
|
|
|||
|
|
@ -49,7 +49,6 @@ class TorchProxy {
|
|||
|
||||
void Backward(
|
||||
void* callback,
|
||||
const std::vector<int64_t>& requires_grads,
|
||||
const std::vector<OrtValue>& tensor_args,
|
||||
const std::vector<int64_t>& tensor_indices,
|
||||
const std::vector<void*>& obj_args,
|
||||
|
|
|
|||
|
|
@ -824,24 +824,24 @@ IMPLEMENT_GRADIENT_BUILDER(GetAddSubGradient) {
|
|||
std::vector<NodeDef> 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<NodeDef> 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<AttributeProto> 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<int64_t> 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<ArgDef> 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<int64_t> 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<size_t>(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<ArgDef> 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());
|
||||
|
|
|
|||
|
|
@ -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<int64_t>(input_type))}});
|
||||
{{"const_zero_casted"},
|
||||
"Cast",
|
||||
{"const_zero_float"},
|
||||
{MakeAttribute("to", static_cast<int64_t>(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<int64_t>(input_type))}});
|
||||
{{"const_one_casted"},
|
||||
"Cast",
|
||||
{"const_one_float"},
|
||||
{MakeAttribute("to", static_cast<int64_t>(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<size_t>(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<int64_t>(ctx.getNumOutputs()); ++i) {
|
||||
updateOutputElemType(ctx, i, static_cast<int32_t>(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<int64_t>(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<size_t>(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<int64_t>(ctx.getNumOutputs()); ++i) {
|
||||
updateOutputElemType(ctx, i, static_cast<int32_t>(output_tensor_types_proto->ints().at(i)));
|
||||
const auto output_tensor_ranks = ctx.getAttribute("output_tensor_ranks");
|
||||
|
|
|
|||
|
|
@ -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"]:
|
||||
|
|
|
|||
|
|
@ -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))
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
@ -28,13 +28,16 @@ std::vector<OrtValue> CreateOrtValueArgs(OpKernelContext* context,
|
|||
|
||||
void PythonOpBase::Init(const OpKernelInfo& info) {
|
||||
ORT_THROW_IF_ERROR(info.GetAttr("name", &name_));
|
||||
inplace_ = info.GetAttrOrDefault("inplace", static_cast<int64_t>(0));
|
||||
ORT_THROW_IF_ERROR(info.GetAttr("inplace", &inplace_));
|
||||
|
||||
is_training_mode_ = static_cast<bool>(info.GetAttrOrDefault("training_mode", static_cast<int64_t>(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<int64_t>());
|
||||
input_tensor_requires_grads_ = info.GetAttrsOrDefault("input_tensor_requires_grads", std::vector<int64_t>());
|
||||
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<int64_t>());
|
||||
|
||||
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<int64_t>());
|
||||
output_tensor_requires_grads_ = info.GetAttrsOrDefault("output_tensor_requires_grads", std::vector<int64_t>());
|
||||
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<int64_t>());
|
||||
output_tensor_requires_grads_ = info.GetAttrsOrDefault("output_tensor_requires_grads", std::vector<int64_t>());
|
||||
ORT_ENFORCE(output_tensor_types_.size() == output_tensor_requires_grads_.size(), "backward tensor output count mismatch");
|
||||
|
||||
SetPositions();
|
||||
}
|
||||
|
||||
void PythonOpGradBase::RunBackward(OpKernelContext* context,
|
||||
std::vector<OrtValue>& 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<Tensor>(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<OrtValue>& returned_ortvalues) const {
|
||||
auto* ctx_internal = reinterpret_cast<onnxruntime::OpKernelContextInternal*>(context);
|
||||
const auto outputs_count = static_cast<size_t>(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<int>(i), returned_ortvalues.at(i)));
|
||||
}
|
||||
|
||||
ORT_ENFORCE(tensor_output_index == ctx_internal->OutputCount(), "backward tensor output count mismatch.");
|
||||
}
|
||||
|
||||
void PythonOpGradBase::SetPositions() {
|
||||
|
|
|
|||
|
|
@ -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<int64_t> 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<int64_t> input_tensor_types_;
|
||||
// input_tensor_types_[i] indicates if the i-th tensor should have gradient.
|
||||
std::vector<int64_t> input_tensor_requires_grads_;
|
||||
|
||||
// Concatenation of all floats from apply(...) 's inputs.
|
||||
std::vector<int64_t> input_int_scalars_;
|
||||
|
|
@ -117,9 +117,10 @@ class PythonOpGradBase {
|
|||
int64_t inplace_;
|
||||
// Input types of MyReLU.backward(...).
|
||||
std::vector<int64_t> input_tensor_types_;
|
||||
|
||||
std::string output_convention_;
|
||||
// Output types of MyReLU.apply(...).
|
||||
std::vector<int64_t> output_tensor_types_;
|
||||
std::vector<int64_t> input_tensor_requires_grads_;
|
||||
std::vector<int64_t> output_tensor_requires_grads_;
|
||||
std::vector<int64_t> arg_positions_;
|
||||
std::vector<int64_t> const_arg_positions_;
|
||||
|
|
|
|||
Loading…
Reference in a new issue