LayerNormGradBeta
函数功能
LayerNormGradBeta接口用于获取反向beta/gmma的数值,和LayerNormGrad共同输出pdx, gmma和beta:
算法公式为:
函数原型
由于该接口的内部实现中涉及复杂的计算,需要额外的临时空间来存储计算过程中的中间变量。临时空间大小BufferSize的获取方法:通过LayerNormGradBeta Tiling中提供的GetLayerNormGradBetaMaxMinTmpSize接口获取所需最大和最小临时空间大小,最小空间可以保证功能正确,最大空间用于提升性能。
临时空间支持接口框架申请和开发者通过sharedTmpBuffer入参传入两种方式,因此LayerNormGradBeta接口的函数原型有两种:
- 通过sharedTmpBuffer入参传入临时空间
1 2
template <typename T, bool isReuseSource = false> __aicore__ inline void LayerNormGradBeta(const LocalTensor<T>& outputPdGamma, const LocalTensor<T>& outputPdBeta, const LocalTensor<T>& resForGamma, const LocalTensor<T>& inputDy, const LocalTensor<uint8_t>& sharedTmpBuffer, const LayerNormGradBetaTiling& tiling)
该方式下开发者需自行申请并管理临时内存空间并管理,并在接口调用完成后,复用该部分内存,内存不会反复申请释放,灵活性较高,内存利用率也较高。
- 接口框架申请临时空间
1 2
template <typename T, bool isReuseSource = false> __aicore__ inline void LayerNormGradBeta(const LocalTensor<T>& outputPdGamma, const LocalTensor<T>& outputPdBeta, const LocalTensor<T>& resForGamma, const LocalTensor<T>& inputDy, LayerNormGradBetaTiling& tiling)
该方式下开发者无需申请,但是需要预留临时空间的大小。
参数说明
参数名 |
描述 |
---|---|
T |
操作数的数据类型。 |
isReuseSource |
是否允许修改源操作数,默认值为false。如果开发者允许源操作数被改写,可以使能该参数,使能后能够节省部分内存空间。 设置为true,则本接口内部计算时复用inputDy的内存空间,节省内存空间;设置为false,则本接口内部计算时不复用inputDy的内存空间。 对于float数据类型输入支持开启该参数,half数据类型输入不支持开启该参数。 isReuseSource的使用样例请参考更多样例。 |
参数名称 |
输入/输出 |
含义 |
---|---|---|
outputPdGamma |
输出 |
目的操作数,类型为LocalTensor,shape为[H],LocalTensor数据结构的定义请参考LocalTensor。尾轴长度需要32B对齐 Atlas A2训练系列产品/Atlas 800I A2推理产品,支持的数据类型为:half/float Atlas推理系列产品AI Core,支持的数据类型为:half/float |
outputPdBeta |
输出 |
目的操作数,类型为LocalTensor,shape为[H],LocalTensor数据结构的定义请参考LocalTensor。尾轴长度需要32B对齐 Atlas A2训练系列产品/Atlas 800I A2推理产品,支持的数据类型为:half/float Atlas推理系列产品AI Core,支持的数据类型为:half/float |
resForGamma |
输入 |
源操作数,类型为LocalTensor,shape为[B, S, H],LocalTensor数据结构的定义请参考LocalTensor。resForGamma的数据类型需要与目的操作数保持一致,尾轴长度需要32B对齐。需提前调用LayerNormGrad接口获取resForGamma参数值。 Atlas A2训练系列产品/Atlas 800I A2推理产品,支持的数据类型为:half/float Atlas推理系列产品AI Core,支持的数据类型为:half/float |
inputDy |
输入 |
源操作数,类型为LocalTensor,shape为[B, S, H],LocalTensor数据结构的定义请参考LocalTensor。inputDy的数据类型需要与目的操作数保持一致,尾轴长度需要32B对齐。 Atlas A2训练系列产品/Atlas 800I A2推理产品,支持的数据类型为:half/float Atlas推理系列产品AI Core,支持的数据类型为:half/float |
sharedTmpBuffer |
输入 |
共享缓冲区,用于存放API内部计算产生的临时数据。该方式开发者可以自行管理sharedTmpBuffer内存空间,并在接口调用完成后,复用该部分内存,内存不会反复申请释放,灵活性较高,内存利用率也较高。 类型为LocalTensor,支持的TPosition为VECIN/VECCALC/VECOUT。 |
tiling |
输入 |
LayerNormGradBeta计算所需Tiling信息,Tiling信息的获取请参考LayerNormGradBeta Tiling。 |
isReuseSource |
输入 |
是否允许修改源操作数,默认值为false。如果开发者允许源操作数被改写,可以使能该参数,使能后能够节省部分内存空间。 设置为true,则本接口内部计算时复用inputDy的内存空间,节省内存空间;设置为false,则本接口内部计算时不复用inputDy的内存空间,本接口内部计算时会额外申请临时内存,接口调用完成后,自动释放临时内存。 对于float数据类型输入支持开启该参数,half数据类型输入不支持开启该参数。 |
返回值
无
支持的型号
Atlas A2训练系列产品/Atlas 800I A2推理产品
Atlas推理系列产品AI Core
约束说明
- 操作数地址偏移对齐要求请参见通用约束。
- src和dst的Tensor空间可以复用。
- 仅支持输入shape为ND格式。
- 输入数据不满足对齐要求时,开发者需要进行补齐,补齐的数据应设置为0,防止出现异常值从而影响网络计算。
- 不支持对尾轴H轴的切分。
调用示例
#include "kernel_operator.h" template <typename T, bool isReuseSource = false> class KernelLayernormGradBeta { public: __aicore__ inline KernelLayernormGradBeta() {} __aicore__ inline void Init(__gm__ uint8_t *resForGammaGm, __gm__ uint8_t *inputDyGm, __gm__ uint8_t *outputPdGammaGm, __gm__ uint8_t *outputPdBetaGm, const LayerNormGradBetaTiling &tiling) { this->bLength = tiling.bLength; this->sLength = tiling.sLength; this->hLength = tiling.hLength; this->tiling = tiling; bshLength = bLength * sLength * hLength; bsLength = bLength * sLength; resForGammaGlobal.SetGlobalBuffer(reinterpret_cast<__gm__ T *>(resForGammaGm), bshLength); inputDyGlobal.SetGlobalBuffer(reinterpret_cast<__gm__ T *>(inputDyGm), bshLength); outputPdGammaGlobal.SetGlobalBuffer(reinterpret_cast<__gm__ T *>(outputPdGammaGm), hLength); outputPdBetaGlobal.SetGlobalBuffer(reinterpret_cast<__gm__ T *>(outputPdBetaGm), hLength); pipe.InitBuffer(inQueueResForGamma, 1, sizeof(T) * bshLength); pipe.InitBuffer(inQueueDy, 1, sizeof(T) * bshLength); pipe.InitBuffer(outQueuePdGamma, 1, sizeof(T) * hLength); pipe.InitBuffer(outQueuePdBeta, 1, sizeof(T) * hLength); } __aicore__ inline void Process() { CopyIn(); Compute(); CopyOut(); } private: __aicore__ inline void CopyIn() { AscendC::LocalTensor<T> resForGammaLocal = inQueueResForGamma.AllocTensor<T>(); AscendC::LocalTensor<T> inputDyLocal = inQueueDy.AllocTensor<T>(); AscendC::DataCopy(resForGammaLocal, resForGammaGlobal, bshLength); AscendC::DataCopy(inputDyLocal, inputDyGlobal, bshLength); inQueueResForGamma.EnQue(resForGammaLocal); inQueueDy.EnQue(inputDyLocal); } __aicore__ inline void Compute() { AscendC::LocalTensor<T> resForGammaLocal = inQueueResForGamma.DeQue<T>(); AscendC::LocalTensor<T> inputDyLocal = inQueueDy.DeQue<T>(); AscendC::LocalTensor<T> outputPdGammaLocal = outQueuePdGamma.AllocTensor<T>(); AscendC::LocalTensor<T> outputPdBetaLocal = outQueuePdBeta.AllocTensor<T>(); AscendC::LayerNormGradBeta<T, isReuseSource>( outputPdGammaLocal, outputPdBetaLocal, resForGammaLocal, inputDyLocal, tiling); outQueuePdGamma.EnQue<T>(outputPdGammaLocal); outQueuePdBeta.EnQue<T>(outputPdBetaLocal); inQueueResForGamma.FreeTensor(resForGammaLocal); inQueueDy.FreeTensor(inputDyLocal); } __aicore__ inline void CopyOut() { AscendC::LocalTensor<T> outputPdGammaLocal = outQueuePdGamma.DeQue<T>(); AscendC::LocalTensor<T> outputPdBetaLocal = outQueuePdBeta.DeQue<T>(); AscendC::DataCopy(outputPdGammaGlobal, outputPdGammaLocal, hLength); AscendC::DataCopy(outputPdBetaGlobal, outputPdBetaLocal, hLength); outQueuePdGamma.FreeTensor(outputPdGammaLocal); outQueuePdBeta.FreeTensor(outputPdBetaLocal); } private: AscendC::TPipe pipe; AscendC::TQue<AscendC::QuePosition::VECIN, 1> inQueueResForGamma, inQueueDy; AscendC::TQue<AscendC::QuePosition::VECOUT, 1> outQueuePdGamma, outQueuePdBeta; AscendC::GlobalTensor<T> resForGammaGlobal; AscendC::GlobalTensor<T> inputDyGlobal; AscendC::GlobalTensor<T> outputPdGammaGlobal; AscendC::GlobalTensor<T> outputPdBetaGlobal; uint32_t bLength; uint32_t sLength; uint32_t hLength; uint32_t bshLength; uint32_t bsLength; LayerNormGradBetaTiling tiling; }; extern "C" __global__ __aicore__ void kernel_layernorm_grad_beta_operator( GM_ADDR outputPdGammaGm, GM_ADDR outputPdBetaGm, GM_ADDR resForGammaGm, GM_ADDR inputDyGm, GM_ADDR tiling) { GET_TILING_DATA(tilingData, tiling); KernelLayernormGradBeta<half, false> op; op.Init(resForGammaGm, inputDyGm, outputPdGammaGm, outputPdBetaGm, tilingData.layerNormGradBetaTiling); op.Process(); }