mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-29 20:14:01 +00:00
Revise pipeline schedule to consider communication ops (#4524)
* Revise pipeline schedule to consider communication ops * Add test * Fix warning * inline some short functions * Fix warnings * Rename a class * Add comment for test * op renamed to task * Fix NVTX wrapper's bug
This commit is contained in:
Wei-Sheng Chin
GitHub
parent
183098e344
commit
21d2728974
7 changed files with 886 additions and 357 deletions
|
|
@ -124,6 +124,7 @@ file(GLOB onnxruntime_test_training_src
|
|||
"${ORTTRAINING_SOURCE_DIR}/test/graph/*.cc"
|
||||
"${ORTTRAINING_SOURCE_DIR}/test/optimizer/*.cc"
|
||||
"${ORTTRAINING_SOURCE_DIR}/test/framework/*.cc"
|
||||
"${ORTTRAINING_SOURCE_DIR}/test/distributed/*.cc"
|
||||
)
|
||||
|
||||
if(WIN32)
|
||||
|
|
|
|||
|
|
@ -204,7 +204,7 @@ Status SequentialExecutor::Execute(const SessionState& session_state, const std:
|
|||
if (node.Description() != "Backward pass" && !forward_range.IsBeginCalled()) {
|
||||
// Start timing forward pass when encountering the first forward node.
|
||||
forward_range.Begin();
|
||||
} else if (node.Description() == "Backward pass" && !backward_range.IsBeginCalled()) {
|
||||
} else if (node.Description() == "Backward pass" && !backward_range.IsBeginCalled() && forward_range.IsBeginCalled()) {
|
||||
// Start timing backward pass when encountering the first backward node.
|
||||
// In the meanwhile, forward range ends.
|
||||
forward_range.End();
|
||||
|
|
|
|||
|
|
@ -3,12 +3,14 @@
|
|||
|
||||
#include "orttraining/models/runner/pipeline.h"
|
||||
|
||||
#include <algorithm>
|
||||
#include <iostream>
|
||||
#include <vector>
|
||||
#include <cstdint>
|
||||
#include <cstdlib>
|
||||
#include <stdexcept>
|
||||
#include <thread>
|
||||
#include <iomanip>
|
||||
|
||||
#include "gsl/gsl"
|
||||
#include "core/framework/ml_value.h"
|
||||
|
|
@ -18,131 +20,198 @@ namespace onnxruntime {
|
|||
namespace training {
|
||||
namespace pipeline {
|
||||
|
||||
Slot::Slot() {
|
||||
type = Empty;
|
||||
batch_id = 0;
|
||||
}
|
||||
|
||||
bool Slot::IsEmpty() const {
|
||||
return type == Empty;
|
||||
}
|
||||
|
||||
bool Slot::IsForward() const {
|
||||
return type == Forward;
|
||||
}
|
||||
|
||||
bool Slot::IsBackward() const {
|
||||
return type == Backward;
|
||||
}
|
||||
|
||||
void Slot::Show() const {
|
||||
const int waited_event_id =
|
||||
waited_events.size() > 0 ? waited_events[0] : -1;
|
||||
const int waited_event_id_after_recv =
|
||||
waited_events.size() > 0 ? waited_events[1] : -1;
|
||||
const int recorded_event_id_before_send =
|
||||
recorded_events.size() > 0 ? recorded_events[0] : -1;
|
||||
const int recorded_event_id =
|
||||
recorded_events.size() > 0 ? recorded_events[1] : -1;
|
||||
switch (type) {
|
||||
case Empty:
|
||||
std::cout << " ";
|
||||
break;
|
||||
case Forward:
|
||||
std::cout << " F" << batch_id << "("
|
||||
<< waited_event_id << "," << waited_event_id_after_recv << ","
|
||||
<< recorded_event_id_before_send << "," << recorded_event_id
|
||||
<< ")"
|
||||
<< " ";
|
||||
break;
|
||||
case Backward:
|
||||
std::cout << " B" << batch_id << "("
|
||||
<< waited_event_id << "," << waited_event_id_after_recv << ","
|
||||
<< recorded_event_id_before_send << "," << recorded_event_id
|
||||
<< ")"
|
||||
<< " ";
|
||||
break;
|
||||
std::ostream& operator<<(std::ostream& stream, const PipelineTask& slot) {
|
||||
if (slot.pass == PipelineTask::Pass::Forward) {
|
||||
switch (slot.type) {
|
||||
case PipelineTask::Type::Compute:
|
||||
stream << "FW";
|
||||
break;
|
||||
case PipelineTask::Type::Send:
|
||||
stream << "FS";
|
||||
break;
|
||||
case PipelineTask::Type::Recv:
|
||||
stream << "FR";
|
||||
break;
|
||||
default:
|
||||
throw std::invalid_argument("Unsupported forwoard action.");
|
||||
break;
|
||||
}
|
||||
} else if (slot.pass == PipelineTask::Pass::Backward) {
|
||||
switch (slot.type) {
|
||||
case PipelineTask::Type::Compute:
|
||||
stream << "BW";
|
||||
break;
|
||||
case PipelineTask::Type::Send:
|
||||
stream << "BS";
|
||||
break;
|
||||
case PipelineTask::Type::Recv:
|
||||
stream << "BR";
|
||||
break;
|
||||
default:
|
||||
throw std::invalid_argument("Unsupported backward action.");
|
||||
break;
|
||||
}
|
||||
} else {
|
||||
throw std::invalid_argument("Unsupported pass type.");
|
||||
}
|
||||
std::ios state(nullptr);
|
||||
state.copyfmt(state);
|
||||
stream << std::setw(2) << std::setfill('0') << slot.batch;
|
||||
stream.copyfmt(state);
|
||||
return stream;
|
||||
}
|
||||
|
||||
PipelineSchedule::PipelineSchedule(int num_stages) {
|
||||
void PipelineSlot::AddSend(const int batch_id, const PipelineTask::Pass pass, const int upstream_time, const int upstream_stage, const int this_rank, const int peer_rank) {
|
||||
tasks_.push_back(PipelineTask{batch_id, PipelineTask::Type::Send, pass, upstream_time, upstream_stage, this_rank, peer_rank});
|
||||
}
|
||||
void PipelineSlot::AddRecv(const int batch_id, const PipelineTask::Pass pass, const int upstream_time, const int upstream_stage, const int this_rank, const int peer_rank) {
|
||||
tasks_.push_back(PipelineTask{batch_id, PipelineTask::Type::Recv, pass, upstream_time, upstream_stage, this_rank, peer_rank});
|
||||
}
|
||||
void PipelineSlot::AddCompute(const int batch_id, const PipelineTask::Pass pass, const int upstream_time, const int upstream_stage) {
|
||||
tasks_.push_back(PipelineTask{batch_id, PipelineTask::Type::Compute, pass, upstream_time, upstream_stage});
|
||||
}
|
||||
|
||||
PipelineTask& PipelineSlot::operator[](int index) {
|
||||
return tasks_[index];
|
||||
}
|
||||
|
||||
const PipelineTask& PipelineSlot::operator[](int index) const {
|
||||
return tasks_[index];
|
||||
}
|
||||
|
||||
PipelineTask& PipelineSlot::GetFrontAction() {
|
||||
return tasks_.front();
|
||||
}
|
||||
|
||||
const PipelineTask& PipelineSlot::GetFrontAction() const {
|
||||
return tasks_.front();
|
||||
}
|
||||
|
||||
PipelineScheduler::PipelineScheduler(const int num_batches, const int num_stages) {
|
||||
num_stages_ = num_stages;
|
||||
num_batches_ = 0;
|
||||
num_batches_ = num_batches;
|
||||
CreateComputeSchedule();
|
||||
|
||||
const size_t num_events_per_slot_compute_side = 2;
|
||||
std::vector<int> compute_default_events(num_events_per_slot_compute_side, -1);
|
||||
InsertEvents(compute_table_, num_events_per_slot_compute_side, compute_default_events);
|
||||
|
||||
CreateFullSchedule();
|
||||
|
||||
const size_t num_events_per_slot_side = 1;
|
||||
std::vector<int> default_events(num_events_per_slot_side, -1);
|
||||
InsertEvents(compute_commute_table_, num_events_per_slot_side, default_events);
|
||||
}
|
||||
|
||||
void PipelineSchedule::Add(int batch_id) {
|
||||
++num_batches_;
|
||||
// Print this structure following a fixed-length format.
|
||||
// It assumes there are at most 2 actions per slot.
|
||||
std::ostream& operator<<(std::ostream& stream, const PipelineSlot& slot) {
|
||||
switch (slot.NumActions()) {
|
||||
case 0:
|
||||
// One empty action.
|
||||
stream << " ";
|
||||
// Another empty action.
|
||||
stream << " ";
|
||||
break;
|
||||
case 1:
|
||||
// Print an action.
|
||||
stream << slot[0];
|
||||
// One empty action.
|
||||
stream << " ";
|
||||
break;
|
||||
case 2:
|
||||
// Print an action.
|
||||
stream << slot[0];
|
||||
// Print another action.
|
||||
stream << slot[1];
|
||||
break;
|
||||
}
|
||||
return stream;
|
||||
}
|
||||
|
||||
// Expand table to accomonadate the new batch.
|
||||
const int required_max_time = 2 * num_stages_ + 2 * (num_batches_ - 1);
|
||||
const int current_max_time = static_cast<int>(table_.size());
|
||||
void PipelineSlot::SetWaitedEvent(const std::vector<int> events) {
|
||||
waited_events_ = std::move(events);
|
||||
}
|
||||
|
||||
for (int t = current_max_time; t < required_max_time; ++t) {
|
||||
table_.push_back(std::vector<Slot>(num_stages_));
|
||||
batch_count_.push_back(0);
|
||||
std::vector<int> PipelineSlot::GetWaitedEvent() const {
|
||||
return waited_events_;
|
||||
}
|
||||
|
||||
void PipelineSlot::SetRecordedEvent(const std::vector<int> events) {
|
||||
recorded_events_ = std::move(events);
|
||||
}
|
||||
|
||||
std::vector<int> PipelineSlot::GetRecordedEvent() const {
|
||||
return recorded_events_;
|
||||
}
|
||||
|
||||
std::ostream& operator<<(std::ostream& stream, PipelineScheduler const& schedule) {
|
||||
// print something from v to str, e.g: Str << v.getX();
|
||||
stream << "-------------View of Compute Schedule-------------" << std::endl;
|
||||
for (int s = 0; s < schedule.num_stages_; ++s) {
|
||||
for (size_t t = 0; t < schedule.compute_table_.size(); ++t) {
|
||||
stream << schedule.compute_table_[t][s];
|
||||
}
|
||||
stream << std::endl;
|
||||
}
|
||||
|
||||
stream << "-------------View of Compute-commute Schedule-------------" << std::endl;
|
||||
for (int s = 0; s < schedule.num_stages_; ++s) {
|
||||
for (size_t t = 0; t < schedule.compute_commute_table_.size(); ++t) {
|
||||
stream << schedule.compute_commute_table_[t][s];
|
||||
}
|
||||
stream << std::endl;
|
||||
}
|
||||
|
||||
return stream;
|
||||
}
|
||||
|
||||
// Return time indexes of a given batch.
|
||||
// i-th returned element is the time of batch_id's forward at stage i.
|
||||
// previous_forward_time[s] is the time that the last forward happens on stage s.
|
||||
std::vector<int> PipelineScheduler::FindForwardComputeTime(const std::vector<int> previous_forward_time) const {
|
||||
// forward_time[i]: i-th stage's forward time of batch_id.
|
||||
std::vector<int> forward_time(num_stages_, 0);
|
||||
|
||||
// Insert forward.
|
||||
for (int s = 0; s < num_stages_; ++s) {
|
||||
for (int t = 0; t < required_max_time; ++t) {
|
||||
if (!table_[t][s].IsEmpty()) {
|
||||
for (int t = previous_forward_time[s]; t < static_cast<int>(compute_table_.size()); ++t) {
|
||||
if (!compute_table_[t][s].IsEmpty()) {
|
||||
// One slot cannot be occupied by two batches.
|
||||
continue;
|
||||
}
|
||||
|
||||
if (s > 0 && t <= forward_time[s - 1]) {
|
||||
// Foward of the s-th stage must happen after Forward of (s-1)-th stage.
|
||||
// Foward of the s-th stage must happen after forward of (s-1)-th stage.
|
||||
// Note that forward_time[s] is the time slot of the s-th stage.
|
||||
continue;
|
||||
}
|
||||
|
||||
if (batch_count_[t] >= num_stages_) {
|
||||
if (compute_batch_count_[t] >= num_stages_) {
|
||||
// At time t, the number of running batches is at maximum,
|
||||
// so we need to put this stage to another time slot.
|
||||
continue;
|
||||
}
|
||||
|
||||
// The s-th stage happens at time t.
|
||||
// For batch_id, its forward happens at time t on the s-th stage.
|
||||
forward_time[s] = t;
|
||||
// This s-th stage is forward pass of the batch_id-th batch.
|
||||
table_[t][s].type = Slot::Type::Forward;
|
||||
table_[t][s].batch_id = batch_id;
|
||||
|
||||
auto events = SearchLastRecordedEvents(t, s);
|
||||
table_[t][s].waited_events = events;
|
||||
|
||||
if (events.size() != 0) {
|
||||
table_[t][s].recorded_events = std::vector<int>{events[events.size() - 1] + 1,
|
||||
events[events.size() - 1] + 2};
|
||||
} else {
|
||||
table_[t][s].recorded_events = std::vector<int>{0, 1};
|
||||
}
|
||||
|
||||
for (int t_ = t + 1; t_ < required_max_time; ++t_) {
|
||||
if (table_[t_][s].IsEmpty()) {
|
||||
continue;
|
||||
}
|
||||
// Find the non-empty slot happens right before this slot.
|
||||
events = SearchLastRecordedEvents(t_, s);
|
||||
// Wait previously recorded events in the slot right before this slot table_[t_][s].
|
||||
table_[t_][s].waited_events = events;
|
||||
// Generate two new unique events; one for RecordEvent before Send and the other one
|
||||
// for RecordEvent after Send.
|
||||
table_[t_][s].recorded_events = std::vector<int>{events[events.size() - 1] + 1,
|
||||
events[events.size() - 1] + 2};
|
||||
}
|
||||
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
// Insert backward.
|
||||
return forward_time;
|
||||
}
|
||||
|
||||
// Return time indexes of a given batch.
|
||||
// i-th returned element is the time of batch_id's backward at stage i.
|
||||
// forward_time[s] is the forward time for the given batch on stage s.
|
||||
std::vector<int> PipelineScheduler::FindBackwardComputeTime(const std::vector<int> forward_time) const {
|
||||
std::vector<int> backward_time(num_stages_, 0);
|
||||
// For a specific batch, the last stage has the earliest backward computation.
|
||||
// Thus, the first loop reversely scans stages.
|
||||
for (int s = num_stages_ - 1; s > -1; --s) {
|
||||
for (int t = 0; t < required_max_time; ++t) {
|
||||
if (!table_[t][s].IsEmpty()) {
|
||||
for (int t = forward_time[s] + 1; t < static_cast<int>(compute_table_.size()); ++t) {
|
||||
if (!compute_table_[t][s].IsEmpty()) {
|
||||
continue;
|
||||
}
|
||||
|
||||
|
|
@ -150,203 +219,419 @@ void PipelineSchedule::Add(int batch_id) {
|
|||
continue;
|
||||
}
|
||||
|
||||
if (t <= forward_time[num_stages_ - 1]) {
|
||||
continue;
|
||||
}
|
||||
|
||||
if (batch_count_[t] >= num_stages_) {
|
||||
if (compute_batch_count_[t] >= num_stages_) {
|
||||
continue;
|
||||
}
|
||||
|
||||
backward_time[s] = t;
|
||||
table_[t][s].type = Slot::Type::Backward;
|
||||
table_[t][s].batch_id = batch_id;
|
||||
|
||||
auto events = SearchLastRecordedEvents(t, s);
|
||||
table_[t][s].waited_events = events;
|
||||
|
||||
if (events.size() != 0) {
|
||||
table_[t][s].recorded_events = std::vector<int>{events[events.size() - 1] + 1,
|
||||
events[events.size() - 1] + 2};
|
||||
} else {
|
||||
table_[t][s].recorded_events = std::vector<int>{/* first event id */ 0};
|
||||
}
|
||||
|
||||
for (int t_ = t + 1; t_ < required_max_time; ++t_) {
|
||||
if (table_[t_][s].IsEmpty()) {
|
||||
continue;
|
||||
}
|
||||
events = SearchLastRecordedEvents(t_, s);
|
||||
table_[t_][s].waited_events = events;
|
||||
table_[t][s].recorded_events = std::vector<int>{events[events.size() - 1] + 1,
|
||||
events[events.size() - 1] + 2};
|
||||
}
|
||||
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
for (int t = forward_time[0]; t <= forward_time[num_stages_ - 1]; ++t) {
|
||||
++batch_count_[t];
|
||||
}
|
||||
|
||||
for (int t = backward_time[num_stages_ - 1]; t <= backward_time[0]; ++t) {
|
||||
++batch_count_[t];
|
||||
}
|
||||
return backward_time;
|
||||
}
|
||||
|
||||
void PipelineSchedule::Add(int batch_id_begin, int batch_id_end) {
|
||||
for (int i = batch_id_begin; i < batch_id_end; ++i) {
|
||||
Add(i);
|
||||
}
|
||||
}
|
||||
int PipelineScheduler::FindSendRecvTime(const int upstream_compute_time, const int upstream_stage, const int stage) const {
|
||||
int good_time = -1;
|
||||
// Search for a time to insert Recv and then Send in full table.
|
||||
// Recv is on slot's stage.
|
||||
// Send is on upstream slot's stage.
|
||||
for (int full_t = static_cast<int>(compute_commute_table_.size()) - 1; full_t > upstream_compute_time; --full_t) {
|
||||
bool is_good_time = true;
|
||||
for (int full_s = 0; full_s < num_stages_; ++full_s) {
|
||||
auto& candidate_slot = compute_commute_table_[full_t][full_s];
|
||||
|
||||
int PipelineSchedule::GetForwardWaitedEventId(int stage_id, int batch_id) const {
|
||||
std::vector<int> events = {-1, -1};
|
||||
for (size_t t = 0; t < table_.size(); ++t) {
|
||||
auto& slot = table_[t][stage_id];
|
||||
if (!slot.IsForward()) {
|
||||
continue;
|
||||
}
|
||||
if (slot.batch_id != batch_id) {
|
||||
continue;
|
||||
}
|
||||
events = slot.waited_events;
|
||||
}
|
||||
return events[0];
|
||||
}
|
||||
if (candidate_slot.HasCompute()) {
|
||||
is_good_time = false;
|
||||
break;
|
||||
}
|
||||
|
||||
int PipelineSchedule::GetForwardWaitedEventIdAfterRecv(int stage_id, int batch_id) const {
|
||||
std::vector<int> events = {-1, -1};
|
||||
for (size_t t = 0; t < table_.size(); ++t) {
|
||||
auto& slot = table_[t][stage_id];
|
||||
if (!slot.IsForward()) {
|
||||
continue;
|
||||
}
|
||||
if (slot.batch_id != batch_id) {
|
||||
continue;
|
||||
}
|
||||
events = slot.waited_events;
|
||||
}
|
||||
return events[1];
|
||||
}
|
||||
if (candidate_slot.HasRecvFrom(upstream_stage)) {
|
||||
is_good_time = false;
|
||||
break;
|
||||
}
|
||||
|
||||
int PipelineSchedule::GetForwardRecordedEventIdBeforeSend(int stage_id, int batch_id) const {
|
||||
std::vector<int> events = {-1, -1};
|
||||
for (size_t t = 0; t < table_.size(); ++t) {
|
||||
auto& slot = table_[t][stage_id];
|
||||
if (!slot.IsForward()) {
|
||||
continue;
|
||||
}
|
||||
if (slot.batch_id != batch_id) {
|
||||
continue;
|
||||
}
|
||||
events = slot.recorded_events;
|
||||
}
|
||||
return events[0];
|
||||
}
|
||||
|
||||
int PipelineSchedule::GetForwardRecordedEventId(int stage_id, int batch_id) const {
|
||||
std::vector<int> events = {-1, -1};
|
||||
for (size_t t = 0; t < table_.size(); ++t) {
|
||||
auto& slot = table_[t][stage_id];
|
||||
if (!slot.IsForward()) {
|
||||
continue;
|
||||
}
|
||||
if (slot.batch_id != batch_id) {
|
||||
continue;
|
||||
}
|
||||
events = slot.recorded_events;
|
||||
}
|
||||
return events[1];
|
||||
}
|
||||
|
||||
int PipelineSchedule::GetBackwardWaitedEventId(int stage_id, int batch_id) const {
|
||||
std::vector<int> events = {-1, -1};
|
||||
for (size_t t = 0; t < table_.size(); ++t) {
|
||||
auto& slot = table_[t][stage_id];
|
||||
if (!slot.IsBackward()) {
|
||||
continue;
|
||||
}
|
||||
if (slot.batch_id != batch_id) {
|
||||
continue;
|
||||
}
|
||||
events = slot.waited_events;
|
||||
}
|
||||
return events[0];
|
||||
}
|
||||
|
||||
int PipelineSchedule::GetBackwardWaitedEventIdAfterRecv(int stage_id, int batch_id) const {
|
||||
std::vector<int> events = {-1, -1};
|
||||
for (size_t t = 0; t < table_.size(); ++t) {
|
||||
auto& slot = table_[t][stage_id];
|
||||
if (!slot.IsBackward()) {
|
||||
continue;
|
||||
}
|
||||
if (slot.batch_id != batch_id) {
|
||||
continue;
|
||||
}
|
||||
events = slot.waited_events;
|
||||
}
|
||||
return events[1];
|
||||
}
|
||||
|
||||
int PipelineSchedule::GetBackwardRecordedEventIdBeforeSend(int stage_id, int batch_id) const {
|
||||
std::vector<int> events = {-1, -1};
|
||||
for (size_t t = 0; t < table_.size(); ++t) {
|
||||
auto& slot = table_[t][stage_id];
|
||||
if (!slot.IsBackward()) {
|
||||
continue;
|
||||
}
|
||||
if (slot.batch_id != batch_id) {
|
||||
continue;
|
||||
}
|
||||
events = slot.recorded_events;
|
||||
}
|
||||
return events[0];
|
||||
}
|
||||
|
||||
int PipelineSchedule::GetBackwardRecordedEventId(int stage_id, int batch_id) const {
|
||||
std::vector<int> events = {-1, -1};
|
||||
for (size_t t = 0; t < table_.size(); ++t) {
|
||||
auto& slot = table_[t][stage_id];
|
||||
if (!slot.IsBackward()) {
|
||||
continue;
|
||||
}
|
||||
if (slot.batch_id != batch_id) {
|
||||
continue;
|
||||
}
|
||||
events = slot.recorded_events;
|
||||
}
|
||||
return events[1];
|
||||
}
|
||||
|
||||
void PipelineSchedule::Show() const {
|
||||
const int num_slots = static_cast<int>(table_.size());
|
||||
|
||||
for (int s = 0; s < num_stages_; ++s) {
|
||||
for (int t = 0; t < num_slots; ++t) {
|
||||
table_[t][s].Show();
|
||||
if (t == num_slots - 1) {
|
||||
std::cout << std::endl;
|
||||
if (candidate_slot.HasRendTo(stage)) {
|
||||
is_good_time = false;
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
std::vector<int> PipelineSchedule::SearchLastRecordedEvents(int time_id, int stage_id) const {
|
||||
std::vector<int> events;
|
||||
for (int t = time_id - 1; t >= 0; --t) {
|
||||
auto& slot = table_[t][stage_id];
|
||||
if (slot.IsEmpty()) {
|
||||
if (!is_good_time) {
|
||||
continue;
|
||||
}
|
||||
|
||||
events = slot.recorded_events;
|
||||
good_time = full_t;
|
||||
break;
|
||||
}
|
||||
return good_time;
|
||||
}
|
||||
|
||||
return events;
|
||||
void PipelineScheduler::CreateFullSchedule() {
|
||||
for (int t = 0; static_cast<size_t>(t) < compute_table_.size(); ++t) {
|
||||
// The last stage is compute, so we append one slot for commute.
|
||||
if (t != 0) {
|
||||
compute_commute_table_.push_back(std::vector<PipelineSlot>(num_stages_));
|
||||
}
|
||||
|
||||
for (int s = 0; s < num_stages_; ++s) {
|
||||
// Read a slot from compute-only schedule.
|
||||
// We will build send-recv pair to connect this slot with its upstream slot.
|
||||
auto slot = compute_table_[t][s];
|
||||
|
||||
if (slot.IsEmpty()) {
|
||||
// No need to insert Send and Recv for empty compute slot.
|
||||
continue;
|
||||
}
|
||||
|
||||
// Ok, this slot is not empty but it must have exactly one action.
|
||||
if (slot.NumActions() != 1) {
|
||||
throw std::invalid_argument("In compute-only schedule, one slot can have only one action, whose type is Compute.");
|
||||
}
|
||||
|
||||
// Get the only action in this slot.
|
||||
const auto action = slot.GetFrontAction();
|
||||
|
||||
const int batch = action.batch;
|
||||
// Stage of upstream compute in compute-only table.
|
||||
const auto upstream_s = action.upstream_stage;
|
||||
// Time of upstream compute in compute-only table.
|
||||
const auto upstream_t = action.upstream_time;
|
||||
|
||||
if (upstream_s < 0 && upstream_t < 0) {
|
||||
// This operator has no upstream, so we don't need to create Send and Recv.
|
||||
continue;
|
||||
}
|
||||
|
||||
if (s == num_stages_ - 1 && action.IsBackward() && action.IsCompute()) {
|
||||
continue;
|
||||
}
|
||||
|
||||
// Upstream of "slot".
|
||||
const auto upstream_slot = compute_table_[upstream_t][upstream_s];
|
||||
|
||||
// Get the only action in upstream slot.
|
||||
const auto upstream_action = upstream_slot.GetFrontAction();
|
||||
|
||||
// Time of upstream compute in full table.
|
||||
const auto upstream_compute_time = upstream_action.full_table_time;
|
||||
|
||||
// Forward/backward information is sent by forward/backward Send.
|
||||
const PipelineTask::Pass recv_pass = action.IsForward() ? PipelineTask::Pass::Forward : PipelineTask::Pass::Backward;
|
||||
// Forward/backward information is received by forward/backward Send.
|
||||
const PipelineTask::Pass send_pass = upstream_action.IsForward() ? PipelineTask::Pass::Forward : PipelineTask::Pass::Backward;
|
||||
|
||||
// Find a time index to insert send-recv pair in the schedule with compute and commute actions.
|
||||
const int good_time = FindSendRecvTime(upstream_compute_time, upstream_s, s);
|
||||
|
||||
// Add Send and Recv to compute-commute schedule.
|
||||
// Send from upstream_s-th stage.
|
||||
// Recv at s-th stage.
|
||||
compute_commute_table_[good_time][upstream_s].AddSend(batch, send_pass, upstream_compute_time, upstream_s, upstream_s, s);
|
||||
compute_commute_table_[good_time][s].AddRecv(batch, recv_pass, good_time, s, s, upstream_s);
|
||||
}
|
||||
|
||||
// Actions in compute_table_[t] are going be copied to full schedule.
|
||||
// For each action, we store its actual time and responsding etage in compute-only schedule.
|
||||
// The stored information may be carried to full schedule.
|
||||
for (int s = 0; s < num_stages_; ++s) {
|
||||
auto slot = compute_table_[t][s];
|
||||
for (int a = 0; a < static_cast<int>(slot.NumActions()); ++a) {
|
||||
auto& task = slot[a];
|
||||
task.full_table_time = static_cast<int>(compute_commute_table_.size());
|
||||
task.full_table_stage = s;
|
||||
}
|
||||
}
|
||||
|
||||
// Copy compute actions from compute-only schedule to compute-commute schedule.
|
||||
compute_commute_table_.push_back(compute_table_[t]);
|
||||
}
|
||||
}
|
||||
|
||||
void PipelineScheduler::InsertEvents(std::vector<std::vector<PipelineSlot>>& schedule, const size_t num_events_per_slot, const std::vector<int> initial_events) {
|
||||
std::vector<std::vector<int>> last_recorded_events(num_stages_, initial_events);
|
||||
|
||||
for (int t = 0; static_cast<size_t>(t) < schedule.size(); ++t) {
|
||||
for (int s = 0; s < num_stages_; ++s) {
|
||||
if (schedule[t][s].IsEmpty()) {
|
||||
continue;
|
||||
}
|
||||
schedule[t][s].SetWaitedEvent(last_recorded_events[s]);
|
||||
|
||||
// Create new recorded events. Their indexes should be greater than those of previous events.
|
||||
const auto max_event = std::max_element(last_recorded_events[s].begin(), last_recorded_events[s].end());
|
||||
std::vector<int> new_recorded_events;
|
||||
for (int i = 0; static_cast<size_t>(i) < num_events_per_slot; ++i) {
|
||||
new_recorded_events.push_back(*max_event + i + 1);
|
||||
}
|
||||
|
||||
schedule[t][s].SetRecordedEvent(new_recorded_events);
|
||||
last_recorded_events[s] = schedule[t][s].GetRecordedEvent();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
void PipelineScheduler::InsertForwardCompute(const int batch_id, const std::vector<int> forward_time) {
|
||||
// Occupy the time slots so that these slots won't be used in later iterations.
|
||||
for (int s = 0; s < num_stages_; ++s) {
|
||||
const auto batch_forward_time = forward_time[s];
|
||||
if (s == 0) {
|
||||
// The first forward compute has no upstream action.
|
||||
compute_table_[batch_forward_time][s].AddCompute(batch_id, PipelineTask::Pass::Forward);
|
||||
} else {
|
||||
// For other cases, forward at stage s happens after forward at stage s - 1.
|
||||
compute_table_[batch_forward_time][s].AddCompute(batch_id, PipelineTask::Pass::Forward, forward_time[s - 1], s - 1);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
void PipelineScheduler::InsertBackwardCompute(const int batch_id, const std::vector<int> forward_time, const std::vector<int> backward_time) {
|
||||
// Occupy the time slots so that these slots won't be used in later iterations.
|
||||
const auto last_stage_index = num_stages_ - 1;
|
||||
for (int s = num_stages_ - 1; s >= 0; --s) {
|
||||
const auto batch_backward_time = backward_time[s];
|
||||
if (s == last_stage_index) {
|
||||
// The first backward (on the last pipeline stage) depends on the a forward on the last pipeline stage.
|
||||
compute_table_[batch_backward_time][s].AddCompute(batch_id, PipelineTask::Pass::Backward, forward_time[s], s);
|
||||
} else {
|
||||
// For other cases, backward at stage s depedns on the backward at stage s + 1.
|
||||
compute_table_[batch_backward_time][s].AddCompute(batch_id, PipelineTask::Pass::Backward, backward_time[s + 1], s + 1);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
void PipelineScheduler::CreateComputeSchedule() {
|
||||
// Expand table to accomonadate the new batch.
|
||||
const int compute_max_time = 2 * num_stages_ + 2 * (num_batches_ - 1);
|
||||
|
||||
compute_table_.resize(compute_max_time, std::vector<PipelineSlot>(num_stages_));
|
||||
compute_batch_count_.resize(compute_max_time);
|
||||
commute_batch_count_.resize(compute_max_time);
|
||||
|
||||
std::vector<int> forward_time(num_stages_, 0);
|
||||
std::vector<int> backward_time(num_stages_, 0);
|
||||
for (int batch_id = 0; batch_id < num_batches_; ++batch_id) {
|
||||
// Find slot to insert forward compute.
|
||||
// The search on stage[s] starts at time forward_time[s].
|
||||
forward_time = FindForwardComputeTime(forward_time);
|
||||
|
||||
// Insert forward Compute's at different stages based on forward_time found.
|
||||
InsertForwardCompute(batch_id, forward_time);
|
||||
|
||||
// Find slot to insert backward compute.
|
||||
// The search on stage[s] starts at time forward_time[s].
|
||||
backward_time = FindBackwardComputeTime(forward_time);
|
||||
|
||||
// Insert backward Compute's at different stages based on backward_time found.
|
||||
InsertBackwardCompute(batch_id, forward_time, backward_time);
|
||||
|
||||
// Increase the batch count for whenever that batch stays in the pipeline.
|
||||
// For a specific batch, its starting/end time is the time it enters/leaves the first pipeline stage.
|
||||
for (int t_compute = forward_time[0]; t_compute <= backward_time[0]; ++t_compute) {
|
||||
++compute_batch_count_[t_compute];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
std::vector<int> PipelineScheduler::TryGetEvent(
|
||||
const bool is_waited_event,
|
||||
const int batch_id,
|
||||
const int stage_id,
|
||||
const PipelineTask::Pass pass,
|
||||
const PipelineTask::Type type,
|
||||
bool& is_found) const {
|
||||
is_found = false;
|
||||
|
||||
// Go through slots of stage stage_id to find the requested action.
|
||||
for (int t = 0; static_cast<size_t>(t) < compute_commute_table_.size(); ++t) {
|
||||
const auto slot = compute_commute_table_[t][stage_id];
|
||||
for (int a = 0; static_cast<size_t>(a) < slot.NumActions(); ++a) {
|
||||
auto task = slot[a];
|
||||
if (task.batch != batch_id) {
|
||||
continue;
|
||||
}
|
||||
if (task.pass != pass) {
|
||||
continue;
|
||||
}
|
||||
if (task.type != type) {
|
||||
continue;
|
||||
}
|
||||
|
||||
// Slot of the asked action is found, so we return its event.
|
||||
is_found = true;
|
||||
return is_waited_event ? slot.GetWaitedEvent() : slot.GetRecordedEvent();
|
||||
}
|
||||
}
|
||||
|
||||
return std::vector<int>();
|
||||
}
|
||||
|
||||
int PipelineScheduler::GetEventOrDefault(
|
||||
const bool is_waited_event,
|
||||
const int batch_id,
|
||||
const int stage_id,
|
||||
const PipelineTask::Pass pass,
|
||||
const PipelineTask::Type type) const {
|
||||
bool is_found = false;
|
||||
auto events = TryGetEvent(is_waited_event, batch_id, stage_id, pass, type, is_found);
|
||||
if (is_found) {
|
||||
return events.front();
|
||||
} else {
|
||||
return -1;
|
||||
}
|
||||
}
|
||||
|
||||
// Forward Compute
|
||||
int PipelineScheduler::GetForwardComputeWaitedEvent(const int batch_id, const int stage_id) const {
|
||||
return GetEventOrDefault(true, batch_id, stage_id, PipelineTask::Pass::Forward, PipelineTask::Type::Compute);
|
||||
}
|
||||
|
||||
// Forward Compute
|
||||
int PipelineScheduler::GetForwardComputeRecordedEvent(const int batch_id, const int stage_id) const {
|
||||
return GetEventOrDefault(false, batch_id, stage_id, PipelineTask::Pass::Forward, PipelineTask::Type::Compute);
|
||||
}
|
||||
|
||||
// Backward Compute
|
||||
int PipelineScheduler::GetBackwardComputeWaitedEvent(const int batch_id, const int stage_id) const {
|
||||
return GetEventOrDefault(true, batch_id, stage_id, PipelineTask::Pass::Backward, PipelineTask::Type::Compute);
|
||||
}
|
||||
|
||||
// Backward Compute
|
||||
int PipelineScheduler::GetBackwardComputeRecordedEvent(const int batch_id, const int stage_id) const {
|
||||
return GetEventOrDefault(false, batch_id, stage_id, PipelineTask::Pass::Backward, PipelineTask::Type::Compute);
|
||||
}
|
||||
|
||||
// Forward Send.
|
||||
int PipelineScheduler::GetForwardSendWaitedEvent(const int batch_id, const int stage_id) const {
|
||||
return GetEventOrDefault(true, batch_id, stage_id, PipelineTask::Pass::Forward, PipelineTask::Type::Send);
|
||||
}
|
||||
|
||||
// Forward Send.
|
||||
int PipelineScheduler::GetForwardSendRecordedEvent(const int batch_id, const int stage_id) const {
|
||||
return GetEventOrDefault(false, batch_id, stage_id, PipelineTask::Pass::Forward, PipelineTask::Type::Send);
|
||||
}
|
||||
|
||||
// Backward Send.
|
||||
int PipelineScheduler::GetBackwardSendWaitedEvent(const int batch_id, const int stage_id) const {
|
||||
return GetEventOrDefault(true, batch_id, stage_id, PipelineTask::Pass::Backward, PipelineTask::Type::Send);
|
||||
}
|
||||
|
||||
// Backward Send.
|
||||
int PipelineScheduler::GetBackwardSendRecordedEvent(const int batch_id, const int stage_id) const {
|
||||
return GetEventOrDefault(false, batch_id, stage_id, PipelineTask::Pass::Backward, PipelineTask::Type::Send);
|
||||
}
|
||||
|
||||
// Forward Recv.
|
||||
int PipelineScheduler::GetForwardRecvWaitedEvent(const int batch_id, const int stage_id) const {
|
||||
return GetEventOrDefault(true, batch_id, stage_id, PipelineTask::Pass::Forward, PipelineTask::Type::Recv);
|
||||
}
|
||||
|
||||
// Forward Recv.
|
||||
int PipelineScheduler::GetForwardRecvRecordedEvent(const int batch_id, const int stage_id) const {
|
||||
return GetEventOrDefault(false, batch_id, stage_id, PipelineTask::Pass::Forward, PipelineTask::Type::Recv);
|
||||
}
|
||||
|
||||
// Backward Recv.
|
||||
int PipelineScheduler::GetBackwardRecvWaitedEvent(const int batch_id, const int stage_id) const {
|
||||
return GetEventOrDefault(true, batch_id, stage_id, PipelineTask::Pass::Backward, PipelineTask::Type::Recv);
|
||||
}
|
||||
|
||||
// Backward Recv.
|
||||
int PipelineScheduler::GetBackwardRecvRecordedEvent(const int batch_id, const int stage_id) const {
|
||||
return GetEventOrDefault(false, batch_id, stage_id, PipelineTask::Pass::Backward, PipelineTask::Type::Recv);
|
||||
}
|
||||
|
||||
std::vector<int> PipelineScheduler::TryGetComputeEvent(
|
||||
const int batch_id,
|
||||
const int stage_id,
|
||||
const PipelineTask::Pass pass,
|
||||
const PipelineTask::Type type,
|
||||
bool& is_found) const {
|
||||
is_found = false;
|
||||
|
||||
// Go through slots of stage stage_id to find the requested action.
|
||||
for (int t = 0; static_cast<size_t>(t) < compute_table_.size(); ++t) {
|
||||
const auto slot = compute_table_[t][stage_id];
|
||||
for (int a = 0; static_cast<size_t>(a) < slot.NumActions(); ++a) {
|
||||
auto task = slot[a];
|
||||
if (task.batch != batch_id) {
|
||||
continue;
|
||||
}
|
||||
if (task.pass != pass) {
|
||||
continue;
|
||||
}
|
||||
if (task.type != PipelineTask::Type::Compute) {
|
||||
// Slots presenting in the table must be either Compute or Empty because it's a compute-only schedule.
|
||||
continue;
|
||||
}
|
||||
|
||||
// Slot of the asked action is found, so we return its event.
|
||||
is_found = true;
|
||||
// Return two Waits' events or two Record's events for
|
||||
// Wait -> Recv -> Wait -> Compute -> Record -> Send -> Record
|
||||
return type == PipelineTask::Type::Recv ? slot.GetWaitedEvent() : slot.GetRecordedEvent();
|
||||
}
|
||||
}
|
||||
|
||||
return std::vector<int>();
|
||||
}
|
||||
|
||||
int PipelineScheduler::GetComputeEventOrDefault(
|
||||
const bool is_before,
|
||||
const int batch_id,
|
||||
const int stage_id,
|
||||
const PipelineTask::Pass pass,
|
||||
const PipelineTask::Type type) const {
|
||||
bool is_found = false;
|
||||
auto events = TryGetComputeEvent(batch_id, stage_id, pass, type, is_found);
|
||||
if (!is_found) {
|
||||
return -1;
|
||||
}
|
||||
|
||||
if (is_before) {
|
||||
return events.front();
|
||||
} else {
|
||||
return events.back();
|
||||
}
|
||||
}
|
||||
|
||||
int PipelineScheduler::GetForwardWaitedEventBeforeRecv(const int batch_id, const int stage_id) const {
|
||||
auto event = GetComputeEventOrDefault(true, batch_id, stage_id, PipelineTask::Pass::Forward, PipelineTask::Type::Recv);
|
||||
return event;
|
||||
}
|
||||
|
||||
int PipelineScheduler::GetForwardWaitedEventAfterRecv(const int batch_id, const int stage_id) const {
|
||||
auto event = GetComputeEventOrDefault(false, batch_id, stage_id, PipelineTask::Pass::Forward, PipelineTask::Type::Recv);
|
||||
return event;
|
||||
}
|
||||
|
||||
int PipelineScheduler::GetForwardRecordedEventBeforeSend(const int batch_id, const int stage_id) const {
|
||||
auto event = GetComputeEventOrDefault(true, batch_id, stage_id, PipelineTask::Pass::Forward, PipelineTask::Type::Send);
|
||||
return event;
|
||||
}
|
||||
|
||||
int PipelineScheduler::GetForwardRecordedEventAfterSend(const int batch_id, const int stage_id) const {
|
||||
auto event = GetComputeEventOrDefault(false, batch_id, stage_id, PipelineTask::Pass::Forward, PipelineTask::Type::Send);
|
||||
return event;
|
||||
}
|
||||
|
||||
int PipelineScheduler::GetBackwardWaitedEventBeforeRecv(const int batch_id, const int stage_id) const {
|
||||
auto event = GetComputeEventOrDefault(true, batch_id, stage_id, PipelineTask::Pass::Backward, PipelineTask::Type::Recv);
|
||||
return event;
|
||||
}
|
||||
|
||||
int PipelineScheduler::GetBackwardWaitedEventAfterRecv(const int batch_id, const int stage_id) const {
|
||||
auto event = GetComputeEventOrDefault(false, batch_id, stage_id, PipelineTask::Pass::Backward, PipelineTask::Type::Recv);
|
||||
return event;
|
||||
}
|
||||
|
||||
int PipelineScheduler::GetBackwardRecordedEventBeforeSend(const int batch_id, const int stage_id) const {
|
||||
auto event = GetComputeEventOrDefault(true, batch_id, stage_id, PipelineTask::Pass::Backward, PipelineTask::Type::Send);
|
||||
return event;
|
||||
}
|
||||
|
||||
int PipelineScheduler::GetBackwardRecordedEventAfterSend(const int batch_id, const int stage_id) const {
|
||||
auto event = GetComputeEventOrDefault(false, batch_id, stage_id, PipelineTask::Pass::Backward, PipelineTask::Type::Send);
|
||||
return event;
|
||||
}
|
||||
|
||||
void PipelineWorkerPool::Join(size_t worker_id) {
|
||||
|
|
|
|||
|
|
@ -16,78 +16,196 @@ namespace onnxruntime {
|
|||
namespace training {
|
||||
namespace pipeline {
|
||||
|
||||
struct Slot {
|
||||
enum Type { Empty,
|
||||
Forward,
|
||||
Backward };
|
||||
// Action in Slot.
|
||||
// One slot can contain multiple actions and these actions can
|
||||
// be computed in parallel.
|
||||
struct PipelineTask {
|
||||
// Types of Action.
|
||||
// Those types are things we can do in a time slot.
|
||||
enum class Type { Empty,
|
||||
Compute,
|
||||
Send,
|
||||
Recv };
|
||||
|
||||
Slot();
|
||||
bool IsEmpty() const;
|
||||
bool IsForward() const;
|
||||
bool IsBackward() const;
|
||||
void Show() const;
|
||||
// Passes of Action.
|
||||
// A pass is a meaningful sub-graph in the training graph.
|
||||
enum class Pass { Unknown,
|
||||
Forward,
|
||||
Backward };
|
||||
|
||||
bool IsForward() const { return pass == Pass::Forward; }
|
||||
bool IsBackward() const { return pass == Pass::Backward; }
|
||||
bool IsCompute() const { return type == Type::Compute; }
|
||||
bool IsSendTo(const int dst_rank) const {
|
||||
if (type != Type::Send) {
|
||||
return false;
|
||||
}
|
||||
return peer_rank == dst_rank;
|
||||
}
|
||||
|
||||
bool IsRecvFrom(const int src_rank) const {
|
||||
if (type != Type::Recv) {
|
||||
return false;
|
||||
}
|
||||
return peer_rank == src_rank;
|
||||
}
|
||||
|
||||
friend std::ostream& operator<<(std::ostream& stream, const PipelineTask& slot);
|
||||
|
||||
// Batch ID this action belongs to.
|
||||
// For the first batch's forward pass (Compute), this value is 0.
|
||||
// A negative value means undefined.
|
||||
int batch;
|
||||
Type type;
|
||||
int batch_id;
|
||||
Pass pass;
|
||||
|
||||
// If type is Forward,
|
||||
// waited_events[0]: The first (before forward Recv)
|
||||
// waited event in forward pass.
|
||||
// waited_events[1]: The second (after forward Recv)
|
||||
// waited event in forward pass.
|
||||
// If type is Backward,
|
||||
// waited_events[0]: The first (before backward Recv)
|
||||
// waited event in backward pass.
|
||||
// waited_events[1]: The second (after backward Recv)
|
||||
// waited event in backward pass.
|
||||
std::vector<int> waited_events;
|
||||
// The last upstream action' time in compute-only schedule.
|
||||
int upstream_time{-1};
|
||||
// The last upstream action' stage in compute-only schedule.
|
||||
int upstream_stage{-1};
|
||||
|
||||
// If type is Forward,
|
||||
// recorded_events[0]: The first (before forward Send)
|
||||
// recorded event in forward pass.
|
||||
// recorded_events[1]: The second (after forward Send)
|
||||
// recorded event in forward pass.
|
||||
// If type is Backward,
|
||||
// recorded_events[0]: The first (before backward Send)
|
||||
// recorded event in backward pass.
|
||||
// recorded_events[1]: The second (after backward Send)
|
||||
// recorded event in backward pass.
|
||||
std::vector<int> recorded_events;
|
||||
// The rank where this Action is executed on.
|
||||
int this_rank{-1};
|
||||
// The remote rank where this Action interact with.
|
||||
int peer_rank{-1};
|
||||
|
||||
// This action's scheduled time in compute-commute schedule.
|
||||
int full_table_time{-1};
|
||||
// This action's scheduled stage in compute-commute schedule.
|
||||
int full_table_stage{-1};
|
||||
};
|
||||
|
||||
class PipelineSchedule {
|
||||
class PipelineSlot {
|
||||
public:
|
||||
PipelineSchedule() = default;
|
||||
PipelineSchedule(int num_stages);
|
||||
void Add(int batch_id);
|
||||
void Add(int batch_id_begin, int batch_id_end);
|
||||
int GetForwardWaitedEventId(int stage_id, int batch_id) const;
|
||||
int GetForwardWaitedEventIdAfterRecv(int stage_id, int batch_id) const;
|
||||
int GetForwardRecordedEventIdBeforeSend(int stage_id, int batch_id) const;
|
||||
int GetForwardRecordedEventId(int stage_id, int batch_id) const;
|
||||
int GetBackwardWaitedEventId(int stage_id, int batch_id) const;
|
||||
int GetBackwardWaitedEventIdAfterRecv(int stage_id, int batch_id) const;
|
||||
int GetBackwardRecordedEventIdBeforeSend(int stage_id, int batch_id) const;
|
||||
int GetBackwardRecordedEventId(int stage_id, int batch_id) const;
|
||||
void Show() const;
|
||||
void AddSend(const int batch_id, const PipelineTask::Pass pass, const int upstream_time = -1, const int upstream_stage = -1, const int this_rank = -1, const int peer_rank = -1);
|
||||
void AddRecv(const int batch_id, const PipelineTask::Pass pass, const int upstream_time = -1, const int upstream_stage = -1, const int this_rank = -1, const int peer_rank = -1);
|
||||
void AddCompute(const int batch_id, const PipelineTask::Pass pass, const int upstream_time = -1, const int upstream_stage = -1);
|
||||
|
||||
bool IsEmpty() const { return tasks_.empty(); };
|
||||
size_t NumActions() const { return tasks_.size(); }
|
||||
bool HasCompute() const {
|
||||
for (auto& task : tasks_) {
|
||||
if (task.IsCompute())
|
||||
return true;
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
bool HasRendTo(const int stage) const {
|
||||
for (auto& task : tasks_) {
|
||||
if (task.IsSendTo(stage)) {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
bool HasRecvFrom(const int stage) const {
|
||||
for (auto& task : tasks_) {
|
||||
if (task.IsRecvFrom(stage)) {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
PipelineTask& operator[](int index);
|
||||
const PipelineTask& operator[](int index) const;
|
||||
PipelineTask& GetFrontAction();
|
||||
const PipelineTask& GetFrontAction() const;
|
||||
|
||||
// Print this structure following a fixed-length format.
|
||||
// It assumes there are at most 2 actions per slot.
|
||||
friend std::ostream& operator<<(std::ostream& stream, const PipelineSlot& slot);
|
||||
void SetWaitedEvent(const std::vector<int> event);
|
||||
std::vector<int> GetWaitedEvent() const;
|
||||
void SetRecordedEvent(const std::vector<int> event);
|
||||
std::vector<int> GetRecordedEvent() const;
|
||||
|
||||
private:
|
||||
std::vector<int> SearchLastRecordedEvents(int time_id, int stage_id) const;
|
||||
// Actions which can be executed in parallel in this time slot.
|
||||
std::vector<PipelineTask> tasks_;
|
||||
|
||||
// 2-D table of pipeline schedule. table_[i][j] is the computation happening in
|
||||
// For MPI PipeDream schedule, it's used to support Wait -> Recv -> Wait -> Compute -> Record -> Send -> Record.
|
||||
// Since Send, Recv, and Compute are stored in the same slot, each slot contains two waited events and two recorded events.
|
||||
//
|
||||
// For NCCL PipDream schedule, it's used to support Wait -> Recv -> Record -> Wait -> Compute -> Record -> Wait -> Send -> Record.
|
||||
|
||||
// Events waited by this slot.
|
||||
std::vector<int> waited_events_;
|
||||
// Events recorded by this slot.
|
||||
std::vector<int> recorded_events_;
|
||||
};
|
||||
|
||||
class PipelineScheduler {
|
||||
public:
|
||||
PipelineScheduler(const int num_batches, const int num_stages);
|
||||
size_t GetScheduleSize() const { return compute_commute_table_.size(); }
|
||||
size_t GetStageSize() const { return num_stages_; }
|
||||
// APIs to get NCCL event for
|
||||
// Wait -> Recv -> Record -> Wait -> Compute -> Record -> Wait -> Send -> Record.
|
||||
int GetForwardComputeWaitedEvent(const int batch_id, const int stage_id) const;
|
||||
int GetForwardComputeRecordedEvent(const int batch_id, const int stage_id) const;
|
||||
int GetBackwardComputeWaitedEvent(const int batch_id, const int stage_id) const;
|
||||
int GetBackwardComputeRecordedEvent(const int batch_id, const int stage_id) const;
|
||||
int GetForwardSendWaitedEvent(const int batch_id, const int stage_id) const;
|
||||
int GetForwardSendRecordedEvent(const int batch_id, const int stage_id) const;
|
||||
int GetBackwardSendWaitedEvent(const int batch_id, const int stage_id) const;
|
||||
int GetBackwardSendRecordedEvent(const int batch_id, const int stage_id) const;
|
||||
int GetForwardRecvWaitedEvent(const int batch_id, const int stage_id) const;
|
||||
int GetForwardRecvRecordedEvent(const int batch_id, const int stage_id) const;
|
||||
int GetBackwardRecvWaitedEvent(const int batch_id, const int stage_id) const;
|
||||
int GetBackwardRecvRecordedEvent(const int batch_id, const int stage_id) const;
|
||||
// APIs to get MPI event event for
|
||||
// Wait -> Recv -> Wait -> Compute -> Record -> Send -> Record.
|
||||
int GetForwardWaitedEventBeforeRecv(const int batch_id, const int stage_id) const;
|
||||
int GetForwardWaitedEventAfterRecv(const int batch_id, const int stage_id) const;
|
||||
int GetForwardRecordedEventBeforeSend(const int batch_id, const int stage_id) const;
|
||||
int GetForwardRecordedEventAfterSend(const int batch_id, const int stage_id) const;
|
||||
int GetBackwardWaitedEventBeforeRecv(const int batch_id, const int stage_id) const;
|
||||
int GetBackwardWaitedEventAfterRecv(const int batch_id, const int stage_id) const;
|
||||
int GetBackwardRecordedEventBeforeSend(const int batch_id, const int stage_id) const;
|
||||
int GetBackwardRecordedEventAfterSend(const int batch_id, const int stage_id) const;
|
||||
// Visualization of this object.
|
||||
friend std::ostream& operator<<(std::ostream& stream, PipelineScheduler const& schedule);
|
||||
|
||||
private:
|
||||
// Return time indexes of a given batch.
|
||||
// i-th returned element is the time of batch_id's forward at stage i.
|
||||
// previous_forward_time[s] is the time that the last forward happens on stage s.
|
||||
std::vector<int> FindForwardComputeTime(const std::vector<int> previous_forward_time) const;
|
||||
// Return time indexes of a given batch.
|
||||
// i-th returned element is the time of batch_id's backward at stage i.
|
||||
// forward_time[s] is the forward time for the given batch on stage s.
|
||||
std::vector<int> FindBackwardComputeTime(const std::vector<int> forward_time) const;
|
||||
void CreateComputeSchedule();
|
||||
void InsertEvents(std::vector<std::vector<PipelineSlot>>& schedule, const size_t num_events_per_slot, const std::vector<int> initial_events);
|
||||
void CreateFullSchedule();
|
||||
int FindSendRecvTime(const int upstream_compute_time, const int upstream_stage, const int stage) const;
|
||||
void InsertForwardCompute(const int batch_id, const std::vector<int> forward_time);
|
||||
void InsertBackwardCompute(const int batch_id, const std::vector<int> forward_time, const std::vector<int> backward_time);
|
||||
|
||||
// Given an action, this function finds its home slot and returns events of that slot.
|
||||
std::vector<int> TryGetEvent(const bool is_waited_event, const int batch_id, const int stage_id, const PipelineTask::Pass pass, const PipelineTask::Type type, bool& is_found) const;
|
||||
// Wrapper over TryGetEvent. It returns -1 when the specified action is not found.
|
||||
int GetEventOrDefault(const bool is_waited_event, const int batch_id, const int stage_id, const PipelineTask::Pass pass, const PipelineTask::Type type) const;
|
||||
|
||||
std::vector<int> TryGetComputeEvent(const int batch_id, const int stage_id, const PipelineTask::Pass pass, const PipelineTask::Type type, bool& is_found) const;
|
||||
int GetComputeEventOrDefault(const bool is_waited_event, const int batch_id, const int stage_id, const PipelineTask::Pass pass, const PipelineTask::Type type) const;
|
||||
|
||||
// Compute-only pipeline schedule as a 2-D table. table_[i][j] is the computation happening in
|
||||
// the i-th time slot at the j-th stage. For example, PipeDream schedule may have
|
||||
// 1. table_[0][0].batch_id is 0 and table_[0][0].type is Forward.
|
||||
// 2. table_[0][1].type is Empty, which means no computation.
|
||||
// 3. table_[1][0].batch_id is 1 and table_[1][0].type is Forward.
|
||||
std::vector<std::vector<Slot>> table_;
|
||||
// Number of active batches per time slot. batch_count_[i] is the number of active
|
||||
// batches at the i-th time slot.
|
||||
std::vector<int> batch_count_;
|
||||
// Total number of stages of this pipeline.
|
||||
// It equals to table_.size().
|
||||
std::vector<std::vector<PipelineSlot>> compute_table_;
|
||||
// Compute + commute pipeline schedule. Its format is the same as compute-only schedule.
|
||||
std::vector<std::vector<PipelineSlot>> compute_commute_table_;
|
||||
// Number of active batches per time slot. compute_batch_count_[i] is the number of active
|
||||
// compute batches at the i-th time slot.
|
||||
std::vector<int> compute_batch_count_;
|
||||
std::vector<int> commute_batch_count_;
|
||||
int num_stages_;
|
||||
// Total number of batches scheduled in this pipeline.
|
||||
// It equals to table_[i].size(), for i = 0, ..., num_stages_ - 1.
|
||||
int num_batches_;
|
||||
};
|
||||
|
||||
|
|
|
|||
|
|
@ -60,7 +60,7 @@ TrainingRunner::TrainingRunner(Parameters params, const Environment& env, Sessio
|
|||
session_options_(session_options),
|
||||
session_(session_options, env),
|
||||
input_allocator_(params.input_allocator ? params.input_allocator : TrainingUtil::GetCpuAllocator()),
|
||||
pipeline_schedule_(params_.pipeline_parallel_size),
|
||||
pipeline_schedule_(params.gradient_accumulation_steps, params_.pipeline_parallel_size),
|
||||
pipeline_worker_pool_(params_.pipeline_parallel_size) {
|
||||
ORT_ENFORCE(!params_.model_path.empty());
|
||||
if (!params.weights_to_train.empty())
|
||||
|
|
@ -281,7 +281,6 @@ Status TrainingRunner::Initialize() {
|
|||
// Configure dimension of this pipeline.
|
||||
pipeline_context_.pipeline_stage_id = config_result.pipeline_config_result.value().pipeline_stage_id;
|
||||
pipeline_context_.num_pipeline_batches = params_.gradient_accumulation_steps;
|
||||
pipeline_schedule_.Add(0, pipeline_context_.num_pipeline_batches);
|
||||
} else {
|
||||
fetch_names = params_.fetch_names;
|
||||
pipeline_context_.pipeline_stage_id = 0;
|
||||
|
|
@ -406,9 +405,9 @@ Status TrainingRunner::PrepareFeedNamesAndFeeds(const SessionMode mode,
|
|||
OrtValue event_id;
|
||||
const int64_t id =
|
||||
(mode == EvaluateStep) ? -1
|
||||
: pipeline_schedule_.GetForwardWaitedEventId(
|
||||
pipeline_context_.pipeline_stage_id,
|
||||
static_cast<int>(step_) % pipeline_context_.num_pipeline_batches);
|
||||
: pipeline_schedule_.GetForwardWaitedEventBeforeRecv(
|
||||
static_cast<int>(step_) % pipeline_context_.num_pipeline_batches,
|
||||
pipeline_context_.pipeline_stage_id);
|
||||
TrainingUtil::CreateCpuMLScalar(
|
||||
id,
|
||||
&event_id,
|
||||
|
|
@ -423,9 +422,9 @@ Status TrainingRunner::PrepareFeedNamesAndFeeds(const SessionMode mode,
|
|||
OrtValue event_id;
|
||||
const int64_t id =
|
||||
(mode == EvaluateStep) ? -1
|
||||
: pipeline_schedule_.GetForwardWaitedEventIdAfterRecv(
|
||||
pipeline_context_.pipeline_stage_id,
|
||||
static_cast<int>(step_) % pipeline_context_.num_pipeline_batches);
|
||||
: pipeline_schedule_.GetForwardWaitedEventAfterRecv(
|
||||
static_cast<int>(step_) % pipeline_context_.num_pipeline_batches,
|
||||
pipeline_context_.pipeline_stage_id);
|
||||
TrainingUtil::CreateCpuMLScalar(
|
||||
id,
|
||||
&event_id,
|
||||
|
|
@ -440,9 +439,9 @@ Status TrainingRunner::PrepareFeedNamesAndFeeds(const SessionMode mode,
|
|||
OrtValue event_id;
|
||||
const int64_t id =
|
||||
(mode == EvaluateStep) ? -1
|
||||
: pipeline_schedule_.GetForwardRecordedEventIdBeforeSend(
|
||||
pipeline_context_.pipeline_stage_id,
|
||||
static_cast<int>(step_) % pipeline_context_.num_pipeline_batches);
|
||||
: pipeline_schedule_.GetForwardRecordedEventBeforeSend(
|
||||
static_cast<int>(step_) % pipeline_context_.num_pipeline_batches,
|
||||
pipeline_context_.pipeline_stage_id);
|
||||
TrainingUtil::CreateCpuMLScalar(
|
||||
id,
|
||||
&event_id,
|
||||
|
|
@ -457,9 +456,9 @@ Status TrainingRunner::PrepareFeedNamesAndFeeds(const SessionMode mode,
|
|||
OrtValue event_id;
|
||||
const int64_t id =
|
||||
(mode == EvaluateStep) ? -1
|
||||
: pipeline_schedule_.GetForwardRecordedEventId(
|
||||
pipeline_context_.pipeline_stage_id,
|
||||
static_cast<int>(step_) % pipeline_context_.num_pipeline_batches);
|
||||
: pipeline_schedule_.GetForwardRecordedEventAfterSend(
|
||||
static_cast<int>(step_) % pipeline_context_.num_pipeline_batches,
|
||||
pipeline_context_.pipeline_stage_id);
|
||||
TrainingUtil::CreateCpuMLScalar(
|
||||
id,
|
||||
&event_id,
|
||||
|
|
@ -474,9 +473,9 @@ Status TrainingRunner::PrepareFeedNamesAndFeeds(const SessionMode mode,
|
|||
OrtValue event_id;
|
||||
const int64_t id =
|
||||
(mode == EvaluateStep) ? -1
|
||||
: pipeline_schedule_.GetBackwardWaitedEventId(
|
||||
pipeline_context_.pipeline_stage_id,
|
||||
static_cast<int>(step_) % pipeline_context_.num_pipeline_batches);
|
||||
: pipeline_schedule_.GetBackwardWaitedEventBeforeRecv(
|
||||
static_cast<int>(step_) % pipeline_context_.num_pipeline_batches,
|
||||
pipeline_context_.pipeline_stage_id);
|
||||
TrainingUtil::CreateCpuMLScalar(
|
||||
id,
|
||||
&event_id,
|
||||
|
|
@ -491,9 +490,9 @@ Status TrainingRunner::PrepareFeedNamesAndFeeds(const SessionMode mode,
|
|||
OrtValue event_id;
|
||||
const int64_t id =
|
||||
(mode == EvaluateStep) ? -1
|
||||
: pipeline_schedule_.GetBackwardWaitedEventIdAfterRecv(
|
||||
pipeline_context_.pipeline_stage_id,
|
||||
static_cast<int>(step_) % pipeline_context_.num_pipeline_batches);
|
||||
: pipeline_schedule_.GetBackwardWaitedEventAfterRecv(
|
||||
static_cast<int>(step_) % pipeline_context_.num_pipeline_batches,
|
||||
pipeline_context_.pipeline_stage_id);
|
||||
TrainingUtil::CreateCpuMLScalar(
|
||||
id,
|
||||
&event_id,
|
||||
|
|
@ -508,9 +507,9 @@ Status TrainingRunner::PrepareFeedNamesAndFeeds(const SessionMode mode,
|
|||
OrtValue event_id;
|
||||
int64_t id =
|
||||
(mode == EvaluateStep) ? -1
|
||||
: pipeline_schedule_.GetBackwardRecordedEventIdBeforeSend(
|
||||
pipeline_context_.pipeline_stage_id,
|
||||
static_cast<int>(step_) % pipeline_context_.num_pipeline_batches);
|
||||
: pipeline_schedule_.GetBackwardRecordedEventBeforeSend(
|
||||
static_cast<int>(step_) % pipeline_context_.num_pipeline_batches,
|
||||
pipeline_context_.pipeline_stage_id);
|
||||
TrainingUtil::CreateCpuMLScalar(
|
||||
id,
|
||||
&event_id,
|
||||
|
|
@ -525,9 +524,9 @@ Status TrainingRunner::PrepareFeedNamesAndFeeds(const SessionMode mode,
|
|||
OrtValue event_id;
|
||||
int64_t id =
|
||||
(mode == EvaluateStep) ? -1
|
||||
: pipeline_schedule_.GetBackwardRecordedEventId(
|
||||
pipeline_context_.pipeline_stage_id,
|
||||
static_cast<int>(step_) % pipeline_context_.num_pipeline_batches);
|
||||
: pipeline_schedule_.GetBackwardRecordedEventAfterSend(
|
||||
static_cast<int>(step_) % pipeline_context_.num_pipeline_batches,
|
||||
pipeline_context_.pipeline_stage_id);
|
||||
TrainingUtil::CreateCpuMLScalar(
|
||||
id,
|
||||
&event_id,
|
||||
|
|
@ -664,12 +663,13 @@ void TrainingRunner::RunWithUpdate(VectorString& feed_names,
|
|||
#endif
|
||||
RunOptions run_options;
|
||||
status = session_.Run(
|
||||
run_options,
|
||||
pipeline_worker_pool_.worker_states[worker_id].feed_names,
|
||||
pipeline_worker_pool_.worker_states[worker_id].feeds,
|
||||
pipeline_worker_pool_.worker_states[worker_id].fetch_names,
|
||||
&(pipeline_worker_pool_.worker_states[worker_id].fetches));
|
||||
}, worker_id, step_);
|
||||
run_options,
|
||||
pipeline_worker_pool_.worker_states[worker_id].feed_names,
|
||||
pipeline_worker_pool_.worker_states[worker_id].feeds,
|
||||
pipeline_worker_pool_.worker_states[worker_id].fetch_names,
|
||||
&(pipeline_worker_pool_.worker_states[worker_id].fetches));
|
||||
},
|
||||
worker_id, step_);
|
||||
|
||||
// Wait all workers to finish this round of pipeline parallelism.
|
||||
// The last batch in a pipeline collects gradient and update the model.
|
||||
|
|
@ -1152,7 +1152,7 @@ Status TrainingRunner::Evaluate(TrainingSession& session, IDataLoader& data_load
|
|||
const std::string training_mode_string = "training_mode";
|
||||
auto input_list = session.GetOverridableInitializers().second;
|
||||
for (auto input : *input_list) {
|
||||
if(input->Name().compare(training_mode_string) == 0) {
|
||||
if (input->Name().compare(training_mode_string) == 0) {
|
||||
feed_names.push_back("training_mode");
|
||||
OrtValue mode_val;
|
||||
TrainingUtil::CreateCpuMLScalar(false, &mode_val, input_allocator_);
|
||||
|
|
|
|||
|
|
@ -247,7 +247,7 @@ class TrainingRunner {
|
|||
// Information for running pipeline.
|
||||
pipeline::PipelineContext pipeline_context_;
|
||||
// Pipeline schedule for deciding when to run batch, forward, or backward.
|
||||
pipeline::PipelineSchedule pipeline_schedule_;
|
||||
pipeline::PipelineScheduler pipeline_schedule_;
|
||||
// Workers to run pipeline stage.
|
||||
pipeline::PipelineWorkerPool pipeline_worker_pool_;
|
||||
};
|
||||
|
|
|
|||
|
|
@ -0,0 +1,125 @@
|
|||
// Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
// Licensed under the MIT License.
|
||||
|
||||
#include "gtest/gtest.h"
|
||||
#include "orttraining/models/runner/pipeline.h"
|
||||
|
||||
namespace onnxruntime {
|
||||
namespace test {
|
||||
|
||||
void TestPipelineScheduler(const int num_batches, const int num_stages, std::vector<std::vector<int>> baseline_events) {
|
||||
onnxruntime::training::pipeline::PipelineScheduler schedule(num_batches, num_stages);
|
||||
for (int s = 0; s < num_stages; ++s) {
|
||||
for (int b = 0; b < num_batches; ++b) {
|
||||
const int forward_wait_before_recv = schedule.GetForwardWaitedEventBeforeRecv(b, s);
|
||||
const int forward_wait_after_recv = schedule.GetForwardWaitedEventAfterRecv(b, s);
|
||||
const int forward_record_before_send = schedule.GetForwardRecordedEventBeforeSend(b, s);
|
||||
const int forward_record_after_send = schedule.GetForwardRecordedEventAfterSend(b, s);
|
||||
const int backward_wait_before_recv = schedule.GetBackwardWaitedEventBeforeRecv(b, s);
|
||||
const int backward_wait_after_recv = schedule.GetBackwardWaitedEventAfterRecv(b, s);
|
||||
const int backward_record_before_send = schedule.GetBackwardRecordedEventBeforeSend(b, s);
|
||||
const int backward_record_after_send = schedule.GetBackwardRecordedEventAfterSend(b, s);
|
||||
|
||||
EXPECT_EQ(forward_wait_before_recv, baseline_events[2 * s + 0][4 * b + 0]) << " batch " << b << " stage " << s;
|
||||
EXPECT_EQ(forward_wait_after_recv, baseline_events[2 * s + 0][4 * b + 1]) << " batch " << b << " stage " << s;
|
||||
EXPECT_EQ(forward_record_before_send, baseline_events[2 * s + 0][4 * b + 2]) << " batch " << b << " stage " << s;
|
||||
EXPECT_EQ(forward_record_after_send, baseline_events[2 * s + 0][4 * b + 3]) << " batch " << b << " stage " << s;
|
||||
|
||||
EXPECT_EQ(backward_wait_before_recv, baseline_events[2 * s + 1][4 * b + 0]) << " batch " << b << " stage " << s;
|
||||
EXPECT_EQ(backward_wait_after_recv, baseline_events[2 * s + 1][4 * b + 1]) << " batch " << b << " stage " << s;
|
||||
EXPECT_EQ(backward_record_before_send, baseline_events[2 * s + 1][4 * b + 2]) << " batch " << b << " stage " << s;
|
||||
EXPECT_EQ(backward_record_after_send, baseline_events[2 * s + 1][4 * b + 3]) << " batch " << b << " stage " << s;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
TEST(Pipeline, ScheduleB8S3) {
|
||||
const int num_batches = 8;
|
||||
const int num_stages = 3;
|
||||
|
||||
// The event baselines at different stages are the same.
|
||||
// The first 4 events are for the first computation on that stage.
|
||||
// Similarly, the last 4 events are for the last computation on that stage.
|
||||
// Below, we add comments to indicate which computation the events associated with.
|
||||
|
||||
// Format per line below:
|
||||
// waited event before Recv, waited event after Recv, recorded event before Send, recorded event after Send.
|
||||
// The value "-1" means a NULL event; RecordEvent and WaitEvent do nothing for NULL events.
|
||||
// Note that the computation pattern is
|
||||
// WaitEvent -> Recv -> WaitEvent -> FW/BW -> RecordEvent -> Send -> RecordEvent.
|
||||
std::vector<int> forward_baseline_events_stage0{
|
||||
-1, -1, 0, 1, // FW00 @ stage 0
|
||||
0, 1, 2, 3, // FW01
|
||||
2, 3, 4, 5, // FW02
|
||||
6, 7, 8, 9, // FW03
|
||||
10, 11, 12, 13, // FW04
|
||||
14, 15, 16, 17, // FW05
|
||||
18, 19, 20, 21, // FW06
|
||||
22, 23, 24, 25, // FW07
|
||||
};
|
||||
|
||||
std::vector<int> backward_baseline_events_stage0{
|
||||
4, 5, 6, 7, // BW00
|
||||
8, 9, 10, 11, // BW01
|
||||
12, 13, 14, 15, // BW02
|
||||
16, 17, 18, 19, // BW03
|
||||
20, 21, 22, 23, // BW04
|
||||
24, 25, 26, 27, // BW05
|
||||
26, 27, 28, 29, // BW06
|
||||
28, 29, 30, 31 // BW07
|
||||
};
|
||||
|
||||
std::vector<int> forward_baseline_events_stage1{
|
||||
-1, -1, 0, 1, // FW00 @ stage 1
|
||||
0, 1, 2, 3, // FW01
|
||||
2, 3, 4, 5, // FW02
|
||||
8, 9, 10, 11, // FW03
|
||||
12, 13, 14, 15, // FW04
|
||||
16, 17, 18, 19, // FW05
|
||||
20, 21, 22, 23, // FW06
|
||||
24, 25, 26, 27, // FW07
|
||||
};
|
||||
|
||||
std::vector<int> backward_baseline_events_stage1{
|
||||
4, 5, 6, 7, // BW00
|
||||
6, 7, 8, 9, // BW01
|
||||
10, 11, 12, 13, // BW02
|
||||
14, 15, 16, 17, // BW03
|
||||
18, 19, 20, 21, // BW04
|
||||
22, 23, 24, 25, // BW05
|
||||
26, 27, 28, 29, // BW06
|
||||
28, 29, 30, 31 // BW07
|
||||
};
|
||||
|
||||
std::vector<int> forward_baseline_events_stage2{
|
||||
-1, -1, 0, 1, // FW00 @ stage 2
|
||||
2, 3, 4, 5, // FW01
|
||||
6, 7, 8, 9, // FW02
|
||||
10, 11, 12, 13, // FW03
|
||||
14, 15, 16, 17, // FW04
|
||||
18, 19, 20, 21, // FW05
|
||||
22, 23, 24, 25, // FW06
|
||||
26, 27, 28, 29, // FW07
|
||||
};
|
||||
|
||||
std::vector<int> backward_baseline_events_stage2{
|
||||
0, 1, 2, 3, // BW00
|
||||
4, 5, 6, 7, // BW01
|
||||
8, 9, 10, 11, // BW02
|
||||
12, 13, 14, 15, // BW03
|
||||
16, 17, 18, 19, // BW04
|
||||
20, 21, 22, 23, // BW05
|
||||
24, 25, 26, 27, // BW06
|
||||
28, 29, 30, 31 // BW07
|
||||
};
|
||||
|
||||
std::vector<std::vector<int>> baseline_events{
|
||||
forward_baseline_events_stage0, backward_baseline_events_stage0,
|
||||
forward_baseline_events_stage1, backward_baseline_events_stage1,
|
||||
forward_baseline_events_stage2, backward_baseline_events_stage2};
|
||||
|
||||
TestPipelineScheduler(num_batches, num_stages, baseline_events);
|
||||
}
|
||||
|
||||
} // namespace test
|
||||
} // namespace onnxruntime
|
||||
Loading…
Reference in a new issue