From d69823f76449bb6b61f8d4052be69b6a6face930 Mon Sep 17 00:00:00 2001 From: cao lei Date: Thu, 2 Mar 2023 07:08:29 -0800 Subject: [PATCH] Do not create Barrier and triggerDownstream steps if the corresponding nodes are split by yield Op in training scenario (#14570) ### Description Do not create Barrier and triggerDownstream steps during execution plan creation if the corresponding nodes are split by yield Op in training scenario. ### Motivation and Context In training scenario, forward and backward processes are running two different partial nodes of a graph. If there are two nodes each in one of the partial graph and separate in two streams, there are still triggerDownstream/barrier steps between them which work quite different from inference process as one of the steps will not be executed due to it is not in the correct range. To make it work, there is a hacky way to trigger the barrier step explicitly for training. This PR is to do some check, and do not create Barrier and triggerDownstream steps if the corresponding nodes are split by yield Op in training scenario. So the hacky way is not needed. --- .../core/framework/allocation_planner.cc | 87 ++++++++++--------- onnxruntime/core/framework/execution_steps.cc | 16 ++-- onnxruntime/core/framework/execution_steps.h | 15 +--- .../partial_graph_execution_state.cc | 9 +- .../framework/sequential_execution_plan.h | 12 ++- .../framework/stream_execution_context.cc | 53 +---------- .../core/framework/stream_execution_context.h | 3 +- .../orttraining/core/agent/training_agent.cc | 3 +- 8 files changed, 72 insertions(+), 126 deletions(-) diff --git a/onnxruntime/core/framework/allocation_planner.cc b/onnxruntime/core/framework/allocation_planner.cc index 8cf07f566a..d32009cd9c 100644 --- a/onnxruntime/core/framework/allocation_planner.cc +++ b/onnxruntime/core/framework/allocation_planner.cc @@ -83,28 +83,26 @@ std::ostream& operator<<(std::ostream& out, std::pair index_to_name; - index_to_name.reserve(name_idx_map.Size()); + std::map index_to_name; // order by Node_Arg index by default out << "Allocation Plan:\n"; out << "(ort_value_idx) output_name : \n"; auto plan_size = plan.allocation_plan.size(); - for (auto& name_index : name_idx_map) { - auto index = name_index.second; - index_to_name[index] = name_index.first; - out << "(" << index << ") " << name_index.first << " : "; + index_to_name[name_index.second] = name_index.first; + } + for (auto it = index_to_name.begin(); it != index_to_name.end(); it++) { + int index = it->first; + out << "(" << index << ")" << it->second << " : "; if (0 <= index && static_cast(index) < plan_size) { auto& elt_plan = plan.allocation_plan[index]; out << elt_plan.alloc_kind; if (elt_plan.alloc_kind == AllocKind::kReuse) out << " " << elt_plan.reused_buffer; - auto& loc = elt_plan.location; out << ", " << loc.ToString(); } else { out << "Index out-of-range!"; } - out << std::endl; } @@ -1788,6 +1786,31 @@ class PlannerImpl { // 2. determing following things: // a. which node need to generate notification // b. which node need to trigger downstream +#ifdef ENABLE_TRAINING + // We will leverage the topological order for the training scenario. + // The nodes before yieldOp in topo order will be executed in RunForward() and nodes after will be executed in RunBackward() + // This partition may not be exactly the same as forward model/gradient model, for example, some nodes in gradient model are + // before yieldOp thus will be executed in RunForward() + // But the final result is still correct, as long as all the nodes will be executed in either RunForward() or RunBackward() + // and no dependency conflict during the execution. + const std::vector& topo_sort = graph_viewer_.GetNodesInTopologicalOrder(context_->GetExecutionOrder()); + plan_.node_index_2_toposort_index.reserve(topo_sort.size()); + size_t yieldOp_index_in_toposort = topo_sort.size(); + for (size_t i = 0; i < topo_sort.size(); i++) { + plan_.node_index_2_toposort_index[topo_sort[i]] = i; + const Node* node = graph_viewer_.GetNode(topo_sort[i]); + if (node->OpType() == "YieldOp") { + ORT_ENFORCE(yieldOp_index_in_toposort == topo_sort.size(), "Two YieldOp in the graph"); + yieldOp_index_in_toposort = i; + } + } + + auto AreNodesSeparatedByYield = [&](NodeIndex producer, NodeIndex consumer) { + size_t producer_topoindex = plan_.node_index_2_toposort_index[producer]; + size_t consumer_topoindex = plan_.node_index_2_toposort_index[consumer]; + return producer_topoindex < yieldOp_index_in_toposort && yieldOp_index_in_toposort < consumer_topoindex; + }; +#endif size_t num_trigger_points = 0; InlinedHashMap node_to_trigger_points; InlinedHashMap node_to_notification; @@ -1797,7 +1820,13 @@ class PlannerImpl { auto* node = graph_viewer_.GetNode(node_index); for (auto it = node->OutputNodesBegin(); it != node->OutputNodesEnd(); ++it) { // if the output node is not in the same stream, generate a trigger point - if (node_stream_map_[it->Index()] != i) { + if (node_stream_map_[it->Index()] != i +#ifdef ENABLE_TRAINING + // Do not insert Barrier/TriggerDownStream step if the producer and consumer are in different sides of yieldOp + // As in this case producer will surely be ready before consumer is running. + && !AreNodesSeparatedByYield(node_index, it->Index()) +#endif + ) { node_to_trigger_points[node_index] = num_trigger_points++; break; } @@ -1869,7 +1898,11 @@ class PlannerImpl { } visited.insert(it->Index()); // check whether we need to add barrier - if (std::find(stream_nodes_[i].begin(), stream_nodes_[i].end(), it->Index()) == stream_nodes_[i].end()) { + if (std::find(stream_nodes_[i].begin(), stream_nodes_[i].end(), it->Index()) == stream_nodes_[i].end() +#ifdef ENABLE_TRAINING + && !AreNodesSeparatedByYield(it->Index(), node_index) +#endif + ) { // find the trigger_point_id auto trigger_point_it = node_to_trigger_points.find(it->Index()); ORT_ENFORCE(trigger_point_it != node_to_trigger_points.end()); @@ -1878,21 +1911,15 @@ class PlannerImpl { size_t barrier_id = plan_.num_barriers++; plan_.downstream_map[trigger_point_index].push_back({i, static_cast(execution_plan[i]->steps_.size())}); - execution_plan[i]->steps_.emplace_back(std::make_unique(barrier_id)); -#ifdef ENABLE_TRAINING - // keep node index first, will turn it to pc index later - execution_plan[i]->step_pc.push_back(node_index); -#endif + execution_plan[i]->steps_.emplace_back(std::make_unique(barrier_id, node_index)); } } auto wait_it = node_to_wait.find(node_index); if (wait_it != node_to_wait.end()) { for (auto wait_param : wait_it->second) { - execution_plan[i]->steps_.emplace_back(std::make_unique(wait_param.second, node_to_notification[wait_param.first])); -#ifdef ENABLE_TRAINING - execution_plan[i]->step_pc.push_back(node_index); -#endif // ENABLE_TRAINING + execution_plan[i]->steps_.emplace_back(std::make_unique(wait_param.second, + node_to_notification[wait_param.first], node_index)); } } @@ -1902,27 +1929,17 @@ class PlannerImpl { } // push launch kernel command execution_plan[i]->steps_.emplace_back(std::make_unique(node_index)); -#ifdef ENABLE_TRAINING - execution_plan[i]->step_pc.push_back(node_index); -#endif // check if any notification generated by this node, if yes, push a activate auto notification_it = node_to_notification.find(node_index); if (notification_it != node_to_notification.end()) { NotificationIndex notification_index = notification_it->second; - execution_plan[i]->steps_.emplace_back(std::make_unique(notification_index)); -#ifdef ENABLE_TRAINING - execution_plan[i]->step_pc.push_back(node_index); -#endif + execution_plan[i]->steps_.emplace_back(std::make_unique(notification_index, node_index)); } // check if any trigger point generated by this node, if yes, push a trigger auto trigger_point_it = node_to_trigger_points.find(node_index); if (trigger_point_it != node_to_trigger_points.end()) { // notify downstreams - execution_plan[i]->steps_.emplace_back(std::make_unique(trigger_point_it->second)); -#ifdef ENABLE_TRAINING - // set the notification step as the triggering part of next node. - execution_plan[i]->step_pc.push_back(node_index); -#endif + execution_plan[i]->steps_.emplace_back(std::make_unique(trigger_point_it->second, node_index)); } } } @@ -2029,14 +2046,6 @@ class PlannerImpl { process_stream(i, -1); } ORT_ENFORCE(plan_.node_execution_order_in_training.size() == num_of_nodes); - // 6. turn the step_node_index to step_pc - for (auto& stream : plan_.execution_plan) { - for (size_t i = 0; i < stream->step_pc.size(); ++i) { - auto it = std::find(plan_.node_execution_order_in_training.begin(), plan_.node_execution_order_in_training.end(), stream->step_pc[i]); - ORT_ENFORCE(it != plan_.node_execution_order_in_training.end()); - stream->step_pc[i] = static_cast(std::distance(plan_.node_execution_order_in_training.begin(), it)); - } - } #endif return Status::OK(); diff --git a/onnxruntime/core/framework/execution_steps.cc b/onnxruntime/core/framework/execution_steps.cc index af3b1b9af3..6eccc45efc 100644 --- a/onnxruntime/core/framework/execution_steps.cc +++ b/onnxruntime/core/framework/execution_steps.cc @@ -3,7 +3,7 @@ #include "core/framework/execution_steps.h" #include "core/framework/sequential_executor.h" namespace onnxruntime { -BarrierStep::BarrierStep(size_t id) : SequentialExecutionPlan::ExecutionStep(), +BarrierStep::BarrierStep(size_t id, NodeIndex node_index) : SequentialExecutionPlan::ExecutionStep(node_index), barrier_id_{id} {} Status BarrierStep::Execute(StreamExecutionContext& ctx, @@ -19,11 +19,9 @@ std::string BarrierStep::ToString() const { return ::onnxruntime::MakeString("Set a barrier with id: ", barrier_id_, ", count: ", 2, "."); } -#ifdef ENABLE_TRAINING -bool BarrierStep::IsBarrier() const { return true; } -#endif + WaitOnEPStep::WaitOnEPStep(WaitNotificationFn handle, - NotificationIndex idx) : SequentialExecutionPlan::ExecutionStep(), + NotificationIndex idx, NodeIndex node_index) : SequentialExecutionPlan::ExecutionStep(node_index), wait_handle_(handle), notification_idx_(idx) {} @@ -48,8 +46,7 @@ std::string WaitOnEPStep::ToString() const { notification_idx_, ". "); } -LaunchKernelStep::LaunchKernelStep(NodeIndex index) : SequentialExecutionPlan::ExecutionStep(), - node_index_{index} {} +LaunchKernelStep::LaunchKernelStep(NodeIndex index) : SequentialExecutionPlan::ExecutionStep(index) {} Status LaunchKernelStep::Execute(StreamExecutionContext& ctx, size_t stream_idx, @@ -74,7 +71,7 @@ std::string LaunchKernelStep::ToString() const { } ActivateNotificationStep::ActivateNotificationStep( - NotificationIndex notification_index) : SequentialExecutionPlan::ExecutionStep(), + NotificationIndex notification_index, NodeIndex node_index) : SequentialExecutionPlan::ExecutionStep(node_index), notification_idx_(notification_index) {} Status ActivateNotificationStep::Execute(StreamExecutionContext& ctx, @@ -96,7 +93,7 @@ std::string ActivateNotificationStep::ToString() const { notification_idx_, ". "); } -TriggerDownstreamStep::TriggerDownstreamStep(size_t trigger_point_index) : SequentialExecutionPlan::ExecutionStep(), +TriggerDownstreamStep::TriggerDownstreamStep(size_t trigger_point_index, NodeIndex node_index) : SequentialExecutionPlan::ExecutionStep(node_index), trigger_point_index_(trigger_point_index) {} Status TriggerDownstreamStep::Execute(StreamExecutionContext& ctx, @@ -113,5 +110,4 @@ std::string TriggerDownstreamStep::ToString() const { return ::onnxruntime::MakeString("TriggerDownstreamStep: trigger downstream of trigger point: ", trigger_point_index_, "."); } - } // namespace onnxruntime diff --git a/onnxruntime/core/framework/execution_steps.h b/onnxruntime/core/framework/execution_steps.h index 47d212dfad..d3eff1a2ee 100644 --- a/onnxruntime/core/framework/execution_steps.h +++ b/onnxruntime/core/framework/execution_steps.h @@ -11,7 +11,7 @@ class SessionScope; class BarrierStep : public SequentialExecutionPlan::ExecutionStep { public: - BarrierStep(size_t id); + BarrierStep(size_t id, NodeIndex node_index); Status Execute(StreamExecutionContext& ctx, size_t /*stream_idx*/, @@ -20,17 +20,13 @@ class BarrierStep : public SequentialExecutionPlan::ExecutionStep { bool& continue_flag) override; std::string ToString() const override; -#ifdef ENABLE_TRAINING - // Only applicable when using PartialExecutor - bool IsBarrier() const override; -#endif private: size_t barrier_id_{0}; }; class WaitOnEPStep : public SequentialExecutionPlan::ExecutionStep { public: - WaitOnEPStep(WaitNotificationFn handle, NotificationIndex idx); + WaitOnEPStep(WaitNotificationFn handle, NotificationIndex idx, NodeIndex node_index); Status Execute(StreamExecutionContext& ctx, size_t stream_idx, @@ -56,14 +52,11 @@ class LaunchKernelStep : public SequentialExecutionPlan::ExecutionStep { bool& continue_flag) override; std::string ToString() const override; - - private: - NodeIndex node_index_{0}; }; class ActivateNotificationStep : public SequentialExecutionPlan::ExecutionStep { public: - ActivateNotificationStep(NotificationIndex notification_index); + ActivateNotificationStep(NotificationIndex notification_index, NodeIndex node_index); Status Execute(StreamExecutionContext& ctx, size_t stream_idx, @@ -79,7 +72,7 @@ class ActivateNotificationStep : public SequentialExecutionPlan::ExecutionStep { class TriggerDownstreamStep : public SequentialExecutionPlan::ExecutionStep { public: - TriggerDownstreamStep(size_t trigger_point_index); + TriggerDownstreamStep(size_t trigger_point_index, NodeIndex node_index); Status Execute(StreamExecutionContext& ctx, size_t /*stream_idx*/, SessionScope& session_scope, diff --git a/onnxruntime/core/framework/partial_graph_execution_state.cc b/onnxruntime/core/framework/partial_graph_execution_state.cc index 7b87c7fb60..4cc53feb26 100644 --- a/onnxruntime/core/framework/partial_graph_execution_state.cc +++ b/onnxruntime/core/framework/partial_graph_execution_state.cc @@ -7,6 +7,7 @@ #include "core/framework/session_state.h" #include "core/framework/stream_execution_context.h" #include "core/framework/execution_frame.h" +#include "core/framework/execution_steps.h" namespace onnxruntime { @@ -28,13 +29,13 @@ ProgramRegion& PartialGraphExecutionState::GetProgramRegions(const SessionState& new_region.stream_pc_range.reserve(plan->execution_plan.size()); for (auto& stream : plan->execution_plan) { size_t cur = 0; - while (cur < stream->step_pc.size() && - stream->step_pc[cur] < new_region.start_pc) { + while (cur < stream->steps_.size() && + plan->node_index_2_toposort_index.at(stream->steps_[cur]->GetNodeIndex()) < new_region.start_pc) { cur++; } size_t start = cur; - while (cur < stream->step_pc.size() && - stream->step_pc[cur] < new_region.end_pc) { + while (cur < stream->steps_.size() && + plan->node_index_2_toposort_index.at(stream->steps_[cur]->GetNodeIndex()) < new_region.end_pc) { cur++; } new_region.stream_pc_range.push_back({start, cur}); diff --git a/onnxruntime/core/framework/sequential_execution_plan.h b/onnxruntime/core/framework/sequential_execution_plan.h index 74c46956c0..468a54b5c7 100644 --- a/onnxruntime/core/framework/sequential_execution_plan.h +++ b/onnxruntime/core/framework/sequential_execution_plan.h @@ -109,6 +109,7 @@ struct SequentialExecutionPlan : public ExecutionPlanBase { // 3. Wait on a notificaiton class ExecutionStep { public: + ExecutionStep(NodeIndex node_index) : node_index_(node_index) {} virtual ~ExecutionStep() {} virtual Status Execute(StreamExecutionContext& ctx, size_t stream_idx, @@ -116,19 +117,15 @@ struct SequentialExecutionPlan : public ExecutionPlanBase { const bool& terminate_flag, bool& continue_flag) = 0; virtual std::string ToString() const = 0; -#ifdef ENABLE_TRAINING - // the partial execution mode for training needs special handling for barrier - virtual bool IsBarrier() const { return false; } -#endif + inline NodeIndex GetNodeIndex() { return node_index_; } + protected: + NodeIndex node_index_; }; // LogicStream is a sequence of execution steps that can be executed independetly. // The steps within a sequence are executed in order, and happened on the same device. struct LogicStream { std::vector> steps_; const OrtDevice& device_; -#ifdef ENABLE_TRAINING - std::vector step_pc; -#endif public: LogicStream(const OrtDevice& device) : device_(device) {} }; @@ -168,6 +165,7 @@ struct SequentialExecutionPlan : public ExecutionPlanBase { #ifdef ENABLE_TRAINING InlinedVector node_execution_order_in_training; + InlinedHashMap node_index_2_toposort_index; #endif const std::vector& GetAllocationPlan() const { diff --git a/onnxruntime/core/framework/stream_execution_context.cc b/onnxruntime/core/framework/stream_execution_context.cc index 8919adac6c..d47f861c29 100644 --- a/onnxruntime/core/framework/stream_execution_context.cc +++ b/onnxruntime/core/framework/stream_execution_context.cc @@ -174,8 +174,7 @@ void StreamExecutionContext::RecycleNodeInputs(onnxruntime::NodeIndex node_index } } -void RunSince(size_t stream_idx, StreamExecutionContext& ctx, SessionScope& session_scope, const bool& terminate_flag, - size_t since, bool is_downstream) { +void RunSince(size_t stream_idx, StreamExecutionContext& ctx, SessionScope& session_scope, const bool& terminate_flag, size_t since) { if (!ctx.TaskStatus().IsOK()) { // already in bad status, terminate it ctx.CompleteTask(); @@ -192,54 +191,6 @@ void RunSince(size_t stream_idx, StreamExecutionContext& ctx, SessionScope& sess end = std::min(end, range->stream_pc_range[stream_idx].second); #endif -#ifdef ENABLE_TRAINING - // this is a special handle for training - // with ORTModule, we are partially execute the graph with a shared context. - // there is a case that in forward pass we want to trigger downstream which - // not in current range. We need to execute one step to consume the Barrier - // 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. - // 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; - } - if (terminate_flag) { - Status status_made = ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Exiting due to terminate flag being set to true."); - ctx.SetStatus(status_made); - ctx.CompleteTask(); - return; - } - bool continue_flag = true; - Status status; - ORT_TRY { - status = logic_stream->steps_[since]->Execute(ctx, stream_idx, session_scope, terminate_flag, continue_flag); - } - ORT_CATCH(const std::exception& ex) { - ORT_HANDLE_EXCEPTION([&]() { - status = ORT_MAKE_STATUS(ONNXRUNTIME, RUNTIME_EXCEPTION, ex.what()); - }); - } - if (!status.IsOK()) { - // terminate it - ctx.SetStatus(status); - ctx.CompleteTask(); - return; - } - if (continue_flag) { - ORT_THROW("Execute the barrier step in backward range passed! this is not expected."); - } - ctx.CompleteTask(); - return; - } -#else - ORT_UNUSED_PARAMETER(is_downstream); -#endif - while (since < end) { if (!ctx.TaskStatus().IsOK()) { ctx.CompleteTask(); @@ -290,7 +241,7 @@ void ScheduleDownstream(StreamExecutionContext& ctx, size_t trigger, bool single // 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, true); + RunSince(downstream.first, ctx, session_scope, terminate_flag, downstream.second); }); } } diff --git a/onnxruntime/core/framework/stream_execution_context.h b/onnxruntime/core/framework/stream_execution_context.h index ede5056209..1815e0d122 100644 --- a/onnxruntime/core/framework/stream_execution_context.h +++ b/onnxruntime/core/framework/stream_execution_context.h @@ -185,8 +185,7 @@ void RunSince(size_t stream_idx, StreamExecutionContext& ctx, SessionScope& session_scope, const bool& terminate_flag, - size_t since, - bool is_downstream = false); + size_t since); // 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/agent/training_agent.cc b/orttraining/orttraining/core/agent/training_agent.cc index 87f5e93211..3b701fa8bf 100644 --- a/orttraining/orttraining/core/agent/training_agent.cc +++ b/orttraining/orttraining/core/agent/training_agent.cc @@ -25,8 +25,7 @@ TrainingAgent::TrainingAgent(InferenceSession& session, std::vector bw_feed_names; size_t break_point = 0; - auto* plan = session_state.GetExecutionPlan(); - auto& training_node_execution_order = plan->node_execution_order_in_training; + auto& training_node_execution_order = session_state.GetGraphViewer().GetNodesInTopologicalOrder(session.GetSessionOptions().execution_order); for (auto node_index : training_node_execution_order) { if (session_state.GetKernel(node_index)->KernelDef().OpName() == "YieldOp") { auto& node = *(session_state.GetGraphViewer().GetGraph().GetNode(node_index));