Cálculo de gradientes gamma y beta en LayerNorm con CANN/asc-devkit

La interfaz LayerNormGradBeta forma parte de AscendC, el lenguaje de desarrollo de operadores que CANN ofrece para los procesadores Ascend. Se utiliza durante la retropropagación de la normalización de capas para obtener los gradientes respecto a los parámetros gamma y beta. Junto con LayerNormGrad, produce los tensores pdx, gamma y beta.

Compatibilidad de productos

Producto	Soportado
Ascend 950PR / Ascend 950DT	√
Atlas A3 serie entrenamiento / Atlas A3 serie inferencia	√
Atlas A2 serie entrenamiento / Atlas A2 serie inferencia	√
Atlas 200I/500 A2 inferencia	×
Atlas serie inferencia AI Core	√
Atlas serie inferencia Vector Core	×
Atlas serie entrenamiento	×

Descripción funcional

Dado que el cálculo interno requiere almacenar variables intermedias, LayerNormGradBeta necesita espacio temporal adicional. El tamaño del búfer se obtiene a través del método GetLayerNormGradBetaMaxMinTmpSize disponible en el Tiling correspondiente. El valor mínimo garantiza la corrección funcional, mientras que el máximo mejora el rendimiento.

El espacio temporal puede gestionarse de dos formas:

Mediante el parámetro sharedTmpBuffer, donde el desarrollo administra y reutiliza la memoria.
Mediante el framework de la interfaz, donde solo es necesario reservar el tamaño indicado.

Prototipos de la función

Versión con búfer compartido gestionado por el usuario:

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);

Versión en la que el framework reserva el espacio temporal:

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);

Parámetros de plantilla

Parámetro	Descripción
`T`	Tipo de datos de los operandos. Soporta `half` y `float`.
`isReuseSource`	Indica si se permite modificar el tensor de origen. Por defecto es `false`. Si se establece en `true`, la interfaz reutiliza la memoria de `inputDy`, lo que ahorra espacio. Solo se admite para `float`; para `half` debe permanecer en `false`.

Parámetros de la interfaz

Parámetro	Dirección	Descripción
`outputPdGamma`	Salida	Gradiente respecto a `gamma`, con forma `[H]`. El eje final debe estar alineado a 32 B. Tipo `LocalTensor`; soporta `VECIN`, `VECCALC` y `VECOUT`.
`outputPdBeta`	Salida	Gradiente respecto a `beta`, con forma `[H]`. El eje final debe estar alineado a 32 B. Tipo `LocalTensor`; soporta `VECIN`, `VECCALC` y `VECOUT`.
`resForGamma`	Entrada	Resultado intermedio de `LayerNormGrad`, con forma `[B, S, H]`. Debe tener el mismo tipo que los operandos de salida y el eje final alineado a 32 B.
`inputDy`	Entrada	Gradiente proveniente de la capa posterior, con forma `[B, S, H]`. Debe coincidir en tipo con los operandos de salida y cumplir la alineación de 32 B en el eje final.
`sharedTmpBuffer`	Entrada	Búfer compartido para datos temporales del cálculo interno. Permite al desarrollador administrar y reutilizar la memoria. El tamaño se obtiene del Tiling de `LayerNormGradBeta`.
`tiling`	Entrada	Información de tiling necesaria para la ejecución. Ver `LayerNormGradBetaTiling`.

Valor de retorno

La interfaz no devuelve ningún valor.

Restricciones

Se debe cumplir el requisito general de alineación de direcciones de memoria.
Los espacios de los tensores de origen y destino pueden solaparse.
Solo se soportan tensores en formato ND.
Si los datos de entrada no cumplen la alineación requerida, deben rellenarse con ceros para evitar valores anómalos.
No se permite dividir el eje H.

Ejemplo de uso

// Variables de ejemplo:
// gradGamma : gradiente respecto a gamma, forma [H]
// gradBeta  : gradiente respecto a beta,  forma [H]
// tempGamma : resultado intermedio de LayerNormGrad, forma [B, S, H]
// upstreamDy: gradiente entrante de la capa posterior, forma [B, S, H]

AscendC::LayerNormGradBeta<T, isReuseSource>(
    gradGamma,   // salida: gradiente de gamma
    gradBeta,    // salida: gradiente de beta
    tempGamma,   // entrada: resultado intermedio
    upstreamDy,  // entrada: gradiente dy
    tiling       // entrada: información de tiling
);

Para un tensor de entrada con forma [1, 8, 8], los resultados numéricos esperados son los siguientes:

// dy (forma [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 ]

// tempGamma (forma [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 ]

// Resultado gradGamma:
// [ 8960,  9416,  9888, 10376, 10880, 11400, 11936, 12488 ]

// Resultado gradBeta:
// [ 224, 232, 240, 248, 256, 264, 272, 280 ]

Etiquetas: CANN AscendC LayerNorm AI-Core Tiling

Publicado el 6-30 00:58

Friki Work