AllGatherVOperation

Applicable Products

Hardware Model	Supported or Not
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

Aggregates data of multiple communication cards in the first dimension according to a sequence of communication numbers and sends data to each card. The length of the data sent by each card can be different.

In the inference scenario, the batch size may not be exactly divided by the DP number. The subsequent compute operators of ReduceScatter need to process data based on the batch dimension and then perform AllGather on the processed data, as shown in Figure 1.

Figure 1 Operator context diagram

Example:

Figure 2 Computational process

Computational process (Python):

# Compute the goldtensor.
gold_outtensor = []
for i in range(len(sendcount)):
    gold_outtensor= gold_outtensor+(tensorafters[i][0:sendcount[i]])
GoldenTensors = (torch.tensor(np.array(gold_outtensor+[0]*(sum*dim[1]-len(gold_outtensor))).reshape(sum,dim[1]), dtype=inTensorDtype))

Application Scenarios

In model parallelism, parameters in forward computation are fully synchronized. Forward computation can be performed only after AllGatherV is used to synchronize the parameters split to different XPUs to one XPU.

>>> rank0 input
tensor([[0,1,2,3],
        [4,5,6,7]], device='npu:0')  shape[2,4]
>>> rank0 sendcount
tensor([4], device='npu:0')  shape[1]
>>> rank1 input
tensor([[3,2,1,0],
        [7,6,5,4],
        [7,6,5,4]], device='npu:1')  shape[3,4]
>>> rank1 sendcount
tensor([2], device='npu:1')  shape[1]
>>> recvout=tensor([4,2])
>>> recvdis=tensor([0,4])
>>> y=tensor([0,1,2,3,4], device='npu:0')
>>> rank0 output
tensor([[0,1,2,3],
        [3,2,0,0],
        [0,0,0,0],
        [0,0,0,0],
        [0,0,0,0]], device='npu:0')  shape[5,4]
>>> rank1 output
tensor([[0,1,2,3],
        [3,2,0,0],
        [0,0,0,0],
        [0,0,0,0],
        [0,0,0,0]], device='npu:1')  shape[5,4]

Definition

struct AllGatherVParam {
    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 Value	Description
rank	int	-1	Communication ID of the current card The value -1 indicates that the value is not passed.
rankSize	int	0	Number of communication cards. The value cannot be 0.
rankRoot	int	0	Primary communication rank
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 scenario where hccl multi-thread is applied, a communicator can only be imported externally. COMM_UNDEFINED: undefined. COMM_MULTI_PROCESS: multi-process communication. COMM_MULTI_THREAD: multi-thread communication.
rankTableFile	std::string	-	Path of the configuration file for cluster information.
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 or Not	Description
x	[dim_0, dim_1, ..., dim_n]	hccl: float16/int8/bfloat16	ND	Yes	Input tensor
sendCount	[1]	int64	ND	Yes	Input tensor, which is the amount of data sent by the current card. The tensor of each card can be different, that is, the data length can be different.
recvCounts	[ranksize]	int64	ND	Yes	Input tensor, which is the amount of data received from the corresponding index card number. The value is the same for each card.
rdispls	[ranksize]	int64	ND	Yes	Offset of the amount of data received from the corresponding index card number. The value is the same for each card. 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.
y	[Sum of the shape of the first dimension of each card]	float16	ND	Yes	The shape is the sum of the first shapes of the merged tensors of all cards, which is used for shape inference.

Output

Parameter	Dimension	Data Type	Format	Mandatory or Not	Description
output	[n, dim_1, ..., dim_n]	hccl: float16/int8/bfloat16	ND	No	Output tensor, which is not written in place. The address of the output tensor is different from that of the input tensor. n is the sum of the first-dimensional shapes of the tensors to be merged on all devices, that is, the shape of y. The data type is the same as that of the input.

Constraints

The length of the array parameter recvCounts or rdispls is equal to rankSize.
sendCount of the card rank is equal to recvCounts[rank].
The shape of the first dimension of the output tensor is equal to the shape of y.
It is supported only by Atlas inference products and Atlas A2 training products/Atlas A2 inference products.
The 0th dimension of the input tensor of each communication card can be different, and other dimensions must be the same.
The sum of recvCounts array elements cannot overflow int64, and the sum of recvCounts arrays must be greater than 0.
The value of sendCount cannot exceed the product of the shapes of all dimensions of the input tensor x.
rank, rankSize, and rankRoot must meet the following conditions:
- 0 ≤ rank < rankSize
- 0 ≤ rankRoot < rankSize

Parent topic: atb/infer_op_params.h