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 ) [source]

Apply Fused RMSNorm kernel. Also support RMSNorm(bias + residual + x) fused pattern.

Parameters

  • x (Tensor) – the input Tensor..

  • norm_weight (Tensor) – the weight Tensor to affine output.

  • norm_bias (Tensor) – the bias Tensor to affine output.

  • epsilon (float) – a small float number to avoid divide 0.

  • begin_norm_axis (int) – the begin axis to normalize.

  • bias (optional|Tensor) – the previous layers’s bias to fused.

  • residual (optional|Tensor) – the residual input to fused.

  • quant_scale (float) – the quant scale.

  • quant_round_type (float) – the quant round type.

  • quant_max_bound (float) – the quant max bound to clip.

  • quant_min_bound (float) – the quant min bound to clip.

Returns

the output Tensor.

Return type

Tensor

Examples

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