AllToAllVV2Operation

Applicable Products

Hardware	Supported
Atlas 350 accelerator card	x
Atlas A3 inference products / Atlas A3 training products	x
Atlas A2 training products / Atlas A2 inference products	√
Atlas training products	x
Atlas inference products	√
Atlas 200I/500 A2 inference products	x

Description

Sends data to all communication cards in the communicator (the data volume can be customized using parameters) and receives data from all communication cards (the data volume can be customized using parameters).

Figure 1 Operator input and output example

Computational process example (Python):

       
            # Compute goldTensor.
gold_outtensor = []
for j in range(len(recvout[rank])):
    gold_outtensor.append(tensorafters[j][senddisp[j][rank]:sendcount[j][rank] + senddisp[j][rank]])
gold_outtensor = [i for arr in gold_outtensor for i in arr]

Operator Context

Figure 2 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.

Scenarios

Send data to each node in many-to-many mode. AllToAll is an extension of AllGather. Compared with AllGather, different AllToAll nodes collect different data from a node.

Figure 3 Comparison between AllToAll and AllGather

It applies to model parallelism, matrix transposition in model parallelism, and matrix transposition from data parallelism to model parallelism.

Example:

>>> rank0 input
tensor([[0,1,2,3],
        [4,5,6,7]], device='npu:0')  shape[2,4]
>>> rank0 sendcount
tensor([2, 4], device='npu:0')  shape[2]
>>> rank0 sdispls
tensor([0, 2])
>>> rank0 recvCounts
tensor([2, 2])
>>> rank0 rdispls
tensor([0, 2])
>>> rank0 tensorForInferShape
tensor([3, 3, 3, 3])

>>> rank1 input
tensor([[0, 1, 2], 
        [3, 4, 6]], device='npu:1')  shape[2,3]
>>> rank1 sendcount
tensor([2, 1], device='npu:1')  shape[2]
>>> rank1 sdispls
tensor([0, 2])
>>> rank1 recvCounts
tensor([4, 1])
>>> rank1 rdispls
tensor([0, 4])
>>> rank1 tensorForInferShape
tensor([3, 3, 3, 3, 3])

>>> rank0 output
tensor([[0, 1, 0, 1]], device='npu:0')  shape[1,4]
>>> rank1 output
tensor([[2, 3, 4, 5, 2]], device='npu:1')  shape[1,5]

Definition

      
           struct AllToAllVV2Param {
    int rank = -1;
    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[64] = {0};
};

Parameters

Member	Type	Default	Description
rank	int32	-1	Communication ID of the current card.
rankSize	int32	0	Number of communication cards. The value cannot be 0.
rankRoot	int32	0	Primary communication ID.
backend	std::string	hccl	Communication computation type. Only hccl is supported.
hcclComm	HcclComm	nullptr	Address pointer obtained by the HCCL communicator API. By default, this parameter is left blank. The ATB is created by users. If the user wants to manage the communicator, the communicator pointer needs to be passed. The ATB uses the passed communicator pointer to execute the communication operator.
commMode	CommMode	COMM_MULTI_PROCESS	Communication mode. For details, see "CommMode." In the case of HCCL multi-threading, a communicator can only be imported externally. enum CommMode : int { COMM_UNDEFINED = -1, //!< Not defined. COMM_MULTI_PROCESS, //!< Specifies multi-process communication. COMM_MULTI_THREAD, //!< Specifies multi-thread communication. };
rankTableFile	std::string	-	Path of the configuration file for cluster information, which is used for multi-node communication. Currently, only the HCCL backend is supported. Path of the configuration file for cluster information, which applies to single-node and multi-node communication scenarios. Currently, only the HCCL backend is supported. If a ranktable is configured for a single-node system, the ranktable is used to initialize the communicator. For details, see .
commDomain	std::string	-	Communicator name used by a communication device group when there are multiple communicators. Currently, only HCCL is supported.
rsv[64]	uint8_t	{0}	Reserved.

Input

Parameter	Dimension	Data Type	Format	Mandatory	Description
x	[dim_0, dim_1, ..., dim_n]	"hccl": float16/int8	ND	Yes	Input tensor.
sendCount	1[rankSize]	int64	ND	No	Array of the sent data volume, which is a tensor on the host. For example, if the send data type is float16, sendCounts[i] = n indicates that the current rank sends n pieces of float16 data to rank_i.
sdispls	1[rankSize]	int64	ND	No	Array of the send offset, which is a tensor on the host. sdispls[i] = n indicates that the current rank sends data to rank_i from a position whose offset relative to the input start position is n.
recvCounts	1[rankSize]	int64	ND	No	Array of the received data volume, which is a tensor on the host. For example, if the send data type is float16, recvCounts[i] = n indicates that the current rank receives n pieces of float16 data from rank_i.
rdispls	1[rankSize]	int64	ND	No	Array of the received offset, which is a tensor on the host. rdispls[i] = n indicates that the current rank starts to receive data of rank_i from a position whose offset relative to the input start position is n.
tensorForInferShape	[recvCountsSum]	int8	ND	No	The shape is the sum of all elements of recvCounts and is used to infer the shape.

Output

Parameter	Dimension	Data Type	Format	Mandatory	Description
output	[1, recvCountsSum]	hccl: float16/int8	ND	No	Output tensor. The shape of the last dimension is the sum of all elements of the recvCounts parameter. The data type is the same as that of the input.

Constraints

In the multi-device scenario, the rank must be specified.
sendCounts, recvCounts, sdispls, and rdispls can be considered as one-dimensional arrays with the length of rankSize, and the values of all elements in the arrays are greater than or equal to 0.
The sum of sendCounts and recvCounts array elements cannot overflow int64. For any element in rdispls, recvCounts[i] + rdispls[i] cannot be greater than recvCountsSum (the last dimension of the output). For any element in sdispls, sendCounts[i] + sdispls[i] cannot be greater than the data volume of the input tensor. For example, if the input tensor has shape [3, 4, 5] and the data amount is 60 (3 × 4 × 5), the input tensor will be considered as a tensor with shape [60] for AllToAllV computation.

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
```

Parent topic: atb/infer_op_params.h