AllGatherVOperation

Applicable Products

Hardware Model	Supported
Atlas 350 accelerator cards	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

Gathers data from multiple devices, concatenates the data along the dim 0 according to the order of their ranks, and then distributes the full result back to each device. Supports varying data lengths across different devices.

In inference tasks, 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

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

Use Cases

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	Rank of the current device within the communicator. -1 indicates that no value is passed.
rankSize	int	0	Number of devices in the communicator. Cannot be 0.
rankRoot	int	0	Root rank for the operation.
backend	std::string	hccl	Communication computation type. Only hccl is supported.
hcclComm	HcclComm	nullptr	Address pointer obtained by the HCCL communicator API. Defaults to null. 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. When HCCL multi-threading 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	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 device. The tensor of each device can be different, that is, the data length can vary.
recvCounts	[ranksize]	int64	ND	Yes	Input tensor, which is the amount of data received from the corresponding index. The value is the same for each device.
rdispls	[ranksize]	int64	ND	Yes	Offset of the amount of data received from the corresponding index. The value is the same for each device. 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 dim 0 size of each device]	float16	ND	Yes	The shape equals the sum of the first dimension's sizes (dim 0) of the concatenated tensor across all devices, which is used for shape inference.

Output

Parameter	Dimension	Data Type	Format	Mandatory	Description
output	[n, dim_1, ..., dim_n]	hccl: float16/int8/bfloat16	ND	No	Output tensor, which is not written in place. Its address differs from that of the input tensor. n is the sum of the first dimension's sizes (dim 0) of the concatenated tensor across 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 dim 0 of the input tensor of each device 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