From 0c3480e56551cbb66339d765fc7c314890200ce2 Mon Sep 17 00:00:00 2001 From: Vincent Wang Date: Tue, 27 Dec 2022 10:18:30 +0800 Subject: [PATCH] [ORTModule] ATen upsample_nearest Gradient Bugfix (#14069) 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. --- .../core/framework/allocation_planner.cc | 3 +- .../framework/stream_execution_context.cc | 28 ++++++++++--------- .../core/framework/stream_execution_context.h | 3 +- .../core/graph/gradient_builder.cc | 4 ++- .../ortmodule/_custom_gradient_registry.py | 22 +++++++++------ 5 files changed, 35 insertions(+), 25 deletions(-) diff --git a/onnxruntime/core/framework/allocation_planner.cc b/onnxruntime/core/framework/allocation_planner.cc index 77638d39d9..03dd59a8c0 100644 --- a/onnxruntime/core/framework/allocation_planner.cc +++ b/onnxruntime/core/framework/allocation_planner.cc @@ -105,7 +105,8 @@ std::ostream& operator<<(std::ostream& out, std::pairdevice_.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; } diff --git a/onnxruntime/core/framework/stream_execution_context.cc b/onnxruntime/core/framework/stream_execution_context.cc index ac58877936..17e6548f9e 100644 --- a/onnxruntime/core/framework/stream_execution_context.cc +++ b/onnxruntime/core/framework/stream_execution_context.cc @@ -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); + }); } } } diff --git a/onnxruntime/core/framework/stream_execution_context.h b/onnxruntime/core/framework/stream_execution_context.h index 1cf57368a1..11432bb66a 100644 --- a/onnxruntime/core/framework/stream_execution_context.h +++ b/onnxruntime/core/framework/stream_execution_context.h @@ -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, diff --git a/orttraining/orttraining/core/graph/gradient_builder.cc b/orttraining/orttraining/core/graph/gradient_builder.cc index 207175f861..668a44d985 100755 --- a/orttraining/orttraining/core/graph/gradient_builder.cc +++ b/orttraining/orttraining/core/graph/gradient_builder.cc @@ -1684,7 +1684,9 @@ IMPLEMENT_GRADIENT_BUILDER(GetExternalGradient) { OpDef op_def(node_def.op_type, node_def.domain); std::vector 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(index))); } else if (input.find("I(") == 0) { diff --git a/orttraining/orttraining/python/training/ortmodule/_custom_gradient_registry.py b/orttraining/orttraining/python/training/ortmodule/_custom_gradient_registry.py index b7461781e7..e1d3d5fcf5 100644 --- a/orttraining/orttraining/python/training/ortmodule/_custom_gradient_registry.py +++ b/orttraining/orttraining/python/training/ortmodule/_custom_gradient_registry.py @@ -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)