fused_rms_norm

paddle.incubate.nn.functional. fused_rms_norm ( x, norm_weight, norm_bias, epsilon, begin_norm_axis, bias=None, residual=None, quant_scale=- 1, quant_round_type=0, quant_max_bound=0, quant_min_bound=0 ) [源代码]

应用 Fused RMSNorm 内核，提供了更高的 GPU 利用率。同时，支持模式融合 RMSNorm(bias + residual + x)。

细节可参考论文 Root Mean Square Layer Normalization 。

fused_rms_norm 算子目前只支持在 GPU 下运行。

参数

  • x (Tensor) - 输入 Tensor。

  • norm_weight (Tensor) - 用于仿射输出的权重张量。

  • norm_bias (Tensor) - 用于仿射输出的偏置张量。

  • epsilon (float) - 一个小的浮点数，用于避免除以零。

  • begin_norm_axis (int) - 归一化的起始轴，默认为 1。

  • bias (Tensor，可选) - 前一层的偏置。

  • residual (Tensor，可选) - 输入的残差。

  • quant_scale (float) - 量化缩放因子。

  • quant_round_type (float) - 量化四舍五入类型。

  • quant_max_bound (float) - 量化裁剪的最大边界值。

  • quant_min_bound (float) - 量化裁剪的最小边界值。

返回

输出 Tensor

形状

Tensor，形状同 x 一致。

代码示例

>>> import paddle
>>> paddle.device.set_device('gpu')

>>> paddle_x = paddle.cast(paddle.randn(shape=[32, 256]), dtype=paddle.float16)
>>> paddle_weight = paddle.cast(paddle.randn(shape=[256]), dtype=paddle.float16)
>>> paddle_bias = paddle.cast(paddle.randn(shape=[256]), dtype=paddle.float16)
>>> epsilon = 1e-6
>>> paddle_rmsnorm = paddle.incubate.nn.functional.fused_rms_norm(paddle_x, paddle_weight, paddle_bias, epsilon, 1)