Lesson 15 · 約 25 分鐘

PyTorch 深入

Tensor 內部 / autograd / hooks / profiler — 用 PyTorch 不再卡

15 ・ PyTorch 深入:Tensor / Autograd / 訓練機制看穿

02 講把模型理論講清楚,本機跑得起小 GPT。但你寫的多半是 PyTorch — 真的搞懂它「在背後做什麼」對 debug、客製化、效能調優都關鍵。這一講拆四件事: tensors、autograd、訓練機制、debugging hooks + profiler。讀完你不會再對著「為什麼這裡 .detach()?」「為什麼這個梯度沒進去?」卡住。

0. 一句話定位

1. Tensor 的真實面目

1.1 內部結構

一個 torch.Tensor 在記憶體裡分兩塊:

Storage:真正的數字 buffer(連續一條 float32 / float16 / int64…)
View metadata:shape / stride / offset(怎麼解讀那條 buffer)

import torch
x = torch.arange(12).reshape(3, 4)
print(x.shape)        # torch.Size([3, 4])
print(x.stride())     # (4, 1) — 跳一列要走 4 個元素,跳一欄走 1
print(x.storage_offset())  # 0

y = x.t()             # 轉置(view operation,不複製)
print(y.shape)        # torch.Size([4, 3])
print(y.stride())     # (1, 4) — stride 翻轉,storage 沒動
print(y.is_contiguous())  # False!

為什麼這重要:.reshape() 通常零成本(改 metadata),但 .contiguous() / .clone() 會複製整條 storage。大模型一不小心多 clone 一次就 OOM。

1.2 dtype 與 device

dtype 決定每元素幾 bytes:float32 = 4, float16 / bfloat16 = 2, int8 = 1。混合精度訓練 = 算的時候用 fp16/bf16,參數本身仍是 fp32。
device 是物理位置:cpu / cuda:0 / mps。跨 device 不能直接運算 — x.to('cuda') 顯式搬。

x = torch.randn(1024, 1024, device='cuda', dtype=torch.float16)
print(x.element_size() * x.numel() / 1024 / 1024, 'MB')  # 2 MB

1.3 view / reshape / transpose / permute 區別

API	行為	複製?
`.view(...)`	改 shape,要求 contiguous	否
`.reshape(...)`	改 shape,需要時 clone	看情況
`.transpose(0, 1)` / `.t()`	兩維互換	否(stride 翻)
`.permute(2, 0, 1)`	任意維重排	否
`.contiguous()`	強制連續(複製到新 storage)	是

2. Autograd:梯度怎麼自動算

2.1 計算圖 = 動態建立

每次 forward,PyTorch 偷偷建一棵「計算圖」:節點是 Tensor,邊是 Function(op)。當你 loss.backward(),它從 loss 反向走一遍,每步用 chain rule 把梯度傳回去。

x = torch.tensor([2.0, 3.0], requires_grad=True)
y = x ** 2          # y.grad_fn = <PowBackward0>
z = y.sum()         # z.grad_fn = <SumBackward0>
z.backward()
print(x.grad)       # tensor([4., 6.]) → dz/dx = 2x

Key 性質:

動態圖:每次 forward 重新建,所以 if/loop/recursion 都能跑(對比 TF 1.x 的靜態圖)
leaf tensors = 起點(requires_grad=True 自己建的),非 leaf 是 op 算出來的中間結果
x.grad 累加(每次 backward +=)→ 訓練迴圈一定要 optimizer.zero_grad()

2.2 `.detach()` / `.no_grad()` / `.requires_grad_(False)` 差別

三個都是「斷開計算圖」但情境不同:

API	用在哪	效果
`tensor.detach()`	拿出一份沒梯度的 view	storage 共享,grad_fn = None
`torch.no_grad():`	整段 block 不建計算圖	推理時必開,省一半記憶體
`tensor.requires_grad_(False)`	凍結某個參數長期不訓	leaf tensor,LoRA / feature extractor 常用

# 推理時的標準寫法
model.eval()
with torch.no_grad():
    logits = model(x)

2.3 backward 進階控制

# 多 loss 場景 — retain_graph 讓你 backward 兩次
loss1.backward(retain_graph=True)
loss2.backward()  # 沒有 retain_graph 第二次會 RuntimeError

# 取得任意 tensor 的梯度(不只是 leaf)
y.retain_grad()  # 訓練時不要亂用,debug 工具
z.backward()
print(y.grad)

3. nn.Module:模型的骨架

3.1 為什麼一定要繼承 nn.Module

import torch.nn as nn

class Net(nn.Module):
    def __init__(self, d):
        super().__init__()
        self.lin = nn.Linear(d, d)      # 自動註冊參數
        self.bn = nn.BatchNorm1d(d)     # 自動註冊 buffer (running_mean 等)

    def forward(self, x):
        return self.bn(self.lin(x))

Module 自動做四件事:

.parameters():遞迴收集所有可訓練張量 → 傳給 optimizer
.train() / .eval():把所有 Module 切到對應 mode(影響 dropout / batchnorm)
.to(device):遞迴搬所有參數 + buffer
.state_dict() / .load_state_dict():存 / 載 checkpoint

3.2 forward 不要直接 call

y = model(x)        # ✓ 正確:走 __call__,觸發 hooks
y = model.forward(x)  # ✗ 錯:跳過 hooks,grad_fn 跟 mode 切換可能壞

4. Hooks:看 / 改中間狀態

forward_hook 跟 backward_hook 讓你不改 model code 就能 inspect 中間 tensor。Debug 時是寶。

# 場景:看 BERT 第 5 層輸出
def print_shape(module, input, output):
    print(f'{module.__class__.__name__} → {output.shape}')

h = model.layers[5].register_forward_hook(print_shape)
model(x)
h.remove()  # 用完拿掉,不然永遠 trigger

進階用法:梯度爆掉時 用 register_full_backward_hook 看 loss 對哪一層輸出最敏感,定位數值不穩的層。

💡 提示 · 看 attention 權重的最簡單辦法

自己寫的模型 forward return logits, attn_weights;Hugging Face 模型則 output_attentions=True:

out = model(input_ids, output_attentions=True)
print(out.attentions[0].shape)  # (B, heads, T, T)

5. Profiler:找瓶頸不靠猜

torch.profiler 比 time.time() 量更準,還告訴你 GPU 哪個 op 最慢、有沒有 host-device 同步停頓。

from torch.profiler import profile, ProfilerActivity

with profile(
    activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
    record_shapes=True
) as prof:
    for _ in range(20):
        loss = model(batch).loss
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()

print(prof.key_averages().table(sort_by='cuda_time_total', row_limit=10))

典型輸出:

Name              Self CUDA   Self CPU
matmul              45 ms       2 ms
softmax              8 ms       1 ms
to                   6 ms      50 ms     ← host→device 太多次,優化點

6. 七個必背的「PyTorch 慣用語」

🏁 里程碑 · 這 7 個你能流暢寫出來 = PyTorch 內功合格

# 1. 訓練迴圈 boilerplate
for x, y in loader:
    x, y = x.to(device), y.to(device)
    optimizer.zero_grad()
    out = model(x)
    loss = criterion(out, y)
    loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    optimizer.step()

# 2. 推理迴圈 boilerplate
model.eval()
with torch.no_grad():
    for x in loader:
        out = model(x.to(device))

# 3. 凍結部分參數
for p in model.encoder.parameters():
    p.requires_grad_(False)

# 4. 存 / 載 checkpoint(只存 state_dict,別存整個 model 物件)
torch.save({'model': model.state_dict(), 'opt': optimizer.state_dict()}, 'ckpt.pt')
ckpt = torch.load('ckpt.pt', map_location=device)
model.load_state_dict(ckpt['model'])

# 5. 抓 mean / std / max(總是配 .item() 拿純值)
print(loss.item(), out.mean().item(), out.std().item())

# 6. 算參數量
n = sum(p.numel() for p in model.parameters() if p.requires_grad)

# 7. random seed
torch.manual_seed(42); torch.cuda.manual_seed_all(42)

7. 自我檢測

答案:看本講對應段落。卡住的時候,Lesson 14 Debugging Playbook 是好搭檔。

接下來

內功有了,下一講 16 · 訓練迴圈深入 — SGD vs Adam vs AdamW、LR scheduler、loss scaling
把這份知識套到自己微調的程式碼上:回去看 code/02_tiny_gpt.py,對著每一行說出它在 PyTorch 裡發生什麼事。能講出 95% = 你真的懂了