WelfordUpdate

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

Welford is a method for calculating the mean and variance online. On one hand, this method can gradually calculate the mean and variance of all samples without needing to store these samples, making it ideal for processing large-scale data. On the other hand, it requires only a single data traversal, reducing memory access and enhancing computational performance. This API is used for pre-processing of the Welford algorithm.

If the reduced axis in the LayerNorm algorithm is large, you can split the axis and use this API and WelfordFinalize together to implement equivalent computation of LayerNorm.

As shown in the following figure, the reduced axis of the data is split. Assume that the shape of each piece of data after splitting is [1, k], and each piece of data is numbered 1, 2, 3, ..., n.

Figure 1 Splitting of the reduced axis

Below is the formula of this API. After data is split, this API is called n times. The following formula is used to compute each piece of split data.

In the formula, the shapes of xi, Meanti, and Mi are all [1, k]. xi indicates the ith data block after splitting. Meanti indicates the mean value of the first i data blocks obtained by calling this API for the ith time. Mi indicates the intermediate result of the variance of the first i data blocks obtained by calling this API for the ith time. (The intermediate result is saved for computing the variance in the following sections.) When this API is called for the first time (i = 1), Meant0 and M0 in the formula are defined by the user as data with the shape of [1, k], and all values of 0.

The following figure shows the computation process of Meantn. After this API is called n times, Meantn and Mn with the shape of [1, k] are obtained. Meantn and Mn are used for subsequent computation by the WelfordFinalize API.

Figure 2 Meantn computation process

Prototype

  • Pass the temporary space through the sharedTmpBuffer input parameter.
    1
    2
    template <typename T, typename U,bool isReuseSource = false, const WelfordUpdateConfig& config = WFUPDATE_DEFAULT_CFG>
    __aicore__ inline void WelfordUpdate(const LocalTensor<U>& outputMean, const LocalTensor<U>& outputVariance, const LocalTensor<U>& inputMean, const LocalTensor<U>& inputVariance, const LocalTensor<T>& inputX, const LocalTensor<uint8_t>& sharedTmpBuffer, const WelfordUpdateParam& para)
    
  • Allocate the temporary space through the API framework.
    1
    2
    template <typename T, typename U,bool isReuseSource = false, const WelfordUpdateConfig& config = WFUPDATE_DEFAULT_CFG>
    __aicore__ inline void WelfordUpdate(const LocalTensor<U>& outputMean, const LocalTensor<U>& outputVariance, const LocalTensor<U>& inputMean, const LocalTensor<U>& inputVariance, const LocalTensor<T>& inputX, const WelfordUpdateParam& para)
    

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 temporary space can be allocated through the API framework or passed by developers through the sharedTmpBuffer input parameter.

  • When the API framework is used for temporary space allocation, developers do not need to allocate the space, but must reserve the required size for the temporary space.
  • When the sharedTmpBuffer input parameter is used for passing the temporary space, the tensor serves as the temporary space. In this case, the API framework is not required for temporary space allocation. 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.

If the API framework is used, developers must reserve the temporary space. If sharedTmpBuffer is used, developers must allocate space for the tensor. The method of obtaining the temporary space size (BufferSize) is as follows: Obtain the required maximum and minimum temporary space sizes using the GetWelfordUpdateMaxMinTmpSize API described in WelfordUpdate Tiling. The minimum space can ensure correct functionality, while the maximum space is used to improve performance.

Parameters

Table 1 Template parameters

Parameter

Description

T

Data type of the inputX operand.

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.

U

Data type of the outputMean, outputVariance, inputMean, and inputVariance operands.

For the Atlas 350 Accelerator Card, the supported data type is float.

For the Atlas A3 training product/Atlas A3 inference product, the supported data type is float.

For the Atlas A2 training product/Atlas A2 inference product, the supported data type is float.

For the Atlas inference product AI Core, the supported data type is 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.

In Atlas inference product AI Core, this parameter is reserved. The default value false is used.

For details about how to use isReuseSource, see Example 4.

config

Reuse relationship between the destination operand and source operand that are not in the specified computation scope range. The WelfordUpdateConfig type is defined as follows:

1
2
3
struct WelfordUpdateConfig {
    bool isInplace = false; // Whether the destination operand reuses the source operand.
};
  • isInplace: The abComputeLength parameter under para in API parameters specifies the calculation length for the inner axis of the input data. The value of the output data beyond the calculation length is determined by the isInplace parameter. The isInplace parameter indicates whether to reuse the source operand for the destination operand beyond the specified calculation length. If yes, the source operand at the corresponding position replaces the destination operand when the output is beyond the specified calculation length. If no, this API does not output the destination operands beyond the calculation length.
    • false (default value): The destination operand does not reuse the source operand.
    • true: The destination operand reuses the source operand. outputMean reuses inputMean, and outputVariance reuses inputVariance.

A configuration example is as follows:

1
constexpr WelfordUpdateConfig WFUPDATE_DEFAULT_CFG = {false};

This parameter is used together with the tiling computation API on the kernel.

Table 2 API parameters

Parameter

Input/Output

Description

outputMean

Output

Destination operand of the mean, corresponding to Meanti in the API formula.

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

The shape must be the same as that of the source operand inputMean.

outputVariance

Output

Destination operand of the intermediate variance result, corresponding to Mi in the API formula.

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

The shape must be the same as that of the source operand inputVariance.

inputMean

Input

Source operand of the mean, corresponding to Meanti-1 in the API formula.

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

inputVariance

Input

Source operand of the intermediate variance result, corresponding to Mi-1 in the API formula.

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

inputX

Input

Source operand, corresponding to xi in the API formula.

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

sharedTmpBuffer

Input

Temporary space.

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

This parameter is used to store intermediate variables during complex internal API computation and is provided by developers.

For details about how to obtain the temporary space size (BufferSize), see WelfordUpdate Tiling.

para

Input

Parameter information required for calculation. The WelfordUpdateParam type is defined as follows:

1
2
3
4
5
6
struct WelfordUpdateParam {
    uint32_t rnLength; 
    uint32_t abLength; 
    uint32_t abComputeLength; 
    float nRec;
};
  • rnLength: This is a reserved parameter, and its value is fixed to 1.
  • abLength indicates the splitting size on the reduced axis.
  • abComputeLength indicates the actual length of the reduced axis computed from the input start address.
  • nRec: The value is 1/i, where i indicates the accumulated number of times that this API is called. The value range of i is [1, n], where n indicates the number of blocks split on the reduced axis of the input data inputX.

The shape of each destination operand and source operand is [rnLength, abLength].

Returns

None

Constraints

  • Currently, the value of para.rnLength can only be 1.
  • The value of para.abLength must be an integer multiple of 32/sizeof (T).
  • The value of para.abComputeLength must be greater than 0.
  • The source operand address must not overlap the destination operand address.
  • The address of sharedTmpBuffer cannot overlap that of the source or destination operand.

Examples

For a complete call example, see WelfordUpdate operator sample.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
// outputMean: updated mean value Meant, with shape [1, abLength]
// outputVariance: updated intermediate variance result Mi, with shape [1, abLength]
// inputMean: previous mean value Meant-1, used as the input
// inputVariance: previous intermediate variance result Mi-1, used as the input
// inputX: input data xi at the current time step, with shape [1, abLength]
// sharedTmpBuffer: temporary space managed by developers, used for complex internal computation
// para: parameter structure that contains the block information and normalization coefficient of the reduced axis

// Use the WelfordUpdate API to update the Welford algorithm online.
struct AscendC::WelfordUpdateParam para = { nLength, rLength, abComputeLength, 0.3 };
AscendC::WelfordUpdate<T, U, false, WELFORD_UPDATE_ENABLE_INPLACE_CFG>(
    outputMean,        // Output: updated mean value
    outputVariance,    // Output: updated intermediate variance result
    inputMean,         // Input: mean value at the previous moment
    inputVariance,     // Input: intermediate variance result at the previous moment
    inputX,            // Input: current input xi
    sharedTmpBuffer,   // Input: temporary buffer (provided by developers)
    para               // Input: Welford update parameters
);
The following is an example:
Input (inputX, shape: [1, 64]):
[1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1.
 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1.
 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. ]
Input (gammaLocal, shape: [64]):
[ 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 (betaLocal, shape: [64]):
[ 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. ]
Output (meanLocal):
[ 0.125  1.     1.875  2.75   3.625  4.5    5.375  6.25   7.125  8.
  8.875  9.75  10.625 11.5   12.375 13.25  14.125 15.    15.875 16.75
 17.625 18.5   19.375 20.25  21.125 22.    22.875 23.75  24.625 25.5
 26.375 27.25  28.125 29.    29.875 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. ]
Output (varianceLocal):
[8.75000e-01 1.00000e+00 2.87500e+00 6.50000e+00 1.18750e+01 1.90000e+01
 2.78750e+01 3.85000e+01 5.08750e+01 6.50000e+01 8.08750e+01 9.85000e+01
 1.17875e+02 1.39000e+02 1.61875e+02 1.86500e+02 2.12875e+02 2.41000e+02
 2.70875e+02 3.02500e+02 3.35875e+02 3.71000e+02 4.07875e+02 4.46500e+02
 4.86875e+02 5.29000e+02 5.72875e+02 6.18500e+02 6.65875e+02 7.15000e+02
 7.65875e+02 8.18500e+02 8.72875e+02 9.29000e+02 9.86875e+02 3.50000e+01
 3.60000e+01 3.70000e+01 3.80000e+01 3.90000e+01 4.00000e+01 4.10000e+01
 4.20000e+01 4.30000e+01 4.40000e+01 4.50000e+01 4.60000e+01 4.70000e+01
 4.80000e+01 4.90000e+01 5.00000e+01 5.10000e+01 5.20000e+01 5.30000e+01
 5.40000e+01 5.50000e+01 5.60000e+01 5.70000e+01 5.80000e+01 5.90000e+01
 6.00000e+01 6.10000e+01 6.20000e+01 6.30000e+01 ]