AllToAllOperation

Description

Sends data of the same amount (the input is divided into ranksize parts) to all communication cards in the communicator, and receives data of the same amount from all communication cards.

Operator Context

Figure 1 Position of AllToAll in communication INT8 quantization

Input the quantized tensor X of int8 and send it to each node through AlltoAll. (In this example, int8 is used for communication, which improves the communication speed.)

Then, the sum operation of reduce is used to sum up x, the int8 data is dequantized to the float16 data, and AllGather is used for communication to transmit the calculation result to each node. (In the requirement diagram, the quantization before and after AllGather may cause accuracy drop, so it is not implemented currently.)

Operator Function Implementation

Sends data of the same amount (the input is divided into ranksize parts) to all communication cards in the communicator, and receives data of the same amount from all communication cards. As shown in the following figure, when ranksize=4, each communication card evenly divides its data into ranksize parts and sends the data to other communication cards.

Figure 2 Functions of the AllToAll operator

Computational process (Python):

goldenTensors = []
for i in range(rankSize):
    golden_out = []
    for j in range(rankSize):
        golden_out_list = intensors[j].reshape(-1).tolist() # intensors[j] indicates the intesor on the jth card, which is considered as a one-dimensional vector during computation.
        split = golden_out_list[i*len(golden_out_list) // rankSize:(i+1)*len(golden_out_list) // rankSize]
        golden_out += split
    golden_out_tensor = torch.tensor(golden_out).reshape(shape) # golden_out_tensor is the output tensor of the ith card. The output shape is the same as the input tensor.
    goldenTensors.append(golden_out_tensor)
return [goldenTensors[rank].cpu()]

Application Scenarios

Send data to each node in many-to-many mode. It is an extension of All-Gather, but different nodes collect different data from a node.

Figure 3 Comparison between AllToAll and AllGather

Applies to model parallelism, matrix transposition in model parallelism, and matrix transposition from data parallelism to model parallelism.

Definition

struct AllToAllParam {
    int rank = 0;
    int rankSize = 0;
    int rankRoot = 0;
    std::string backend = "hccl";
    HcclComm hcclComm = nullptr;
    CommMode commMode = COMM_MULTI_PROCESS;
    std::string rankTableFile;
    std::string commDomain;
    uint8_t rsv[16] = {0};
};

Parameters

Input

Output

Restrictions

• Currently, only the Atlas A2 inference products is supported.
• The input tensor dimensions of each communication card are the same.
• During AllToAll computation, the input tensor is considered as a 1D vector, and the output shape is the same as that of the input tensor.
• rank, rankSize, and rankRoot must meet the following conditions:
  • 0 ≤ rank < rankSize
  • 0 ≤ rankRoot < rankSize

• If there are multiple users, ATB_SHARE_MEMORY_NAME_SUFFIX (see ATB Environment Variables) needs to be used to distinguish the shared memory and synchronize initialization information.
• When the communication operator of the ATB exits abnormally, run the following commands to clear the residual data to avoid affecting subsequent operations.

  rm -rf /dev/shm/sem.lccl*
  rm -rf /dev/shm/sem.hccl*
  ipcrm -a