saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-30 20:18:08 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

[ORTModule] ATen upsample_nearest Gradient Bugfix (#14069)

Browse source 
PyTorch removed upsample_nearest related backward functions with "vec"
overload name since 1.13. The functions without overload name are
available for all versions, though they are not that convienent to use.
This PR changes the gradient builder code to use functions without
overload name for ATen upsample_nearest nodes.

This PR also fixed a bug for ORTModule's corner case introduced by the
multi-stream PR. There is some code to execute the barrier step for
triggered downsteam is the barrier is out of range. But this should be
applied to triggered downstream only. If it's a normal run with start
step as a barrier step but out of range, we should not apply the logic.
For example, for ORTModule, if the barrier is the 1st step of whole CPU
plan, and the forward part is empty, then the forward normal run will
run step from start-0 to end-0 (actually nothing), and step-0 is the
barrier, then we should not execute the barrier in such case.
This commit is contained in:
Vincent Wang 2022-12-27 10:18:30 +08:00 committed by GitHub
parent 770fb9649b
commit 0c3480e565
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
5 changed files with 35 additions and 25 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -105,7 +105,8 @@ std::ostream& operator<<(std::ostream& out, std::pair<const SequentialExecutionP
out << "\nExecution Plan:\n";
for (size_t i = 0; i < plan.execution_plan.size(); ++i) {
auto& execution_plan = plan.execution_plan[i];
out << " Start logic stream : " << i << "on device: " << execution_plan->device_.Type() << std::endl;
out << "Start logic stream: " << i << " on device: " << std::to_string(execution_plan->device_.Type())
<< std::endl;
for (auto& step : execution_plan->steps_) {
out << step->ToString() << std::endl;
}

28
onnxruntime/core/framework/stream_execution_context.cc
View file

@ -109,8 +109,8 @@ StreamExecutionContext ::StreamExecutionContext(const SessionState& sess_state,
}
}
synchronize::Notification* StreamExecutionContext ::GetNotification(size_t /*idx*/) {
ORT_THROW("Try to get notification in a build which doesn't enable Stream!");
synchronize::Notification* StreamExecutionContext ::GetNotification(size_t /*idx*/) {
ORT_THROW("Try to get notification in a build which doesn't enable Stream!");
}
bool StreamExecutionContext ::DecCountDownBarrier(size_t /*barrier_id*/) {
@ -166,7 +166,8 @@ void StreamExecutionContext ::RecycleNodeInputs(onnxruntime::NodeIndex node_inde
}
}
void RunSince(size_t stream_idx, StreamExecutionContext& ctx, SessionScope& session_scope, const bool& terminate_flag, size_t since) {
void RunSince(size_t stream_idx, StreamExecutionContext& ctx, SessionScope& session_scope, const bool& terminate_flag,
size_t since, bool is_downstream) {
if (!ctx.TaskStatus().IsOK()) {
// already in bad status, terminate it
ctx.CompleteTask();
@ -191,7 +192,10 @@ void RunSince(size_t stream_idx, StreamExecutionContext& ctx, SessionScope& sess
// counter otherwise later in backward the downstream won't execute correctly.
// this is ugly, hopefully we won't need to worry about if deprecate ORTModule
// by Torch Dynamo.
if (since >= end && since < logic_stream->steps_.size() && logic_stream->steps_[since]->IsBarrier()) {
// We only need to do this on a triggered downstream. For example if the barrier is the first step of whole CPU plan,
// and the forward part is empty, the normal run of the forward part will not do this extra barrier handling.
if (is_downstream && since >= end && since < logic_stream->steps_.size() &&
logic_stream->steps_[since]->IsBarrier()) {
if (!ctx.TaskStatus().IsOK()) {
ctx.CompleteTask();
return;
@ -224,6 +228,8 @@ void RunSince(size_t stream_idx, StreamExecutionContext& ctx, SessionScope& sess
ctx.CompleteTask();
return;
}
#else
ORT_UNUSED_PARAMETER(is_downstream);
#endif
while (since < end) {
@ -265,11 +271,8 @@ void RunSince(size_t stream_idx, StreamExecutionContext& ctx, SessionScope& sess
return;
}
void ScheduleDownstream(StreamExecutionContext& ctx,
size_t trigger,
bool single_thread_mode,
const bool& terminate_flag,
SessionScope& session_scope) {
void ScheduleDownstream(StreamExecutionContext& ctx, size_t trigger, bool single_thread_mode,
const bool& terminate_flag, SessionScope& session_scope) {
auto* plan = ctx.GetSessionState().GetExecutionPlan();
auto& downstream_map = plan->downstream_map;
auto* tp = single_thread_mode ? nullptr : ctx.GetSessionState().GetInterOpThreadPool();
@ -278,10 +281,9 @@ void ScheduleDownstream(StreamExecutionContext& ctx,
for (auto downstream : it->second) {
// increase the task count before schedule down-stream
ctx.AddTask();
concurrency::ThreadPool::Schedule(tp,
[&ctx, downstream, &terminate_flag, &session_scope]() {
RunSince(downstream.first, ctx, session_scope, terminate_flag, downstream.second);
});
concurrency::ThreadPool::Schedule(tp, [&ctx, downstream, &terminate_flag, &session_scope]() {
RunSince(downstream.first, ctx, session_scope, terminate_flag, downstream.second, true);
});
}
}
}

3
onnxruntime/core/framework/stream_execution_context.h
View file

@ -184,7 +184,8 @@ void RunSince(size_t stream_idx,
StreamExecutionContext& ctx,
SessionScope& session_scope,
const bool& terminate_flag,
size_t since);
size_t since,
bool is_downstream = false);
// Schedule the downstream jobs from other streams at 'trigger' step, based on the execution plan.
void ScheduleDownstream(StreamExecutionContext& ctx,

4
orttraining/orttraining/core/graph/gradient_builder.cc
View file

@ -1684,7 +1684,9 @@ IMPLEMENT_GRADIENT_BUILDER(GetExternalGradient) {
OpDef op_def(node_def.op_type, node_def.domain);
std::vector<ArgDef> input_args;
for (const auto& input : node_def.inputs) {
if (input.find("GO(") == 0) {
if (input == "") {
input_args.emplace_back(ArgDef());
} else if (input.find("GO(") == 0) {
int index = std::stoi(input.substr(3, input.length() - 4));
input_args.emplace_back(GO(static_cast<size_t>(index)));
} else if (input.find("I(") == 0) {

22
orttraining/orttraining/python/training/ortmodule/_custom_gradient_registry.py
View file

@ -237,31 +237,35 @@ def native_group_norm_gradient():
]
def _upsample_nearest_gradient(backward_fn):
# PyTorch removed related backward functions with "vec" overload name since 1.13. The functions with no overload name
# are available for all versions, though they are not that convienent to use.
def _upsample_nearest_gradient(backward_fn, dims):
scales = ["" for _ in range(dims)]
return [
("Shape", ["I(0)"], ["Shape_X"]),
("Shape", ["O(0)"], ["Shape_Y"]),
("Constant", [], ["Const_Start"], {"value": {"value": [2], "dtype": "int", "is_tensor": True}}),
("Constant", [], ["Const_End"], {"value": {"value": [2 + dims], "dtype": "int", "is_tensor": True}}),
("Slice", ["Shape_Y", "Const_Start", "Const_End"], ["Sliced_Shape_Y"]),
(
("ATen", "org.pytorch.aten"),
["GO(0)", "I(1)", "Shape_X", "I(2)"],
["GO(0)", "Sliced_Shape_Y", "Shape_X"] + scales,
["GI(0)"],
{
"operator": {"value": backward_fn, "dtype": "string"},
"overload_name": {"value": "vec", "dtype": "string"},
},
{"operator": {"value": backward_fn, "dtype": "string"}},
),
]
@register_gradient("org.pytorch.aten", "ATen", "upsample_nearest1d", "vec")
def upsample_nearest1d_gradient():
return _upsample_nearest_gradient("upsample_nearest1d_backward")
return _upsample_nearest_gradient("upsample_nearest1d_backward", 1)
@register_gradient("org.pytorch.aten", "ATen", "upsample_nearest2d", "vec")
def upsample_nearest2d_gradient():
return _upsample_nearest_gradient("upsample_nearest2d_backward")
return _upsample_nearest_gradient("upsample_nearest2d_backward", 2)
@register_gradient("org.pytorch.aten", "ATen", "upsample_nearest3d", "vec")
def upsample_nearest3d_gradient():
return _upsample_nearest_gradient("upsample_nearest3d_backward")
return _upsample_nearest_gradient("upsample_nearest3d_backward", 3)