【llm-leetcode】 TASK01

• Python 基础：列表推导、字典操作、装饰器等 LLM 代码中的常见语法
• NumPy 操作：数组广播、einsum 符号，为理解 Attention 机制做准备
• PyTorch 基础：Tensor 操作、自动求导、模块定义，深度学习的核心工具
• 性能分析：Profiling 工具、显存优化、调试技巧，工程实践的必备技能

Python Essentials for LLM

学习重点：

• 列表推导：生成序列、过滤数据
• 字典操作：管理模型配置、参数字典
• 装饰器：实现缓存、计时等功能
• 类与继承：定义模型基类

!info- 列表推导：

比如 Transformer 有很多层：

layers = [TransformerBlock(config) for _ in range(num_layers)]

这句话的意思是：

创建 num_layers 个 TransformerBlock，放进一个列表里。

等价于：

layers = []for _ in range(num_layers):    layers.append(TransformerBlock(config))

在 PyTorch 里经常会写成：

self.layers = nn.ModuleList([    TransformerBlock(config) for _ in range(config.num_layers)])

带条件过滤的列表推导

例如筛选偶数：

evens = [x for x in range(10) if x % 2 == 0]

结果：

[0, 2, 4, 6, 8]

意思是：

遍历 0 到 9，只保留偶数。

常见错误：条件位置写错

正确：

[x for x in nums if x > 0]

错误：

[x if x > 0 for x in nums]

因为如果你只是在过滤，if 应该放在最后。

但是如果你要写"三元表达式"，可以这样：

["positive" if x > 0 else "non-positive" for x in nums]

这个不是过滤，而是对每个元素做判断替换。

区别很重要：

[x for x in nums if x > 0]

是 筛选。

[x if x > 0 else 0 for x in nums]

是 替换。

字典操作

字典是 LLM 代码里非常常见的数据结构，尤其用于保存配置。

例如：

config = {
    "vocab_size": 30522,
    "hidden_size": 768,
    "num_layers": 12,
    "num_heads": 12,
}

访问参数：

hidden_size = config["hidden_size"]

结果：

768

config"key" 和 config.get("key") 的区别

写法一：config["hidden_size"]

config["hidden_size"]

如果 key 存在，正常返回。

如果 key 不存在，会报错：

KeyError: 'dropout'

写法二：config.get("dropout")

config.get("dropout")

如果 key 不存在，不会报错，而是返回：

None

也可以设置默认值：

dropout = config.get("dropout", 0.1)

意思是：

如果 config 里有 dropout，就用它；如果没有，就默认用 0.1。
!info- 函数

函数用于封装重复逻辑。

例如 Attention 里面可能要写一个 softmax：

def softmax(x):
    exp_x = np.exp(x)

   return exp_x / exp_x.sum()

但是这个版本只适合一维数组。

如果是矩阵，需要指定维度：

def softmax(x, axis=-1):

exp_x = np.exp(x - np.max(x, axis=axis, keepdims=True))

return exp_x / np.sum(exp_x, axis=axis, keepdims=True)

这里有几个关键点：

axis=-1

表示沿着最后一个维度做 softmax。

keepdims=True

表示求和之后保留维度，方便广播。

例如：

x.shape = (batch, seq_len, seq_len)

如果对最后一维 softmax：

softmax(x, axis=-1)

就是对每个 query 对所有 key 的分数做归一化。

这正是 Attention 里的操作。

!info- 装饰器

装饰器本质是：

在不修改原函数代码的情况下，给函数额外加功能。

例如你有一个函数：

def add(a, b):    return a + b

你想统计它运行时间，不想改函数内部，就可以用装饰器。

简单计时装饰器

import time
def timer(func):    
	def wrapper(*args, **kwargs):        
		start = time.time()        
		result = func(*args, **kwargs)        
		end = time.time()        
		print(f"{func.__name__} took {end - start:.4f}s")        
		return result    
	return wrapper

使用：

@timerdef matrix_multiply(A, B):    return A @ B

等价于：

matrix_multiply = timer(matrix_multiply)

!info- 继承

例如定义一个 Transformer 模型：

class Transformer(Model):    
	def __init__(self, config):        
		super().__init__(config)        
		self.layers = [TransformerBlock(config) for _ in range(config["num_layers"])]    
	def forward(self, x):        
		for layer in self.layers:            
			x = layer(x)        
		return x

这里：

class Transformer(Model):

表示 Transformer 继承 Model。

super().__init__(config)

表示调用父类的初始化函数。
!info- 三维数组

LLM 里最常见的是三维张量：

x.shape = (batch_size, seq_len, hidden_dim)

例如：

x = np.random.randn(2, 4, 8)

表示：

batch_size = 2
seq_len = 4
hidden_dim = 8

也就是：

有 2 个样本，每个样本有 4 个 token，每个 token 是 8 维向量。

这是 Transformer 输入最常见的形状。
!info- 索引和切片

假设：

x.shape = (2, 4, 8)

那么：

x[0]

取第 0 个 batch，形状是：

(4, 8)

表示一个句子的所有 token embedding。

---

x[0, 0]

取第 0 个 batch 的第 0 个 token，形状是：

(8,)

表示一个 token 的 hidden vector。

---

x[:, 0, :]

取所有 batch 的第 0 个 token，形状是：

(2, 8)

---

x[:, :, 0]

取所有 batch、所有 token 的第 0 个特征，形状是：

(2, 4)

!info- 广播机制

广播机制是 NumPy 自动扩展维度做运算的规则。

例如：

A = np.ones((2, 3))
b = np.array([10, 20, 30])

A 是：

shape = (2, 3)

b 是：

shape = (3,)

执行：

A + b

NumPy 会自动把 b 看成：

[[10, 20, 30], [10, 20, 30]]

所以结果是：

[[11, 21, 31], [11, 21, 31]]

广播规则

从最后一个维度开始对齐。

两个维度可以广播，当且仅当：

1. 两个维度相等
2. 其中一个维度是 1
3. 其中一个数组缺少这个维度

LLM 里广播在哪里常见？

比如给每个 token 加 bias：

x.shape = (batch, seq_len, hidden_dim)
bias.shape = (hidden_dim,)

执行：

x + bias

NumPy 会自动把 bias 广播到：

(batch, seq_len, hidden_dim)

这就是线性层里的 bias 加法。

再比如 LayerNorm：

x.shape = (batch, seq_len, hidden_dim)
mean.shape = (batch, seq_len, 1)
std.shape = (batch, seq_len, 1)

执行：

(x - mean) / std

这里 mean 和 std 会沿着 hidden_dim 维度广播。
!info- einsum 的核心思想

先看：

np.einsum("ij,jk->ik", A, B)

假设：

A.shape = (i, j)B.shape = (j, k)

输出是：

(i, k)

中间的 j 出现在输入里，但没有出现在输出里，所以对 j 求和。

数学上就是：

C[i, k] = sum_j A[i, j] * B[j, k]

这正是矩阵乘法。

einsum 规则

1. 出现在输出里的下标会保留

"ij,jk->ik"

输出里有 i 和 k，所以结果保留这两个维度。

---

2. 没出现在输出里的重复下标会被求和

"ij,jk->ik"

j 在输入中出现，但是输出中没有，所以对 j 求和。

---

3. 下标顺序决定输出维度顺序

np.einsum("ij->ji", A)

就是转置。

如果：

A.shape = (2, 3)

那么输出：

(3, 2)

!info- Batch 矩阵乘法

LLM 里经常不是一个矩阵乘法，而是每个 batch 都做一次矩阵乘法。

C = np.einsum("bij,bjk->bik", A, B)

假设：

A.shape = (batch, i, j)
B.shape = (batch, j, k)

输出：

C.shape = (batch, i, k)

其中：

b = batch 维度，保留j = 中间维度，被求和i, k = 输出矩阵维度

数学上：

C[b, i, k] = sum_j A[b, i, j] * B[b, j, k]

---

具体例子

A = np.random.randn(2, 3, 4)
B = np.random.randn(2, 4, 5)
C = np.einsum("bij,bjk->bik", A, B)

shape：

A: (2, 3, 4)
B: (2, 4, 5)
C: (2, 3, 5)

含义：

对 batch 里的每个样本，分别做一个 (3,4) @ (4,5) 的矩阵乘法。
!info- Attention 里的 QKᵀ

这是最重要的一个公式：

scores = np.einsum("bqd,bkd->bqk", Q, K)

假设：

Q.shape = (batch, query_len, dim)
K.shape = (batch, key_len, dim)

也就是：

Q: (b, q, d)
K: (b, k, d)

输出：

scores.shape = (batch, query_len, key_len)

也就是：

scores: (b, q, k)

---

这是什么意思？

Attention 要计算每个 query token 和每个 key token 的相似度。

对于一个 query 向量：

Q[b, q, :]

和一个 key 向量：

K[b, k, :]

它们的点积是：

sum_d Q[b, q, d] * K[b, k, d]

所以：

scores = np.einsum("bqd,bkd->bqk", Q, K)

表示：

scores[b, q, k] = sum_d Q[b, q, d] * K[b, k, d]

这里 d 没有出现在输出里，所以对 d 求和。

---

为什么是 QKᵀ？

普通矩阵形式是：

Q @ K.T

如果单个样本：

Q.shape = (query_len, dim)
K.shape = (key_len, dim)

那么：

K.T.shape = (dim, key_len)

所以：

Q @ K.T

得到：

(query_len, key_len)

对应 einsum：

np.einsum("qd,kd->qk", Q, K)

如果加上 batch 维度：

np.einsum("bqd,bkd->bqk", Q, K)

---

9. Attention 完整流程

Self-Attention 输入：

X.shape = (batch, seq_len, hidden_dim)

通过三个线性层得到：

Q = X @ Wq
K = X @ Wk
V = X @ Wv

形状通常是：

Q.shape = (batch, seq_len, head_dim)
K.shape = (batch, seq_len, head_dim)
V.shape = (batch, seq_len, head_dim)

然后计算分数：

scores = np.einsum("bqd,bkd->bqk", Q, K)

结果：

scores.shape = (batch, seq_len, seq_len)

再缩放：

scores = scores / np.sqrt(head_dim)

再 softmax：

attn = softmax(scores, axis=-1)

最后加权求和：

out = np.einsum("bqk,bkd->bqd", attn, V)

解释：

attn: (batch, query_len, key_len)
V:    (batch, key_len, dim)
out:  (batch, query_len, dim)

数学上：

out[b, q, d] = sum_k attn[b, q, k] * V[b, k, d]

这就是 Attention 的核心。
!info- PyTorch 写法的对应关系

NumPy einsum：

scores = np.einsum("bqd,bkd->bqk", Q, K)

PyTorch 写法：

scores = torch.einsum("bqd,bkd->bqk", Q, K)

等价于：

scores = Q @ K.transpose(-1, -2)

---

输出：

out = np.einsum("bqk,bkd->bqd", attn, V)

PyTorch 写法：

out = torch.einsum("bqk,bkd->bqd", attn, V)

等价于：

out = attn @ V

多头 Attention 里的 einsum

真正的 Transformer 里还有 head 维度。

形状通常是：

Q.shape = (batch, num_heads, query_len, head_dim)
K.shape = (batch, num_heads, key_len, head_dim)
V.shape = (batch, num_heads, key_len, head_dim)

可以写成：

scores = np.einsum("bhqd,bhkd->bhqk", Q, K)

输出：

scores.shape = (batch, num_heads, query_len, key_len)

解释：

b: batch
h: head
q: query position
k: key position
d: head dimension

d 被求和，所以得到每个 head 里的 attention scores。

然后：

out = np.einsum("bhqk,bhkd->bhqd", attn, V)

输出：

out.shape = (batch, num_heads, query_len, head_dim)

这就是多头 Attention 的核心运算。

einops 是什么？

你的进阶方向提到了 einops。

einops 主要用于更清楚地做 reshape、transpose、repeat。

比如多头 Attention 里，需要把：

x.shape = (batch, seq_len, hidden_dim)

拆成：

(batch, num_heads, seq_len, head_dim)

普通写法：

x = x.reshape(batch, seq_len, num_heads, head_dim)
x = x.transpose(0, 2, 1, 3)

用 einops：

from einops import rearrange
x = rearrange(x, "b s (h d) -> b h s d", h=num_heads)

这个表达式非常直观：

b = batch
s = seq_len
h = num_heads
d = head_dim
hidden_dim = h * d

LLM 代码里越来越常见这种写法。

这部分是从 NumPy 张量理解 过渡到 真正训练神经网络。0A 主要是让你理解 Python 和张量运算；0B 开始进入 PyTorch 的四个核心能力：

题号	核心问题	本质
02	Tensor 怎么创建和操作？	数据如何表示
03	Autograd 怎么算梯度？	模型如何学习
04	nn.Module 怎么定义模型？	模型如何组织
05	Loss + Optimizer 怎么训练？	参数如何更新

一句话概括：

PyTorch = NumPy 风格张量操作 + 自动求导 + 神经网络模块 + 优化器训练循环。

Tensor 有数据类型。

x = torch.randn(2, 3)print(x.dtype)

一般是：

torch.float32

神经网络参数通常是 float32 或 float16/bfloat16。
!info- masked_fill

masked_fill 在 Attention 里非常常见。

例如：

scores = torch.randn(2, 4, 4)
mask = torch.tensor([    [False, False, True, True],    [False, False, False, True],])

如果要把 mask 为 True 的位置填成一个极小值：

scores = scores.masked_fill(mask[:, None, :], -1e9)

这里常用于 Attention mask：

scores = scores.masked_fill(attention_mask == 0, -1e9)

为什么填 -1e9？

因为后面会做 softmax：

attn = torch.softmax(scores, dim=-1)

极小值经过 softmax 后几乎变成 0，相当于不关注这些位置。

masked_fill() 是 PyTorch Tensor 内置方法，作用：根据布尔掩码（mask），把张量中掩码为 True 的位置全部填充成指定数值 。

常用在注意力机制遮挡 padding token、上三角掩码屏蔽未来 token、无效数值屏蔽等场景。

tensor.masked_fill(mask, value)

参数说明

1. mask：布尔型张量，形状必须和原张量广播兼容

   • mask[i][j] = True：对应位置替换为 value
   • mask[i][j] = False：原数值保持不变

2. value：要填充的标量数值（float/int）

关键特性

• 原地操作区分

  • masked_fill()：返回新张量，不修改原张量
  • masked_fill_()：下划线结尾，原地修改原张量

• 掩码维度支持广播，无需严格同 shape

• Transformer 里最经典用法：用 -inf 遮挡无效位置，后续 softmax 后权重趋近 0

view vs reshape vs permute

!info- view()

view() 用来改变 Tensor 的形状。

x = torch.randn(2, 3, 4)y = x.view(6, 4)

原来：

x.shape = (2, 3, 4)

变成：

y.shape = (6, 4)

本质是把原来的数据重新解释成新的形状。

但是 view() 有一个要求：

Tensor 必须是内存连续的。

如果 Tensor 不连续，view() 会报错。

---

!info- reshape()

reshape() 也能改变形状：

y = x.reshape(6, 4)

它比 view() 更宽松。

如果内存连续，它可能直接返回 view；如果不连续，它会自动复制一份数据。

所以初学时可以优先用：

reshape()

更安全。

---

!info- permute()

permute() 用来改变维度顺序。

例如：

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

现在：

x.shape = (batch, seq_len, hidden_dim)

如果想变成：

(batch, hidden_dim, seq_len)

可以写：

y = x.permute(0, 2, 1)

结果：

y.shape = (2, 8, 4)

意思是：

原来的第 0 维还在第 0 维原来的第 2 维放到第 1 维原来的第 1 维放到第 2 维

---

!info- transpose()

transpose() 只交换两个维度。

x = torch.randn(2, 4, 8)y = x.transpose(1, 2)

结果：

y.shape = (2, 8, 4)

对于二维矩阵：

A = torch.randn(3, 4)A.T

结果是：

(4, 3)

device = "cuda" if torch.cuda.is_available() else "cpu"  
  
x = x.to(device)  
model = model.to(device)

!info- squeeze() 和 unsqueeze() - 增减维度

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

# 移除所有大小为 1 的维度
x_squeezed = x.squeeze()
print(x_squeezed.shape)  # torch.Size([2, 3, 4])

# 移除指定维度（必须大小为 1）
x_squeezed = x.squeeze(1)
print(x_squeezed.shape)  # torch.Size([2, 3, 1, 4])

# 在指定位置增加维度
x = torch.randn(2, 3)
x_unsqueezed = x.unsqueeze(0)  # 在第 0 维增加
print(x_unsqueezed.shape)  # torch.Size([1, 2, 3])

x_unsqueezed = x.unsqueeze(1)  # 在第 1 维增加
print(x_unsqueezed.shape)  # torch.Size([2, 1, 3])

!info- 索引和切片

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

# 基础索引
print(x[0])  # 第一个元素（4x5 的 Tensor）
print(x[0, 1])  # 第一个元素的第二行（长度为 5 的 Tensor）
print(x[0, 1, 2])  # 单个标量

# 切片
print(x[:, 0, :])  # 所有 batch 的第一行
print(x[0, :2, :])  # 第一个 batch 的前两行
print(x[..., -1])  # 所有元素的最后一列（... 表示所有维度）

# 高级索引
indices = torch.tensor([0, 2])
print(x[indices])  # 选择第 0 和第 2 个元素

# 布尔索引
mask = x > 0
print(x[mask])  # 选择所有大于 0 的元素（返回一维 Tensor）

# masked_fill - 根据 mask 填充值
x_masked = x.masked_fill(mask, 0)  # 将所有大于 0 的元素设为 0

!info- 设备转移（CPU ↔ GPU）

# 创建 CPU Tensor
x = torch.randn(2, 3)
print(x.device)  # cpu

# 转移到 GPU（如果可用）
if torch.cuda.is_available():
    x_gpu = x.to('cuda')  # 或 x.cuda()
    print(x_gpu.device)  # cuda:0
    
    # 转回 CPU
    x_cpu = x_gpu.to('cpu')  # 或 x_gpu.cpu()
    print(x_cpu.device)  # cpu
    
    # 指定 GPU 设备
    x_gpu1 = x.to('cuda:1')  # 转到第二块 GPU

# 通用写法（自动检测）
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
x = x.to(device)

!info- 数据类型转换

x = torch.randn(2, 3)
print(x.dtype)  # torch.float32

# 转换数据类型
x_int = x.to(torch.int32)  # 或 x.int()
x_long = x.to(torch.int64)  # 或 x.long()
x_float = x.to(torch.float32)  # 或 x.float()
x_double = x.to(torch.float64)  # 或 x.double()
x_half = x.to(torch.float16)  # 或 x.half()

# 常用数据类型
# torch.float32 (float) - 默认，4 字节
# torch.float16 (half) - 半精度，2 字节
# torch.int64 (long) - 长整型，8 字节
# torch.int32 (int) - 整型，4 字节
# torch.bool - 布尔型

!info- 内存连续性

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

# 检查是否内存连续
print(x.is_contiguous())  # True

# 转置后内存不连续
x_t = x.transpose(0, 1)
print(x_t.is_contiguous())  # False

# 使内存连续
x_t_contiguous = x_t.contiguous()
print(x_t_contiguous.is_contiguous())  # True

# 为什么需要连续内存？
# - view() 要求内存连续
# - 某些操作在连续内存上更快
# - 与 C/C++ 代码交互时需要连续内存

!info- Tensor 与 NumPy 互转【注意共享内存的问题】

import numpy as np

# Tensor -> NumPy
x = torch.randn(2, 3)
x_np = x.numpy()  # 共享内存，修改一个会影响另一个
print(type(x_np))  # <class 'numpy.ndarray'>

# NumPy -> Tensor
arr = np.array([1, 2, 3])
x = torch.from_numpy(arr)  # 共享内存
print(type(x))  # <class 'torch.Tensor'>

# ⚠️ 注意：共享内存意味着修改会相互影响
x = torch.randn(2, 3)
x_np = x.numpy()
x_np[0, 0] = 999
print(x[0, 0])  # 999.0（也被修改了）

# 如果不想共享内存，使用 clone()
x_np = x.clone().numpy()  # 不共享内存

!info- torch.triu

• triu = upper triangle 上三角

• diagonal>0：往右缩，只截更远的右上区域

• diagonal<0：往左扩，包含一小部分下三角

  torch.triu(input, diagonal=0)

功能：提取矩阵上三角部分，其余位置填充 0。

• input：输入张量（二维矩阵）
• diagonal：控制从哪一条对角线开始保留上三角，默认 0
• 返回同形状新张量

diagonal 参数取值规则

设行索引 i，列索引 j：

• diagonal = k：保留满足 j >= i + k 的位置，其余置 0

1. k = 0：主对角线及右上方全部保留
2. k = 1：主对角线右侧第一条斜线开始保留（不含主对角线）
3. k = -1：主对角线往左下移一行，包含部分下三角

def create_causal_mask(seq_len):

    """

    创建 Causal Mask

    Args:

        seq_len: 序列长度

    Returns:

        mask: 形状为 (seq_len, seq_len) 的布尔 Tensor

              mask[i, j] = True 表示位置 i 可以看到位置 j

    示例:

        seq_len = 4

        返回:

        [[True, False, False, False],

         [True,  True, False, False],

         [True,  True,  True, False],

         [True,  True,  True,  True]]

    """

    # TODO: 实现 Causal Mask

    # 提示: 使用 torch.triu() 创建上三角矩阵

    # pass

    upper_tri = torch.triu(torch.ones(seq_len,seq_len),diagonal=1)

    mask = (upper_tri == 0)

    return mask

  
  

# 测试

mask = create_causal_mask(4)

print(mask)

!info- 多头注意力维度变换

在多头注意力中，需要将 (batch, seq, hidden) 变换为 (batch, num_heads, seq, head_dim)。

def split_heads(x, num_heads):

    """

    将隐藏维度拆分为多个头

    Args:

        x: 形状为 (batch, seq, hidden) 的 Tensor

        num_heads: 头的数量

    Returns:

        形状为 (batch, num_heads, seq, head_dim) 的 Tensor

        其中 head_dim = hidden // num_heads

    示例:

        x.shape = (2, 10, 512), num_heads = 8

        返回 shape = (2, 8, 10, 64)

    """

    batch, seq, hidden = x.shape

    head_dim = hidden // num_heads

    # TODO: 实现维度变换

    # 提示: 先 reshape 再 permute

    x = x.reshape(batch,seq,num_heads,head_dim)

    x = x.permute(0,2,1,3)

    return x

!info- 批量矩阵乘法

实现批量矩阵乘法，用于计算 Attention 的 QK^T。

def batch_matmul(a, b):

    """

    批量矩阵乘法

    Args:

        a: 形状为 (batch, n, m) 的 Tensor

        b: 形状为 (batch, m, p) 的 Tensor

    Returns:

        形状为 (batch, n, p) 的 Tensor

    示例:

        a.shape = (2, 3, 4), b.shape = (2, 4, 5)

        返回 shape = (2, 3, 5)

    """

    # TODO: 实现批量矩阵乘法

    # 提示: 使用 torch.bmm() 或 @

    torch.bmm(a,b)

  

# 测试

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

b = torch.randn(2, 4, 5)

c = batch_matmul(a, b)

print(c.shape)  # 应该是 torch.Size([2, 3, 5])
 #等价写法
 def batch_matmul(a, b): 
	 return a @ b

!info- Pytorch Autograd and Backward

什么是自动求导？

自动求导（Automatic Differentiation, Autograd）是深度学习框架的核心功能，它能够自动计算函数的梯度，无需手动推导和编写梯度计算代码。

计算图（Computational Graph）

PyTorch 使用动态计算图（Dynamic Computational Graph）来追踪操作：

• 前向传播：构建计算图，记录每个操作

• 反向传播：沿着计算图反向计算梯度

  前向传播: x → f(x) → y
  反向传播: ∂L/∂x ← ∂L/∂y

• L L L：损失函数（Loss），模型最终要最小化的误差值，标量数字；

• y y y：模型前向输出；

• x x x：模型输入 / 中间变量 / 权重参数；

• ∂ L ∂ y \displaystyle \frac{\partial L}{\partial y} ∂y∂L：偏导数，损失 L 对 y 的梯度，代表 y 变化一点点，L 会变化多少；

• ∂ L ∂ x \displaystyle \frac{\partial L}{\partial x} ∂x∂L：损失 L 对 x 的梯度，反向传播最终要求的东西，用来更新参数。

• 前向：数据流正向走 x → f ( x ) → y x \rightarrow f(x) \rightarrow y x→f(x)→y

  拿输入 x，套函数f，算出输出y，同时 PyTorch 默默搭建计算图，存下来每一步做了什么运算

• 反向：梯度回流反向走 ∂ L ∂ x ← ∂ L ∂ y \frac{\partial L}{\partial x} \leftarrow \frac{\partial L}{\partial y} ∂x∂L←∂y∂L

  箭头是梯度传递方向 ：先拿到上层梯度 ∂ L ∂ y \frac{\partial L}{\partial y} ∂y∂L，再链式求导算出下层梯度 ∂ L ∂ x \frac{\partial L}{\partial x} ∂x∂L。

!info- 链式求导

y = f ( x ) y = f(x) y=f(x)，多元复合函数求导法则：

∂ L ∂ x = ∂ L ∂ y ⋅ ∂ y ∂ x \frac{\partial L}{\partial x} = \frac{\partial L}{\partial y} \cdot \frac{\partial y}{\partial x} ∂x∂L=∂y∂L⋅∂x∂y

翻译成人话：

1. 反向传播最开始，我们已经知道损失对输出的梯度 ∂ L ∂ y \frac{\partial L}{\partial y} ∂y∂L；
2. 再补上当前层自己的导数 ∂ y ∂ x \frac{\partial y}{\partial x} ∂x∂y；
3. 相乘，就得到损失对输入 x 的梯度 ∂ L ∂ x \frac{\partial L}{\partial x} ∂x∂L；
4. 再把这个梯度继续往前传给更靠前的变量，一路回溯完整张计算图。

多步计算

x = torch.tensor([2.0], requires_grad=True)

# 多步计算
y = x ** 2      # y = x^2
z = y * 3       # z = 3y = 3x^2
out = z.mean()  # out = z

# 反向传播
out.backward()

# 梯度: dout/dx = d(3x^2)/dx = 6x = 12
print(x.grad)  # tensor([12.])

多变量求导

x = torch.tensor([2.0], requires_grad=True)
y = torch.tensor([3.0], requires_grad=True)

# 计算 z = x^2 + y^2
z = x ** 2 + y ** 2

# 反向传播
z.backward()

# 查看梯度
print(x.grad)  # tensor([4.])  因为 dz/dx = 2x = 4
print(y.grad)  # tensor([6.])  因为 dz/dy = 2y = 6

!info- requires_grad

如果你希望 PyTorch 记录某个 Tensor 的计算过程，就设置：

x = torch.tensor(2.0, requires_grad=True)

例如：

y = x ** 2

数学上：

y = x²

当 x = 2 时：

y = 4

对 x 求导：

dy/dx = 2x = 4

调用：

y.backward()

然后看梯度：

print(x.grad)

输出：

tensor(4.)

!info- grad_fn - 记录操作历史

x = torch.tensor([2.0], requires_grad=True)
y = x ** 2

print(x.grad_fn)  # None（叶子节点）
print(y.grad_fn)  # <PowBackward0>（记录了 ** 操作）

z = y * 3
print(z.grad_fn)  # <MulBackward0>（记录了 * 操作）

叶子节点（Leaf Tensor）

x = torch.tensor([2.0], requires_grad=True)
y = x ** 2

print(x.is_leaf)  # True（用户创建的 Tensor）
print(y.is_leaf)  # False（由操作产生的 Tensor）

# 只有叶子节点会保留梯度
y.backward()
print(x.grad)  # tensor([4.])
print(y.grad)  # None（非叶子节点的梯度会被释放）

!info- 梯度清零

x = torch.tensor([2.0], requires_grad=True)

# 第一次
y1 = x ** 2
y1.backward()
print(x.grad)  # tensor([4.])

# 清零梯度
x.grad.zero_()

# 第二次
y2 = x ** 3
y2.backward()
print(x.grad)  # tensor([12.])（不再累加）

• 梯度累积

# 梯度累积

x = torch.tensor([2.0], requires_grad=True)

# 第一次反向传播

y1 = x ** 2

y1.backward()

print(x.grad)  # tensor([4.])

  

y1.backward()

print(x.grad)  # tensor([4.])

# 第二次反向传播（梯度会累加！）

y2 = x ** 3

y2.backward()

print(x.grad)  # tensor([16.])  = 4 + 12

---

RuntimeError Traceback (most recent call last)

Cell In33, line 9

6 y1.backward()

7 print(x.grad) # tensor(4.)

----> 9 y1.backward()

10 print(x.grad) # tensor(4.)

11 # 第二次反向传播（梯度会累加！）

File ~/.conda/envs/AcademicAgentStudio/lib/python3.11/site-packages/torch/_tensor.py:522, in Tensor.backward(self, gradient, retain_graph, create_graph, inputs)

512 if has_torch_function_unary(self):

513 return handle_torch_function(

514 Tensor.backward,

515 (self,),

(...)

520 inputs=inputs,

521 )

--> 522 torch.autograd.backward(

523 self, gradient, retain_graph, create_graph, inputs=inputs

524 )

File ~/.conda/envs/AcademicAgentStudio/lib/python3.11/site-packages/torch/autograd/init .py:266, in backward(tensors, grad_tensors, retain_graph, create_graph, grad_variables, inputs)

261 retain_graph = create_graph

263 # The reason we repeat the same comment below is that

264 # some Python versions print out the first line of a multi-line function

265 # calls in the traceback and some print out the last line

--> 266 Variable.execution_engine.run_backward( # Calls into the C++ engine to run the backward pass
267 tensors,
268 grad_tensors ,

269 retain_graph,

270 create_graph,

271 inputs,

272 allow_unreachable=True,

273 accumulate_grad=True,

274 )

RuntimeError: Trying to backward through the graph a second time (or directly access saved tensors after they have already been freed). Saved intermediate values of the graph are freed when you call .backward() or autograd.grad(). Specify retain_graph=True if you need to backward through the graph a second time or if you need to access saved tensors after calling backward.

默认情况下不能对同一个计算图连续两次调用 backward() ，第一次反向传播结束后，计算图中间缓存被释放了，第二次

找不到运算节点就报错；你示例里想连续两次 y1.backward() 必然报错。

一、报错原因拆解

1. y1 = x**2 构建了一张独立计算图；
2. 第一次 y1.backward()：
   • 正常计算梯度、累加到 x.grad；
   • 默认 retain_graph=False，反向走完立刻销毁这张计算图的中间变量，节省显存；
3. 再一次执行 y1.backward()：计算图已经没了，无法再次回溯求导，抛出 RuntimeError。

方案 1：第二次反向保留计算图 retain_graph=True

只需要第一次反向时开启保留图：

import torch
x = torch.tensor([2.0], requires_grad=True)

y1 = x ** 2
y1.backward(retain_graph=True)  # 保留计算图不销毁
print(x.grad)  # tensor([4.])

# 现在可以再次反向y1
y1.backward()
print(x.grad)  # 梯度累加 4+4 = tensor([8.])

# 再执行新的y2反向
y2 = x ** 3
y2.backward()
print(x.grad)  # 8+12 = tensor([20.])

方案 2：分开构建两张独立计算图（推荐）

不要重复复用同一个 y1，重新前向计算一遍，各自反向，互不干扰，不用开 retain_graph：

import torch
x = torch.tensor([2.0], requires_grad=True)

# 第一轮前向+反向
y1 = x ** 2
y1.backward()
print(x.grad)  # tensor([4.])

# 重新前向，全新计算图
y1_new = x ** 2
y1_new.backward()
print(x.grad)  # 4+4 = tensor([8.])

# 第三轮
y2 = x ** 3
y2.backward()
print(x.grad)  # 8+12 = tensor([20.])

三、梯度累积的真实逻辑（你原本想演示的）

PyTorch 只要不手动清零 x.grad，多次反向梯度会自动累加，这就是梯度累积训练的原理。

正确演示示例（无报错版）：

import torch
x = torch.tensor([2.0], requires_grad=True)

# 第1次
y1 = x**2
y1.backward()
print(x.grad)  # 4

# 梯度不清零，第2次全新前向
y2 = x**3
y2.backward()
print(x.grad)  # 4+12=16

四、补充关键知识点

4. retain_graph=True 代价：计算图一直驻留在显存里，多次循环会显存暴涨，训练循环里尽量少用；
5. 梯度清零：x.grad.zero_()，每个 batch 迭代前必须清零，不然梯度持续累积错乱；
6. 梯度累积训练标准写法：多个小 batch 前向 + 反向，不清梯度，累计 N 步后再更新参数，等价大 batch 训练。
   !info- 为什么要在同一张计算图上多次 backward ()

为什么要在同一张计算图上多次 backward () ，同时附上代码示例，顺带说明 retain_graph=True 的必要性。

前置常识

默认 backward() 执行完毕后：

1. 梯度正常累加到 .grad；
2. 计算图中间节点、缓存全部释放；
3. 无法再次反向，必须加 retain_graph=True 锁住计算图。

场景 1：多任务联合损失（最常用）

一个网络同时输出多个任务分支，Loss1、Loss2 都依赖同一套主干网络特征。

• 正向只跑一次前向，只构建一张计算图；

• 分别对 loss1.backward ()、loss2.backward ()；

• 两个损失梯度自动累加，一次性更新主干权重。
  示例：

  import torch

  x = torch.tensor([2.0], requires_grad=True)

  共享计算图，两个输出

  loss1 = x ** 2
  loss2 = x ** 3

  第一次反向，保留计算图

  loss1.backward(retain_graph=True)
  print(x.grad) # 4

  同一张图再次反向

  loss2.backward()
  print(x.grad) # 4+12=16

优势：只前向一次，节省算力，不用重复跑网络。

场景 2：二阶导数、Hessian 矩阵、梯度正则

需要对梯度本身再求导 （高阶微分）。

y = f ( x ) y=f(x) y=f(x)，先求 d y d x \frac{dy}{dx} dxdy，再求 d 2 y d x 2 \frac{d^2y}{dx^2} dx2d2y。

第一次 backward 拿到一阶梯度，还要在原图上继续反向求二阶导，必须锁住计算图。

x = torch.tensor(3.0, requires_grad=True)
y = x ** 3

# 一阶导数 dy/dx = 3x²
grad1, = torch.autograd.grad(y, x, retain_graph=True, create_graph=True)
# 对一阶梯度再求导，得到二阶导数 d²y/dx² = 6x
grad2, = torch.autograd.grad(grad1, x)

print(grad1)  # 27
print(grad2)  # 18

科研里牛顿法、曲率分析、GAN 梯度惩罚、Hessian 特征值计算都要多次反向。

场景 3：梯度惩罚（WGAN-GP 标准用法）

WGAN-GP 需要在插值样本处约束梯度范数，步骤：

1. 前向计算判别器输出；
2. 第一次反向求插值样本梯度；
3. 用梯度范数构造惩罚 Loss，再一次反向传回判别器权重 ；
   两张损失共用同一判别器计算图，必须 retain_graph。

场景 4：循环依赖、强化学习（RL）Actor-Critic

Actor 策略网络、Critic 价值网络共享特征提取层：

4. Critic Loss 先反向更新特征层；

5. 再用 Policy Loss 继续在同一张图反向；

共享主干不重复前向，梯度叠加更新。

场景 5：多次采样蒙特卡洛梯度估计

同一个输入前向一次，多次从输出端施加不同梯度权重、多次回传，做期望梯度数值估计。

不需要重复前向，只复用一张计算图多次反向。

场景 6：同一个 Loss 分多段回传、分步梯度裁剪

超大模型无法一次性把完整梯度传回，拆分多次反向，每次传回一部分梯度，逐步叠加。

重要注意事项

6. retain_graph=True 会一直占用显存，循环多次极易 OOM，用完务必手动释放；
7. 绝大多数普通单任务分类 / 检测模型：永远不需要多次反向 ，一个 loss 只调用一次 backward()；
8. 梯度累积训练（多张小批次累加梯度）：
   每一步都是独立前向、独立计算图 ，不属于 "同一张图多次反向"，不需要 retain_graph，别混淆。

一句话总结适用边界

只要满足：

只跑一次前向（一张计算图），但是要从输出端执行≥2 次梯度回传

就必须连续多次 backward + retain_graph=True。

普通单 Loss 深度学习训练完全碰不到这个需求。
!info- 训练循环中的梯度管理

import torch.nn as nn

import torch.optim as optim

  

model = nn.Linear(10, 1)

optimizer = optim.SGD(model.parameters(), lr=0.01)

  

for epoch in range(10):

    # 前向传播

    output = model(input)

    loss = criterion(output, target)

    # 清零梯度（重要！）

    optimizer.zero_grad()

    # 反向传播

    loss.backward()

    # 更新参数

    optimizer.step()

!info- torch.no_grad() - 禁用梯度计算

x = torch.tensor([2.0], requires_grad=True)

# 正常计算（会追踪梯度）
y = x ** 2
print(y.requires_grad)  # True

# 禁用梯度计算
with torch.no_grad():
    y = x ** 2
    print(y.requires_grad)  # False

# 用于推理阶段，节省内存
model.eval()
with torch.no_grad():
    output = model(input)

!info- detach() - 分离计算图

x = torch.tensor([2.0], requires_grad=True)
y = x ** 2

# 分离 y，不再追踪梯度
y_detached = y.detach()
print(y_detached.requires_grad)  # False

# y_detached 的操作不会影响 x 的梯度
z = y_detached * 3
z.backward()  # 报错！因为 z 不需要梯度

!info- @torch.no_grad() 装饰器

@torch.no_grad()
def inference(model, input):
    """推理函数，不计算梯度"""
    return model(input)

# 等价于
def inference(model, input):
    with torch.no_grad():
        return model(input)

!info- grad_outputs - 指定输出梯度

x = torch.tensor([2.0, 3.0], requires_grad=True)
y = x ** 2  # y = [4, 9]

# 指定 dy = [1, 2]
y.backward(torch.tensor([1.0, 2.0]))

# dx = dy * dy/dx = [1, 2] * [4, 6] = [4, 12]
print(x.grad)  # tensor([4., 12.])

!info- create_graph - 计算高阶导数

x = torch.tensor([2.0], requires_grad=True)
y = x ** 3

# 一阶导数
grad_y = torch.autograd.grad(y, x, create_graph=True)[0]
print(grad_y)  # tensor([12.])  因为 dy/dx = 3x^2 = 12

# 二阶导数
grad2_y = torch.autograd.grad(grad_y, x)[0]
print(grad2_y)  # tensor([12.])  因为 d^2y/dx^2 = 6x = 12

!info- 自定义 autograd.Function

class MyReLU(torch.autograd.Function):

@staticmethod

def forward(ctx, input):

    """

    前向传播

    Args:

        ctx: 上下文对象，用于保存信息供反向传播使用

        input: 输入 Tensor

    Returns:

        output: 输出 Tensor

    """

    # 保存输入，供反向传播使用

    ctx.save_for_backward(input)

    # 计算输出

    output = input.clamp(min=0)

    return output

@staticmethod

def backward(ctx, grad_output):

    """

    反向传播

    Args:

        ctx: 上下文对象

        grad_output: 输出的梯度 (dL/dy)

    Returns:

        grad_input: 输入的梯度 (dL/dx)

    """

    # 获取保存的输入

    input, = ctx.saved_tensors

    # 计算输入的梯度

    grad_input = grad_output.clone()

    grad_input[input < 0] = 0  # ReLU 的导数

    return grad_input

使用自定义函数

relu = MyReLU.apply

x = torch.tensor(-1.0, 2.0, -3.0, 4.0, requires_grad=True)

y = relu(x)

print(y) # tensor(0., 2., 0., 4.)

y.sum().backward()

print(x.grad) # tensor(0., 1., 0., 1.)
!info- Sigmoid 函数

Sigmoid 函数定义：

σ ( x ) = 1 1 + e − x \sigma(x) = \frac{1}{1+e^{-x}} σ(x)=1+e−x1

求导计算：

σ ′ ( x ) = d d x ( 1 1 + e − x ) = 0 ⋅ ( 1 + e − x ) − 1 ⋅ ( − e − x ) ( 1 + e − x ) 2 = e − x ( 1 + e − x ) 2 \begin{aligned} \sigma'(x) &= \frac{d}{dx}\left( \frac{1}{1+e^{-x}} \right) \\ &= \frac{0\cdot(1+e^{-x}) - 1\cdot (-e^{-x})}{(1+e^{-x})^2} \\ &= \frac{e^{-x}}{(1+e^{-x})^2} \end{aligned} σ′(x)=dxd(1+e−x1)=(1+e−x)20⋅(1+e−x)−1⋅(−e−x)=(1+e−x)2e−x

做恒等变形，凑回 σ ( x ) \sigma(x) σ(x)：

σ ( x ) ( 1 − σ ( x ) ) = 1 1 + e − x ⋅ ( 1 − 1 1 + e − x ) = 1 1 + e − x ⋅ e − x 1 + e − x = e − x ( 1 + e − x ) 2 \begin{aligned} \sigma(x)\big(1-\sigma(x)\big) &= \frac{1}{1+e^{-x}} \cdot \left(1-\frac{1}{1+e^{-x}}\right) \\ &= \frac{1}{1+e^{-x}} \cdot \frac{e^{-x}}{1+e^{-x}} \\ &= \frac{e^{-x}}{(1+e^{-x})^2} \end{aligned} σ(x)(1−σ(x))=1+e−x1⋅(1−1+e−x1)=1+e−x1⋅1+e−xe−x=(1+e−x)2e−x

于是得到经典结论：

σ ′ ( x ) = σ ( x ) ⋅ ( 1 − σ ( x ) ) \boldsymbol{\sigma'(x) = \sigma(x)\cdot\big(1-\sigma(x)\big)} σ′(x)=σ(x)⋅(1−σ(x))

二、链式法则接入反向传播（grad_output）

设：

本层输出： y = σ ( x ) y = \sigma(x) y=σ(x)

上游传来梯度： ∂ L ∂ y = grad_output \displaystyle \frac{\partial L}{\partial y} = \text{grad\_output} ∂y∂L=grad_output

我们需要求： ∂ L ∂ x \displaystyle \frac{\partial L}{\partial x} ∂x∂L

根据多元链式求导：

∂ L ∂ x = ∂ L ∂ y ⋅ ∂ y ∂ x \frac{\partial L}{\partial x} = \frac{\partial L}{\partial y}\cdot \frac{\partial y}{\partial x} ∂x∂L=∂y∂L⋅∂x∂y

代入导数公式：

∂ L ∂ x = grad_output    ⋅    σ ( x ) ( 1 − σ ( x ) ) \frac{\partial L}{\partial x} = \text{grad\_output}\;\cdot\; \sigma(x)\big(1-\sigma(x)\big) ∂x∂L=grad_output⋅σ(x)(1−σ(x))

代码里：

output 就是 σ ( x ) \sigma(x) σ(x)

grad_input 就是 ∂ L ∂ x \displaystyle \frac{\partial L}{\partial x} ∂x∂L

对应代码：

grad_input = grad_output * output * (1 - output)

三、代入测试值验算

x = 0 x=0 x=0：

σ ( 0 ) = 1 1 + e 0 = 0.5 \sigma(0) = \frac{1}{1+e^{0}} = 0.5 σ(0)=1+e01=0.5

σ ′ ( 0 ) = 0.5 × ( 1 − 0.5 ) = 0.25 \sigma'(0) = 0.5 \times (1-0.5) = 0.25 σ′(0)=0.5×(1−0.5)=0.25

y.backward() 时，标量 loss 默认上游梯度 g r a d o u t p u t = 1.0 grad_{output}=1.0 gradoutput=1.0：

∂ L ∂ x = 1 × 0.25 = 0.25 \frac{\partial L}{\partial x}=1 \times 0.25 = 0.25 ∂x∂L=1×0.25=0.25

和打印输出 tensor([0.25]) 完全对上。

四、补充 ctx 保存张量的原因

正向计算只算出了 σ ( x ) \sigma(x) σ(x)，没有保留原始 x x x；

而求导只需要 σ ( x ) \sigma(x) σ(x) 就能算出导数值，不需要 x x x，所以：

ctx.save_for_backward(output)

把正向结果存起来，反向直接取出复用，不用重复计算，节省算力。

class MySigmoid(torch.autograd.Function):

    @staticmethod

    def forward(ctx, input):

        output = 1 / (1 + torch.exp(-input))

        ctx.save_for_backward(output)

        return output

    @staticmethod

    def backward(ctx, grad_output):

        output, = ctx.saved_tensors

        # sigmoid 的导数: σ'(x) = σ(x) * (1 - σ(x))

        grad_input = grad_output * output * (1 - output)

        return grad_input

  

# 测试

sigmoid = MySigmoid.apply

x = torch.tensor([0.0], requires_grad=True)

y = sigmoid(x)

y.backward()

print(x.grad)  # tensor([0.25])  因为 σ'(0) = 0.5 * 0.5 = 0.25

• 解析梯度（理论梯度） ：数学求导公式直接算出精确梯度，PyTorch backward() 自动求解，速度快、无误差；
• 数值梯度（近似梯度）：微小扰动输入，用有限差分近似求导，是数值近似值，速度慢，但可以用来验证手写反向传播代码是否写错；
• gradcheck ：PyTorch 内置工具，自动对比自定义 autograd.Function 的解析梯度和数值梯度，一键校验反向传播实现正确性。

1）数值梯度数学原理（中心差分）

一元函数导数定义：

f ′ ( x ) ≈ f ( x + ε ) − f ( x − ε ) 2 ε f'(x)\approx \frac{f(x+\varepsilon)-f(x-\varepsilon)}{2\varepsilon} f′(x)≈2εf(x+ε)−f(x−ε)

不用求导公式，只给自变量加、减一个极小值 ε \varepsilon ε，两次前向算函数值，相除近似导数。

• 优势：不用推导求导公式；
• 劣势：多次前向计算，慢，有微小浮点误差。

2）numerical_gradient 代码逐行解释

def numerical_gradient(f, x, eps=1e-5):
    grad = torch.zeros_like(x)
    # 遍历x里每一个元素，逐个求偏导
    for i in range(x.numel()):
        # x第i个元素+eps，其余不变
        x_plus = x.clone()
        x_plus.view(-1)[i] += eps
        f_plus = f(x_plus)

        # x第i个元素-eps，其余不变
        x_minus = x.clone()
        x_minus.view(-1)[i] -= eps
        f_minus = f(x_minus)

        # 中心差分计算该位置偏导
        grad.view(-1)[i] = (f_plus - f_minus) / (2 * eps)
    return grad

• x.numel()：张量总元素个数，循环逐个维度单独扰动；
• view(-1)：任意形状张量展平成一维，方便按索引修改单个元素；
• 为什么用中心差分 ？比单侧差分 ( f ( x + ε ) − f ( x ) ) / ε (f(x+\varepsilon)-f(x))/\varepsilon (f(x+ε)−f(x))/ε 精度更高。

!info- 数值梯度 vs 解析梯度

def numerical_gradient(f, x, eps=1e-5):

    """

    使用数值方法计算梯度

    Args:

        f: 函数
        x: 输入
        eps: 扰动大小

    Returns:
        数值梯度

    """

    grad = torch.zeros_like(x)

    for i in range(x.numel()):

        # f(x + eps)

        x_plus = x.clone()

        x_plus.view(-1)[i] += eps

        f_plus = f(x_plus)

        # f(x - eps)

        x_minus = x.clone()

        x_minus.view(-1)[i] -= eps

        f_minus = f(x_minus)

        # 数值梯度

        grad.view(-1)[i] = (f_plus - f_minus) / (2 * eps)

    return grad

  

# 测试

def f(x):

    return (x ** 2).sum()

  

x = torch.tensor([2.0, 3.0], requires_grad=True)

  

# 解析梯度

y = f(x)

y.backward()

analytical_grad = x.grad.clone()

  

# 数值梯度

x.grad.zero_()

numerical_grad = numerical_gradient(f, x)

  

# 对比

print("Analytical:", analytical_grad)  # tensor([4., 6.])
print("Numerical:", numerical_grad)    # tensor([4., 6.])
print("Difference:", (analytical_grad - numerical_grad).abs().max())  # 很小

!info- 梯度检查（Gradient Checking）

自动校验你手写的反向传播导数实现是否正确 ：

它内部自动算出数值梯度（有限差分近似） ，再调用你自定义算子的 backward 得到解析梯度 ，逐元素比对两者误差；误差在阈值内返回 True（导数写对了），否则 False（求导公式代码写错）。

torch.autograd.gradcheck

from torch.autograd import gradcheck

# 定义函数
class MyFunction(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input):
        ctx.save_for_backward(input)
        return input ** 2
    
    @staticmethod
    def backward(ctx, grad_output):
        input, = ctx.saved_tensors
        return grad_output * 2 * input

# 测试
input = torch.randn(3, 4, dtype=torch.double, requires_grad=True)
test = gradcheck(MyFunction.apply, input, eps=1e-6, atol=1e-4)
print("Gradient check:", "Passed" if test else "Failed")

当你用 torch.autograd.Function 手动重写 forward、backward 时：

• 手动推导导数公式极易笔误、链式法则漏乘上游梯度、维度写错；
• 直接扔进网络训练，会出现 loss 不下降、参数不更新、梯度爆炸 / 消失，但很难定位是算子导数写错了；
• gradcheck 做单元测试，提前锁定自定义算子正确性，不用跑完整训练。

完整工作流程（内部自动执行）

1. 传入自定义算子 func.apply 和输入张量；
2. 微小扰动输入，用中心有限差分批量计算数值梯度（标准答案）；
3. 执行你的 backward()，得到你手写代码算出的解析梯度；
4. 逐元素计算两组梯度绝对差，和设定的 atol/rtol 容忍阈值对比；
5. 返回布尔标记：通过 / 不通过。

1. input

就是你传给自定义算子的原始输入张量，示例里就是 x，形状 (batch, num_classes)。

• 前向：用来计算 softmax 输出；
• 反向不能直接用原始 input，要在前向用 ctx.save_for_backward() 存下来，反向取出。

2. ctx（context，上下文对象）

是 torch.autograd.Function 专属的上下文容器，只在前向、反向两个静态方法之间传递数据，核心用途只有两个：

1. ctx.save_for_backward(t1, t2, ...)：把前向计算出来的中间张量存起来，反向里取用；
2. 附带标记开关：ctx.needs_input_grad、ctx.mark_dirty() 等高级用法。

⚠️ 关键特性：

• 不能直接存普通数字、列表、numpy 数组，只能存张量；

• 反向通过 ctx.saved_tensors 按顺序解包取出保存的张量。
  !info- 梯度裁剪

  import torch

  def train_with_gradient_accumulation(model, data_loader, optimizer, accumulation_steps=4):
  """
  使用梯度累积模拟大 batch size

    Args:
        model: 模型
        data_loader: 数据加载器
        optimizer: 优化器
        accumulation_steps: 累积步数
    """
    model.train()
    optimizer.zero_grad()  # 初始梯度清零
    
    for idx, (inputs, targets) in enumerate(data_loader):
        # 1. 前向传播
        outputs = model(inputs)
        loss = torch.nn.functional.cross_entropy(outputs, targets)
        
        # 2. 损失归一化：除以累积步数，等价大batch平均
        loss = loss / accumulation_steps
        
        # 3. 反向传播，梯度自动累加，不触发参数更新
        loss.backward()
  
        # 4. 累积满指定步数，更新参数+梯度清零
        if (idx + 1) % accumulation_steps == 0:
            optimizer.step()
            optimizer.zero_grad()

!info- backward() 做了什么？

y.backward()

意思是：

从 y 开始，沿着计算图反向传播，把每个需要梯度的变量的梯度算出来。

例如：

x = torch.tensor(3.0, requires_grad=True)
y = x * x + 2 * x + 1
y.backward()
print(x.grad)

数学上：

y = x² + 2x + 1dy/dx = 2x + 2

当 x = 3：

dy/dx = 8

所以：

x.grad# tensor(8.)

!info- 参数管理

定义了：

self.linear = nn.Linear(4, 2)

PyTorch 会自动把它的参数注册到模型里。

查看参数：

for name, param in model.named_parameters():    print(name, param.shape)

可能输出：

linear.weight torch.Size([2, 4])linear.bias torch.Size([2])

注意 nn.Linear(in_features, out_features) 的 weight shape 是：

(out_features, in_features)

所以：

nn.Linear(4, 2)

里面的 weight 是：

(2, 4)

parameters()

优化器需要知道要更新哪些参数：

optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

这里的：

model.parameters()

会返回所有可训练参数。

ModuleList

如果模型有很多层，比如 Transformer 有 N 层 block，不能直接用普通 list：

不推荐：

self.layers = [Block() for _ in range(12)]

因为 PyTorch 可能不会正确注册这些子模块。

应该写：

self.layers = nn.ModuleList([    Block() for _ in range(12)])

然后：

def forward(self, x):    for layer in self.layers:        x = layer(x)    return x

这就是 Transformer 堆叠多层 block 的基本形式。
!info- Sequential

如果模型只是简单按顺序执行，可以用：

self.net = nn.Sequential(    nn.Linear(10, 32),    nn.ReLU(),    nn.Linear(32, 2))

forward 里：

def forward(self, x):    return self.net(x)

适合简单 MLP。
!info- 自定义 MLP 示例

class MLP(nn.Module):    def __init__(self, input_dim, hidden_dim, output_dim):        super().__init__()        self.net = nn.Sequential(            nn.Linear(input_dim, hidden_dim),            nn.ReLU(),            nn.Linear(hidden_dim, output_dim)        )    def forward(self, x):        return self.net(x)

使用：

model = MLP(input_dim=8, hidden_dim=16, output_dim=3)x = torch.randn(4, 8)logits = model(x)print(logits.shape)

输出：

torch.Size([4, 3])

含义：

4 个样本每个样本输出 3 个类别的 logits

!info- Optimizer 优化器

优化器负责根据梯度更新参数。

最常见：

optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

或者：

optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

LLM / Transformer 训练里更常用 AdamW：

optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)

完整训练流程。

02 Tensor
    ↓
输入数据、参数、模型输出都是 Tensor

03 Autograd
    ↓
loss.backward() 自动计算梯度

04 nn.Module
    ↓
用类组织模型参数和 forward 逻辑

05 Loss + Optimizer
    ↓
计算损失并更新参数

---

自动求导（03）

核心概念：

• 计算图：PyTorch 自动构建的有向无环图
• 梯度累积：多次 backward() 会累加梯度
• 梯度清零 ：optimizer.zero_grad() 或 tensor.grad.zero_()

自定义 autograd.Function：

class MyFunction(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input):
        ctx.save_for_backward(input)
        return output
    
    @staticmethod
    def backward(ctx, grad_output):
        input, = ctx.saved_tensors
        return grad_input

nn.Module（04）

核心要点：

• 继承 nn.Module：所有自定义模块都要继承
• init() 中定义子模块 ：使用 self.layer = nn.Linear(...)
• forward() 中定义前向传播：不要手动调用 forward()
• 参数管理 ：parameters(), named_parameters(), state_dict()

常见错误：

• ❌ 列表推导中的条件语句位置错误
• ❌ 字典的 get() 和 \[\] 访问的区别
• ❌ 装饰器的参数传递

进阶方向：

• 理解 Python 的闭包和作用域
• 学习 functools 模块的高级用法

01: NumPy and Einsum

学习重点：

• 数组创建与索引：理解多维数组的操作
• 广播机制：自动扩展维度进行运算
• einsum 符号：简洁地表达复杂的张量运算

核心公式：

• 矩阵乘法：np.einsum('ij,jk->ik', A, B)
• Attention 的 QK^T：np.einsum('bqd,bkd->bqk', Q, K)
• Batch 矩阵乘法：np.einsum('bij,bjk->bik', A, B)

常见错误：

• ❌ einsum 下标顺序错误
• ❌ 广播维度不匹配
• ❌ 忘记指定输出维度

进阶方向：

• 理解 einsum 的性能优化
• 学习 einops 库的使用

---

📗 0B: PyTorch 基础（02-05）

🎯 学习目标

• ✅ 掌握 Tensor 的创建、操作、设备转移
• ✅ 理解自动求导的原理和使用
• ✅ 能够定义自定义的 nn.Module
• ✅ 掌握损失函数和优化器的使用

📚 题目列表

题号	题目	难度	核心知识点
02	PyTorch Tensor Fundamentals	Easy	Tensor 创建、操作、设备转移、数据类型
03	PyTorch Autograd and Backward	Medium	自动求导、梯度计算、反向传播
04	PyTorch nn.Module Basics	Medium	模块定义、前向传播、参数管理
05	PyTorch Optimizers and Loss	Medium	损失函数、优化器、学习率

核心概念解析

Tensor 操作（02）

关键操作：

• 创建 ：torch.randn(), torch.zeros(), torch.ones()
• 形状变换 ：view(), reshape(), permute(), transpose()
• 设备转移 ：tensor.to('cuda'), tensor.cpu()
• 索引切片 ：tensor[0, :, 1:3], tensor.masked_fill()

view vs reshape vs permute：

• view()：要求内存连续，速度快
• reshape()：自动处理内存不连续的情况
• permute()：改变维度顺序

---

自动求导（03）

核心概念：

• 计算图：PyTorch 自动构建的有向无环图
• 梯度累积：多次 backward() 会累加梯度
• 梯度清零 ：optimizer.zero_grad() 或 tensor.grad.zero_()

自定义 autograd.Function：

class MyFunction(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input):
        ctx.save_for_backward(input)
        return output
    
    @staticmethod
    def backward(ctx, grad_output):
        input, = ctx.saved_tensors
        return grad_input

---

nn.Module（04）

核心要点：

• 继承 nn.Module：所有自定义模块都要继承
• init() 中定义子模块 ：使用 self.layer = nn.Linear(...)
• forward() 中定义前向传播：不要手动调用 forward()
• 参数管理 ：parameters(), named_parameters(), state_dict()

---

📗 0C: 深度学习基础（06-09）

🎯 学习目标

• ✅ 掌握完整的训练循环
• ✅ 理解常见激活函数的原理
• ✅ 掌握归一化技术
• ✅ 理解 Attention 机制的基础

📚 题目列表

题号	题目	难度	核心知识点
06	Simple Neural Network Training	Medium	训练循环、验证、保存模型
07	Activation Functions	Easy	ReLU、GELU、SiLU 的实现与对比
08	Normalization Techniques	Medium	BatchNorm、LayerNorm 的原理与实现
09	Attention Mechanism Intro	Medium	Scaled Dot-Product Attention 基础

核心概念解析

训练循环（06）

标准训练流程：

for epoch in range(num_epochs):
    # 训练阶段
    model.train()
    for batch in train_loader:
        optimizer.zero_grad()
        output = model(batch)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()
    
    # 验证阶段
    model.eval()
    with torch.no_grad():
        for batch in val_loader:
            output = model(batch)
            val_loss = criterion(output, target)

关键技巧：

• model.train() vs model.eval()：影响 Dropout、BatchNorm 的行为
• torch.no_grad()：验证时不计算梯度，节省显存
• 早停 (Early Stopping)：防止过拟合

---

激活函数（07）

常见激活函数对比：

激活函数	公式	优点	缺点
ReLU	max(0, x)	简单、快速	死神经元问题
GELU	x * Φ(x)	平滑、性能好	计算稍慢
SiLU (Swish)	x * sigmoid(x)	平滑、自门控	计算稍慢

为什么 LLM 使用 GELU/SiLU？

• 平滑的梯度，训练更稳定
• 自门控机制，表达能力更强

---

归一化技术（08）

BatchNorm vs LayerNorm：

特性	BatchNorm	LayerNorm
归一化维度	Batch 维度	Feature 维度
适用场景	CNN、大 batch	RNN、Transformer、小 batch
依赖 batch size	是	否

为什么 Transformer 使用 LayerNorm？

• 不依赖 batch size，适合序列长度不固定的场景
• 每个样本独立归一化，适合自回归生成

---

Attention 机制（09）

Scaled Dot-Product Attention：

Attention(Q, K, V) = softmax(QK^T / √d_k) V

关键要点：

• 缩放因子 √d_k：防止 softmax 饱和
• Causal Mask：自回归生成时屏蔽未来信息
• 多头注意力：并行计算多个 Attention（Chapter 2 详细讲解）

---

📗 0D: 工具与调试（10-13）

🎯 学习目标

• ✅ 掌握 PyTorch Profiler 的使用
• ✅ 学会显存分析和优化
• ✅ 掌握常见的调试技巧
• ✅ 了解 Jupyter 和 Git 的基础使用

📚 题目列表

题号	题目	难度	核心知识点
10	PyTorch Profiling Basics	Medium	torch.profiler、性能分析、瓶颈定位
11	Memory Profiling and Optimization	Medium	内存分析、显存优化、梯度累积
12	Debugging Techniques	Medium	梯度检查、NaN 调试、断点调试
13	Jupyter and Git Basics	Easy	Notebook 使用、版本控制基础

核心概念解析

PyTorch Profiler（10）

核心功能：

• 性能分析：测量每个操作的 CPU/GPU 时间
• 内存分析：追踪显存使用情况
• 可视化：导出 Chrome Trace 或 TensorBoard

关键指标：

• Self CPU time：操作本身的 CPU 时间
• CUDA time：GPU 执行时间
• CPU Mem / CUDA Mem：内存/显存使用

使用示例：

with torch.profiler.profile(
    activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
    record_shapes=True,
    profile_memory=True
) as prof:
    output = model(input)
    loss.backward()

print(prof.key_averages().table(sort_by="cuda_time_total"))

---

显存优化（11）

常用技巧：

1. 梯度累积：模拟大 batch size
2. 混合精度训练：使用 FP16 节省显存
3. 梯度检查点：用计算换显存
4. 及时释放 ：del tensor, torch.cuda.empty_cache()

显存分析命令：

torch.cuda.memory_allocated()      # 已分配的显存
torch.cuda.max_memory_allocated()  # 峰值显存
torch.cuda.memory_summary()        # 显存摘要

---

调试技巧（12）

常见问题排查：

• Loss 不下降：检查学习率、梯度、数据标准化
• Loss 变成 NaN：检查学习率过大、数值溢出
• 显存溢出：减小 batch size、使用梯度累积
• 训练速度慢：使用 profiler 定位瓶颈

梯度检查：

for name, param in model.named_parameters():
    if param.grad is not None:
        if torch.isnan(param.grad).any():
            print(f"NaN in {name}")
        if torch.isinf(param.grad).any():
            print(f"Inf in {name}")

---

💡 学习建议

学习方法

1. 动手实践：每个题目都要自己实现，不要只看答案
2. 对比验证：用 PyTorch 官方实现验证你的代码
3. 循序渐进：按照学习组顺序学习，不要跳跃
4. 记录笔记：记录关键概念和常见错误

关键数字速查表

PyTorch 数据类型：

• torch.float32 (FP32)：4 Bytes
• torch.float16 (FP16)：2 Bytes
• torch.bfloat16 (BF16)：2 Bytes
• torch.int8：1 Byte

常用 Tensor 操作：

• view() vs reshape()：view 要求内存连续
• permute() vs transpose()：permute 可以任意维度，transpose 只能两个维度
• contiguous()：使 Tensor 内存连续

常见问题

Q: 没有 GPU 能学吗？

• A: 可以！Chapter 0 的所有内容都可以在 CPU 上运行

Q: 需要多长时间完成 Chapter 0？

• A: 快速通关 1-3 天，系统学习 1 周左右

Q: Chapter 0 和 Chapter 2 的 00 题有什么区别？

• A: Chapter 0 更基础，Chapter 2 的 00 题是 PyTorch Warmup，假设你已经掌握了 Chapter 0 的内容

Q: 可以跳过 Chapter 0 直接学习后续章节吗？

• A: 如果你已经熟悉 PyTorch，可以跳过 0A-0C，但建议学习 0D（Profiling 和调试）

---

einops

einops 是轻量张量维度操作库，名字源自爱因斯坦求和记法（Einstein notation），ICLR 2022 Oral 论文工具，专门简化深度学习里各种复杂张量变形、转置、压缩、复制操作。

支持：PyTorch / JAX / TensorFlow / NumPy / MLX / Paddle 等，一套语法全框架通用。

张量维度变换 (Tensor Reshaping)、嵌入层查表 (Embedding Lookup) 以及链式法则的反向传播 (Backpropagation)。

import torch
import torch.nn as nn
import torch.nn.functional as F
import einops

原生 PyTorch 维度操作需要一堆 permute / view / reshape / transpose / squeeze / unsqueeze，数字下标抽象、极易改错：

# 传统写法：把 [b,c,h,w] 转 [b,h,w,c]，看不懂数字含义
x = x.permute(0,2,3,1)

einops 用命名维度字符串，输入输出维度一目了然：

from einops import rearrange
# 一眼看懂：batch,channel,height,width → batch,height,width,channel
x = rearrange(x, 'b c h w -> b h w c')

然而，在实际的工业级代码中（尤其是 Transformer 的多头注意力机制等高维张量操作），原生方法往往缺乏可读性且极易出错。 举个典型的例子：将形状为 [batch, heads, seq_len, head_dim] 的多头张量合并为 [batch, seq_len, hidden_dim]：

• 原生实现：x.permute(0, 2, 1, 3).reshape(batch, seq_len, -1) ------ 开发者必须在脑海中硬记数字索引 (0,2,1,3) 的物理含义，代码维护成本极高。
• einops 实现：rearrange(x, 'b h s d -> b s (h d)') ------ 维度变换的语义直接写在字符串中，代码即文档（Self-documenting）。
  这正是为什么现代深度学习框架和开源模型广泛拥抱 einops 库，它能让复杂的张量操作变得语义清晰、安全可防错。

def embedding_warmup(input_ids: torch.Tensor, vocab_size: int, hidden_dim: int):
    """
    演示 Embedding 查表的过程，并用纯 Tensor 索引模拟它。
    
    Args:
        input_ids: 形状 [batch_size, seq_len]，包含整数类型的 Token IDs
    """
    # ==========================================
    # TODO 2.1: 实例化一个官方的 nn.Embedding，并用其进行前向传播
    # ==========================================
    # emb_layer = ???
    # emb_layer.weight.data.normal_(0, 0.1)  # 随便初始化一下
    # out_official = ???
    emb_layer = nn.Embedding(vocab_size, hidden_dim)
    emb_layer.weight.data.normal_(0, 0.1)# 初始化
    out_official = emb_layer(input_ids)

    
    # ==========================================
    # TODO 2.2: 使用纯 PyTorch 张量索引 (Advanced Indexing)，不使用 nn.Embedding，
    # 达到和上面官方 API 完全一模一样的输出。
    # 提示: Embedding 的本质是查表，思考如何用索引从权重矩阵中提取向量
    # ==========================================
    # out_manual = ???
    emb_weight = emb_layer.weight
    out_manual = emb_weight[input_ids]
    
    # return out_official, out_manual
    return out_official, out_manual

嵌入层 (Embedding Layer) 的本质

文本是离散的（Token IDs，如 [10, 42, 99]）。神经网络只能处理连续的稠密向量（Dense Vectors）。

Embedding 层的本质： 就是一个大规模的查表（Lookup Table）。给定一个 ID 列表，它直接把对应的行向量抽出来拼在一起。

它在数学上等价于：把离散的 ID 转换成 One-hot 向量，然后去乘以一个全连接层（Linear）。

def tensor_warmup(x: torch.Tensor):
    """
    假设 x 是一批图像的特征 (例如在多模态大模型中)，形状为 [batch_size, channels, height, width]
    我们需要将其展平为序列 (Sequence)，以输入给 Transformer。
    目标形状: [batch_size, height * width, channels]
    """
    
    # ==========================================
    # TODO 1.1: 使用原生的 PyTorch 方法 (permute + reshape/flatten) 完成变换
    # 提示: 先调整维度顺序，再合并空间维度
    # ==========================================
    b,c,h,w = x.shape
    x_native = x.permute(0,2,3,1).reshape()

    
    # ==========================================
    # TODO 1.2: 使用 einops.rearrange 优雅地完成完全相同的操作
    # 提示: 使用括号表示要合并的维度
    # ==========================================
    x_einops = einops.rearrange(x, 'b c h w -> b (h w) c')
    
    # return x_native, x_einops
    return x_native,x_einops

if __name__ == "__main__":
    ids = torch.randn(4, 128, 16, 16)  # [batch_size, channels, height, width]
    out1, out2 = tensor_warmup(ids)
    print(out1.shape, out2.shape)
    print(torch.allclose(out1, out2))  # True
    print(torch.equal(out1,out2))

前向传播与反向传播 (Forward & Backward)

为什么要理解前向和反向传播？

大模型的训练机制完全建立在反向传播算法 (Backpropagation) 与 链式法则 (Chain Rule) 之上。

前向传播 (Forward Pass)： 数据从输入层流向输出层，经过一系列的线性变换和非线性激活函数。在这个过程中，我们需要保存中间结果（如激活值、mask 等），供反向传播使用。

反向传播 (Backward Pass)： 梯度从输出层反向流向输入层，利用链式法则逐层计算每个参数的梯度。这是深度学习训练的核心机制。

在日常使用中，我们只需要写前向传播，然后调用 loss.backward()，PyTorch 的 Autograd 会自动帮我们算梯度。但为了真正理解底层原理，我们需要手动实现一个包含 Linear 和 ReLU 的自定义算子的完整前向和反向逻辑。

本节目标： 实现一个 LinearReLU 算子，公式为 y = relu(x @ W^T + b)，并手动推导其梯度。这将帮助你深入理解：

• 前向传播如何计算输出并保存中间结果
• 反向传播如何利用链式法则计算梯度
• 为什么需要在前向传播时保存某些张量（如 mask）、

---

01. RMSNorm Tutorial | 均方根层归一化 (RMSNorm)

难度： Easy | 标签： 基础架构, PyTorch | 目标人群： 模型微调与工程部署

本节我们将实现大语言模型（如 LLaMA、Gemma）中最常用的归一化技术：RMSNorm (Root Mean Square Normalization)。相比于传统的 LayerNorm，它能带来可观的训练加速，同时几乎不损失模型表现。

相关阅读 :

本节使用纯 PyTorch 实现了算法逻辑与数学推导。

如果你想学习工业界如何打破该算子的 Memory Bound (访存瓶颈)，请前往 Triton 篇：

03_Triton_Fused_RMSNorm

Step 1: 核心思想与痛点

为什么抛弃了 LayerNorm？
标准的 LayerNorm 需要计算均值（Mean）和方差（Variance）。
RMSNorm 的本质：
假设输入的均值已经接近 0（在大型网络中通常成立），那么我们直接去掉减去均值的操作，只用均方根（RMS）去归一化特征。这减少了同步开销，显著提升了前向和反向传播的计算速度。

Step 2: 核心公式与张量维度

给定输入向量 x ∈ R d x \in \mathbb{R}^d x∈Rd，RMSNorm 的输出 y y y 为：

1. 计算均方根 (RMS)：

   RMS ( x ) = 1 d ∑ i = 1 d x i 2 + ϵ \text{RMS}(x) = \sqrt{\frac{1}{d} \sum_{i=1}^d x_i^2 + \epsilon} RMS(x)=d1i=1∑dxi2+ϵ 其中 ϵ \epsilon ϵ 是为了防止除以 0 的极小值（如 1e-6）。

2. 归一化并缩放 (Scale)：

   y = x RMS ( x ) ⊙ γ y = \frac{x}{\text{RMS}(x)} \odot \gamma y=RMS(x)x⊙γ

   其中 γ ∈ R d \gamma \in \mathbb{R}^d γ∈Rd 是可学习的权重参数（Weight）。RMSNorm 没有偏置项 (Bias)。

对比项	标准 LayerNorm	RMSNorm（LLaMA/Gemma 同款）
核心假设	无前置假设，严格标准化	假设输入特征均值≈0，省略均值中心化
归一化分母	1 d ∑ ( x i − μ ) 2 + ϵ \sqrt{\frac{1}{d}\sum (x_i-\mu)^2+\epsilon} d1∑(xi−μ)2+ϵ μ = 1 d ∑ x i \mu=\frac{1}{d}\sum x_i μ=d1∑xi	1 d ∑ x i 2 + ϵ \sqrt{\frac{1}{d}\sum x_i^2+\epsilon} d1∑xi2+ϵ
输出计算	y = x − μ Var + ϵ ⊙ γ + β y=\frac{x-\mu}{\sqrt{\text{Var}+\epsilon}} \odot \gamma + \beta y=Var+ϵ x−μ⊙γ+β	y = x RMS ( x ) ⊙ γ y=\frac{x}{\text{RMS}(x)} \odot \gamma y=RMS(x)x⊙γ
可学习参数	γ \gamma γ（缩放）+ β \beta β（偏置）2组参数	仅 γ \gamma γ 缩放参数，无bias
计算步骤	1.求均值 μ \mu μ 2.每个元素减均值 3.平方求和求方差 4.开根号归一化 5.缩放+偏置	1.直接平方求和 2.开根号得到RMS 3.归一化+缩放 省去均值、中心化、bias三步

计算、显存、速度工程对比

维度	LayerNorm	RMSNorm	优势方
浮点计算量	更高，多一轮均值、减法运算	更少，约减少30%算术运算	RMSNorm
反向传播梯度计算	链式求导复杂，中间缓存多	求导公式简洁，中间临时张量更少	RMSNorm
显存读写次数	多次中间结果缓存，访存压力大	访存操作更少，更容易做算子融合	RMSNorm
多卡/分布式同步开销	均值聚合存在额外同步	无均值计算，同步开销极低	RMSNorm
算子融合潜力	中等	极高，Triton/CUDA可完全融合消除访存瓶颈	RMSNorm
参数量开销	2×隐藏层维度（γ+β）	1×隐藏层维度（仅γ）	RMSNorm

三、模型效果与适用场景对比

对比项	LayerNorm	RMSNorm
模型精度损失	基准	几乎无下降，大LLM场景精度持平
训练速度提升	基准	前向/反向速度提升10%~25%（大模型更明显）
主流使用模型	BERT、GPT1/2、T5、老款Transformer	LLaMA1/2/3、Gemma、Mistral、Qwen、现代开源大模型
适用场景	小模型、CV、输入分布偏移大的任务	大语言模型、超长文本、推理部署、微调加速
数值稳定性	强，不受均值偏移影响	依赖激活分布均值接近0；预训练激活偏移大时微调易波动

四、PyTorch 代码逻辑极简对比

LayerNorm 关键逻辑

mu = x.mean(-1, keepdim=True)
var = ((x - mu) ** 2).mean(-1, keepdim=True)
x_norm = (x - mu) / torch.sqrt(var + eps)
out = x_norm * self.weight + self.bias

RMSNorm 关键逻辑

rms = torch.sqrt((x ** 2).mean(-1, keepdim=True) + eps)
x_norm = x / rms
out = x_norm * self.weight  # 无bias

五、优缺点总结汇总表

类型	LayerNorm	RMSNorm
优点	数学严谨，不依赖输入分布；稳定适配CV、小NLP模型	计算更快、显存占用更低、参数量更少、大LLM训练加速明显
缺点	计算冗余、访存多、训练慢、参数量翻倍	依赖激活均值趋近0；部分小模型/分布偏移任务精度轻微下滑
工业选型结论	传统模型、计算机视觉首选	现代大语言模型预训练、微调、推理部署标配

代码实现与混合精度 (AMP) 陷阱

在 PyTorch 中，我们需要通过 torch.mean 计算均方，加上一个极小的 eps 防止除以零，最后乘以可学习的参数 weight。

在代码实现时，有一个非常关键的工程细节需要处理：数值溢出 (Numerical Overflow)。

工程经验：为什么要强制转换精度？

现代大模型训练与推理几乎都会使用混合精度 (AMP) 或半精度格式 (FP16) 以节省显存。但我们需要注意，FP16 的最大安全数值仅为 65504。

在计算 RMSNorm 时，第一步是求输入张量的平方 ( x 2 x^2 x2)。如果输入特征中某个值大于 256 256 256（由于 256 2 = 65536 > 65504 256^2 = 65536 > 65504 2562=65536>65504），该位置计算后就会溢出变为 inf（无穷大），进而导致损失函数出现 NaN，引发训练崩溃。
标准处理方案 (Upcasting)：

无论模型输入是什么精度格式，在执行平方和均值操作前，通常需要显式地将其转换为 float32 计算。待归一化计算完毕后，再将结果转换回原有精度。这是深度学习框架中处理该算子的标准做法。

class RMSNorm(nn.Module):

    def __init__(self, hidden_size: int, eps: float = 1e-6):

        super().__init__()

        self.eps = eps

        # ==========================================

        # TODO 1: 定义可学习参数 weight，并初始化为全 1

        # 形状: [hidden_size]

        # 提示: 使用 nn.Parameter 包装张量使其可学习

        # ==========================================

        # self.weight = ???

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

  

  

    def _norm(self, x: torch.Tensor) -> torch.Tensor:

        # ==========================================

        # TODO 2: 实现 RMSNorm 核心计算逻辑

        # 提示:

        # 1. 为防止 FP16 溢出，需要在高精度下计算

        # 2. 计算输入的均方值（平方后求均值），注意保持维度以便广播

        # 3. 使用均方根的倒数进行归一化，torch.rsqrt 比 1/sqrt 更快

        # 4. 返回归一化后的结果（保持高精度，便于后续操作）

        # ==========================================

        # variance = ???

        # return ???

        x = x.float() # 转换为高精度

        variance = x.pow(2).mean(dim=-1,keepdim=True) # 计算均方值

        return x * torch.rsqrt(variance + self.eps) # 归一化

  
  

    def forward(self, x: torch.Tensor) -> torch.Tensor:

        # ==========================================

        # TODO 3: 组合归一化与权重缩放

        # 提示: 调用 _norm 进行归一化，乘以可学习的 weight，最后转回输入精度

        # ==========================================

        # output = ???

        # return ???

        datatype = x.dtype # 记录输入精度

        x = self._norm(x) * self.weight # 归一化并缩放

        return x.to(datatype) # 转回输入精度

# 运行此单元格以测试你的实现

def test_rmsnorm():

    try:

        # 构造输入

        hidden_size = 512

        x = torch.randn(2, 16, hidden_size, dtype=torch.float16)  # FP16 输入模拟大模型

        # 测试你的实现

        my_norm = RMSNorm(hidden_size)

        # 将模型参数也转换为 FP16，对齐真实的工业半精度运行环境，防止发生隐式的 Type Promotion

        my_norm.to(x.dtype)

        my_out = my_norm(x)

        assert my_out.dtype == torch.float16, "输出类型必须与输入一致 (FP16)"

        assert my_out.shape == x.shape, "输出形状改变了！"

        # LLaMA 原版实现作为标准答案 (HuggingFace 提取)

        def hf_rmsnorm(hidden_states, weight, eps):

            input_dtype = hidden_states.dtype

            hidden_states = hidden_states.to(torch.float32)

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

            hidden_states = hidden_states * torch.rsqrt(variance + eps)

            return weight.to(torch.float32) * hidden_states.to(input_dtype)

        hf_out = hf_rmsnorm(x, my_norm.weight, my_norm.eps)

        # 检查容差

        assert torch.allclose(my_out.float(), hf_out.float(), rtol=1e-3, atol=1e-4), "计算结果与 HuggingFace 不一致！"

        print("\n✅ All Tests Passed! RMSNorm 实现通过测试。")

    except NotImplementedError:

        print("请先完成 TODO 部分的代码！")

    except AttributeError:

        print("代码未完成，无法找到 Parameter")

    except Exception as e:

        print(f"\n❌ 测试失败: {e}")

  

test_rmsnorm()

解析

1. TODO 1 (可学习参数)

• 参数定义： RMSNorm 的 weight（论文中称为 γ \gamma γ）是逐元素乘以归一化结果的，形状应与特征维度 hidden_size 一致，初始化为全 1。

2. TODO 2 (核心计算逻辑)

• 防溢出： 大模型特征的平方和极易越界（超过 FP16 的 65504 上限），因此在计算均方值前，必须将输入强制转换为 float32。
• 张量广播： 使用 .mean(dim=-1, keepdim=True) 保留维度数量（形状变为 (batch_size, seq_len, 1)），以便与 x_fp32 正确广播相乘。
• 指令优化： 使用 torch.rsqrt(x)（相当于 1 / x 1/\sqrt{x} 1/x ）而非 1.0 / torch.sqrt(x)，前者直接映射为 CUDA 快速倒数平方根指令，速度更快且数值更稳定。
• 精度保持： 返回 float32 结果，不要急着转换精度。

3. TODO 3 (类型恢复与权重缩放)

• 精度一致性： 必须确保最终输出的精度与输入一致。在经过 FP32 的归一化计算后，将其与 weight 相乘，最后统一通过 .to(x.dtype) 转换回原生精度（如 float16）。
• 进阶思考： 为什么最后的乘法敢在低精度做（真实场景下 weight 也是低精度），不怕溢出吗？因为 _norm(x) 计算完毕后，数值的均方根为 1，绝大多数值落在 -3, 3 区间（3σ 原则），乘以 weight（通常接近 1）后仍远低于 FP16 的溢出上限 65504。而 weight（初始为 1）通常在 [0.5, 2.0] 附近波动。两者的乘积一般在 [-6, 6] 之间，距离 FP16 的溢出红线 65504 差了一万倍，因此发生溢出的概率极低。

假设一组数据服从正态分布（高斯分布） ，( σ \sigma σ) 是这组数据的标准差，代表数据离散程度；( μ \mu μ) 是均值。3σ 原则描述数值落在均值 ±1σ、±2σ、±3σ 区间的概率：

1. ( μ − 1 σ , μ + 1 σ \\mu-1\\sigma,\\mu+1\\sigma μ−1σ,μ+1σ)：约 68.27% 的数据
2. ( μ − 2 σ , μ + 2 σ \\mu-2\\sigma,\\mu+2\\sigma μ−2σ,μ+2σ)：约 95.45% 的数据
3. ( μ − 3 σ , μ + 3 σ \\mu-3\\sigma,\\mu+3\\sigma μ−3σ,μ+3σ)：约 99.73% 的数据

简单说：99.73% 的样本都不会超出均值 ±3 倍标准差范围，只有不到 0.3% 的极端值会落在区间外，工程上基本可以忽略。

文中的 _norm(x) 是RMS 归一化 / LayerNorm 类归一化，归一化后满足两个条件：

• 均值 ( μ ≈ 0 \mu \approx 0 μ≈0)
• 均方根 RMS = 1，等价于标准差 ( σ ≈ 1 \sigma \approx 1 σ≈1)

代入 3σ 原则：

绝大多数归一化后的 x 取值落在 ( − 3 , 3 -3, 3 −3,3)，只有极少极端点超出。

---