layer_norm.h

LayerNorm 函数。

函数

void nvte_layernorm_fwd(const NVTETensor x, const NVTETensor gamma, const NVTETensor beta, const float epsilon, NVTETensor z, NVTETensor mu, NVTETensor rsigma, cudaStream_t stream, const int multiprocessorCount, NVTETensor workspace, NVTETensor barrier)

计算输入的 LayerNorm。

使用的公式

\[ y = \frac{x - E[x]}{\sqrt{Var[x] + \varepsilon}}\gamma + \beta \]

如果将 workspace 和 barrier 设置为空张量来调用此函数，将不会执行该操作，而是将 workspace 和 barrier 张量的形状和类型设置为所需的值。

参数

• x – [in] 形状为 [N, H] 的输入张量。

• gamma – [in] 形状为 [H] 的 Gamma 张量。

• beta – [in] 形状为 [H] 的 Beta 张量。

• epsilon – [in] 添加到分母以提高数值稳定性的值。

• z – [inout] 形状为 [N, H] 的输出张量。

• mu – [out] 输入在最后一个维度上计算的均值。形状：[N]。

• rsigma – [out] 输入在最后一个维度上计算的方差的倒数。形状：[N]。

• stream – [in] 用于操作的 CUDA 流。

• multiprocessorCount – [in] 设备中 SM 的数量。

• workspace – [out] 工作区张量。

• barrier – [out] Barrier 张量。

void nvte_layernorm1p_fwd(const NVTETensor x, const NVTETensor gamma, const NVTETensor beta, const float epsilon, NVTETensor z, NVTETensor mu, NVTETensor rsigma, cudaStream_t stream, const int multiprocessorCount, NVTETensor workspace, NVTETensor barrier)

计算对输入进行零中心化 gamma 的 LayerNorm。

使用的公式

\[ y = \frac{x - E[x]}{\sqrt{Var[x] + \varepsilon}}(1 + \gamma) + \beta \]

如果将 workspace 和 barrier 设置为空张量来调用此函数，将不会执行该操作，而是将 workspace 和 barrier 张量的形状和类型设置为所需的值。

参数

• x – [in] 形状为 [N, H] 的输入张量。

• gamma – [in] 形状为 [H] 的 Gamma 张量。

• beta – [in] 形状为 [H] 的 Beta 张量。

• epsilon – [in] 添加到分母以提高数值稳定性的值。

• z – [inout] 形状为 [N, H] 的输出张量。

• mu – [out] 输入在最后一个维度上计算的均值。形状：[N]。

• rsigma – [out] 输入在最后一个维度上计算的方差的倒数。形状：[N]。

• stream – [in] 用于操作的 CUDA 流。

• multiprocessorCount – [in] 设备中 SM 的数量。

• workspace – [out] 工作区张量。

• barrier – [out] Barrier 张量。

void nvte_layernorm_bwd(const NVTETensor dz, const NVTETensor x, const NVTETensor mu, const NVTETensor rsigma, const NVTETensor gamma, NVTETensor dx, NVTETensor dgamma, NVTETensor dbeta, NVTETensor dgamma_part, NVTETensor dbeta_part, cudaStream_t stream, const int multiprocessorCount, NVTETensor workspace, NVTETensor barrier)

计算 LayerNorm 的反向传播。

此函数计算函数的梯度

\[ y = \frac{x - E[x]}{\sqrt{Var[x] + \varepsilon}}\gamma + \beta \]

关于 \(x\), \(\gamma\) 和 \(\beta\)。

调用此函数时，将 workspace、barrier、dgamma_part 和 dbeta_part 设置为空张量将不会执行操作，而是将这些张量的形状和类型设置为所需的值。

参数

• dz – [输入] 形状为 [N, H] 的传入梯度张量。

• x – [输入] 形状为 [N, H] 的前向输入张量。

• mu – [输入] 在最后一个维度上计算的输入的平均值。 形状：[N]。

• rsigma – [输入] 在最后一个维度上计算的输入的方差的倒数。 形状：[N]。

• gamma – [in] 形状为 [H] 的 Gamma 张量。

• dx – [输出] 形状为 [N, H] 的输出梯度。

• dgamma – [输出] 形状为 [H] 的 gamma 张量的梯度。

• dbeta – [输出] 形状为 [H] 的 beta 张量的梯度。

• dgamma_part – [输出] 用于存储部分 gamma 梯度的存储。

• dbeta_part – [输出] 用于存储部分偏差梯度的存储。

• stream – [in] 用于操作的 CUDA 流。

• multiprocessorCount – [in] 设备中 SM 的数量。

• workspace – [out] 工作区张量。

• barrier – [out] Barrier 张量。

void nvte_layernorm1p_bwd(const NVTETensor dz, const NVTETensor x, const NVTETensor mu, const NVTETensor rsigma, const NVTETensor gamma, NVTETensor dx, NVTETensor dgamma, NVTETensor dbeta, NVTETensor dgamma_part, NVTETensor dbeta_part, cudaStream_t stream, const int multiprocessorCount, NVTETensor workspace, NVTETensor barrier)

计算具有零中心化 gamma 的 LayerNorm 的反向传播。

此函数计算函数的梯度

\[ y = \frac{x - E[x]}{\sqrt{Var[x] + \varepsilon}}(1 + \gamma) + \beta \]

关于 \(x\), \(\gamma\) 和 \(\beta\)。

调用此函数时，将 workspace、barrier、dgamma_part 和 dbeta_part 设置为空张量将不会执行操作，而是将这些张量的形状和类型设置为所需的值。

参数

• dz – [输入] 形状为 [N, H] 的传入梯度张量。

• x – [输入] 形状为 [N, H] 的前向输入张量。

• mu – [输入] 在最后一个维度上计算的输入的平均值。 形状：[N]。

• rsigma – [输入] 在最后一个维度上计算的输入的方差的倒数。 形状：[N]。

• gamma – [in] 形状为 [H] 的 Gamma 张量。

• dx – [输出] 形状为 [N, H] 的输出梯度。

• dgamma – [输出] 形状为 [H] 的 gamma 张量的梯度。

• dbeta – [输出] 形状为 [H] 的 beta 张量的梯度。

• dgamma_part – [输出] 用于存储部分 gamma 梯度的存储。

• dbeta_part – [输出] 用于存储部分偏差梯度的存储。

• stream – [in] 用于操作的 CUDA 流。

• multiprocessorCount – [in] 设备中 SM 的数量。

• workspace – [out] 工作区张量。

• barrier – [out] Barrier 张量。