NCCL Communication C++ Example (Part 4): Implementation and Analysis of AlltoAll_Split

In NCCL cluster communication operations, there is no direct alltoall communication API, nor is there an alltoall_split method. This article mainly introduces the principle of alltoall_split and one implementation method, along with profiling of the send/recv operator calls in the alltoall operation.

Other references:

• NCCL Communication C++ Example (Part 1): Basic Use Case Interpretation and Execution

• NCCL Communication C++ Example (Part 2): Establishing Multi-Machine Connections with Sockets

• NCCL Communication C++ Example (Part 3): Concurrent Communication with Multiple Streams (Non-blocking)

01

Principle of Operation

First, let’s look at the general process of an alltoall operation: suppose there are n ranks (cluster communication processes). For a single rank, it needs to divide its data into n equal parts and send them sequentially to each rank while receiving data sent from other ranks. The following diagram illustrates the alltoall operation with 4 ranks, where the data length is set to 1 for simplicity. The operation process in the diagram is somewhat similar to a matrix transposition operation.

Of course, the number of data items transmitted in an alltoall operation does not need to equal the number of ranks; it can be a multiple of the ranks. The following is an example with a data length that is double the number of ranks:

If the data is split not according to equal parts but according to a specified array, it constitutes an alltoall_split operation. Still with n=4 and a data length of 8, if rank 0 splits the data according to the lengths [2, 3, 2, 1] and sends this data to ranks 0 to 3, as shown below, it achieves non-equal length splitting and sending.

We construct a split array for ranks 0 to 1, where each rank splits according to the corresponding index and sends the data to the specified rank, thus completing the entire alltoall process.

02

Code Implementation and Execution

• Main Function:

Implementation idea: Complete data sending and receiving through point-to-point communication APIs, while using group operations to “package” the data send/receive operations. The key code is as follows:

ncclResult_t AlltoAllSplit(const void *sendbuff, void *recvbuff, const size_t *sendSplitList,                           const size_t *recvSplitList, ncclDataType_t type, ncclComm_t comm, cudaStream_t stream){    int nRanks;    NCCLCHECK(ncclCommCount(comm, &nRanks));    size_t sendOffset = 0;    size_t recvOffset = 0;    NCCLCHECK(ncclGroupStart());    for (int r = 0; r < nRanks; r++) {        NCCLCHECK(ncclSend(((char *)sendbuff) + sendOffset, sendSplitList[r], type, r, comm, stream));        NCCLCHECK(ncclRecv(((char *)recvbuff) + recvOffset, recvSplitList[r], type, r, comm, stream));        sendOffset += wordSize(type) * sendSplitList[r];        recvOffset += wordSize(type) * recvSplitList[r];    }    NCCLCHECK(ncclGroupEnd());    return ncclSuccess;}

Here, sendOffset/recvOffset are the accumulated pointer indices. sendSplitList and recvSplitList are arrays allocated in advance based on the previous data splits. For example, with 4 ranks, suppose:

The split design is as follows: sendArray[4][4] = {{1, 2, 3, 4}, {4, 2, 3, 1}, {3, 2, 1, 4}, {2, 3, 4, 1}};

Then the shape of the received data would be: recvArray[4][4] = {{1, 4, 3, 2}, {2, 2, 2, 3}, {3, 3, 1, 4}, {4, 1, 4, 1}};

For the complete example, please refer to:https://github.com/CalvinXKY/BasicCUDA/blob/master/nccl/alltoall.cu

• Compilation and Debugging

For debugging convenience, we use an image that includes the NCCL library, specifically the NVIDIA ngc-pytorch image.

• Image Pull:

docker pull nvcr.io/nvidia/pytorch:24.07-py3

• Startup Reference:

sudo docker run  --net=host  --name megatron_test  --gpus=all -it    -e UID=root    --ipc host --shm-size="32g" 
-v /home/xky/:/home/xky 
-u 0 
--name=nccl nvcr.io/nvidia/pytorch:24.07-py3 bash

• Compilation:

It is recommended to compile within the image container, at the project location:BasicCUDA/nccl at master · CalvinXKY/BasicCUDA

cd BasicCUDA/nccl/make

To obtain the alltoall_test executable file.

• Execution Method:

./alltoall_test

You can see the following output:

====== AlltoAllSplit case begin =====
GPU:3 input data: 3 3 3 3 3 3 3 3 3 3
GPU:2 input data: 2 2 2 2 2 2 2 2 2 2
GPU:1 input data: 1 1 1 1 1 1 1 1 1 1
GPU:0 input data: 0 0 0 0 0 0 0 0 0 0
GPU:2 output data: 0 0 0 1 1 1 2 3 3 3 3
GPU:1 output data: 0 0 1 1 2 2 3 3 3
GPU:3 output data: 0 0 0 0 1 2 2 2 2 3
GPU:0 output data: 0 1 1 1 1 2 2 2 3 3
====== AlltoAllSplit case end =====

If you want to observe profiling, run it as follows:

nvprof --csv -o profile_output.csv ./alltoall_test

03

Profiling Observation

The kernel execution situation of alltoall is shown below. The following is the profiling data from running alltoall four times, where you can see:

• The send and recv operations within alltoall have been integrated into a single operation, rather than being scattered operations;

• There are time gaps between alltoall operations, mainly due to stream synchronization operations and the time taken for new kernel dispatch.

References:

• This article’s code: https://github.com/CalvinXKY/BasicCUDA/tree/master/nccl

• NCCL Use Cases:https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/examples.html

• NCCLtest:https://github.com/NVIDIA/nccl-tests

Scan the code to follow us and learn more about AI infrastructure basics.