saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-05 04:17:53 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Disable reuse for YieldOp's inputs (FW partial graph's output) (#7767)

Browse source 
* Disable reuse for YieldOp's input

Co-authored-by: Sherlock Huang <bahuang@OrtTrainingDev3.af05slrtruoetgaxwwjv5nsq5e.px.internal.cloudapp.net>
This commit is contained in:
Sherlock 2021-05-20 21:39:36 -07:00 committed by GitHub
parent c2435d24ec
commit 2a02871157
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 19 additions and 6 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

14
onnxruntime/core/framework/allocation_planner.cc
View file

@ -257,6 +257,20 @@ class PlannerImpl {
// Find if there exists some input tensor that we can use in-place for output_arg_num-th input in the node.
bool FindReusableInput(const onnxruntime::Node& node, int output_arg_num, OrtValueIndex* reusable_input) {
#ifdef ENABLE_TRAINING
// Inputs of Yields are essentially the outputs for FW partial subgraph
// Thses tensors will be pass back to pytorch, thus cannot share the buffer with other tensors
// Unhandled corner case:
// If FW output tensor is consumed by BW graph, and pytorch performs an inplace operation on th returned tensor,
// we will run into a buffer corruption problem.
// One potential fix is returning a copy of output tensor, if it has downstream dependency
auto p_next_node = node.OutputNodesBegin();
if (p_next_node != node.OutputNodesEnd() && p_next_node->OpType() == "YieldOp") {
return false;
}
#endif //ENABLE_TRAINING
auto p_output_arg = node.OutputDefs()[output_arg_num];
const KernelCreateInfo& ci = GetKernelCreateInfo(kernel_create_info_map_, node.Index());

11
orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py
View file

@ -46,6 +46,8 @@ class NeuralNetMultiplePositionalArgumentsMultiOutputsWithoutDependency(torch.nn
self.fc2 = torch.nn.Linear(input_size, hidden_size)
self.softmax1 = torch.nn.Softmax(dim=1)
self.softmax2 = torch.nn.Softmax(dim=1)
self.relu1 = torch.nn.ReLU()
self.relu2 = torch.nn.ReLU()
def forward(self, input1, input2):
model_input = input1 + input2
@ -53,11 +55,8 @@ class NeuralNetMultiplePositionalArgumentsMultiOutputsWithoutDependency(torch.nn
out2 = self.fc2(model_input)
out1 = self.softmax1(out1)
out2 = self.softmax2(out2)
# TODO: Using relu here will cause the forward prediction error
# ORT's Relu output is sharing the same buffer as input,
# and this buffer is returned as ORTModule's output to Pytorch
# out1 = self.relu1(out1)
# out2 = self.relu2(out2)
out1 = self.relu1(out1)
out2 = self.relu2(out2)
return out1, out2
class NeuralNetMultiplePositionalArgumentsMultiOutputsWithDependency(torch.nn.Module):
@ -593,7 +592,7 @@ def test_gradient_correctness_conv1d(use_fp16, input_requires_grad):
if use_fp16:
_test_helpers.assert_values_are_close(ort_prediction, pt_prediction, atol=1e-3, rtol=1e-3)
_test_helpers.assert_gradients_match_and_reset_gradient(ort_model, pt_model, rtol=1e-2, atol=1e-2)
_test_helpers.assert_gradients_match_and_reset_gradient(ort_model, pt_model, rtol=1e-2, atol=1.1e-2)
else:
_test_helpers.assert_values_are_close(ort_prediction, pt_prediction, atol=1e-5)
_test_helpers.assert_gradients_match_and_reset_gradient(ort_model, pt_model, rtol=5e-3, atol=4e-3)