saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-24 19:43:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Do not create Barrier and triggerDownstream steps if the corresponding nodes are split by yield Op in training scenario (#14570)

Browse source 
### 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.
This commit is contained in:
cao lei 2023-03-02 07:08:29 -08:00 committed by GitHub
parent c66af46fc1
commit d69823f764
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
8 changed files with 72 additions and 126 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -83,28 +83,26 @@ std::ostream& operator<<(std::ostream& out, std::pair<const SequentialExecutionP
const SessionState& session_state = *planinfo.second;
const auto& name_idx_map = session_state.GetOrtValueNameIdxMap();
InlinedHashMap<int, std::string_view> index_to_name;
index_to_name.reserve(name_idx_map.Size());
std::map<int, std::string_view> index_to_name; // order by Node_Arg index by default
out << "Allocation Plan:\n";
out << "(ort_value_idx) output_name : <allocation plan>\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<size_t>(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<NodeIndex>& 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<NodeIndex, size_t> node_to_trigger_points;
InlinedHashMap<NodeIndex, NotificationIndex> 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<int>(execution_plan[i]->steps_.size())});
execution_plan[i]->steps_.emplace_back(std::make_unique<BarrierStep>(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<BarrierStep>(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<WaitOnEPStep>(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<WaitOnEPStep>(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<LaunchKernelStep>(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<ActivateNotificationStep>(notification_index));
#ifdef ENABLE_TRAINING
execution_plan[i]->step_pc.push_back(node_index);
#endif
execution_plan[i]->steps_.emplace_back(std::make_unique<ActivateNotificationStep>(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<TriggerDownstreamStep>(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<TriggerDownstreamStep>(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<int>(std::distance(plan_.node_execution_order_in_training.begin(), it));
}
}
#endif
return Status::OK();

16
onnxruntime/core/framework/execution_steps.cc
View file

@ -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

15
onnxruntime/core/framework/execution_steps.h
View file

@ -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,

9
onnxruntime/core/framework/partial_graph_execution_state.cc
View file

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

12
onnxruntime/core/framework/sequential_execution_plan.h
View file

@ -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<std::unique_ptr<ExecutionStep>> steps_;
const OrtDevice& device_;
#ifdef ENABLE_TRAINING
std::vector<NodeIndex> step_pc;
#endif
public:
LogicStream(const OrtDevice& device) : device_(device) {}
};
@ -168,6 +165,7 @@ struct SequentialExecutionPlan : public ExecutionPlanBase {
#ifdef ENABLE_TRAINING
InlinedVector<NodeIndex> node_execution_order_in_training;
InlinedHashMap<NodeIndex, size_t> node_index_2_toposort_index;
#endif
const std::vector<AllocPlanPerValue>& GetAllocationPlan() const {

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

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

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

@ -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,

3
orttraining/orttraining/core/agent/training_agent.cc
View file

@ -25,8 +25,7 @@ TrainingAgent::TrainingAgent(InferenceSession& session,
std::vector<std::string> 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));