KJT a2a redesign and input dist fusion (#797)

Summary:
Pull Request resolved: #797

# tldr:
* make it easier for users add new variants of the KeyedJaggedTensor (KJT)
* make it easier to automate communication for these new KJT variants
* improve QPS when models contain a high number of KJTs.

We expanded kjt interface to include these new calls:
* `dist_labels` -> names of tensors to transmit
* `dist_splits` -> shapes of internal tensors, split by key
* `dist_tensors` -> the actual tensor data to transmit
* `dist_init` -> builds a new kjt from raw collective output

These methods define the data the a2a collective will transmit from the KJT, and so the actual KJT a2a will work for KJT variants that properly implement these methods.

Next we changed the KJT a2a such that its first awaitable is to transmit the tensor splits so each rank will know the size of the tensors it will receive (this collective call is blocking). The second awaitable is asynchronous and it transmits the provided tensor data to the correct ranks. Previously it was an explicit awaitable for the lengths which was blocking, followed by an awaitable for the indices (values).

Finally we enabled input dist fusion by delaying calling wait() on the first awaitable in the KJT a2a until gathering all KJTs then calling the first wait in succession (becoming the cpu blocking part), and then calling the second awaitable to transmit the actual tensor data completely asynchronously.

Reviewed By: dstaay-fb, RenfeiChen-FB

Differential Revision:
D39520093

LaMa Project: L1138451

fbshipit-source-id: b7892a3af4e26c99ac01a88dd2c696e79e89cf42

Author: joshuadeng

torchrec/distributed/comm_ops.py

@@ -355,7 +355,7 @@ def alltoall_sequence(
             AlltoAll.
         input_splits (Tensor): input splits.
         output_splits (Tensor): output splits.
-        variable_batch_size (bool): whether varibale batch size is enabled
+        variable_batch_size (bool): whether variable batch size is enabled
         group (Optional[dist.ProcessGroup]): The process group to work on. If None, the
             default process group will be used.
         codecs: Optional[QuantizedCommCodecs]: Quantized communication codecs