LayerNorm
Applicability
|
Product |
Supported |
|---|---|
|
Atlas 350 Accelerator Card |
√ |
|
|
√ |
|
|
√ |
|
|
x |
|
|
√ |
|
|
x |
|
|
x |
Function Usage
This section describes the following two LayerNorm APIs based on the API output.
- For the input data with the shape of [B, S, H], the normalized result, mean value, and variance are output.
During the training process of a deep neural network, updating the training parameters in the earlier layers can cause changes in the input data distribution for subsequent layers, resulting in unbalanced weight updates and reduced learning efficiency. Implementing the normalization policy to scale the input data of network layers to the [0, 1] range standardizes the distributions of both input and output data across network layers. This expedites the convergence of training parameters and ensures more stable improvements in learning efficiency. LayerNorm is one of many normalization methods.
This API implements LayerNorm normalization for input data with a shape of [B, S, H]. The calculation formula is as follows, where γ is the scale coefficient, β is the translation coefficient, and ε is the weight coefficient for preventing division by zero:

The following two parameters respectively represent the mean and variance of the input on the H axis:

- For the input data with the shape of [A, R], the variance, normalized result, mean value, and reciprocal of the standard deviation are output.
This API implements LayerNorm normalization for input data with a shape of [A, R]. The calculation formula is as follows, where γ is the scale coefficient, β is the translation coefficient, and ε is the weight coefficient for preventing division by zero:

The following three parameters respectively represent the mean, variance, and reciprocal of the standard deviation of the input on the R axis:

Principles
- For the input data with the shape of [B, S, H], the normalized result, mean value, and variance are output.
The figure below illustrates the internal algorithm block diagram of LayerNorm high-level APIs, taking the float type, ND format, and inputs inputX[B, S, H], gamma[H], and beta[H] as examples.
Figure 1 LayerNorm algorithm block diagram
The computation process is divided into the following steps, all of which are performed on vectors (m indicates the length of the last axis H):
- Calculate the mean: Muls calculates the value of x x 1/m, and then calculates the accumulated value ReduceSum to obtain the mean outputMean.
- Calculate the variance: Sub calculates the difference between input x and the mean, uses Mul to square the difference, multiplies Muls by 1/m, and calculates the accumulated value to obtain the variance outputVariance.
- Process gamma and beta: Obtain gamma and beta in the BSH dimension by broadcasting.
- Compute the output: Broadcast (or Duplicate) the variance to obtain the BSH-dimension tensor, which passes through Adds(outputVariance, eps), Ln, Muls, and Exp (or Sqrt) in sequence and is then multiplied by (x – mean). The obtained result is multiplied by gamma and added with beta to obtain the output result.
- For the input data with the shape of [A, R], the variance, normalized result, mean value, and reciprocal of the standard deviation are output.
The figure below illustrates the internal algorithm block diagram of LayerNorm high-level APIs, taking the float type, ND format, and inputs inputX[A, R], gamma[R], and beta[R] as examples.
Figure 2 LayerNorm-Rstd algorithm block diagram
The computation process is divided into the following steps, all of which are performed on vectors with the A axis being considered as the outermost loop.
- Compute the mean: Multiply each element of x by 1/(2k + m) to prevent overflow of subsequent accumulations. Then, sum up the data in binary accumulation mode: Split the data into a whole block and a tail block. The whole block contains 2k elements, and the tail block contains m elements. Add the data of the tail block to the data of the whole block. For ease of description, Vnum is defined as the number of elements participating in a single computation. Perform Vadd on odd and even bits in the whole block by Vnum to obtain a result of Vnum length. Perform WholeReduceSum on the result to obtain the output mean.
- Calculate rstd: Sub calculates the difference between input x and the mean, and Mul squares the difference. To prevent overflow, use the same dichotomy accumulation method to calculate Variance of the squared result. Add the variance to the coefficient for preventing division by zero ε, and calculate the output rstd by using Rsqrt.
- Calculate the output: Sub calculates the difference between input x and the mean. Multiply the difference by rstd and gamma, and add the obtained result with beta to obtain the output result.
Prototype
Due to the complex computation involved in the internal implementation of this API, additional temporary space is required to store intermediate variables generated during computation. The method of obtaining the temporary space size (BufferSize) is as follows: Obtain the required maximum and minimum temporary space sizes using the GetLayerNormMaxMinTmpSize API provided in LayerNorm Tiling. The minimum space can ensure correct functionality, while the maximum space is used to improve performance.
The temporary space can be allocated through the API framework or passed by developers through the sharedTmpBuffer input parameter. Therefore, there are two types of function prototypes for the LayerNorm API.
- For the input data with the shape of [B, S, H], the normalized result, mean value, and variance are output.
- Pass the temporary space through the sharedTmpBuffer input parameter.
1 2
template <typename T, bool isReuseSource = false> __aicore__ inline void LayerNorm(const LocalTensor<T>& output, const LocalTensor<T>& outputMean, const LocalTensor<T>& outputVariance, const LocalTensor<T>& inputX, const LocalTensor<T>& gamma, const LocalTensor<T>& beta, const LocalTensor<uint8_t>& sharedTmpBuffer, const T epsilon, LayerNormTiling& tiling)
This method enables developers to allocate and manage the temporary buffer space on their own, and reuse the buffer after calling the API, so that the buffer is not repeatedly allocated or deallocated, improving the flexibility and buffer utilization.
- Allocate the temporary space through the API framework.
1 2
template <typename T, bool isReuseSource = false> __aicore__ inline void LayerNorm(const LocalTensor<T>& output, const LocalTensor<T>& outputMean, const LocalTensor<T>& outputVariance, const LocalTensor<T>& inputX, const LocalTensor<T>& gamma, const LocalTensor<T>& beta, const T epsilon, LayerNormTiling& tiling)
When using this method, developers do not need to allocate the space, but must reserve the required size for the space.
- Pass the temporary space through the sharedTmpBuffer input parameter.
- For the input data with the shape of [A, R], the variance, normalized result, mean value, and reciprocal of the standard deviation are output.
- Pass the temporary space through the sharedTmpBuffer input parameter.
1 2
template <typename U, typename T, bool isReuseSource = false, const LayerNormConfig& config = LNCFG_NORM> __aicore__ inline void LayerNorm(const LocalTensor<T>& output, const LocalTensor<float>& outputMean, const LocalTensor<float>& outputRstd, const LocalTensor<T>& inputX, const LocalTensor<U>& gamma, const LocalTensor<U>& beta, const float epsilon, const LocalTensor<uint8_t>& sharedTmpBuffer, const LayerNormPara& para, const LayerNormSeparateTiling& tiling)
This method enables developers to allocate and manage the temporary buffer space on their own, and reuse the buffer after calling the API, so that the buffer is not repeatedly allocated or deallocated, improving the flexibility and buffer utilization.
- Allocate the temporary space through the API framework.
1 2
template <typename U, typename T, bool isReuseSource = false, const LayerNormConfig& config = LNCFG_NORM> __aicore__ inline void LayerNorm(const LocalTensor<T>& output, const LocalTensor<float>& outputMean, const LocalTensor<float>& outputRstd, const LocalTensor<T>& inputX, const LocalTensor<U>& gamma, const LocalTensor<U>& beta, const float epsilon, const LayerNormPara& para, const LayerNormSeparateTiling& tiling)
When using this method, developers do not need to allocate the space, but must reserve the required size for the space.
- Pass the temporary space through the sharedTmpBuffer input parameter.
Parameters
- API for outputting the normalized result, mean value, and variance for the input data with the shape of [B, S, H]
Table 1 Template parameters Parameter
Description
T
Data type of the operand.
For the Atlas 350 Accelerator Card, the supported data types are half and float.
For the
Atlas A3 training product /Atlas A3 inference product , the supported data types are half and float.For the
Atlas A2 training product /Atlas A2 inference product , the supported data types are half and float.For the
Atlas inference product AI Core, the supported data types are half and float.isReuseSource
Whether the source operand can be modified. The default value is false. If developers allow the source operand to be modified, enable this parameter, to reduce memory space usage.
If this parameter is set to true, the inputX memory space is reused during internal computation of this API to reduce memory space usage. If this parameter is set to false, the inputX memory space is not reused during internal computation of this API.
This parameter can be enabled for float data inputs but cannot be enabled for half data inputs.
For details about how to use isReuseSource, see Example 4.
Table 2 API parameters Parameter
Input/Output
Description
output
Output
Destination operand, with a shape of [B, S, H]. For details about the definition of the LocalTensor data structure, see LocalTensor.
The type is LocalTensor, and TPosition can be VECIN, VECCALC, or VECOUT.
outputMean
Output
Mean, with a shape of [B, S]. For details about the definition of the LocalTensor data structure, see LocalTensor.
The type is LocalTensor, and TPosition can be VECIN, VECCALC, or VECOUT.
outputVariance
Output
Variance, with a shape of [B, S]. For details about the definition of the LocalTensor data structure, see LocalTensor.
The type is LocalTensor, and TPosition can be VECIN, VECCALC, or VECOUT.
inputX
Input
Source operand, with a shape of [B, S, H]. For details about the definition of the LocalTensor data structure, see LocalTensor. The data type of inputX must be the same as that of the destination operand, and the last axis length must be 32-byte aligned.
The type is LocalTensor, and TPosition can be VECIN, VECCALC, or VECOUT.
gamma
Input
Scaling coefficient, with a shape of [H]. For details about the definition of the LocalTensor data structure, see LocalTensor. The data type of gamma must be the same as that of the destination operand, and the last axis length must be 32-byte aligned.
The type is LocalTensor, and TPosition can be VECIN, VECCALC, or VECOUT.
beta
Input
Translation coefficient, with a shape of [H]. For details about the definition of the LocalTensor data structure, see LocalTensor. The data type of beta must be the same as that of the destination operand, and the last axis length must be 32-byte aligned.
The type is LocalTensor, and TPosition can be VECIN, VECCALC, or VECOUT.
sharedTmpBuffer
Input
Shared buffer, which is used to store temporary data generated during internal API computation. This enables developers to manage the sharedTmpBuffer space and reuse the buffer after calling the API, so that the buffer is not repeatedly allocated or deallocated, improving the flexibility and buffer utilization. For details about how to obtain the size of the shared buffer, see LayerNorm Tiling.
The type is LocalTensor, and TPosition can be VECIN, VECCALC, or VECOUT.
epsilon
Input
Weight coefficient for preventing division by zero.
tiling
Input
Tiling information required for LayerNorm computation. For details about how to obtain the tiling information, see LayerNorm Tiling.
- API for outputting the variance, normalized result, mean value, or reciprocal of the standard deviation for the input data with the shape of [A, R]
Table 3 Template parameters Parameter
Description
U
Data type of the beta and gamma operands.
For the Atlas 350 Accelerator Card, the supported data types are half, bfloat16_t, and float.
For the
Atlas A3 training product /Atlas A3 inference product , the supported data types are half and float.For the
Atlas A2 training product /Atlas A2 inference product , the supported data types are half and float.For the
Atlas inference product AI Core, the supported data types are half and float.T
Data type of the output and inputX operands.
For the Atlas 350 Accelerator Card, the supported data types are half, bfloat16_t, and float.
For the
Atlas A3 training product /Atlas A3 inference product , the supported data types are half and float.For the
Atlas A2 training product /Atlas A2 inference product , the supported data types are half and float.For the
Atlas inference product AI Core, the supported data types are half and float.isReuseSource
This parameter is reserved. Pass the default value false.
config
A parameter used to configure the input and output information of the LayerNorm API. The LayerNormConfig type is defined as follows:
1 2 3 4 5 6
struct LayerNormConfig { bool isNoBeta = false; bool isNoGamma = false; bool isOnlyOutput = false; bool isOutputRstd = true; };
- isNoBeta: Whether to use beta in computation.
- false: Default value. The input beta is used in LayerNorm computation.
- true: The input beta is not used in LayerNorm computation. In this case, computation related to beta in the formula is omitted.
- isNoGamma: Whether the optional input gamma is used.
- false: Default value. The optional input gamma is used in LayerNorm computation.
- true: The input gamma is not used in LayerNorm computation. In this case, computation related to gamma in the formula is omitted.
- isOnlyOutput: Whether only y is output, and mean and reciprocal of the standard deviation rstd are not output. Currently, this parameter can only be set to false, indicating that all y, mean, and rstd results are output.
- isOutputRstd: indicates whether to output the reciprocal of the standard deviation (rstd) or the variance. Only the Atlas 350 Accelerator Card supports this option.
- true: The reciprocal of the standard deviation is output. This is the default value.
- false: The variance is output.
Table 4 API parameters Parameter
Input/Output
Description
output
Output
Destination operand, with a shape of [A, R]. For details about the definition of the LocalTensor data structure, see LocalTensor.
The type is LocalTensor, and TPosition can be VECIN, VECCALC, or VECOUT.
outputMean
Output
Mean, with a shape of [A]. For details about the definition of the LocalTensor data structure, see LocalTensor.
The type is LocalTensor, and TPosition can be VECIN, VECCALC, or VECOUT.
outputRstd
Output
If isOutputRstd in the template parameter config is set to true, outputRstd is the reciprocal of the standard deviation. If isOutputRstd is false, outputRstd is the variance. The shape is [A]. For details about the definition of the LocalTensor data structure, see LocalTensor.
Note that only the Atlas 350 Accelerator Card supports the output variance.
The type is LocalTensor, and TPosition can be VECIN, VECCALC, or VECOUT.
inputX
Input
Source operand, with a shape of [A, R]. For details about the definition of the LocalTensor data structure, see LocalTensor. The data type of inputX must be the same as that of the destination operand, and the last axis length must be 32-byte aligned.
The type is LocalTensor, and TPosition can be VECIN, VECCALC, or VECOUT.
gamma
Input
Scaling coefficient, with a shape of [R]. For details about the definition of the LocalTensor data structure, see LocalTensor. The data type precision of gamma must be greater than or equal to that of the source operand.
The type is LocalTensor, and TPosition can be VECIN, VECCALC, or VECOUT.
beta
Input
Translation coefficient, with a shape of [R]. For details about the definition of the LocalTensor data structure, see LocalTensor. The data type precision of beta must be greater than or equal to that of the source operand.
The type is LocalTensor, and TPosition can be VECIN, VECCALC, or VECOUT.
epsilon
Input
Weight coefficient for preventing division by zero.
sharedTmpBuffer
Input
Shared buffer, which is used to store temporary data generated during internal API computation. This enables developers to manage the sharedTmpBuffer space and reuse the buffer after calling the API, so that the buffer is not repeatedly allocated or deallocated, improving the flexibility and buffer utilization. For details about how to obtain the size of the shared buffer, see LayerNorm Tiling.
The type is LocalTensor, and TPosition can be VECIN, VECCALC, or VECOUT.
para
Input
Parameter information required for LayerNorm computation. The LayerNormPara type is defined as follows:
1 2 3 4 5
struct LayerNormPara { uint32_t aLength; uint32_t rLength; uint32_t rLengthWithPadding; };
- aLength: Specifies the length of inputX along the A axis.
- rLength: Specifies the length of the to-be-processed data of inputX on the R axis.
- rLengthWithPadding: Specifies the 32-byte aligned length of inputX along the R axis.
tiling
Input
Tiling information required for LayerNorm computation. For details about how to obtain the tiling information, see LayerNorm Tiling.
- isNoBeta: Whether to use beta in computation.
Returns
None
Constraints
- For details about the operand address alignment requirements, see General Address Alignment Restrictions.
- Constraints on the API that outputs the normalization result, mean value, and variance for the input data with the shape of [B, S, H]:
- The space of output and inputX can be reused. The space of other outputs and inputs cannot be reused.
- If the last axis H in the input data does not meet the alignment requirements, you need to pad the data. The padded data should be set to 0 to prevent abnormal values from affecting network computation.
- The last axis (H axis) cannot be split.
- The H axis lengths of inputX, output, gamma, and beta must be the same.
- The B axis lengths and S axis lengths of inputX, output, outputMean, and outputVariance must be the same.
- Constraints on the API that outputs the variance, normalized result, mean value, or reciprocal of the standard deviation for the input data with the shape of [A, R]:
- The data type precision of gamma and beta must be greater than or equal to that of the source operand. For example, for the Atlas 350 Accelerator Card, if the data type of inputX is bfloat16_t, the data types of gamma and beta can be bfloat16_t or float, and their data type precision is not lower than that of inputX.
- The tensor space of src and dst cannot be reused.
- The R axis cannot be split.
Examples
- Example of calling the API to output the normalized result, mean value, and variance for the input data with shape [B, S, H]
For details about the complete call example, see LayerNorm operator sample.
1 2 3 4 5 6 7 8 9 10
AscendC::LayerNorm<float, false>( output, // [Output] normalized result y, with shape [B, S, H] mean, // [Output] mean value of the H dimension at each (B, S) position, with shape [B, S] variance, // [Output] variance of the H dimension at each (B, S) position, with shape [B, S] inputX, // [Input] original input data x, with shape [B, S, H], to be normalized gamma, // [Input] scaling coefficient γ, with shape [H], used to scale the normalized data beta, // [Input] translation coefficient β, with shape [H], used to offset the normalized data (float)epsilon, // [Input] coefficient ε for preventing division by zero, used to prevent division by zero when the variance is 0. tiling // [Input] tiling information, including the hardware computing block policy (such as block and thread). );
The following is an example:Input (inputX, shape: [1, 8, 8]): [ 0. 1. 2. 3. 4. 5. 6. 7. 8. 9. 10. 11. 12. 13. 14. 15. 16. 17. 18. 19. 20. 21. 22. 23. 24. 25. 26. 27. 28. 29. 30. 31. 32. 33. 34. 35. 36. 37. 38. 39. 40. 41. 42. 43. 44. 45. 46. 47. 48. 49. 50. 51. 52. 53. 54. 55. 56. 57. 58. 59. 60. 61. 62. 63. ] Input (gamma, shape: [8]): [ 0. 1. 2. 3. 4. 5. 6. 7. ] Input (beta, shape: [8]): [ 0. 1. 2. 3. 4. 5. 6. 7. ] Output (output): [ 0. -0.09107912 0.69070506 2.3453526 4.8728633 8.273237 12.546474 17.692575 0. -0.09107912 0.69070506 2.3453526 4.8728633 8.273237 12.546474 17.692575 0. -0.09107912 0.69070506 2.3453526 4.8728633 8.273237 12.546474 17.692575 0. -0.09107912 0.69070506 2.3453526 4.8728633 8.273237 12.546474 17.692575 0. -0.09107912 0.69070506 2.3453526 4.8728633 8.273237 12.546474 17.692575 0. -0.09107912 0.69070506 2.3453526 4.8728633 8.273237 12.546474 17.692575 0. -0.09107912 0.69070506 2.3453526 4.8728633 8.273237 12.546474 17.692575 0. -0.09107912 0.69070506 2.3453526 4.8728633 8.273237 12.546474 17.692575 ] Output (mean): [ 3.5 11.5 19.5 27.5 35.5 43.5 51.5 59.5 ] Output (variance): [ 5.25 5.25 5.25 5.25 5.25 5.25 5.25 5.25 ]
- Example of calling the API to output the variance, normalized result, mean, or reciprocal of the standard deviation for the input data with shape [A, R]
For details about the complete call example, see LayerNormV2 operator sample.
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17
// config: compile-time constant, which defines the behavior configuration of LayerNorm constexpr auto config = AscendC::LayerNormConfig{false, false, false, true}; // para: runtime parameter, which describes the dimension information of the input tensor AscendC::LayerNormPara para = {aLength, rLength, rLengthWithPadding}; // Call the LayerNorm API. AscendC::LayerNorm<float, float, false, config>( output, // [Output] normalized result y, with shape [A, R] mean, // [Output] mean value of the R dimension at each A position, with shape [A] output1, // [Output] reciprocal of the standard deviation rstd (or variance), with shape [A] inputX, // [Input] original input data x, with shape [A, R] gamma, // [Input] scaling coefficient γ, with shape [R] beta, // [Input] translation coefficient β, with shape [R] (float)epsilon, // [Input] coefficient ε for preventing division by zero para, // [Input] structure containing information such as the lengths of axes A and R tiling // [Input] tiling policy information );