大模型实战：深入解析 LLaMA 核心算子 RMSNorm 开发

前言

在前期课程中，我们打下了坚实的算子开发基础。从本期开始，我们将正式进入 '大模型（LLM）算子实战' 阶段。

为什么首选 RMSNorm？在 LLaMA、ChatGLM、Baichuan 等主流大模型中，传统的 LayerNorm 已经被 RMSNorm 全面取代。

1. 计算更简：去掉了均值（Mean）计算，减少了计算量。
2. 效果相当：在收敛性和稳定性上与 LayerNorm 几乎无异。
3. 实战典型：它包含 Element-wise（平方、乘法）和 Reduce（求和）两类操作，是练习 Ascend C 混合指令调用的绝佳案例。

本期文章，我们将手把手带你用 Ascend C 实现这个 LLM 的'标准零件'。

一、核心原理：给数据'去油解腻'

Normalization（归一化） 的作用是把数据拉回到一个标准的分布，防止训练过程中梯度爆炸或消失。

传统的 LayerNorm 公式：

$$y = \frac{x - \text{Mean}(x)}{\sqrt{\text{Var}(x) + \epsilon}} * \gamma + \beta$$

它需要算均值（Mean）和方差（Var）。

RMSNorm 认为减去均值没必要，直接除以均方根（RMS）就行，公式更简单：

$$y = \frac{x}{\text{RMS}(x)} * \gamma$$

其中：

$$\text{RMS}(x) = \sqrt{\frac{1}{n} \sum_{i=1}^{n} x_i^2 + \epsilon}$$

简单来说，就是三步走：平方求和 -> 开根号求倒数 -> 乘回原值。

二、关键挑战：Reduce 操作与精度控制

RMSNorm 开发中最大的难点在于 ReduceSum 和 精度保持。

1. Tiling 策略：通常采用 '按行切分'。即保证每一行数据（Row）完整地在一个 AI Core 上被处理。
   • 如果 Hidden Size ($N$) 较小（如 4096），可以一次搬进 UB，效率最高。
   • 如果 $N$ 很大（如 32K），UB 放不下，就需要分段计算平方和，最后汇总。为了简化演示，本期我们假设一行能放入 UB。
2. 精度陷阱：输入通常是 FP16。如果直接用 FP16 做平方和（Square + Sum），极易溢出或损失精度。黄金法则：在做 ReduceSum 时，必须转换成 FP32 进行累加。

三、代码实战：Ascend C 实现

3.1 Kernel 类定义

我们需要三个 Tensor：输入 x，参数 gamma，输出 y。

class KernelRMSNorm {
 public:
  __aicore__ inline void  {
    
    
    ->rowLength = rowLength;
    ->totalRows = totalRows;
    
  }

  {
    
    
     ( i = ; i < totalRows; i++) {
      (i);
      (i);
      (i);
    }
  }
  
};