Pytorch调用#

实际上PyTorch的调用接口是高度统一的，它们都继承自nn.Module，核心都是沿着某个维度计算统计量，之后标准化，可选仿射变换。

BN、LN、RMS的区别

区别在于沿着哪个轴做归一化，以及计算哪些统计量

理解了这个本质，剩下的就是查找参数，看维度。

BatchNorm#

1
torch.nn.BatchNorm1d(num_features, eps=1e-05, momentum=0.1, affine=True, ...)
2
torch.nn.BatchNorm2d(num_features, eps=1e-05, momentum=0.1, affine=True, ...)
3
torch.nn.BatchNorm3d(num_features, eps=1e-05, momentum=0.1, affine=True, ...)

对于BatchNorm来说，主要有上述三个版本，区别在于它们对于输入的维度布局不同，我们来看一下基本用法：

1
def batchnorm_example():
2
    x = torch.rand(100, 16, 784)
3
    layer = nn.BatchNorm1d(16)
4
    out = layer(x)
5

6
    y = torch.rand(1, 16, 7, 7)
7
    layer = nn.BatchNorm2d(16) # 传入通道数C
8
    out = layer(y)

Pytorch要求显式传入特征通道数，以便可以在内部初始化可学习参数weight和bias。它们的形状都是(C,)，实际上BatchNorm会根据自身是否为training来自动选择后续行为，对于调用者是完全透明的。

学习了官方的调用之后，我们来看一下具体是如何实现的：

1
def batchnorm(x, running_mean, running_var, weight, bias, training, momentum=0.1, eps=1e-5):
2

3

4
    """
5
        x.dim() 返回张量的维度数，例如对于一个形状为 (3, 4, 5) 的张量，x.dim() 将返回 3。
6
        range(x.dim()) 生成一个从 0 到 x.dim() - 1 的序列，例如生成 [0, 1, 2]。
7
        [d for d in range(x.dim()) if d != 1] 是列表推导式，它遍历上述生成的序列，并将不等于 1 的维度索引添加到列表中。
8
    """
9
    dims = [d for d in range(x.dim()) if d != 1]
10

11
    """
12
        [1: -1]是一个列表，第一个元素是1，第二个元素是-1
13
        [1] * (x.dim() - 2) 是一个列表，包含 x.dim() - 2 个元素，每个元素都是1。
14
        shape表示前两个维度是1和-1，后边的维度全都是1
15
    """
16
    shape = [1, -1] + [1] * (x.dim() - 2)
17

18
    # 训练模式
19
    if training:
20
        mean = x.mean(dim=dims, keepdim=True)                       # shape: (1, C, 1, 1)
21
        var_biased = x.var(dim=dims, keepdim=True, correction=0)    # 有偏估计，用于当前批次的归一化
22
        var_unbiased = x.var(dim=dims, keepdim=True, correction=1)  # 无偏估计，用于更新全局统计量
23

24
        with torch.no_grad():
25
            # running_mean 和running_var 形状是(C,)， 用squeeze()挤掉多余的维度
26
            running_mean.data = (1 - momentum) * running_mean.data + momentum * mean.squeeze()
27
            running_var.data  = (1 - momentum) * running_var.data  + momentum * var_unbiased.squeeze()
28
        var = var_biased
29
    else:
30
        mean = running_mean.view(shape)
31
        var = running_var.view(shape)
32

33
    # 归一化
34
    x_norm = (x - mean) / torch.sqrt(var + eps)
35
    # 仿射变换：y=γx+β
36
    if weight is not None:
37
        w = weight.view(shape)
38
        b = bias.view(shape)
39
        return x_norm * w + b
40
    return x_norm

具体内容如上，关键点我都已经写出来了，大家自行查阅即可。

优化思维时时刻刻都要有，我还写了一个优化版本的，mean、var_unbiased可以同时计算出来：

1
def manual_batchnorm_v2(x, running_mean, running_var, weight, bias, training, momentum=0.1, eps=1e-5):
2
    dims = [d for d in range(x.dim()) if d != 1]
3
    shape = [1, -1] + [1] * (x.dim() - 2)
4

5
    if training:
6
        # 计算参与归约的元素总数
7
        n = 1
8
        for d in dims:
9
            n *= x.shape[d]
10

11
        # 一次性计算均值和方差（无偏方差）
12
        var_unbiased, mean = torch.var_mean(x, dim=dims, keepdim=True, correction=1)
13
        # 转为有偏方差用于归一化
14
        var_biased = var_unbiased * ((n - 1) / n)
15

16
        with torch.no_grad():
17
            running_mean.data = (1 - momentum) * running_mean.data + momentum * mean.squeeze()
18
            running_var.data  = (1 - momentum) * running_var.data  + momentum * var_unbiased.squeeze()
19
        var = var_biased
20
    else:
21
        mean = running_mean.view(shape)
22
        var = running_var.view(shape)
23

24
    x_norm = (x - mean) / torch.sqrt(var + eps)
25
    if weight is not None:
26
        return x_norm * weight.view(shape) + bias.view(shape)
27
    return x_norm

LayerNorm#

LayerNorm不依赖batch统计量，在NLP和Transformer是绝对主力。

1
torch.nn.LayerNorm(normalized_shape, eps=1e-05, elementwise_affine=True, ...)
2
# shape可以是int或者tuple，表示输入最后几个维度的大小，归一化在这些维度进行
3
# eps，同BatchNorm
4
# elementwise_affine，是否学习逐元素的weight和bias

调用示例如下：

1
ln = nn.LayerNorm(512)
2
x = torch.randn(32, 10, 512)
3
out = ln.(x) # 形状不变，每个token的512维向量被归一化

我们手动实现的版本：

1
import torch
2
import torch.nn as nn
3

4

5

6
class layernorm(nn.Module):
7
    def __init__(self, embed_size, eps=1e-5):
8
        super().__init__()
9
        self.gamma = nn.Parameter(torch.ones(embed_size))
10
        self.beta = nn.Parameter(torch.zeros(embed_size))
11
        self.eps = eps
12

13
    def forward(self, x):
14
        # correction=0表示有偏估计
15
        var, mean = torch.var_mean(x, dim=-1, keepdim=True, correction=0)
16

17
        x_norm = (x - mean) / torch.sqrt(var + self.eps)
18
        return x_norm * self.gamma + self.beta
19

20
if __name__ == "__main__":
21
    x = torch.randn(2, 4, 3)
22
    embed_size = x.size(-1)
23
    official = nn.LayerNorm(embed_size, eps=1e-5)
24
    my  = layernorm(embed_size, eps=1e-5)
25

26
    # 复制相同权重
27
    my.gamma.data = official.weight.data.clone()
28
    my.beta.data  = official.bias.data.clone()
29

30
    diff = (official(x) - my(x)).abs().max().item()
31
    print(f"Max difference: {diff:.2e}")

RMSNorm#

RMSNorm是LayerNorm的一个简化变体，直接减去均值计算的步骤，只除以均方根，在最新的大模型里广泛使用。

1
torch.nn.RMSNorm(normalized_shape, eps=1e-06, elementwise_affine=True, ...)

使用方法完全同LayerNorm一致。

特性	BatchNorm	LayerNorm	RMSNorm
归一化轴	沿 batch 和空间维度 (每个通道独立)	沿特征维度 (每个样本独立)	沿特征维度 (每个样本独立)
依赖 batch	是	否	否
可学习参数	`weight` + `bias` (通道数)	`weight` + `bias` (归一化形状)	仅 `weight`
适用场景	CNN、CV	NLP、Transformer	大模型
PyTorch 接口	`BatchNorm1d/2d/3d`	`LayerNorm`	`RMSNorm`

Triton实现#

下边我们核心先讲解LN和RMS的Triton版本，BN就忽略掉，大家可以自行学习。

归一化的视角#

无论是LN还是RMS，输入都是一个二维矩阵 $X \in R^{N \times K}$ ，其中N是所有token的总行数，K是特征维度。

每一行的计算是完全独立的。对于第n行，算子聚合该行内所有K个元素，算出指定的统计量，再用统计量去逐元素变换该行。

Tip

这是经典的Reduce+Broadcast模式

具体来说，LayerNorm需要两步聚合，分别计算均值和方差；RMSNorm只进行依次平方和聚类。

在GPU上，Reduce是代价最高的部分，Broadcast是element wise操作，因此很明显，归一化操作是带宽密集型。

LayerNOrm#

我在这里参考了flash-attenion官方的实现: flash-attention

关于反向传播的代码

实际上，除了前向传播，在写绝大多数算子的时候，也要考虑反向传播，后续我会写一个完整的macroTorch，来完整的学习是如何进行前向传播与反向传播的。

关于自动调优，可以写一个函数来囊括进去常用的内容。同时我在这里保存了mean和rstd，这样能够给未来的反向传播直接复用，避免重新进行规约计算，因此可以看到我显式的进行了两次保存。

并且还有一个设计上的巧思，即步长参数化，而不是假设内存严格连续，这可以让kernel实现更加复杂的张量布局。

1
import torch
2
import torch.nn as nn
3
import triton
4
import triton.language as tl
5
from triton.testing import do_bench
6

7

8
def get_autotune_configs():
9
    warp_size = 32
10
    max_threads_per_block = 1024
11
    configs = []
12
    for num_warps in [1,2, 4, 8, 16, 32]:
13
        if num_warps * warp_size <= max_threads_per_block:
14
            configs.append(triton.Config({}, num_warps=num_warps))
15
    return configs
16

17
@triton.autotune(
18
        configs=get_autotune_configs(),
19
        key=["N"],
20
)
21
@triton.jit
22
def layer_norm_fwd_kernel(
23
    X_ptr, Y_ptr, W_ptr, B_ptr, Mean_ptr, Rstd_ptr,
24
    stride_x_row, stride_y_row,     # 传入行步长即可灵活索引，Triton编程常见模式
25
    N, eps, BLOCK_N: tl.constexpr
26
):
27
    # 每个program处理1行
28
    row_idx = tl.program_id(0)
29
    X_row_ptr = X_ptr + row_idx * stride_x_row
30
    Y_row_ptr = Y_ptr + row_idx * stride_y_row
31

32

33
    cols = tl.arange(0, BLOCK_N)
34
    mask = cols < N
35

36
    # 加载该行的所有元素，在计算中提升为FP32
37
    x = tl.load(X_row_ptr + cols, mask=mask, other=0.0).to(tl.float32)
38
    w = tl.load(W_ptr + cols, mask=mask).to(tl.float32)
39
    b = tl.load(B_ptr + cols, mask=mask).to(tl.float32)
40

41
    # 规约：均值 and 方差
42
    mean = tl.sum(x, axis=0) / N
43
    tl.store(Mean_ptr + row_idx, mean)
44

45
    x_bar = tl.where(mask, x - mean, 0.0)
46
    var = tl.sum(x_bar * x_bar, axis=0) / N
47
    rstd = 1.0 / tl.sqrt(var + eps)
48
    tl.store(Rstd_ptr + row_idx, rstd)
49

50
    # 广播：归一化+仿射变换
51
    y = (x - mean) * rstd * w + b
52
    tl.store(Y_row_ptr + cols, y, mask=mask)
53

54

55
def layer_norm_fwd(x, weight, bias=None, eps=1e-5):
56
    M, N = x.shape
57
    y = torch.empty_like(x)
58
    mean = torch.empty(M, device=x.device, dtype=torch.float32)
59
    rstd = torch.empty(M, device=x.device, dtype=torch.float32)
60

61
    if bias is None:
62
        bias = torch.zeros(N, device=x.device, dtype=weight.dtype)
63
    if weight is None:
64
        weight = torch.ones(N, device=x.device, dtype=x.dtype)
65

66
    BLOCK_N = triton.next_power_of_2(N)
67
    layer_norm_fwd_kernel[(M,)](
68
        x, y, weight, bias, mean, rstd,
69
        x.stride(0), y.stride(0),
70
        N, eps, BLOCK_N=BLOCK_N
71
    )
72
    return y, mean, rstd
73

74
def test_correctness(shapes=[(128, 256), (512, 1024)]):
75
    for M, N in shapes:
76
        x = torch.randn(M, N, device='cuda', dtype=torch.float32)
77
        weight = torch.randn(N, device='cuda', dtype=torch.float32)
78
        bias = torch.randn(N, device='cuda', dtype=torch.float32)
79
        eps = 1e-5
80

81
        ln = nn.LayerNorm(N, eps=eps).to('cuda')
82
        ln.weight.data = weight
83
        ln.bias.data = bias
84
        y_ref = ln(x)
85

86
        y_tri, _, _ = layer_norm_fwd(x, weight, bias, eps)
87

88
        max_diff = (y_tri - y_ref).abs().max().item()
89
        print(f"Shape ({M}, {N}): max diff = {max_diff:.6e}")
90

91

92
bench_perf_report = None
93
try:
94
    from triton.testing import Benchmark, perf_report
95

96
    @perf_report(
97
        Benchmark(
98
            x_names=["N"],
99
            x_vals=[256, 512, 1024, 2048, 4096, 8192],
100
            line_arg="provider",
101
            line_vals=["triton", "pytorch"],
102
            line_names=["Triton", "PyTorch"],
103
            styles=[("blue", "-"), ("red", "-")],
104
            ylabel="Latency (ms)",
105
            plot_name="LayerNorm Fwd Performance",
106
            args={"M": 1024, "eps": 1e-5, "dtype": torch.float32},
107
        )
108
    )
109
    def _bench_perf_report(M, N, eps, dtype, provider):
110
        device = 'cuda'
111
        x = torch.randn(M, N, device=device, dtype=dtype)
112
        weight = torch.randn(N, device=device, dtype=dtype)
113
        bias = torch.randn(N, device=device, dtype=dtype)
114

115
        if provider == "triton":
116
            def run():
117
                return layer_norm_fwd(x, weight, bias, eps)
118
        else:
119
            ln = nn.LayerNorm(N, eps=eps).to(device)
120
            ln.weight.data = weight
121
            ln.bias.data = bias
122
            def run():
123
                return ln(x)
124
        return do_bench(run, quantiles=[0.5, 0.2, 0.8])
125

126
    bench_perf_report = _bench_perf_report
127
except ImportError:
128
    print("当前 Triton 版本不支持 perf_report，跳过高阶绘图功能\n")
129

130

131
if __name__ == "__main__":
132
    test_correctness()
133
    if bench_perf_report is not None:
134
        bench_perf_report.run(show_plots=True, print_data=True, save_path="./layer_norm_fwd_perf.png")

性能对比如下：

RMSNorm#

LayerNorm需要计算两个统计量：mean和std。 mean就得涉及到规约，而RMSNorm只计算一个：

1
rms = sqrt(E[x^2] + eps)
2
y = x / rms * γ

核心代码如下：

1
@triton.autotune(
2
    configs=autotune_configs(),
3
    key=["N"],
4
)
5
@triton.jit
6
def rms_norm_fwd_kernel(
7
    X_ptr, Y_ptr, W_ptr,
8
    stride_x_row, stride_y_row,
9
    N, eps,
10
    BLOCK_N: tl.constexpr,
11
):
12
    row_idx = tl.program_id(0)
13
    X_row_ptr = X_ptr + row_idx * stride_x_row
14
    Y_row_ptr = Y_ptr + row_idx * stride_y_row
15

16
    cols = tl.arange(0, BLOCK_N)
17
    mask = cols < N
18

19
    x = tl.load(X_row_ptr + cols, mask=mask, other=0.0).to(tl.float32)
20
    w = tl.load(W_ptr + cols, mask=mask, other=0.0).to(tl.float32)
21

22
    # RMSNorm: rstd = rsqrt(mean(x^2) + eps)
23
    x2 = x * x
24
    mean_x2 = tl.sum(x2, axis=0) / N
25
    rstd = tl.rsqrt(mean_x2 + eps)
26
    # 广播：缩放+仿射变换
27
    y = x * rstd * w
28
    tl.store(Y_row_ptr + cols, y, mask=mask)
29

30

31
def rms_norm_fwd(x, weight, eps=1e-5):
32
    M, N = x.shape
33
    y = torch.empty_like(x)
34

35
    if weight is None:
36
        weight = torch.ones(N, device=x.device, dtype=x.dtype)
37

38
    BLOCK_N = triton.next_power_of_2(N)
39
    rms_norm_fwd_kernel[(M,)](
40
        x, y, weight,
41
        x.stride(0), y.stride(0),
42
        N, eps,
43
        BLOCK_N=BLOCK_N,
44
    )
45
    return y

我在特征维度从256到8192进行了测试，Triton版本始终优于Pytorch，并且自动调优保证了在不同尺寸下的均衡表现。

RMS在FP16下的优势更加明显，因为Triton对混合精度路径的控制更加精细，避免了不必要的类型转换开销。

1
====================================================================================================
2
RMSNorm Performance: dtype=torch.float16
3
====================================================================================================
4
     M       N       Triton(ms)     PyTorch (ms)    my vs PT
5
----------------------------------------------------------------------------------------------------
6
   128     256         0.003816         0.021189       5.55x
7
   128     512         0.004489         0.021973       4.89x
8
   128    1024         0.004901         0.023449       4.78x
9
   128    2048         0.005038         0.026430       5.25x
10
   128    4096         0.006675         0.032510       4.87x
11
   128    8192         0.008356         0.038082       4.56x
12
   512     256         0.004816         0.022823       4.74x
13
   512     512         0.005085         0.025100       4.94x
14
   512    1024         0.006658         0.029662       4.45x
15
   512    2048         0.008573         0.038389       4.48x
16
   512    4096         0.011734         0.058861       5.02x
17
   512    8192         0.022541         0.093368       4.14x
18
  1024     256         0.005728         0.025211       4.40x
19
  1024     512         0.006578         0.029643       4.51x
20
  1024    1024         0.009162         0.037848       4.13x
21
  1024    2048         0.011643         0.057440       4.93x
22
  1024    4096         0.020953         0.093154       4.45x
23
  1024    8192         0.041657         0.170544       4.09x
24
  2048     256         0.006578         0.029480       4.48x
25
  2048     512         0.007969         0.037949       4.76x
26
  2048    1024         0.011659         0.057081       4.90x
27
  2048    2048         0.021078         0.092297       4.38x
28
  2048    4096         0.041632         0.169116       4.06x
29
  2048    8192         0.077019         0.681771       8.85x
30
  4096     256         0.007757         0.037739       4.87x
31
  4096     512         0.012584         0.057165       4.54x
32
  4096    1024         0.020957         0.091967       4.39x
33
  4096    2048         0.040349         0.168362       4.17x
34
  4096    4096         0.077127         0.669127       8.68x
35
  4096    8192         0.149073         1.469917       9.86x
36

37
====================================================================================================
38
RMSNorm Performance: dtype=torch.float32
39
====================================================================================================
40
     M       N       Triton(ms)     PyTorch (ms)    my vs PT
41
----------------------------------------------------------------------------------------------------
42
   128     256         0.004033         0.016027       3.97x
43
   128     512         0.004791         0.016718       3.49x
44
   128    1024         0.005044         0.018569       3.68x
45
   128    2048         0.006306         0.021332       3.38x
46
   128    4096         0.009204         0.026772       2.91x
47
   128    8192         0.013090         0.030454       2.33x
48
   512     256         0.005348         0.017874       3.34x
49
   512     512         0.006673         0.019975       2.99x
50
   512    1024         0.008565         0.023732       2.77x
51
   512    2048         0.012434         0.031180       2.51x
52
   512    4096         0.020615         0.048329       2.34x
53
   512    8192         0.043017         0.077455       1.80x
54
  1024     256         0.006282         0.019679       3.13x
55
  1024     512         0.008333         0.023323       2.80x
56
  1024    1024         0.012538         0.030153       2.40x
57
  1024    2048         0.020061         0.046295       2.31x
58
  1024    4096         0.043250         0.076708       1.77x
59
  1024    8192         0.079363         0.145954       1.84x
60
  2048     256         0.008264         0.023329       2.82x
61
  2048     512         0.011833         0.030187       2.55x
62
  2048    1024         0.023103         0.046089       1.99x
63
  2048    2048         0.042146         0.075980       1.80x
64
  2048    4096         0.078298         0.144876       1.85x
65
  2048    8192         0.150213         0.479396       3.19x
66
  4096     256         0.012040         0.030257       2.51x
67
  4096     512         0.020640         0.045097       2.18x
68
  4096    1024         0.040291         0.075542       1.87x
69
  4096    2048         0.079232         0.144678       1.83x
70
  4096    4096         0.150072         0.465174       3.10x
71
  4096    8192         0.296848         1.028277       3.46x

FP32下加速比在2-4倍之间，FP16下可达4-9倍。

CUDA实现#

二者对于每一行N，都会遍历该行的所有元素，计算统计量（平方和、均值和方差等），并且使用这些统计量做归一化变换。

实际上，Norm的核心分为两个阶段：

1
阶段 1 — Reduce（归约）:
2
  每一行的 K 个元素 →  汇总成 1 个统计量（标量）
3

4
  例如：RMSNorm 把一行 8192 个 float 聚合成一个平方和
5
       LayerNorm 把一行 8192 个 float 先聚合成均值，再聚合成方差
6

7
阶段 2 — Broadcast（广播）:
8
  用这 1 个标量去变换该行的每一个元素
9

10
  y[n][k] = f(x[n][k], stat[n])

用 C++ 伪代码表达就是：

1
for (int n = 0; n < N; n++) {         // N 行，行间独立
2
    float stat = 0;
3
    for (int k = 0; k < K; k++) {     // Reduce: K 个元素 → 1 个值
4
        stat += compute(x[n][k]);
5
    }
6
    stat = finalize(stat);             // 例如 rsqrt(stat/K + eps)
7

8
    for (int k = 0; k < K; k++) {     // Broadcast: 1 个值 → K 个元素
9
        y[n][k] = transform(x[n][k], stat, gamma[k], beta[k]);
10
    }
11
}

Important

所有归一化的算子，本质上都是Reduce + Broadcast

音乐

Pytorch调用#

BatchNorm#

LayerNorm#

RMSNorm#

Triton实现#

归一化的视角#

LayerNOrm#

RMSNorm#

CUDA实现#

支持与分享

音乐

目录

音乐

LeetGPU习题07：Norm系列代码实现

Pytorch调用#

BatchNorm#

LayerNorm#

RMSNorm#

Triton实现#

归一化的视角#

LayerNOrm#

RMSNorm#

CUDA实现#

支持与分享

音乐

目录