Documentation
Download

LayerNormGrad

Applicability

Product

Supported

Atlas 350 Accelerator Card

√

Atlas A3 training product/Atlas A3 inference product

√

Atlas A2 training product/Atlas A2 inference product

√

Atlas 200I/500 A2 inference product

x

Atlas inference product AI Core

√

Atlas inference product Vector Core

x

Atlas training product

x

Function Usage

LayerNormGrad is a function used to calculate the backpropagation gradient of LayerNorm. If this API is used alone, it generates outputs x and resForGamma. When combined with LayerNormGradBeta, resForGamma that this API produces can be passed to LayerNormGradBeta to generate the gamma and beta outputs. Therefore, you can simultaneously obtain x, gamma, and beta by using these two APIs.

The formulas are as follows:

1
2
3
4
5
pd_xl(BSH) = data_dy x data_gamma
pd_var(H) = np.sum(((-0.5) * pd_xl * (data_x - data_mean) * np.power((data_variance + EPSILON), (-1.5))), reduce_axis, keepdims=True)
pd_mean(BS1) = np.sum(((-1.0) * pd_xl * np.power((data_variance + EPSILON), (-0.5))), reduce_axis, keepdims=True) + pd_var * (1.0 / H) * np.sum(((-2.0) * (data_x - data_mean)), reduce_axis, keepdims=True)
pd_x(BSH) = pd_xl * np.power((data_variance + EPSILON), (-0.5)) + pd_var * (2.0 / H) * (data_x - data_mean) + pd_mean * (1.0 / H)
res_for_gamma(BSH) = (data_x - data_mean) * np.power((data_variance + EPSILON), (-0.5))

Principles

The figure below illustrates the internal algorithm block diagram of LayerNormGrad high-level APIs, taking the float type, ND format, and inputs inputDy[B, S, H], inputX[B, S, H], inputVariance[B, S], inputMean[B, S], and inputGamma[H] as examples.

Figure 1 LayerNormGrad algorithm block diagram

The computation process is divided into the following steps, all of which are performed on vectors:

  1. ComputePdX1: Calculate inputDy x inputGamma and store the result to x1Tensor.
  2. ComputePdX2: Extend the shape of inputMean to [B, S, H] through Brcb, calculate inputX-inputMean, and store the result to x2Tensor.
  3. ComputePdVar: Implement the computation of the np.sum(((–0.5) x x1Tensor x x2Tensor x np.power((inputVariance + EPSILON), (–1.5)))) formula, with the power method implemented by combining the Sqrt, Div, and Mul basic APIs. Store the result to pdVarTensor.
  4. ComputePdMean: Implement the computation of the formula np.sum(((–1.0) x x1Tensor x np.power((inputVariance + EPSILON), (–0.5)))) + pd_var x (1.0/H) x np.sum(((–2.0) x (x2Tensor))), with the power method implemented by combining the Sqrt and Div basic APIs. Store the result to pdMeanTensor. At the same time, use the intermediate computation result to calculate the formula resForGamma result according to the x2Tensor x np.power((inputVariance + EPSILON), (–0.5)).
  5. ComputePdX: Implement the computation of the formula x1Tensor x np.power((inputVariance + EPSILON), (–0.5)) + pd_var x (2.0/H) x (x2Tensor) + pd_mean x (1.0/H) and store the result to outputPdX.

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 GetLayerNormGradMaxMinTmpSize API provided in LayerNormGrad 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 LayerNormGrad API.

  • Pass the temporary space through the sharedTmpBuffer input parameter.
    1
2
    template <typename T, bool isReuseSource = false>
__aicore__ inline void LayerNormGrad(const LocalTensor<T>& outputPdX, const LocalTensor<T>& resForGamma, const LocalTensor<T>& inputDy, const LocalTensor<T>& inputX, const LocalTensor<T>& inputVariance, const LocalTensor<T>& inputMean, const LocalTensor<T>& inputGamma, LocalTensor<uint8_t>& sharedTmpBuffer, T epsilon, LayerNormGradTiling &tiling, const LayerNormGradShapeInfo& shapeInfo = {})

    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 LayerNormGrad(const LocalTensor<T>& outputPdX, const LocalTensor<T>& resForGamma, const LocalTensor<T>& inputDy, const LocalTensor<T>& inputX, const LocalTensor<T>& inputVariance, const LocalTensor<T>& inputMean, const LocalTensor<T>& inputGamma, T epsilon, LayerNormGradTiling& tiling, const LayerNormGradShapeInfo& shapeInfo = {})

    When using this method, developers do not need to allocate the space, but must reserve the required size for the space.

Parameters

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 More Examples.
Table 2 API parameters

Parameter

Input/Output

Description

outputPdX

Output

Destination operand, with a shape of [B, S, H]. For details about the definition of the LocalTensor data structure, see LocalTensor. The last axis length must be 32-byte aligned.

The type is LocalTensor, and TPosition can be VECIN, VECCALC, or VECOUT.

resForGamma

Output

Destination operand, with a shape of [B, S, H]. For details about the definition of the LocalTensor data structure, see LocalTensor. The last axis length must be 32-byte aligned.

The type is LocalTensor, and TPosition can be VECIN, VECCALC, or VECOUT.

inputDy

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 inputDy 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.

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.

inputVariance

Input

Variance, with a shape of [B, S]. For details about the definition of the LocalTensor data structure, see LocalTensor. The data type of inputVariance must be the same as that of the destination operand, and the last axis length must be 32-byte aligned. The LayerNorm API needs to be called in advance to obtain the variance.

The type is LocalTensor, and TPosition can be VECIN, VECCALC, or VECOUT.

inputMean

Input

Mean, with a shape of [B, S]. For details about the definition of the LocalTensor data structure, see LocalTensor. The data type of inputMean must be the same as that of the destination operand, and the last axis length must be 32-byte aligned. The LayerNorm API needs to be called in advance to obtain the mean.

The type is LocalTensor, and TPosition can be VECIN, VECCALC, or VECOUT.

inputGamma

Input

Source operand, with a shape of [H]. For details about the definition of the LocalTensor data structure, see LocalTensor. The data type of inputGamma 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 LayerNormGrad 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 LayerNormGrad computation.

shapeInfo

Input

Layout format of each input data of LayerNormGrad. The default value indicates that the input format is ND. The value can be DataFormat::ND. LayerNormGradShapeInfo type. The specific definition is as follows:

1
2
3
struct LayerNormGradShapeInfo {
    DataFormat dataFormat = DataFormat::ND;
};

Returns

None

Constraints

  • For details about the operand address alignment requirements, see General Address Alignment Restrictions.
  • The tensor space of the source operand and destination operand can be reused.
  • The input shape must be in ND format.
  • If 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.

Examples

In this sample, the shape of inputX and inputDy is [2, 32, 16], the shape of inputVariance and inputMean is [2, 32], and the shape of inputGamma is [16]. The shape of outputPdX and resForGamma is [2, 32, 16]. The data layout is in ND format and the data type is float. The memory space of the source operand is not reused.

For a complete call example, see LayerNormGrad operator sample.

 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
// outputPdX: output gradient of input X, that is, dX, with shape [B, S, H]
// resForGamma: output intermediate result (for example, dy x normalized_x) used to calculate the gamma and beta gradients, with shape [B, S, H]
// inputDy: input upper-layer gradient dy, with shape [B, S, H]
// inputX: input X during forward propagation, with shape [B, S, H]
// inputVariance: variance calculated by forward LayerNorm, with shape [B, S]
// inputMean: mean calculated by forward LayerNorm, with shape [B, S]
// inputGamma: scaling parameter gamma in LayerNorm, with shape [H]
// sharedTmpBuffer: temporary buffer managed by developers, used to store intermediate variables during internal computation
// epsilon: small value to prevent division by zero, for example, 1e-5
// tiling: structure containing tiling information required for computation (such as block and thread)
// shapeInfo: (optional) data layout format of the input tensor. Currently, only the ND format is supported.

// Use the LayerNormGrad API to perform backward propagation computation for layer normalization.
AscendC::LayerNormGrad<float, isReuseSource>(
    outputPdX,        // Output: input gradient dX, with shape [B, S, H]
    resForGamma,      // Output: intermediate result, used to calculate dgamma/dbeta
    inputDy,          // Input: upper-layer gradient dy, with shape [B, S, H]
    inputX,           // Input: original input X, with shape [B, S, H]
    inputVariance,    // Input: variance calculated by forward propagation, with shape [B, S]
    inputMean,        // Input: mean calculated by forward propagation, with shape [B, S]
    inputGamma,       // Input: scaling parameter gamma, with shape [H]
    sharedTmpBuffer,  // Input: temporary space provided by developers (the size needs to be obtained by calling GetLayerNormGradMaxMinTmpSize)
    epsilon,          // Input: coefficient ε for preventing division by zero
    tiling,           // Input: tiling information, generated by the tiling tool
    {DataFormat::ND}  // Input: shapeInfo, which defaults to DataFormat::ND
);
The following is an example:
Input (inputDy, 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 (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 (inputMean, shape: [8]):
[ 3.5 11.5 19.5 27.5 35.5 43.5 51.5 59.5 ]
Input (inputVariance, shape: [8]):
[ 5.25 5.25 5.25 5.25 5.25 5.25 5.25 5.25 ]
Input (inputGamma, shape: [8]):
[ 0. 1. 2. 3. 4. 5. 6. 7. ]
Output (outputPdX):
[ 3.0548172 0.4362857 -1.3093826 -2.182187 -2.1821284 -1.309207 0.4365778 3.0552254 3.0545845 0.4361186 -1.3094826 -2.1822214 -2.1820965 -1.3091087 0.4367447 3.055458 3.0543518 0.4359522 -1.3095818 -2.1822548 -2.182064 -1.3090096 0.43690872 3.055687 3.054119 0.43578815 -1.309679 -2.1822853 -2.182026 -1.3089066 0.437088 3.0559235 3.0538864 0.43562222 -1.3097801 -2.1823158 -2.1819916 -1.3088074 0.43724823 3.05616 3.0536423 0.4354477 -1.3098869 -2.1823578 -2.181961 -1.3087158 0.43740845 3.0563965 3.0534134 0.43528175 -1.3099861 -2.1823883 -2.1819305 -1.308609 0.43756104 3.0566254 3.0531921 0.43511963 -1.3100777 -2.1824188 -2.1818848 -1.3085022 0.43774414 3.0568542 ]
Output (resForGamma):
[ -1.5275106 -1.091079 -0.6546474 -0.21821581 0.21821581 0.6546474 1.091079 1.5275106 -1.5275106 -1.091079 -0.6546474 -0.21821581 0.21821581 0.6546474 1.091079 1.5275106 -1.5275106 -1.091079 -0.6546474 -0.21821581 0.21821581 0.6546474 1.091079 1.5275106 -1.5275106 -1.091079 -0.6546474 -0.21821581 0.21821581 0.6546474 1.091079 1.5275106 -1.5275106 -1.091079 -0.6546474 -0.21821581 0.21821581 0.6546474 1.091079 1.5275106 -1.5275106 -1.091079 -0.6546474 -0.21821581 0.21821581 0.6546474 1.091079 1.5275106 -1.5275106 -1.091079 -0.6546474 -0.21821581 0.21821581 0.6546474 1.091079 1.5275106 -1.5275106 -1.091079 -0.6546474 -0.21821581 0.21821581 0.6546474 1.091079 1.5275106 ]
Parent topic: Normalization
This site uses cookies to improve your browsing experience. By continuing to browse this site, you agree to our use of cookies.View Details