Fix interleaved 1f1b race

[H-Huang]

Current fixes:

• We need to delay the send in the warmup phase right before 1f1b phase so that it is sent along with the recvs in the next step. This is to align the communication between ranks to prevent hang
• Sorted_batch_isend_irecv() is not working for interleaved 1f1b since we need to ensure that recv / send across different peers are all part of the same coalesced op. So we remove its usage from interleaved 1f1b.

---

Test added:

When I add sleep for certain ranks then there is a failure on rank 0 [Rank 0] Some NCCL operations have failed or timed out. Due to the asynchronous nature of CUDA kernels, subsequent GPU operations might run on corrupted/incomplete data.

Run with:
PIPPY_VERBOSITY=DEBUG pytest test/test_pipeline_schedule.py -vsk test_interleaved_1f1b_with_model_sleep

Stack from ghstack (oldest at bottom):

• -> Fix interleaved 1f1b race #1088

[kwen2501]

Thanks for the fix! I haven't run through the entire change but one small comment:
maybe try not to reuse the name ops too much. It could lead to bugs if other devs did not use it carefully.

[kwen2501]

Nice work.

[kwen2501]

What does group_size mean here? Number of PP ranks? If so, is [len(stages) - 1] * group_size a square of # PP ranks?
And what does "hop" mean?

[H-Huang]

Yep group_size is pp ranks. And the warmup steps is going to be a multiple of the pp ranks. Hop means for every rank less than the highest. The comment is a bit old so i will update it.

[kwen2501]

oh, so len(stages) means the number of local stages?

[kwen2501]

I like this comment. To actuate it, we can add a has_skip_conn filed in pipe_info, so that the schedule can check this field and error out here. (Rather than silently hang)