WelfordFinalize

功能说明

Welford计算是一种在线计算均值和方差的方法。一方面，它可以在不存储所有样本的情况下，逐步计算所有样本的均值和方差，更适合处理海量数据；另一方面，它只需要对数据进行一次遍历，能减少访存次数，提高计算性能。本接口为Welford算法的后处理。

LayerNorm算法中reduce轴较大的场景，可以通过切分reduce轴，联合使用本接口与WelfordUpdate，能够实现等效计算LayerNorm。本接口的计算公式如下：

• 不带尾块/不带counts参数场景：

  

  其中，为均值输出，为方差输出。

  代表输入的第个均值，代表输入的第个方差。代表Reduce轴拆分后一次计算的大小，代表Reduce轴按拆分的次数(只有整除的情况下使用此公式)。

• 带尾块/带counts参数场景：

  

  其中，除上述参数含义外，代表对应的系数，代表未切分的原始Reduce轴长度。

函数原型

• 通过sharedTmpBuffer入参传入临时空间
  • 不带counts参数场景

    1 2

    template <bool isReuseSource = false> __aicore__ inline void WelfordFinalize(const LocalTensor<float>& outputMean, const LocalTensor<float>& outputVariance, const LocalTensor<float>& inputMean, const LocalTensor<float>& inputVariance, const LocalTensor<uint8_t>& sharedTmpBuffer, WelfordFinalizePara& para)

  • 带counts参数场景

    1 2

    template <bool isReuseSource = false> __aicore__ inline void WelfordFinalize(const LocalTensor<float>& outputMean, const LocalTensor<float>& outputVariance, const LocalTensor<float>& inputMean, const LocalTensor<float>& inputVariance, const LocalTensor<int32_t>& counts, const LocalTensor<uint8_t>& sharedTmpBuffer, WelfordFinalizePara& para)

• 接口框架申请临时空间
  • 不带counts参数场景

    1 2	template <bool isReuseSource = false> __aicore__ inline void WelfordFinalize(const LocalTensor<float>& outputMean, const LocalTensor<float>& outputVariance, const LocalTensor<float>& inputMean, const LocalTensor<float>& inputVariance, WelfordFinalizePara& para)

  • 带counts参数场景

    1 2	template <bool isReuseSource = false> __aicore__ inline void WelfordFinalize(const LocalTensor<float>& outputMean, const LocalTensor<float>& outputVariance, const LocalTensor<float>& inputMean, const LocalTensor<float>& inputVariance, const LocalTensor<int32_t>& counts, WelfordFinalizePara& para)

由于该接口的内部实现中涉及复杂的计算，需要额外的临时空间来存储计算过程中的中间变量。临时空间支持接口框架申请和开发者通过sharedTmpBuffer入参传入两种方式。

• 接口框架申请临时空间，开发者无需申请，但是需要预留临时空间的大小。

• 通过sharedTmpBuffer入参传入，使用该tensor作为临时空间进行处理，接口框架不再申请。该方式开发者可以自行管理sharedTmpBuffer内存空间，并在接口调用完成后，复用该部分内存，内存不会反复申请释放，灵活性较高，内存利用率也较高。

接口框架申请的方式，开发者需要预留临时空间；通过sharedTmpBuffer传入的情况，开发者需要为tensor申请空间。临时空间大小BufferSize的获取方式如下：通过WelfordFinalize Tiling中提供的GetWelfordFinalizeMaxMinTmpSize接口获取所需最大和最小临时空间大小，最小空间可以保证功能正确，最大空间用于提升性能。

参数说明

 1
 2
 3
 4
 5
 6
 7
 8
 9
10

struct WelfordFinalizePara {
    uint32_t rnLength; // Reduce轴按abLength拆分的次数（向下取整）
    uint32_t abLength; // Reduce轴拆分的大小
    uint32_t headCount; // 不带counts接口中的headCount值
    uint32_t headCountLength; // 不带counts接口中，headCount值对应的长度
    uint32_t tailCount; // 不带counts接口中的tailCount值
    uint32_t tailCountLength; // 不带counts接口中，tailCount值对应的长度
    float abRec; // 表示1/abLength的值
    float rRec; // 当没有尾块时，表示1/(rnLength*abLength), 有尾块时，表示1/abLength的值
};

返回值

无

支持的型号

Atlas 推理系列产品AI Core

Atlas A2 训练系列产品/Atlas 800I A2 推理产品

注意事项

• 接口参数para.abLength的取值必须为32/sizeof(float)的整数倍。
• 接口参数para.headCountLength与para.tailCountLength的加和必须等于参数para.ablength。
• 接口处理逻辑以参数para中设置的具体参数值为准，不依赖源操作数的shape信息。
• 接口参数para.tailCount为0时，禁止配置para.tailCountLength为非0值。

调用示例

1
2
3
4

pipe.InitBuffer(sharedTmpBuffer, stackBufferSize);        
AscendC::LocalTensor<uint8_t> tmpLocalTensor = sharedTmpBuffer.Get<uint8_t>();         
struct AscendC::WelfordFinalizePara para = {rnLength, abLength, head, headLength, tail, tailLength, abRec, rRec};
AscendC::WelfordFinalize<false>(meanLocal, varianceLocal, inputMeanLocal, inputVarianceLocal, inputCountsLocal, tmpLocalTensor, para);