第4章：PyTorch 框架核心

PyTorch 是当今 AI 训练与推理的事实标准。无论你后面做 CUDA 算子、分布式训练还是推理优化,几乎都要在 PyTorch 之上落地。本章不讲”怎么搭模型”,而是聚焦 AI Infra 工程师真正用得到的:Tensor 内部机制、autograd 反向传播原理、Module 的 hook 与状态管理、显存与性能 profiling、以及从零实现一个 mini GPT 把这些串起来。

📑 目录

• 1. Tensor 的本质:Storage + View
• 2. autograd:动态图与反向传播
• 3. nn.Module:状态、参数与 hook
• 4. 训练循环模板
• 5. 显存与性能 profiling
• 6. 常见坑与最佳实践
• 7. 实战:从零写一个 mini GPT
• 自我检验清单
• 参考资料

---

1. Tensor 的本质:Storage + View

PyTorch 的 Tensor 不是一块独立的内存,而是 (Storage, 形状元信息) 的二元组。理解这点对调试性能问题至关重要。

1.1 Storage 与 View

x = torch.arange(12)          # Storage: [0,1,2,...,11], shape=(12,)
y = x.view(3, 4)              # 共享同一 Storage,只是改了元信息
z = x.reshape(2, 6)           # 同上(连续时不复制)
y[0, 0] = 99                  # x 也变了!
print(x)                      # tensor([99, 1, 2, ..., 11])

view / reshape / transpose 都不复制数据,只改 strides。一旦 Tensor 不再 contiguous,某些 op 会自动 .contiguous() 复制——这是隐藏的性能开销来源。

1.2 stride 与 contiguous

x = torch.randn(2, 3, 4)
print(x.stride())             # (12, 4, 1)
y = x.transpose(0, 2)         # shape (4, 3, 2)
print(y.stride())             # (1, 4, 12)  ← 非连续
print(y.is_contiguous())      # False

非连续 Tensor 不能直接做某些操作,比如 view,需要先 .contiguous()。这一步会复制全部数据——大 Tensor 上是个隐形的性能杀手。

1.3 device 管理

x = torch.randn(1000, 1000)                       # CPU
x = x.to('cuda:0', non_blocking=True)             # GPU,异步
x = x.cuda()                                      # 等价
y = x.cpu()                                       # 拉回 CPU(同步)

# 多 GPU 之间
x_on_gpu1 = x.to('cuda:1')                        # P2P 传输

🌟 重要:tensor.to(device) 是异步的(在 stream 上排队),但 tensor.cpu() 默认会同步(因为 CPU 端要立即读)。这是为什么循环里的 loss.item() 慢的原因——每次都触发 GPU→CPU 同步。

---

2. autograd:动态图与反向传播

2.1 计算图的构建

每个 requires_grad=True 的 Tensor 在被 op 操作时,会动态记录一个 grad_fn:

x = torch.tensor([2.0], requires_grad=True)
y = x ** 3
z = y + 5
print(z.grad_fn)              # <AddBackward0>
print(z.grad_fn.next_functions)  # ((<PowBackward0>, 0), (None, 0))

整个计算过程形成一个 DAG(有向无环图),backward() 沿这个 DAG 从 loss 倒推到所有参数。

2.2 backward 的关键点

loss.backward()               # 计算梯度,累加到 .grad

# 等价于
torch.autograd.backward(loss)

梯度累加:每次 backward() 都会把新梯度累加到 .grad 上,而不是覆盖。所以训练前必须 optimizer.zero_grad()。这个特性可以用来做梯度累积(gradient accumulation):

optimizer.zero_grad()
for i, batch in enumerate(loader):
    loss = model(batch) / accum_steps
    loss.backward()
    if (i + 1) % accum_steps == 0:
        optimizer.step()
        optimizer.zero_grad()

2.3 关闭梯度的两种方式

# 1. 局部禁用
with torch.no_grad():
    output = model(x)         # 推理时省显存,不构建计算图

# 2. detach
y = x.detach()                # 切断 y 的反向链路,但仍共享 storage

2.4 自定义 autograd Function

写自定义 CUDA 算子时,需要手动定义反向:

class MyReLU(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x):
        ctx.save_for_backward(x)
        return x.clamp(min=0)

    @staticmethod
    def backward(ctx, grad_out):
        x, = ctx.saved_tensors
        grad_in = grad_out.clone()
        grad_in[x < 0] = 0
        return grad_in

y = MyReLU.apply(x)

---

3. nn.Module:状态、参数与 hook

3.1 Module 的核心三件套

class MyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(128, 64)        # 子 Module(自动注册)
        self.weight = nn.Parameter(torch.randn(64))  # 可训练参数
        self.register_buffer('mean', torch.zeros(64))  # 不训练但属于状态

    def forward(self, x):
        return self.linear(x) + self.weight - self.mean

类型	是否参与训练	是否随 state_dict 保存	用途
nn.Parameter	✅	✅	可学习权重
register_buffer	❌	✅	不训练但要保存(如 BN 的 running mean)
普通 attribute	❌	❌	临时变量、配置

3.2 train / eval 模式

model.train()                # Dropout 启用,BN 用 batch 统计
model.eval()                 # Dropout 关闭,BN 用 running 统计

评估和推理时一定要 model.eval() + torch.no_grad(),否则 Dropout 会引入随机性,BN 用错统计量。

3.3 forward / backward Hook

debug 显存或检查中间值时极其有用:

def hook(module, input, output):
    print(f"{module.__class__.__name__}: out shape = {output.shape}")
    print(f"  mean={output.mean():.4f}, std={output.std():.4f}")

handle = model.linear.register_forward_hook(hook)
model(x)
handle.remove()

---

4. 训练循环模板

工业界通用骨架——可以原样套用:

import torch
from torch import nn
from torch.utils.data import DataLoader
from torch.cuda.amp import autocast, GradScaler

model = MyModel().cuda()
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=0.1)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs)
scaler = GradScaler()         # 混合精度
loader = DataLoader(dataset, batch_size=64, num_workers=4, pin_memory=True)

for epoch in range(epochs):
    model.train()
    for batch in loader:
        x, y = batch[0].cuda(non_blocking=True), batch[1].cuda(non_blocking=True)

        optimizer.zero_grad(set_to_none=True)    # set_to_none 更快
        with autocast(dtype=torch.bfloat16):     # 混合精度
            logits = model(x)
            loss = F.cross_entropy(logits, y)

        scaler.scale(loss).backward()            # 缩放后反传
        scaler.unscale_(optimizer)               # 还原梯度
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)  # 防梯度爆炸
        scaler.step(optimizer)
        scaler.update()
    scheduler.step()

    # checkpoint
    if epoch % save_interval == 0:
        torch.save({
            'epoch': epoch,
            'model': model.state_dict(),
            'optimizer': optimizer.state_dict(),
            'scaler': scaler.state_dict(),
        }, f'ckpt-{epoch}.pt')

🍎 关键细节:

• pin_memory=True + non_blocking=True 允许数据加载和 GPU 计算重叠
• set_to_none=True 比清零省一次内存写
• BF16 训练通常不需要 GradScaler(动态范围足够),可以简化掉

---

5. 显存与性能 profiling

5.1 显存查询

# 当前 GPU 已分配
torch.cuda.memory_allocated() / 1e9       # GB
# 历史峰值
torch.cuda.max_memory_allocated() / 1e9
# 完整摘要
print(torch.cuda.memory_summary())

# 重置峰值统计(分阶段测量)
torch.cuda.reset_peak_memory_stats()

5.2 torch.profiler

from torch.profiler import profile, ProfilerActivity, schedule

with profile(
    activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
    schedule=schedule(wait=1, warmup=1, active=3),
    on_trace_ready=torch.profiler.tensorboard_trace_handler('./logs'),
    record_shapes=True,
    profile_memory=True,
    with_stack=True,
) as prof:
    for step in range(5):
        train_step()
        prof.step()

# 终端查看
print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=20))

trace 文件可以用 chrome://tracing 或 Perfetto 打开,可视化每个 op 的时间线。

5.3 显存可视化

torch.cuda.memory._record_memory_history(max_entries=100000)
# ... run training ...
torch.cuda.memory._dump_snapshot("mem.pickle")
# 上传到 https://pytorch.org/memory_viz 看时序

---

6. 常见坑与最佳实践

坑	症状	解决
忘记 optimizer.zero_grad()	loss 不收敛 / 梯度乱跳	每次 step 前清零
在 hot loop 里 .cpu() / .item()	GPU 利用率忽高忽低	累积在 GPU,一次性同步
non_blocking 没配 pin_memory	数据传输不能异步	两者要一起用
推理忘记 eval()	输出有随机性	model.eval() + torch.no_grad()
多个 optimizer 共享参数	梯度被覆盖	用 param_groups 分组
Tensor 在 CPU 但模型在 GPU	RuntimeError: expected device	数据 .to(device)
Inplace op 破坏 autograd	RuntimeError: a leaf Variable	避免 _ 后缀 op 在带梯度的 leaf 上
梯度累积时 BN 统计错乱	val acc 暴跌	累积期间 BN running stats 不要更新,或用 SyncBN

---

7. 实战:从零写一个 mini GPT

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

class CausalSelfAttention(nn.Module):
    def __init__(self, dim, n_head):
        super().__init__()
        self.n_head = n_head
        self.head_dim = dim // n_head
        self.qkv = nn.Linear(dim, 3 * dim, bias=False)
        self.proj = nn.Linear(dim, dim, bias=False)

    def forward(self, x):
        B, S, D = x.shape
        qkv = self.qkv(x).reshape(B, S, 3, self.n_head, self.head_dim)
        q, k, v = qkv.unbind(dim=2)                       # (B, S, H, D)
        q, k, v = [t.transpose(1, 2) for t in (q, k, v)]   # (B, H, S, D)
        # PyTorch 2.0+ 内置 SDPA,自动用 FlashAttention
        out = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        out = out.transpose(1, 2).reshape(B, S, D)
        return self.proj(out)

class Block(nn.Module):
    def __init__(self, dim, n_head):
        super().__init__()
        self.norm1 = nn.RMSNorm(dim)
        self.attn = CausalSelfAttention(dim, n_head)
        self.norm2 = nn.RMSNorm(dim)
        self.ffn = nn.Sequential(
            nn.Linear(dim, 4 * dim, bias=False),
            nn.GELU(),
            nn.Linear(4 * dim, dim, bias=False),
        )

    def forward(self, x):
        x = x + self.attn(self.norm1(x))
        x = x + self.ffn(self.norm2(x))
        return x

class MiniGPT(nn.Module):
    def __init__(self, vocab=50257, dim=384, n_head=6, n_layer=6, max_len=1024):
        super().__init__()
        self.tok = nn.Embedding(vocab, dim)
        self.pos = nn.Embedding(max_len, dim)
        self.blocks = nn.ModuleList([Block(dim, n_head) for _ in range(n_layer)])
        self.norm = nn.RMSNorm(dim)
        self.head = nn.Linear(dim, vocab, bias=False)
        # 权重共享:embedding 和 lm_head
        self.head.weight = self.tok.weight

    def forward(self, idx, targets=None):
        B, S = idx.shape
        pos_idx = torch.arange(S, device=idx.device)
        x = self.tok(idx) + self.pos(pos_idx)
        for blk in self.blocks:
            x = blk(x)
        x = self.norm(x)
        logits = self.head(x)
        if targets is None:
            return logits
        loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))
        return logits, loss

# 训练
model = MiniGPT().cuda()
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)
for step in range(1000):
    x = torch.randint(0, 50257, (32, 128), device='cuda')
    y = torch.randint(0, 50257, (32, 128), device='cuda')
    _, loss = model(x, y)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    if step % 100 == 0:
        print(f"step {step}: loss = {loss.item():.4f}")

50 行代码,涵盖了 LLM 的核心结构。后续的所有优化(FlashAttention、量化、分布式…)都是在这个骨架上展开。

---

✅ 自我检验清单

• Storage / View:能解释 view / reshape / transpose 何时不复制数据,何时必须 .contiguous()
• autograd 流程:能口头讲清 loss.backward() 内部做了什么、为什么要 zero_grad()
• 梯度累积:能解释 gradient accumulation 的代码,以及为什么 loss /= accum_steps
• train / eval:能解释为什么推理一定要 model.eval(),具体改变了哪些层的行为
• 混合精度:能写出 BF16 训练循环,并说明为什么 BF16 不需要 GradScaler
• 训练模板:能默写一个含 AMP + CosineLR + ClipGrad + Checkpoint 的标准训练循环
• profiler 实战:能用 torch.profiler 抓一次 trace,在 chrome://tracing 中找到耗时最长的算子
• OOM 排查:能用 torch.cuda.memory_summary() 找出哪一步显存激增
• mini GPT 复现:能从零写一个 50-100 行的 mini GPT,在小数据集上跑通 loss 下降

---

📚 参考资料

• PyTorch 官方教程:https://pytorch.org/tutorials/
• PyTorch Internals (ezyang):http://blog.ezyang.com/2019/05/pytorch-internals/
• Karpathy:nanoGPT:https://github.com/karpathy/nanoGPT —— 最简 GPT 实现的事实标准
• PyTorch Profiler 文档:https://pytorch.org/tutorials/recipes/recipes/profiler_recipe.html
• PyTorch Memory Visualization:https://pytorch.org/blog/understanding-gpu-memory-1/
• HuggingFace Performance Tuning:https://huggingface.co/docs/transformers/performance