PyTorch for LLM：动手之前的必备工具箱

本章定位：第二部分「构建 (Build)」的起点。前面 10 章解决了「懂原理」，从本章开始解决「写得出来」。目标：读完本章，你能读懂并自己写出 small LLM 的全部代码。本章不引入新模型概念，只补齐动手所需的 PyTorch 工具箱。

11.1 为什么单独练 PyTorch

结论先行——一个 LLM 工程师 90% 的时间只用这 12 个工具：

11.2 Tensor：shape / dtype / device

import torch

x = torch.tensor([[1.0, 2.0], [3.0, 4.0]])   # 从 Python 列表创建
print(x.shape)    # torch.Size([2, 2])
print(x.dtype)    # torch.float32
print(x.device)   # cpu

y = torch.randn(2, 3)                 # 标准正态
z = torch.zeros(2, 3)                 # 全 0
w = torch.arange(0, 10, 2)            # [0, 2, 4, 6, 8]
e = torch.empty(4, 4)                 # 未初始化（省时间，但值随机）

x_gpu = x.to("cuda")                  # 搬到 GPU
x_half = x.to(torch.bfloat16)         # 改精度
x_reshaped = x.view(1, 4)             # 改形状

三种最常打交道的 dtype

11.3 形状变换四件套（LLM 代码的重灾区）

对比表（背下来）

最小数值例子

x = torch.arange(12).reshape(2, 3, 2)     # [2, 3, 2]
print(x.shape)                            # torch.Size([2, 3, 2])

a = x.transpose(1, 2)                     # 交换第 1、2 维 → [2, 2, 3]
b = x.permute(2, 0, 1)                    # 任意重排 → [2, 2, 3]
c = x.reshape(2, 6)                       # → [2, 6]
d = x.view(6, 2)                          # → [6, 2]（x 连续，可 view）

print(a.shape, b.shape)                   # 都是 [2, 2, 3]，但值不同！

x = torch.arange(12).reshape(2, 3, 2)
xt = x.transpose(1, 2)          # [2, 2, 3]，stride 已改变
xt.view(2, 6)                   # ❌ RuntimeError: view size is not compatible ...
xt.reshape(2, 6)                # ✅ reshape 自动拷贝一份再变形
xt.contiguous().view(2, 6)      # ✅ 显式拷贝，然后 view 合法

B, S, d_model, H = x.size(0), x.size(1), x.size(2), 8
D = d_model // H

# 1) 投影 → [B, S, d_model]，拆成多头 → [B, S, H, D]
q = Wq(x).view(B, S, H, D)
# 2) 把 head 维换到前面 → [B, H, S, D]，方便批量算注意力
q = q.transpose(1, 2)          # 或 permute(0, 2, 1, 3)
k = Wk(x).view(B, S, H, D).transpose(1, 2)
v = Wv(x).view(B, S, H, D).transpose(1, 2)

# 3) 注意力 → [B, H, S, S] @ [B, H, S, D] = [B, H, S, D]
attn = (q @ k.transpose(-2, -1)) / (D ** 0.5)
out = attn.softmax(-1) @ v

# 4) 拼回 → [B, S, d_model]
out = out.transpose(1, 2).contiguous().view(B, S, d_model)
out = Wo(out)

11.4 广播：不加循环做批量运算

三条规则

1. 从最后一维开始向左对齐；
2. 两维相等，或其中一维是 1 → 兼容；
3. 不兼容 → 报错。

a = torch.randn(3, 4)         # [3, 4]
b = torch.randn(4)            # [4]      → 广播成 [3, 4]
c = a + b                     # ✅

d = torch.randn(3, 1)         # [3, 1]   → 广播成 [3, 4]
e = a + d                     # ✅

f = torch.randn(3)            # [3]      → 与 [3, 4] 不兼容
# a + f  →  RuntimeError: The size of tensor a (4) ...

LLM 三个经典用法

# ① 因果掩码：scores [B, S, S] + mask [S, S]
scores = scores + causal_mask                # 广播自动处理 B 维

# ② padding mask：scores [B, S, S] + mask [B, 1, S]
scores = scores.masked_fill(pad_mask[:, None, :] == 0, float("-inf"))

# ③ softmax 数值稳定：每行减最大值
scores = scores - scores.max(dim=-1, keepdim=True).values
# keepdim=True 保留 [B, S, 1] 的形状，才能正确广播回 [B, S, S]

11.5 矩阵乘法：matmul / @ / einsum

三种写法对照

A = torch.randn(2, 3, 4)      # [B=2, m=3, k=4]
B_ = torch.randn(2, 4, 5)     # [B=2, k=4, n=5]

C1 = A @ B_                   # [2, 3, 5]（批量矩阵乘）
C2 = torch.matmul(A, B_)      # 同上
C3 = torch.einsum("bmk,bkn->bmn", A, B_)   # 同上，下标显式写出

# 只对最后两维做矩阵乘、前面维度广播：
A4 = torch.randn(2, 3, 4, 5)
B4 = torch.randn(4, 5, 6)
C4 = A4 @ B4                  # [2, 3, 4, 6]

用 einsum 重写 Attention（检查 shape 的神器）

# 第 7 章的公式：Attention(Q,K,V) = softmax(QKᵀ/√d)V
# Q,K,V: [B, H, S, D]

scores = torch.einsum("bhsd,bhtd->bhst", Q, K) / (D ** 0.5)   # [B,H,S,S]
attn = scores.softmax(dim=-1)
out = torch.einsum("bhst,bhtd->bhsd", attn, V)                 # [B,H,S,D]

表格：LLM 里常见的矩阵乘

11.6 nn.Module 与 Parameter：组织模型

手写一个 Linear（理解内部）

import torch, torch.nn as nn
import math

class MyLinear(nn.Module):
    def __init__(self, in_features, out_features):
        super().__init__()                                # 必须调用
        self.weight = nn.Parameter(torch.empty(out_features, in_features))
        self.bias = nn.Parameter(torch.zeros(out_features))
        nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))

    def forward(self, x):                                 # x: [..., in_features]
        return x @ self.weight.t() + self.bias            # [..., out_features]

lin = MyLinear(768, 256)
print(sum(p.numel() for p in lin.parameters()))           # 768*256 + 256 = 196,864

nn.Module 的三条规则

class Block(nn.Module):
    def __init__(self, d, h):
        super().__init__()
        self.norm1 = nn.LayerNorm(d)
        self.attn = nn.MultiheadAttention(d, h, batch_first=True)
        self.norm2 = nn.LayerNorm(d)
        self.ffn = nn.Sequential(
            nn.Linear(d, 4 * d), nn.GELU(), nn.Linear(4 * d, d)
        )

    def forward(self, x):
        h = x + self.attn(self.norm1(x), self.norm1(x), self.norm1(x))[0]
        out = h + self.ffn(self.norm2(h))
        return out

model = nn.Sequential(*[Block(256, 8) for _ in range(4)])
print(sum(p.numel() for p in model.parameters()))

11.7 autograd：把第 1 章的手算交给框架

import torch

x = torch.tensor(2.0)
y = torch.tensor(5.0)
w = torch.tensor(1.0, requires_grad=True)   # 需要求导

z = w * x            # 前向（自动记录计算图）
L = (z - y) ** 2
L.backward()         # 反向：自动微分

print(w.grad)        # tensor(-12.)   ← 与第 1 章手算完全一致

手动对照：计算图长这样

w ──×── z = wx ──(z−y)²── L
∂L/∂w ← ∂L/∂w = ∂L/∂z · ∂z/∂w
     = 2(z−y) · x = (−6)·2 = −12

三条必须记住的规则

11.8 Dataset / DataLoader：喂数据

from torch.utils.data import Dataset, DataLoader

class TokenDataset(Dataset):
    """data.bin：把整个语料存成一串 uint16 token id"""
    def __init__(self, path, seq_len):
        import numpy as np
        self.data = np.memmap(path, dtype=np.uint16, mode="r")
        self.seq_len = seq_len

    def __len__(self):
        # 合法起点 i 满足 i+1+seq_len <= N，即 i ∈ [0, N-seq_len-1]，共 N-seq_len 个样本
        return len(self.data) - self.seq_len

    def __getitem__(self, i):
        x = torch.from_numpy(self.data[i:i + self.seq_len].astype("int64"))
        y = torch.from_numpy(self.data[i + 1:i + 1 + self.seq_len].astype("int64"))
        return x, y                      # 输入与「右移一位」的标签

# 自测：N=10、S=4 时应恰好 6 个样本，最后一个样本是 x=data[5:9], y=data[6:10]
import numpy as np
arr = np.arange(10, dtype=np.uint16)
arr.tofile("/tmp/tiny.bin")
tiny = TokenDataset("/tmp/tiny.bin", seq_len=4)
assert len(tiny) == 6, len(tiny)
x_last, y_last = tiny[5]
assert x_last.tolist() == [5, 6, 7, 8] and y_last.tolist() == [6, 7, 8, 9]

ds = TokenDataset("data.bin", seq_len=512)
loader = DataLoader(ds, batch_size=32, shuffle=True, num_workers=4, pin_memory=True)

x, y = next(iter(loader))
print(x.shape, y.shape)                  # [32, 512] [32, 512]

11.9 训练循环：标准骨架

model.train()
for step, (x, y) in enumerate(loader):
    x, y = x.to(device), y.to(device)

    logits = model(x)                          # ① 前向：  [B,S,V]
    loss = F.cross_entropy(                       # ② 算损失
        logits.view(-1, logits.size(-1)),
        y.view(-1),
    )

    optimizer.zero_grad(set_to_none=True)      # ③ 清梯度
    loss.backward()                            # ④ 反向：算梯度
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)  # 梯度裁剪
    optimizer.step()                           # ⑤ 更新参数

    if step % 50 == 0:
        print(f"step {step}  loss {loss.item():.4f}")

五个零件各自的角色

scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)  # warmup+cosine

accum_steps = 8
for step, (x, y) in enumerate(loader):
    logits = model(x)
    loss = F.cross_entropy(logits.view(-1, V), y.view(-1)) / accum_steps
    loss.backward()                              # 梯度在 .grad 里累积

    if (step + 1) % accum_steps == 0:            # 攒够 8 个 micro-batch 才更新
        clip_grad_norm_(model.parameters(), 1.0)
        optimizer.step()
        scheduler.step()
        optimizer.zero_grad(set_to_none=True)

11.10 no_grad / eval / inference_mode

model.eval()
with torch.no_grad():
    logits = model(x_val)
    val_loss = F.cross_entropy(logits.view(-1, V), y_val.view(-1))

11.11 device / dtype / autocast：进入 GPU 世界

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

# 混合精度：前向/反向用 bf16，主权重仍是 fp32（第 17 章细讲）
with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
    logits = model(x)
    loss = F.cross_entropy(logits.view(-1, V), y.view(-1))
loss.backward()

# 查看显存
print(torch.cuda.memory_allocated() / 1e9, "GB")
print(torch.cuda.max_memory_allocated() / 1e9, "GB")

11.12 state_dict / checkpoint：保存与续训

# 只保存模型权重（推理用）
torch.save(model.state_dict(), "model.pt")
model.load_state_dict(torch.load("model.pt", map_location="cpu"))

# 断点续训 checkpoint（训练用）
ckpt = {
    "model": model.state_dict(),
    "optimizer": optimizer.state_dict(),
    "scheduler": scheduler.state_dict(),
    "step": step,
}
torch.save(ckpt, f"ckpt_{step}.pt")

# 恢复
ckpt = torch.load("ckpt_1000.pt", map_location=device)
model.load_state_dict(ckpt["model"])
optimizer.load_state_dict(ckpt["optimizer"])
scheduler.load_state_dict(ckpt["scheduler"])
start_step = ckpt["step"]

11.13 调试清单：shape 断言与最小复现

# ① shape 断言：把期望写出来，让错误尽早暴露
assert x.shape == (B, S), f"expected [B,S], got {x.shape}"
assert q.shape == (B, H, S, D)

# ② device / dtype 断言
assert x.device == model.device or x.device.type == "cuda"
assert y.dtype == torch.long

# ③ 最小复现：造 2 个样本、3 个 token、8 维的小输入
x = torch.randint(0, 100, (2, 3)).to(device)

常见报错对照表（收藏）