4、模型加载优化：模型并行策略、张量并行与流水线并行、混合精度训练(AMP)、梯度检查点
模型加载这块，说白了就是怎么把大模型塞进显存里跑起来。我刚开始做大模型压测时，经常遇到「OOM」——模型刚加载完，显存就爆了。后来才明白，光靠单卡硬扛是不行的，得用点策略。
4.1 模型并行策略：为什么单卡扛不住？
一个 70B 的模型，光参数就要 140GB 显存（FP16）。你想想看，现在最贵的 A100 也就 80GB。怎么办？
答案就是：把模型拆开，分到多张卡上。
我个人习惯把模型并行分为两类：
- 层间并行（流水线并行）：按层切分，每张卡负责几层
- 层内并行（张量并行）：按矩阵切分，每张卡负责一部分计算
嗯，这里要注意：不是切得越细越好。切太碎，通信开销会吃掉你的性能提升。
4.2 张量并行：把矩阵拆开算
张量并行，说白了就是把一个大的矩阵乘法拆成多个小矩阵，分到不同 GPU 上并行计算。
举个例子，Transformer 里的 FFN 层：
# 伪代码：张量并行前
output = Linear(input, weight) # 单卡计算，显存压力大
# 张量并行后（以列切分为例）
weight1 = weight[:, :half_dim] # GPU 0
weight2 = weight[:, half_dim:] # GPU 1
output1 = Linear(input, weight1) # GPU 0 算一半
output2 = Linear(input, weight2) # GPU 1 算一半
output = concat([output1, output2], dim=-1) # 通信合并我在项目中遇到过一个问题：张量并行虽然降低了单卡显存，但引入了 all-reduce 通信。如果跨机网络带宽不够，反而比单卡还慢。所以我的建议是——机内用张量并行，机间用流水线并行。
核心要点：张量并行适合单机多卡场景，通信延迟低。跨机时慎用。
4.3 流水线并行：把模型切成几段
流水线并行更直观——把 Transformer 的 40 层切成 4 段，每张卡管 10 层。
但有个坑：气泡问题。什么意思？
假设 4 张卡，每张卡算一层。卡 0 算完第 1 层传给卡 1，卡 1 算第 2 层时，卡 0 闲着。这就产生了气泡。
我曾经踩过这个坑——一开始用 naive 流水线，GPU 利用率不到 30%。后来改用 1F1B（一次前向一次后向） 策略，气泡才降下来。
# 1F1B 策略示意
# 卡0: F1 -> F2 -> B2 -> B1
# 卡1: F1 -> F2 -> B2 -> B1
# 气泡被交错的前向后向填满了小技巧：流水线并行时，尽量让每段计算量均衡。我习惯用 profiling 工具先跑一遍，看看每层耗时，再决定怎么切。
4.4 混合精度训练（AMP）：用更少的显存跑更大的模型
混合精度训练，说白了就是：该用 FP16 的地方用 FP16，该用 FP32 的地方用 FP32。
为什么能省显存？FP16 占 2 字节，FP32 占 4 字节，直接省一半。
但有个问题：精度损失。我刚开始用 AMP 时，模型直接不收敛了。后来发现是梯度下溢——FP16 能表示的最小正数是 6e-8，比这个小的梯度直接变 0。
解决方案：Loss Scaling。
from torch.cuda.amp import autocast, GradScaler
scaler = GradScaler()
for data, label in dataloader:
with autocast():
output = model(data)
loss = loss_fn(output, label)
# 先放大 loss，再反传，防止梯度下溢
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()避坑指南：我曾经在压测时发现，AMP 虽然省显存，但某些算子（如 LayerNorm）在 FP16 下精度不够。建议用 torch.cuda.amp.custom_fwd 手动指定这些算子用 FP32。
4.5 梯度检查点：用时间换空间
梯度检查点，是我压测时最常用的「救命稻草」。
原理很简单：前向时不保存中间激活值，反向时重新算一遍。
你想想看，一个 70B 的模型，中间激活值可能占几十 GB。用梯度检查点，可以省掉 60%-80% 的显存。
from torch.utils.checkpoint import checkpoint
class TransformerLayer(nn.Module):
def forward(self, x):
# 常规写法：保存中间激活
# y = self.attention(x)
# z = self.ffn(y)
# 梯度检查点写法：不保存中间激活
y = checkpoint(self.attention, x)
z = checkpoint(self.ffn, y)
return z代价是什么？训练时间增加约 20%-30%。因为反向时要重新算一遍前向。
我个人习惯：只在显存瓶颈时用梯度检查点。如果显存够，就别用，毕竟时间也是成本。
实战建议：在压测场景下，我通常这样组合使用：
- 单机多卡：张量并行 + AMP
- 多机多卡：流水线并行 + 梯度检查点
- 显存极度紧张：三者全上
4.6 知识体系总览
下面这张图，是我压测时常用的决策流程。你可以把它当作一个「显存优化 checklist」：
嗯，这张图基本概括了我压测时的思考路径。先判断显存够不够，不够再根据硬件环境选策略。
个人经验：我建议你在压测前，先用 torch.cuda.max_memory_allocated() 看看峰值显存。如果超过 80%，就该上优化策略了。
最后说一句：这些策略不是互斥的。我在实际项目中，经常是张量并行 + 流水线并行 + AMP 一起上。关键是要找到那个「显存够用、速度不慢」的平衡点。