[QUESTION]Pipe-offload: Why is B phase scheduled between OFFLOAD RECV PREP & START ops in schedule?

[Butterfingrz]

When I run Pipe-offload and print the scheduling graph, I don't understand. Why, when running Pipe-offload, can a B phase be scheduled between the OFFLOAD RECV PREP and OFFLOAD RECV START operations in the generated scheduling graph?

[ufotalent]

Hi @2395959141 , Thanks for the interest in our work. Theoretically B has dependency on offload_recv_start (which does resume). I assume in the schedule you print the passes are from different microbatch. Maybe can print the microbach number and have a check?
Thanks

[Butterfingrz]

Hello, I'm still a bit confused. On the scheduling diagram for a single stage that was printed, both OFFLOAD_RECV_PREP and OFFLOAD_RECV_START have a start time of 83.19, and OFFLOAD_RECV_PREP comes before B. Does this mean I can also place OFFLOAD_RECV_START before B?

[ufotalent]

Hello, I'm still a bit confused. On the scheduling diagram for a single stage that was printed, both OFFLOAD_RECV_PREP and OFFLOAD_RECV_START have a start time of 83.19, and OFFLOAD_RECV_PREP comes before B. Does this mean I can also place OFFLOAD_RECV_START before B?

@2395959141 My suspect is that the B you printed and OFFLOAD_RECV_PREP OFFLOAD_RECV_START are passes from different microbatch. Say OFFLOAD_RECV_PREP and OFFLOAD_RECV_START are from microbatch 2 but B is from microbatch 1. Then that's totally possible for a B for a earlier microbatch happens before offload operations for a following microbatch.

[Lhongpei]

Hi, I have a similar question. In my opinion, each device process ScheduleNode one by one from schedule order of this device. However, when I print the schedule in device 0, I also find that OFFLOAD_RECV_PREP and OFFLOAD_RECV_START are not adjacent? And I check the code in offloading.py, I find that the order of OFFLOAD_RECV_PREP are compared with the completeion time of schedule nodes while OFFLOAD_RECV_START are compared with start time of nodes. Shouldn't I pass OFFLOAD_RECV_PREP immediately after OFFLOAD_RECV_START?

[ufotalent]

Hi, I have a similar question. In my opinion, each device process ScheduleNode one by one from schedule order of this device. However, when I print the schedule in device 0, I also find that OFFLOAD_RECV_PREP and OFFLOAD_RECV_START are not adjacent? And I check the code in offloading.py, I find that the order of OFFLOAD_RECV_PREP are compared with the completeion time of schedule nodes while OFFLOAD_RECV_START are compared with start time of nodes. Shouldn't I pass OFFLOAD_RECV_PREP immediately after OFFLOAD_RECV_START?

@Lhongpei OFFLOAD_RECV_START does not necessarily come immediately after OFFLOAD_RECV_PREP. The purpose of OFFLOAD_RECV_PREP is to allocate memory for resuming GPU tensors and happens on the compute stream while OFFLOAD_RECV_START happens on the offload stream. It's possible that there're some time difference between OFFLOAD_RECV_PREP and OFFLOAD_RECV_START. So that's why we generally order by the time of each passes.

For example, in this case, the scheduling order can be RECV_PREP, B, SEND_START, RECV_START

Another thing is that in my understanding, whether we schedule B or RECV_START first is actually not very important in this case, as there're on different streams and there're no dependency between this B and RECV_START.

Hope this answers your questions well.

[Lhongpei]

Hi, I have a similar question. In my opinion, each device process ScheduleNode one by one from schedule order of this device. However, when I print the schedule in device 0, I also find that OFFLOAD_RECV_PREP and OFFLOAD_RECV_START are not adjacent? And I check the code in offloading.py, I find that the order of OFFLOAD_RECV_PREP are compared with the completeion time of schedule nodes while OFFLOAD_RECV_START are compared with start time of nodes. Shouldn't I pass OFFLOAD_RECV_PREP immediately after OFFLOAD_RECV_START?

@Lhongpei OFFLOAD_RECV_START does not necessarily come immediately after OFFLOAD_RECV_PREP. The purpose of OFFLOAD_RECV_PREP is to allocate memory for resuming GPU tensors and happens on the compute stream while OFFLOAD_RECV_START happens on the offload stream. It's possible that there're some time difference between OFFLOAD_RECV_PREP and OFFLOAD_RECV_START. So that's why we generally order by the time of each passes.

For example, in this case, the scheduling order can be RECV_PREP, B, SEND_START, RECV_START

Another thing is that in my understanding, whether we schedule B or RECV_START first is actually not very important in this case, as there're on different streams and there're no dependency between this B and RECV_START.

Hope this answers your questions well.

Thank you for your reply! I understand offload should on different stream, but I notice that all operation are processed in TrainingIteration.run() in runtime.py, where a for-loop are used to iterate schedule node and schedule_b is still on the main stream in this class, TrainingIteration, I think. Isn't it means that OFFLOAD_RECV_START can only be processed after B?

[ufotalent]

Hi, I have a similar question. In my opinion, each device process ScheduleNode one by one from schedule order of this device. However, when I print the schedule in device 0, I also find that OFFLOAD_RECV_PREP and OFFLOAD_RECV_START are not adjacent? And I check the code in offloading.py, I find that the order of OFFLOAD_RECV_PREP are compared with the completeion time of schedule nodes while OFFLOAD_RECV_START are compared with start time of nodes. Shouldn't I pass OFFLOAD_RECV_PREP immediately after OFFLOAD_RECV_START?

@Lhongpei OFFLOAD_RECV_START does not necessarily come immediately after OFFLOAD_RECV_PREP. The purpose of OFFLOAD_RECV_PREP is to allocate memory for resuming GPU tensors and happens on the compute stream while OFFLOAD_RECV_START happens on the offload stream. It's possible that there're some time difference between OFFLOAD_RECV_PREP and OFFLOAD_RECV_START. So that's why we generally order by the time of each passes.

For example, in this case, the scheduling order can be RECV_PREP, B, SEND_START, RECV_START
Another thing is that in my understanding, whether we schedule B or RECV_START first is actually not very important in this case, as there're on different streams and there're no dependency between this B and RECV_START.
Hope this answers your questions well.

Thank you for your reply! I understand offload should on different stream, but I notice that all operation are processed in TrainingIteration.run() in runtime.py, where a for-loop are used to iterate schedule node and schedule_b is still on the main stream in this class, TrainingIteration, I think. Isn't it means that OFFLOAD_RECV_START can only be processed after B?

in this case, the scheduling time of OFFLOAD_RECV_START is after B. But however, because they are on different streams and there're no dependency (implemented as cuda events) between them, there're no strict requirement on who runs first on GPU. They might run concurrently.

[Lhongpei]

Hi, I have a similar question. In my opinion, each device process ScheduleNode one by one from schedule order of this device. However, when I print the schedule in device 0, I also find that OFFLOAD_RECV_PREP and OFFLOAD_RECV_START are not adjacent? And I check the code in offloading.py, I find that the order of OFFLOAD_RECV_PREP are compared with the completeion time of schedule nodes while OFFLOAD_RECV_START are compared with start time of nodes. Shouldn't I pass OFFLOAD_RECV_PREP immediately after OFFLOAD_RECV_START?

@Lhongpei OFFLOAD_RECV_START does not necessarily come immediately after OFFLOAD_RECV_PREP. The purpose of OFFLOAD_RECV_PREP is to allocate memory for resuming GPU tensors and happens on the compute stream while OFFLOAD_RECV_START happens on the offload stream. It's possible that there're some time difference between OFFLOAD_RECV_PREP and OFFLOAD_RECV_START. So that's why we generally order by the time of each passes.

For example, in this case, the scheduling order can be RECV_PREP, B, SEND_START, RECV_START
Another thing is that in my understanding, whether we schedule B or RECV_START first is actually not very important in this case, as there're on different streams and there're no dependency between this B and RECV_START.
Hope this answers your questions well.

Thank you for your reply! I understand offload should on different stream, but I notice that all operation are processed in TrainingIteration.run() in runtime.py, where a for-loop are used to iterate schedule node and schedule_b is still on the main stream in this class, TrainingIteration, I think. Isn't it means that OFFLOAD_RECV_START can only be processed after B?

in this case, the scheduling time of OFFLOAD_RECV_START is after B. But however, because they are on different streams and there're no dependency (implemented as cuda events) between them, there're no strict requirement on who runs first on GPU. They might run concurrently.

So, in the inplementation, OFFLOAD_RECV_START can be launched, immediately after B is launched? But I'm not sure where you use CUDA Event in function schedule_b .

[ufotalent]

I think general guidance here is that there's no requirement on the launch order between two streams which has no dependencies. E.g. the B and OFFLOAD_RECV_START we discuss here which are from two different microbatches.
On the other hand, to synchronize dependencies in the same microbatch, we use cuda event. For example, in B, when autograd asks for an activation tensor, we'll wait on the event that makes sure the corresponding resume is complete.

zero-bubble-pipeline-parallelism/megatron/core/pipeline_parallel/offload.py

Line 217 in 39f2c18

self._resume_event.wait()