Fix paddle.incubate.nn.functional.fused_rms_norm big Tensor

[xingmingyyj]

PR Category

Operator Mechanism

PR Types

Bug fixes

Description

• 修复fused_rms_norm 大Tensor 访存越界问题
• fused_rms_norm的kernel计算流程是将每一行加载到shared mem中，所以当col的取值过大时会导致kernel launch失败，在kernel中未做强制检查，导致kernel未launch就直接退出，输出结果变为全0。这里补充检查。
• 当输入数据类型为float16时，fused_rms_norm中会将数据cast成float32参与norm计算，以提升精度。在float16下，可以和下面的torch实现对齐精度。

class RMSNormFunction(torch.autograd.Function): @staticmethod def forward( ctx, x: torch.Tensor, norm_weight: torch.Tensor, norm_bias: Optional[torch.Tensor], epsilon: float, begin_norm_axis: int, bias: Optional[torch.Tensor] = None, residual: Optional[torch.Tensor] = None, quant_scale: float = -1.0, quant_round_type: int = 0, quant_max_bound: float = 0.0, quant_min_bound: float = 0.0, ) -> torch.Tensor: """Forward pass of RMSNorm.""" def _flatten_from_axis(tensor: torch.Tensor, axis: int) -> torch.Tensor: """Flatten tensor starting from given axis.""" rows = torch.prod(torch.tensor(tensor.shape[:axis])).item() cols = torch.prod(torch.tensor(tensor.shape[axis:])).item() return tensor.reshape(rows, cols) # Save original dtype and shape for later origin_dtype = x.dtype origin_shape = x.shape # Convert inputs to float32 if needed x = x.float() if x.dtype == torch.float16 else x norm_weight = norm_weight.float() if norm_weight.dtype == torch.float16 else norm_weight residual = residual.float() if residual is not None and residual.dtype == torch.float16 else residual bias = bias.float() if bias is not None and bias.dtype == torch.float16 else bias norm_bias = norm_bias.float() if norm_bias is not None and norm_bias.dtype == torch.float16 else norm_bias # Apply residual and bias if provided output = x if residual is not None: output = output + residual if bias is not None: output = output + bias # Normalization output = _flatten_from_axis(output, begin_norm_axis) output_sq = output.pow(2) mean_output_sq = output_sq.mean(dim=-1, keepdim=True) rms = torch.sqrt(mean_output_sq + epsilon) invvar = 1.0 / rms output_norm = output * invvar output = output_norm * norm_weight # Add norm_bias if provided if norm_bias is not None: output = output + _flatten_from_axis(norm_bias, begin_norm_axis) # Quantization if enabled if quant_scale > 0: output = output / quant_scale if quant_round_type == 0: output = torch.round(output) elif quant_round_type == 1: output = torch.where( output >= 0, torch.ceil(output - 0.5), torch.floor(output + 0.5), ) else: raise ValueError(f"Unsupported quant_round_type: {quant_round_type}") output = output * quant_scale output = torch.clamp(output, min=quant_min_bound, max=quant_max_bound) # Convert back to original dtype if no quantization if origin_dtype == torch.float16 and quant_scale <= 0: output = output.to(origin_dtype) norm_weight = norm_weight.to(origin_dtype) # Save tensors and metadata for backward ctx.save_for_backward(x, norm_weight, invvar) ctx.epsilon = epsilon ctx.exist_residual = residual is not None ctx.exist_bias = bias is not None ctx.exist_norm_bias = norm_bias is not None ctx.quant_scale = quant_scale ctx.begin_norm_axis = begin_norm_axis ctx.origin_shape = origin_shape return output @staticmethod def backward(ctx, grad_output: torch.Tensor) -> Tuple[Optional[torch.Tensor], ...]: """Backward pass of RMSNorm.""" def _flatten_from_axis(tensor: torch.Tensor, axis: int) -> torch.Tensor: """Flatten tensor starting from given axis.""" rows = torch.prod(torch.tensor(tensor.shape[:axis])).item() cols = torch.prod(torch.tensor(tensor.shape[axis:])).item() return tensor.reshape(rows, cols) exist_bias = ctx.exist_bias exist_residual = ctx.exist_residual exist_norm_bias = ctx.exist_norm_bias quant_scale = ctx.quant_scale if quant_scale > 0 or exist_norm_bias or exist_bias or exist_residual: raise NotImplementedError # Retrieve saved tensors and metadata x, weight, invvar = ctx.saved_tensors origin_shape = ctx.origin_shape origin_dtype = grad_output.dtype # Flatten tensors for computation grad_output = _flatten_from_axis(grad_output.float(), ctx.begin_norm_axis) x = _flatten_from_axis(x.float(), ctx.begin_norm_axis) weight = weight.float() # Gradient w.r.t. weight (gamma) x_norm = x * invvar grad_weight = (grad_output * x_norm).sum(dim=tuple(range(grad_output.dim() - 1)), keepdim=False) grad_weight = grad_weight.to(origin_dtype) # Gradient w.r.t. input (x) D = x.size(-1) S = (grad_output * weight * x * invvar).sum(dim=1, keepdim=True) term1 = invvar / D grad_x = (D * grad_output * weight - x * invvar * S) * term1 grad_x = grad_x.to(origin_dtype).reshape(origin_shape) # Return gradients (order matches forward inputs) return ( grad_x, # x grad_weight, # norm_weight None, # norm_bias None, # epsilon None, # begin_norm_axis None, # bias None, # residual None, # quant_scale None, # quant_round_type None, # quant_max_bound None, # quant_min_bound )

Pcard-73263

[paddle-bot]

你的PR提交成功，感谢你对开源项目的贡献!
请关注后续CI自动化测试结果，详情请参考Paddle-CI手册。
Your PR has been submitted. Thanks for your contribution!
Please wait for the result of CI firstly. See Paddle CI Manual for details.

[wanghuancoder]

const int64_t read_idx = static_cast<int64_t>(threadIdx.y) * blockDim.x + threadIdx.x;

[xingmingyyj]

done

[xingmingyyj]

/re-run all-failed

[wanghuancoder]

LGTM