ReduceScatterVOperation

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

Computes the data on multiple communication cards, including addition, maximum, and minimum calculations, and sends the result to each card.

Application Scenarios

In the inference scenario, the batch size may not be exactly divided by the TP 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

For the original ReduceScatter, the amount of data on each card must be the same. For ReduceScatterV, assuming that there are five copies of data, three copies can be allocated to one card and two copies to another card.

Example:

Input

[[[0], [1], [2], [3], [4]], [[0], [1], [2], [3], [4]]]
[[2, 3], [2, 3]]
[[0, 2], [0, 2]]
[2, 3]

Output
Output of NPU 0:
```
[[0], [2]]
```
Output of NPU 1:
```
[[4], [6], [8]]
```

Definition

struct ReduceScatterVParam {
    int rank = 0;
    int rankSize = 0;
    int rankRoot = 0;
    std::vector<int64_t> sendCounts;
    std::vector<int64_t> sdispls;
    std::int64_t recvCount = 0;
    std::string reduceType = "sum";
    HcclComm hcclComm = nullptr;
    CommMode commMode = COMM_MULTI_PROCESS;
    std::string backend = "hccl";
    std::string rankTableFile;
    std::string commDomain;
    uint8_t rsv[64] = {0};
};

Parameters

Member	Type	Default Value	Value Range	Mandatory or Not	Description
rank	int	0	[0, rankSize-1]	Yes	Communication ID of the current card
rankSize	int	0	-	Yes	Number of communication cards
rankRoot	int	0	[0, rankSize-1]	Yes	Primary communication rank
reduceType	string	"sum"	sum max min	Yes	Communication computation type. The value can be sum, max, or min. Atlas inference products support only sum.
backend	string	"hccl"	hccl	Yes	Communication computation type. Only hccl is supported.
hcclComm	HcclComm	nullptr	hccl	No	Pointer to the HCCL communicator. Currently, only hccl is supported. This parameter is reserved. 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	COMM_MULTI_PROCESS/COMM_MULTI_THREAD	No	Communication mode. For details, see "CommMode."
rankTableFile	string	None	-	No	Path of the configuration file for cluster information.
commDomain	string	None	-	No	Communicator name of a communication device group.
rsv[64]	uint8_t	{0}	[0]	No	Reserved

Input

Parameter	Dimension	Data Type	Format	Mandatory or Not	Description
x	[dim_0, dim_1]	Atlas inference products: float16/int8 Atlas A2 training products/Atlas A2 inference products: float16/int8/bf16	ND	Yes	Input tensor. The sizes of dim_0 and dim_1 are not limited.
sendCounts	[rankSize]	int64	ND	Yes	1D tensor. The length is equal to the number of rankSize cards. The value of each index in the tensor indicates the amount of data allocated to each card. For example, sendCounts[0] indicates the amount of data allocated to rank 0. The sum of values in the tensor is equal to dim_0 of x.
sdispls	[rankSize]	int64	ND	Yes	1D tensor. The length is equal to the number of rankSize cards. The value of each index in the tensor indicates the offset of the amount of data received from the corresponding index card number. sdispls [0] = n indicates that rank 0 starts to receive the amount of data specified by sendCounts[0] from the position whose offset relative to the input start position is n. The value in the tensor must be less than dim_0 of x.
recvCount	[1]	int64	ND	Yes	1D tensor. The dimension size is fixed to 1. recvCount indicates the amount of data received by the card. The value is the same as that of sendCounts.
y	[dim_0]	float16	ND	Yes	1D tensor, which is used to derive the shape of the output. The value is the same as that of recvCount.

Output

Parameter	Dimension	Data Type	Format	Mandatory or Not	Description
output	[dim_0, dim_1]	"hccl": float16/int8	ND	Yes	Output tensor. dim_0 is the shape of y. dim_0 is equal to recvCount.shape[rank], and dim_1 is equal to x.shape[1].

Constraints

rank, rankSize, and rankRoot must meet the following conditions:
- 0 ≤ rank < rankSize
- 0 ≤ rankRoot < rankSize

Parent topic: atb/infer_op_params.h