FastSoftMaxGradOperation

Applicable Products

Product	Supported
Atlas 350 accelerator cards	x
Atlas A3 inference products / Atlas A3 training products	√
Atlas A2 training products / Atlas A2 inference products	√
Atlas training products	x
Atlas inference products	x
Atlas 200I/500 A2 inference products	x

Description

Training operator, which is used to calculate the reverse result of FastSoftMax. The data arrangement structure is the same as that of FastSoftMax.

$\text{[math]}$

Application Scenarios

Figure 1 Model flowchart

The FastSoftMaxGrad operator is used to calculate the backward result of FastSoftMax. The data arrangement structure is the same as that of FastSoftMax.

Figure 2 Arrangement of input data

Definition

      
           struct FastSoftMaxGradParam {
    int32_t headNum = 0;
    std::vector<int32_t> qSeqLen;
    uint8_t rsv[8] = {0};
};

Parameters

Member	Type	Default Value	Description
headNum	int32_t	0	Number of attention heads. The value must be greater than 0.
qSeqLen	std::vector<int32_t>	-	Actual input length of each batch. The number of elements is the same as the batch size. The maximum value is 32.
rsv[8]	uint8_t	{0}	Reserved

Input

Parameter	Dimension	Data Type	Format	Description
yInput	[nSquareTokens]	float16	ND	Input tensor. The forward computation result of the SoftMax operator is the result of the matrix of the size (headNum, qSeqLen[i], qSeqLen[i]), with the batch value range of (0, 1) and the sum of the last axis being 1. The result is concatenated into a one-dimensional result based on the ND layout.
yGrad	[nSquareTokens]	float16	ND	Gradient data transferred by the next operator. The layout is the same as that of yInput.

Output

Parameter	Dimension	Data Type	Format	Description
output	[nSquareTokens]	float16	ND	The layout is the same as that of yInput.

Constraints

The length of the qSeqLen array cannot exceed 32, and each element must be greater than 0.
The dimension size nSquareTokens of the input tensors (yInput and yGrad) and output tensor (outTensor) is related to headNum and qSeqLen in the parameters.
$\text{[math]}$

: indicates the kth element in qSeqLen.

Functions

When the parameter list contains

$\text{[math]}$ :

The shape size of the corresponding yInput is as follows:

$\text{[math]}$ .

For example,

$\text{[math]}$

In this case, the shape size of yInput is:

$\text{[math]}$

The shape of the input tensors yInput and yGrad, and output tensor is [5528144].

Parent topic: atb/train_op_params.h