Ring-based decomposition for Allgather+GEMM overlap ATen implementation

[nsarka]

Implementation using ATen of https://docs.google.com/document/d/1Fzr9Zs2Dqfj3e4yR8LKxFrRqC1EkMUfQczJMYQGJQUI/edit?tab=t.0#heading=h.5x7hptdjzhet

[wujingyue]

Did you nsys and did that show overlapping?

[wujingyue]

LGTM otherwise

[nsarka]

Did you nsys and did that show overlapping?

@wujingyue I ran an experiment that showed some overlap, but it seems like I need to spend more time (and with more detailed profiling instrumentation) to understand it fully.

[samnordmann]

Did you nsys and did that show overlapping?

@wujingyue I ran an experiment that showed some overlap, but it seems like I need to spend more time (and with more detailed profiling instrumentation) to understand it fully.

When ready, please share a screenshot of the night profile here for the record. Even if the observed overlap is not perfect, just seeing some overlap might be enough to validate the algo and merge the pr ; then, it will be non-trivial to fully understand how to get the best performance

[nsarka]

Here is a screenshot. This is with M=K=N=2^12 on sjc-dc2-dgx1-2 (4x V100)

[nsarka]

This is a better screenshot:

The nameless ...'s are nccl sendrecv, then gemm. You can see the sendrecv is fully overlapped with the gemm in a few cases. The profile still looks different than the profile in the Nemo Megatron Parallelization Techniques document, though, which is not what I was expecting

[nsarka]

After swapping the order that comm and compute are posted, the overlap looks almost perfect:

[wujingyue]

After swapping the order that comm and compute are posted, the overlap looks almost perfect:

It's hard to tell from the figure which blocks are gemm and which are allgather -- they all look blue boxes starting with "void". But I believe what you said!

[nsarka]

After swapping the order that comm and compute are posted, the overlap looks almost perfect:

It's hard to tell from the figure which blocks are gemm and which are allgather -- they all look blue boxes starting with "void". But I believe what you said!

Thanks! Sorry--the ones that say "void" are GEMM kernels, the smaller boxes are all ncclSendrecv

[nsarka]

!build