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));