normalization.h

LayerNorm 和 RMSNorm 函数。

函数

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

计算输入的 LayerNorm。

使用的公式

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

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

参数:

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

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

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

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

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

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

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

• workspace – [out] Workspace 张量。

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

• zero_centered_gamma – [in] 将归一化值乘以 \( \gamma+1 \) 而不是 \( \gamma \)

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

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 workspace, const int multiprocessorCount, const bool zero_centered_gamma, cudaStream_t stream)

计算 LayerNorm 的反向传播。

此函数计算函数

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

关于 \(x\)、\(\gamma\) 和 \(\beta\) 的梯度。

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

参数:

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

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

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

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

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

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

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

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

• workspace – [out] Workspace 张量。

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

• zero_centered_gamma – [in] 将归一化值乘以 \( \gamma+1 \) 而不是 \( \gamma \)

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

void nvte_rmsnorm_fwd(const NVTETensor x, const NVTETensor gamma, const float epsilon, NVTETensor z, NVTETensor rsigma, NVTETensor workspace, const int multiprocessorCount, const bool zero_centered_gamma, cudaStream_t stream)

计算 RMSNorm。

使用的公式

\[ y = \frac{x}{RMS_\varepsilon(x)}\gamma \]

其中

\[ RMS_\varepsilon(x) = \sqrt{\frac{1}{n}\sum_{i=0}^{n-1} x_i^2 + \varepsilon} \]

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

参数:

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

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

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

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

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

• workspace – [out] Workspace 张量。

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

• zero_centered_gamma – [in] 将归一化值乘以 \( \gamma+1 \) 而不是 \( \gamma \)

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

void nvte_rmsnorm_bwd(const NVTETensor dz, const NVTETensor x, const NVTETensor rsigma, const NVTETensor gamma, NVTETensor dx, NVTETensor dgamma, NVTETensor workspace, const int multiprocessorCount, const bool zero_centered_gamma, cudaStream_t stream)

计算 RMSNorm 的反向传播。

此函数计算函数

\[ y = \frac{x}{RMS_\varepsilon(x)}\gamma \]

其中

\[ RMS_\varepsilon(x) = \sqrt{\frac{1}{n}\sum_{i=0}^{n-1} x_i^2 + \varepsilon} \]

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

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

参数:

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

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

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

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

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

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

• workspace – [out] Workspace 张量。

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

• zero_centered_gamma – [in] 将归一化值乘以 \( \gamma+1 \) 而不是 \( \gamma \)

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

void nvte_enable_cudnn_norm_fwd(bool enable)

启用 normalization 的 cuDNN 后端的辅助函数。

参数:

bool – [in] 如果为 True 则启用

void nvte_enable_cudnn_norm_bwd(bool enable)