Revise pipeline schedule to consider communication ops by wschin · Pull Request #4524 · microsoft/onnxruntime

wschin · 2020-07-15T20:53:37Z

NCCL is not as flexible as MPI. It requires user to schedule the communication to happen at the right time, so we extend our pipeline-parallel schedule to cover Send and Recv. The following two tables are visualizations for schedules w/wo Send/Recv.

Symbols in the tables:

FW00: forward compute of batch 00
FS00: forward send of batch 00
FR00: forward recv of batch 00
BW00: backward compute of batch 00
BS00: backward send of batch 00
BR00: backward recv of batch 00.

The class PipelineSchedule now includes two timeline tables, compute_table_ and compute_commute_table_. compute_table_ stores existing schedule for backward combability (we haven't replaced MPI with NCCL P2P yet) and compute_compute_table_ will be used later when NCCL P2P is introduced. Each schedule is a 2-D table. The slot at time t and stage s locates at table[t][s], which is a PipelineSlot. A PipelineSlot contains parallelly-computed operations (PipelineTasks) such as Forward and Send.

ke1337 · 2020-07-15T21:52:14Z

orttraining/orttraining/models/runner/pipeline.h

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


would it better to use one function to return a struct at given batch_id and stage_id? We might make it a template function and use struct to differentiate between NCCL and MPI. Or maybe have MPIPipelineScheduler/NCCLPipelineScheduler that inherit from the basic PipelineScheduler. #Resolved

Having MpiPipelineScheduler and NcclPipelineScheduler is doable (MPI scheduler is a special case of NCCL scheduler) but I'd like to first have NCCL P2P working and then come back to refactorize these classes again. Would this be ok to you?

In reply to: 455384433 [](ancestors = 455384433)

Sure.

In reply to: 455584382 [](ancestors = 455584382,455384433)

xzhu1900 · 2020-07-15T22:46:31Z

orttraining/orttraining/models/runner/pipeline.cc

-Slot::Slot() {
-  type = Empty;
-  batch_id = 0;
+bool Action::IsForward() const {


For small function, maybe we can try use inline? #Resolved

Good idea.

In reply to: 455407917 [](ancestors = 455407917)

xzhu1900 · 2020-07-16T22:18:56Z

orttraining/orttraining/test/distributed/pipeline_schedule_test.cc

+
+  std::vector<int> forward_baseline_events_stage2{
+      -1, -1, 0, 1,    // FW00 @ stage 2
+      2, 3, 4, 5,      // FW01


2, 3 [](start = 6, length = 4)

Does the event distinguish fw pass or bw pass? I guess I am not clear if this '2' comes from the forward pass in stage 1 (FW01) or the backward pass in stage 2 (BW00)

the variable name seems to indicate it's FW. event id itself is just an int and does not carry FW/BW semantics I think.

In reply to: 456110213 [](ancestors = 456110213)

Right. This Wait '2' is in the FW pass and it's waiting a Record '2' in other place to sent it the signal. But the Record '2' can be from the FW pass in stage 1 (FW01) or the BW pass in stage 2(BW00). So not sure which signal it is listening to.

In reply to: 456111879 [](ancestors = 456111879,456110213)

Yes, event ID doesn't define forward or backward. It's just a format of mine; I choose to split events based on forward/backward after events have been generated.

Wait 2 is on stage 2, which will be triggered by Record 2 on stage 2 (in the first row of backward_baseline_events_stage2).

First of all, different stages are run by different processes, so their events are independent event they have the same ID. That is, Event 0 on stage 0 has nothing to do with Event 0 on stage 1. The format of these 4 ints is

// Format per line below: // waited event before Recv, waited event after Recv, recorded event before Send, recorded event after Send.

It means 2 is the event waited before Recv in FW01 on stage 2. Similarly, 3, 4, and 5 sequentially maps to waited event after Recv, recorded event before Send, and recorded event after Send.

In reply to: 456124585 [](ancestors = 456124585,456111879,456110213)

ke1337

xzhu1900 · 2020-07-16T22:31:52Z

orttraining/orttraining/models/runner/pipeline.h

-  std::vector<int> recorded_events;
+  Pass pass;
+
+  // The last upstream action' time in compute-only schedule.


' [](start = 29, length = 1)

Nit: remove?

Same as below.

xzhu1900 · 2020-07-16T22:37:59Z

orttraining/orttraining/models/runner/pipeline.h

+  // 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.


Record [](start = 70, length = 6)

Out of curiosity, why NCCl needs an extra Record after each Recv and Send? #Resolved

For NCCL, we need more control on the timing of communication components. It means we need to know the start and end of each of Send, Recv, Forward, Backward.

In reply to: 456117125 [](ancestors = 456117125)

xzhu1900 · 2020-07-16T22:43:13Z

orttraining/orttraining/models/runner/pipeline.cc

-                << recorded_event_id_before_send << "," << recorded_event_id
-                << ")"
-                << " ";
+    case 2:


case 2: [](start = 4, length = 7)

To make it a bit more robust, maybe add a default and assert it will never have more than 2 actions?

Thank you. I will do it when refactorizing this class.

In reply to: 456118801 [](ancestors = 456118801)

Revise pipeline schedule to consider communication ops

359c249

wschin added the training issues related to ONNX Runtime training; typically submitted using template label Jul 15, 2020

wschin requested a review from a team as a code owner July 15, 2020 20:53

ke1337 reviewed Jul 15, 2020

View reviewed changes

xzhu1900 reviewed Jul 15, 2020

View reviewed changes

Add test

3b837df

wschin force-pushed the wechi/revise-pipeline-schedule branch from fa60ad0 to 3b837df Compare July 16, 2020 00:42

wschin added 2 commits July 16, 2020 01:27

Fix warning

a3ce298

inline some short functions

a4ec543

wschin force-pushed the wechi/revise-pipeline-schedule branch from a3b543e to a4ec543 Compare July 16, 2020 07:04

wschin added 5 commits July 16, 2020 07:51

Fix warnings

5a53330

Rename a class

36ee7fe

Add comment for test

e44ca70

op renamed to task

835b3f4

Fix NVTX wrapper's bug

81c90ac

xzhu1900 reviewed Jul 16, 2020

View reviewed changes

ke1337 approved these changes Jul 16, 2020

View reviewed changes

xzhu1900 reviewed Jul 16, 2020

View reviewed changes

wschin merged commit 21d2728 into master Jul 17, 2020

wschin deleted the wechi/revise-pipeline-schedule branch July 17, 2020 17:04

Conversation

wschin commented Jul 15, 2020 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

ke1337 Jul 15, 2020 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

wschin Jul 16, 2020 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

xzhu1900 Jul 15, 2020 • edited by wschin Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ke1337 left a comment

Choose a reason for hiding this comment

Uh oh!

xzhu1900 Jul 16, 2020 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

xzhu1900 Jul 16, 2020 • edited by wschin Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

wschin Jul 17, 2020 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

wschin commented Jul 15, 2020 •

edited

Loading

ke1337 Jul 15, 2020 •

edited

Loading

wschin Jul 16, 2020 •

edited

Loading

xzhu1900 Jul 15, 2020 •

edited by wschin

Loading

xzhu1900 Jul 16, 2020 •

edited

Loading

xzhu1900 Jul 16, 2020 •

edited by wschin

Loading

wschin Jul 17, 2020 •

edited

Loading