diff --git a/cmake/onnxruntime_unittests.cmake b/cmake/onnxruntime_unittests.cmake index 3467db2895..13298f984b 100644 --- a/cmake/onnxruntime_unittests.cmake +++ b/cmake/onnxruntime_unittests.cmake @@ -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) diff --git a/onnxruntime/core/framework/sequential_executor.cc b/onnxruntime/core/framework/sequential_executor.cc index f9f71d8a07..805f669ea3 100644 --- a/onnxruntime/core/framework/sequential_executor.cc +++ b/onnxruntime/core/framework/sequential_executor.cc @@ -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(); diff --git a/orttraining/orttraining/models/runner/pipeline.cc b/orttraining/orttraining/models/runner/pipeline.cc index 793ddefb98..9f69447513 100644 --- a/orttraining/orttraining/models/runner/pipeline.cc +++ b/orttraining/orttraining/models/runner/pipeline.cc @@ -3,12 +3,14 @@ #include "orttraining/models/runner/pipeline.h" +#include #include #include #include #include #include #include +#include #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 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 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(table_.size()); +void PipelineSlot::SetWaitedEvent(const std::vector events) { + waited_events_ = std::move(events); +} - for (int t = current_max_time; t < required_max_time; ++t) { - table_.push_back(std::vector(num_stages_)); - batch_count_.push_back(0); +std::vector PipelineSlot::GetWaitedEvent() const { + return waited_events_; +} + +void PipelineSlot::SetRecordedEvent(const std::vector events) { + recorded_events_ = std::move(events); +} + +std::vector 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 PipelineScheduler::FindForwardComputeTime(const std::vector previous_forward_time) const { + // forward_time[i]: i-th stage's forward time of batch_id. std::vector 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(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{events[events.size() - 1] + 1, - events[events.size() - 1] + 2}; - } else { - table_[t][s].recorded_events = std::vector{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{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 PipelineScheduler::FindBackwardComputeTime(const std::vector forward_time) const { std::vector 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(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{events[events.size() - 1] + 1, - events[events.size() - 1] + 2}; - } else { - table_[t][s].recorded_events = std::vector{/* 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{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(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 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 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 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 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 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 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 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 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(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 PipelineSchedule::SearchLastRecordedEvents(int time_id, int stage_id) const { - std::vector 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(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(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(slot.NumActions()); ++a) { + auto& task = slot[a]; + task.full_table_time = static_cast(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>& schedule, const size_t num_events_per_slot, const std::vector initial_events) { + std::vector> last_recorded_events(num_stages_, initial_events); + + for (int t = 0; static_cast(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 new_recorded_events; + for (int i = 0; static_cast(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 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 forward_time, const std::vector 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(num_stages_)); + compute_batch_count_.resize(compute_max_time); + commute_batch_count_.resize(compute_max_time); + + std::vector forward_time(num_stages_, 0); + std::vector 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 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(t) < compute_commute_table_.size(); ++t) { + const auto slot = compute_commute_table_[t][stage_id]; + for (int a = 0; static_cast(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 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 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(t) < compute_table_.size(); ++t) { + const auto slot = compute_table_[t][stage_id]; + for (int a = 0; static_cast(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 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) { diff --git a/orttraining/orttraining/models/runner/pipeline.h b/orttraining/orttraining/models/runner/pipeline.h index 6464a78906..9d40bc8d3f 100644 --- a/orttraining/orttraining/models/runner/pipeline.h +++ b/orttraining/orttraining/models/runner/pipeline.h @@ -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 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 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 event); + std::vector GetWaitedEvent() const; + void SetRecordedEvent(const std::vector event); + std::vector GetRecordedEvent() const; private: - std::vector SearchLastRecordedEvents(int time_id, int stage_id) const; + // Actions which can be executed in parallel in this time slot. + std::vector 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 waited_events_; + // Events recorded by this slot. + std::vector 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 FindForwardComputeTime(const std::vector 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 FindBackwardComputeTime(const std::vector forward_time) const; + void CreateComputeSchedule(); + void InsertEvents(std::vector>& schedule, const size_t num_events_per_slot, const std::vector 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 forward_time); + void InsertBackwardCompute(const int batch_id, const std::vector forward_time, const std::vector backward_time); + + // Given an action, this function finds its home slot and returns events of that slot. + std::vector 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 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> 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 batch_count_; - // Total number of stages of this pipeline. - // It equals to table_.size(). + std::vector> compute_table_; + // Compute + commute pipeline schedule. Its format is the same as compute-only schedule. + std::vector> 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 compute_batch_count_; + std::vector 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_; }; diff --git a/orttraining/orttraining/models/runner/training_runner.cc b/orttraining/orttraining/models/runner/training_runner.cc index a6c1f61fb2..f97b72d0aa 100644 --- a/orttraining/orttraining/models/runner/training_runner.cc +++ b/orttraining/orttraining/models/runner/training_runner.cc @@ -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(step_) % pipeline_context_.num_pipeline_batches); + : pipeline_schedule_.GetForwardWaitedEventBeforeRecv( + static_cast(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(step_) % pipeline_context_.num_pipeline_batches); + : pipeline_schedule_.GetForwardWaitedEventAfterRecv( + static_cast(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(step_) % pipeline_context_.num_pipeline_batches); + : pipeline_schedule_.GetForwardRecordedEventBeforeSend( + static_cast(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(step_) % pipeline_context_.num_pipeline_batches); + : pipeline_schedule_.GetForwardRecordedEventAfterSend( + static_cast(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(step_) % pipeline_context_.num_pipeline_batches); + : pipeline_schedule_.GetBackwardWaitedEventBeforeRecv( + static_cast(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(step_) % pipeline_context_.num_pipeline_batches); + : pipeline_schedule_.GetBackwardWaitedEventAfterRecv( + static_cast(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(step_) % pipeline_context_.num_pipeline_batches); + : pipeline_schedule_.GetBackwardRecordedEventBeforeSend( + static_cast(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(step_) % pipeline_context_.num_pipeline_batches); + : pipeline_schedule_.GetBackwardRecordedEventAfterSend( + static_cast(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_); diff --git a/orttraining/orttraining/models/runner/training_runner.h b/orttraining/orttraining/models/runner/training_runner.h index 65432fe5c2..51ef943ba0 100644 --- a/orttraining/orttraining/models/runner/training_runner.h +++ b/orttraining/orttraining/models/runner/training_runner.h @@ -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_; }; diff --git a/orttraining/orttraining/test/distributed/pipeline_schedule_test.cc b/orttraining/orttraining/test/distributed/pipeline_schedule_test.cc new file mode 100644 index 0000000000..568a7f4fe1 --- /dev/null +++ b/orttraining/orttraining/test/distributed/pipeline_schedule_test.cc @@ -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> 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 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 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 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 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 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 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> 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 \ No newline at end of file