llm-algo-leetcode Task 01:PyTorch Warmup & RMSNorm

PyTorch Warmup & RMSNorm 学习笔记

来源章节：

• Datawhale LLM Algorithm Practice Lab：00. PyTorch Warmup
• Datawhale LLM Algorithm Practice Lab：01. RMSNorm Tutorial

---

0. 今天这两章的主线

今天学的内容可以分成两条线：

第一条线是 PyTorch 基础操作：

• Tensor 到底是什么
• Tensor shape 怎么变换
• Embedding 为什么本质是查表
• nn.Module 为什么可以像函数一样被调用
• 自定义 autograd.Function 里的 ctx 是干什么的
• 梯度是什么，为什么要算 grad_x / grad_weight / grad_bias

第二条线是 RMSNorm：

• RMSNorm 的本质是什么
• 为什么它比 LayerNorm 更简单
• 为什么 LLM 里常用 RMSNorm
• dim=-1、keepdim=True、torch.rsqrt 这些实现细节怎么理解

这两章看起来一个是基础，一个是归一化层，但其实底层都围绕一个核心能力：

看懂 Tensor 的形状变化，以及每一步数学操作在神经网络训练中的意义。

---

Part 1：Tensor 与 shape

1.1 Tensor 的本质

Tensor 可以先理解成：

一堆数字 + shape + dtype + device + 是否需要梯度

在 PyTorch 中，输入、参数、中间结果、输出和梯度，基本都是 Tensor。

例如：

x = torch.randn(2, 3, 4)

它不是一个普通三层列表，而是一个三维 Tensor：

shape = [2, 3, 4]

理解 Tensor 最重要的不是先看具体数值，而是先看 shape。

常见 shape 含义：

[B]             一批标量
[B, D]          B 个样本，每个样本 D 个特征
[B, S]          B 个序列，每个序列 S 个 token id
[B, S, H]       B 个序列，每个 token 是 H 维向量
[B, C, H, W]    图像 batch，C 个通道，高 H，宽 W

今天后面所有问题，本质都在问：

这个 Tensor 的 shape 是什么？这个操作沿哪个维度做？结果 shape 会变成什么？

---

1.2 图像 Tensor 为什么要从 [B, C, H, W] 变成 [B, H*W, C]

在图像或多模态模型中，图像特征常见 shape 是：

[B, C, H, W]

含义是：

B：batch size
C：channel / feature 维度
H：height
W：width

但 Transformer 更喜欢吃“序列”：

[B, seq_len, hidden_dim]

所以需要把图像的每个空间位置当成一个 token：

seq_len = H * W
hidden_dim = C

目标 shape：

[B, H*W, C]

原生 PyTorch 写法：

B, C, H, W = x.shape
x_native = x.permute(0, 2, 3, 1).reshape(B, H * W, C)

逐步理解：

原始： [B, C, H, W]
permute(0, 2, 3, 1)
变成： [B, H, W, C]
reshape(B, H*W, C)
变成： [B, H*W, C]

permute 是调整维度顺序，reshape 是合并维度。

等价的 einops 写法：

from einops import rearrange

x_einops = rearrange(x, "b c h w -> b (h w) c")

这个写法更像在直接描述意图：

b c h w -> b (h w) c

也就是：保留 batch，把 h 和 w 合并，把 c 放到最后作为特征维。

---

Part 2：Embedding 的本质

2.1 为什么 token id 不能直接给神经网络用

文本会先被 tokenizer 转成整数 id：

input_ids = torch.tensor([10, 42, 99])

这些数字只是词表编号，不代表大小关系。

例如：

token 99 并不比 token 42 “更大”或“更强”

神经网络真正需要的是连续向量，所以要用 Embedding 把离散 id 变成 dense vector。

---

2.2 Embedding 本质是查表

Embedding 层内部有一个权重矩阵：

emb_layer.weight.shape = [vocab_size, hidden_dim]

可以理解成一张大表：

第 0 行：token 0 的向量
第 1 行：token 1 的向量
第 2 行：token 2 的向量
...

当输入是：

input_ids = torch.tensor([2, 0, 3])

那么：

emb_layer.weight[input_ids]

就是取：

[weight[2], weight[0], weight[3]]

所以：

out_official = emb_layer(input_ids)
out_manual = emb_layer.weight[input_ids]

在普通默认情况下，二者前向结果是一样的。

---

2.3 为什么 emb_layer(input_ids) 能这样写

因为 emb_layer 是一个 nn.Module。

PyTorch 里模块可以像函数一样调用：

out = emb_layer(input_ids)

本质上会进入：

emb_layer.forward(input_ids)

Embedding 的 forward 逻辑就是：把 input_ids 当行号，去 weight 中查对应行。

需要注意：

input_ids 必须是整数 Tensor，通常 dtype 是 torch.long
input_ids 里的值必须在 [0, vocab_size - 1] 范围内

---

2.4 为什么 Tensor 可以放在 [] 里当索引

这一句：

out_manual = emb_layer.weight[input_ids]

用到了 PyTorch 的 Tensor indexing / advanced indexing。

input_ids 是一个整数 Tensor，里面每个值都是要取的行号。

例如：

weight = torch.tensor([
    [0.1, 0.2],
    [0.3, 0.4],
    [0.5, 0.6],
    [0.7, 0.8],
])

ids = torch.tensor([2, 0, 3])
out = weight[ids]

结果是：

[
  weight[2],
  weight[0],
  weight[3]
]

也就是：

[
  [0.5, 0.6],
  [0.1, 0.2],
  [0.7, 0.8]
]

所以 Tensor 放在 [] 里时，不是当普通数值计算，而是当一批索引。

---

Part 3：Linear 层、参数数量与矩阵乘法

3.1 Linear 层不是只有一个 w

一个 Linear 层：

linear = nn.Linear(in_features, out_features)

它的参数形状是：

weight.shape = [out_features, in_features]
bias.shape   = [out_features]

权重数量是：

in_features × out_features

只有当：

1 个输入 → 1 个输出

才只有一个 w。

---

3.2 例子：1 个输入 → 1 个输出

linear = nn.Linear(1, 1)

公式：

z = x * w + b

参数：

weight.shape = [1, 1]
bias.shape   = [1]

这时确实只有一个 w。

---

3.3 例子：2 个输入 → 1 个输出

如果输入是：

x = [x1, x2]

输出只有一个：

z1

公式是：

z1 = x1*w1 + x2*w2 + b1

这时需要两个权重：

w1 连接 x1
w2 连接 x2

PyTorch 中：

linear = nn.Linear(2, 1)

参数 shape：

weight.shape = [1, 2]
bias.shape   = [1]

注意，PyTorch 里 weight 是 [out_features, in_features]，所以是 [1, 2]，不是 [2, 1]。

---

3.4 例子：2 个输入 → 2 个输出

如果输入是：

x = [x1, x2]

输出是：

z = [z1, z2]

那么每个输出神经元都要看所有输入：

z1 = x1*w11 + x2*w12 + b1
z2 = x1*w21 + x2*w22 + b2

所以有 4 个权重：

w11, w12, w21, w22

写成矩阵：

weight = [
  [w11, w12],
  [w21, w22]
]

PyTorch 中：

linear = nn.Linear(2, 2)

参数 shape：

weight.shape = [2, 2]
bias.shape   = [2]

这里的“2×2”不是说只能处理一个 2×2 的输入，而是说：

输入最后一维必须是 2
输出最后一维会变成 2

所以这些输入都可以：

x = torch.randn(1, 2)
x = torch.randn(5, 2)
x = torch.randn(32, 2)

因为最后一维都是 2。

---

3.5 为什么前向是 x @ weight.T

PyTorch 的 Linear 前向等价于：

z = x @ weight.T + bias

假设：

x.shape      = [B, in_features]
weight.shape = [out_features, in_features]

那么：

weight.T.shape = [in_features, out_features]

所以：

[B, in_features] @ [in_features, out_features]
= [B, out_features]

这正是 Linear 想要的输出 shape。

@ 表示矩阵乘法，等价于：

torch.matmul(a, b)

---

Part 4：梯度与反向传播

4.1 梯度到底是什么

梯度可以先理解成：

某个变量稍微变一点，结果会怎么变

在神经网络训练里，最关心的是：

某个参数稍微变一点，loss 会怎么变

所以：

梯度 = loss 对某个变量的变化率

如果某个参数的梯度是正的，说明参数增大时 loss 会增大，所以应该让参数减小。

如果某个参数的梯度是负的，说明参数增大时 loss 会减小，所以应该让参数增大。

统一更新公式是：

param = param - learning_rate * param.grad

---

4.2 为什么还要计算对 x 的梯度

在一层 Linear 里：

z = x @ weight.T + bias

反向传播要计算：

grad_x
grad_weight
grad_bias

其中：

grad_weight / grad_bias：用于更新当前层参数
grad_x：用于把梯度继续传给上一层

x 本身通常不是参数，但它可能是上一层的输出。

例如：

input → layer1 → x → layer2 → loss

对 layer2 来说，x 是输入。

但对 layer1 来说，x 是它的输出。

如果 layer2 不计算 grad_x，梯度就传不回 layer1，前面的参数就没法更新。

所以：

grad_x 不是为了更新 x 本身，而是为了继续反传。

---

4.3 Linear 的反向传播公式

前向：

z = x @ weight.T + bias

反向：

grad_x = grad_z @ weight
grad_weight = grad_z.T @ x
grad_bias = grad_z.sum(dim=0)

维度关系：

x.shape          = [B, in_features]
weight.shape     = [out_features, in_features]
z.shape          = [B, out_features]
grad_z.shape     = [B, out_features]

---

4.4 为什么 grad_x = grad_z @ weight

前向：

z = x @ weight.T

假设有 2 个输入、2 个输出：

x = [x1, x2]

weight = [
  [w11, w12],
  [w21, w22]
]

前向展开：

z1 = x1*w11 + x2*w12
z2 = x1*w21 + x2*w22

现在反向时，假设上游给了：

grad_z = [dz1, dz2]

意思是：

loss 对 z1 的梯度是 dz1
loss 对 z2 的梯度是 dz2

x1 同时影响 z1 和 z2：

x1 → z1，系数是 w11
x1 → z2，系数是 w21

所以：

grad_x1 = dz1*w11 + dz2*w21

同理：

grad_x2 = dz1*w12 + dz2*w22

写成矩阵形式就是：

grad_x = grad_z @ weight

这一步的直觉是：

输出梯度 grad_z 按照 weight 记录的连接关系，汇总回每个输入维度。

维度也正好匹配：

[B, out_features] @ [out_features, in_features]
= [B, in_features]

而 grad_x 必须和 x 的 shape 一样。

---

4.5 为什么 grad_bias = grad_z.sum(dim=0)

前向中：

z = x @ weight.T + bias

bias.shape = [out_features]。

但它会被广播加到 batch 里的每一个样本上。

例如 batch size 为 3：

z[0] = ... + bias
z[1] = ... + bias
z[2] = ... + bias

同一个 bias 被用了 3 次。

如果：

grad_z = [
  [1, 10],
  [2, 20],
  [3, 30]
]

那么：

grad_b1 = 1 + 2 + 3 = 6
grad_b2 = 10 + 20 + 30 = 60

所以：

grad_bias = grad_z.sum(dim=0)

因为 dim=0 是 batch 维度，沿 batch 求和后，shape 变成 [out_features]，正好和 bias.shape 一样。

---

4.6 ctx 是什么

在自定义 autograd 函数里：

class LinearReLUFunction(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, weight, bias):
        ...

ctx 是 context，上下文对象。

作用是：

forward 阶段保存反向传播需要的信息
backward 阶段再取出来用

例如：

ctx.save_for_backward(x, weight, mask)

表示保存：

x：计算 grad_weight 要用
weight：计算 grad_x 要用
mask：ReLU 反向传播要用

在 backward 里：

x, weight, mask = ctx.saved_tensors

把 forward 保存的张量取出来。

可以把 ctx 理解成 forward 和 backward 之间的临时小背包。

---

4.7 ReLU 的反向传播

ReLU 前向：

y = relu(z)

规则：

z > 0：输出 y = z，梯度可以通过
z <= 0：输出 y = 0，梯度被截断

所以 forward 里保存：

mask = z > 0

backward 里：

grad_z = grad_output * mask

意思是：

只有前向时 z > 0 的位置，梯度才能继续往前传。

---

4.8 Linear + ReLU 自定义 Function 完整代码

import torch
import torch.nn.functional as F


class LinearReLUFunction(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, weight, bias):
        z = F.linear(x, weight, bias)
        y = F.relu(z)
        mask = z > 0
        ctx.save_for_backward(x, weight, mask)
        return y

    @staticmethod
    def backward(ctx, grad_output):
        x, weight, mask = ctx.saved_tensors

        grad_z = grad_output * mask

        grad_x = grad_z @ weight

        x_2d = x.reshape(-1, x.shape[-1])
        grad_z_2d = grad_z.reshape(-1, grad_z.shape[-1])

        grad_weight = grad_z_2d.T @ x_2d
        grad_bias = grad_z_2d.sum(dim=0)

        return grad_x, grad_weight, grad_bias

---

Part 5：RMSNorm

5.1 RMSNorm 的本质

RMSNorm，全称 Root Mean Square Normalization。

它的本质是：

把每个 token 的 hidden vector 按整体大小缩放到稳定范围。

核心公式：

RMSNorm(x) = x / sqrt(mean(x^2) + eps) * weight

它做两件事：

1. 用 RMS 控制 hidden vector 的整体尺度
2. 用可学习 weight 调整每个 hidden 维度的重要性

---

5.2 RMSNorm 和 LayerNorm 的区别

LayerNorm：

先减均值，再除标准差

公式直觉：

LayerNorm(x) = (x - mean(x)) / std(x) * gamma + beta

RMSNorm：

不减均值，只除以均方根

公式：

RMSNorm(x) = x / sqrt(mean(x^2) + eps) * weight

对比：

项目	LayerNorm	RMSNorm
是否减均值	是	否
是否控制尺度	是	是
参数	gamma + beta	weight
计算量	更大	更小
直觉	居中 + 缩放	只缩放

RMSNorm 不是绝对更好，而是在很多 LLM 中：

计算更少，速度更快，训练稳定性通常够用。

---

5.3 为什么 RMSNorm 在 LLM 里常用

LLM 中每个 token 都有一个很长的 hidden vector，例如：

hidden_size = 4096

Transformer 每层都可能有 Norm：

RMSNorm → Attention → RMSNorm → MLP

层数很多、token 很多、hidden_size 很大，所以 Norm 的计算成本会不断累计。

RMSNorm 省掉了 LayerNorm 中的减均值和中心化方差计算，因此更轻量。

它主要控制的是：

hidden vector 的整体大小不要爆炸或消失

在很多大模型结构里，只控制尺度已经足够稳定。

---

5.4 RMSNorm 代码实现

import torch
import torch.nn as nn


class RMSNorm(nn.Module):
    def __init__(self, hidden_size: int, eps: float = 1e-6):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(hidden_size))

    def _norm(self, x: torch.Tensor) -> torch.Tensor:
        x_float = x.float()
        variance = x_float.pow(2).mean(dim=-1, keepdim=True)
        return x_float * torch.rsqrt(variance + self.eps)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        input_dtype = x.dtype
        output = self._norm(x) * self.weight
        return output.to(input_dtype)

---

5.5 为什么 weight 是 [hidden_size]

RMSNorm 是对每个 token 的 hidden vector 做归一化。

如果：

x.shape = [batch, seq_len, hidden_size]

那么每个 token 有 hidden_size 个维度。

归一化之后，模型仍然需要决定每个维度是否应该放大或缩小，所以有一个可学习参数：

self.weight = nn.Parameter(torch.ones(hidden_size))

它的 shape 是：

[hidden_size]

例如：

hidden_size = 4
weight = [w1, w2, w3, w4]

每个 hidden 维度一个缩放系数。

初始化为全 1，是因为一开始不额外改变归一化结果的大小。

---

5.6 为什么要 x.float()

RMSNorm 要计算：

x.pow(2).mean(...)

如果输入是 FP16，平方时容易有数值溢出或精度不足。

所以通常先转成 float32：

x_float = x.float()

计算完成后，再转回原始 dtype：

return output.to(input_dtype)

这样既保证 RMSNorm 的计算稳定，又不破坏后续模型的混合精度流程。

---

5.7 dim=-1 怎么理解

dim=-1 表示最后一个维度。

如果：

x.shape = [batch, seq_len, hidden_size]

那么：

dim=0  是 batch
dim=1  是 seq_len
dim=2  是 hidden_size
dim=-1 也是 hidden_size

所以：

x_float.pow(2).mean(dim=-1, keepdim=True)

意思是：

对每个 token 的 hidden vector 内部求 mean(x^2)

例如一个 token：

x = [1, 2, 3, 4]

平方：

x^2 = [1, 4, 9, 16]

求均值：

mean(x^2) = 7.5

RMSNorm 就是用这个值来缩放整个 token 的 hidden vector。

---

5.8 为什么要 keepdim=True

假设：

x.shape = [2, 3, 4]

执行：

variance = x.pow(2).mean(dim=-1, keepdim=True)

结果：

variance.shape = [2, 3, 1]

最后一维从 4 个数变成 1 个数，但这个维度被保留了。

这样后面可以广播：

x.shape        = [2, 3, 4]
variance.shape = [2, 3, 1]

相乘时：

[2, 3, 1] 会自动广播成 [2, 3, 4]

也就是每个 token 用自己的一个缩放系数去缩放自己的所有 hidden 值。

---

5.9 torch.rsqrt 是什么

torch.rsqrt(a)

表示：

1 / sqrt(a)

所以：

x_float * torch.rsqrt(variance + self.eps)

等价于：

x_float / torch.sqrt(variance + self.eps)

使用 rsqrt 的好处是表达更直接，也通常更适合底层优化。

---

Part 6：今天最容易混的点

6.1 weight[input_ids] 里的 Tensor 不是参与计算，而是当索引

emb_layer.weight[input_ids]

不是把 input_ids 当普通数字去乘，而是把它当行号列表。

---

6.2 Linear 层不是只有一个 w

nn.Linear(in_features, out_features)

权重数量是：

in_features × out_features

一个输出神经元需要一组权重。

多个输出神经元就有多组权重。

---

6.3 grad_x 不是为了更新 x

grad_x 是为了把梯度传给上一层。

grad_weight / grad_bias：当前层参数更新用
grad_x：上一层继续反向传播用

---

6.4 bias 的梯度要沿 batch 求和

因为同一个 bias 被加到了每个 batch 样本上。

grad_bias = grad_z.sum(dim=0)

dim=0 是 batch 维度。

---

6.5 RMSNorm 里的 dim=-1 是 hidden 维度

RMSNorm 对每个 token 的 hidden vector 做归一化，所以沿最后一维求均值。

mean(dim=-1, keepdim=True)

意思是：

每个 token 自己算自己的 RMS

---