Lesson 16 · 約 22 分鐘

訓練迴圈深入

Optimizer 怎麼挑 · LR Scheduler 怎麼排 · Mixed Precision 必開

16 ・訓練迴圈深入:Optimizer / LR Scheduler / Mixed Precision

大部分人「會跑訓練迴圈」但「不知道為什麼調這些超參數」。這一講把訓練迴圈的三個核心齒輪拆開:選哪個 optimizer、怎麼排 learning rate、怎麼用混合精度多跑 2× 模型。配 lesson 15 的 PyTorch 基礎一起讀。

0. 從一行公式長出來的整個訓練

🔑 重點 · 一切都是這條更新規則的變體

θ ← θ - lr · g

θ = 參數,g = 梯度,lr = 學習率。SGD / Adam / AdamW / LAMB 都是在「怎麼算 g」跟「怎麼算 lr」上動手腳。Scheduler 是「怎麼讓 lr 隨時間變」。Mixed precision 是「怎麼讓 g 算得更快」。理論上就這四個變數,讓現代 AI 動起來。

1. Optimizer:從 SGD 到 AdamW

1.1 SGD with momentum(最古典)

optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9, weight_decay=1e-4)

更新規則:

v ← β·v + g           # momentum 累積過去的梯度方向
θ ← θ - lr · v        # 沿累積方向走

特性:

β = 0.9 是 GD 的「移動平均」,讓更新方向比較穩、不被單一 batch 震盪帶歪
簡單、可解釋,理論收斂分析最完整
缺點:每個參數共用同一個 lr → 維度量級差很大時(像 embeddings vs final layer)調起來很煩

1.2 Adam(深度學習的工作馬)

optimizer = torch.optim.Adam(model.parameters(), lr=3e-4)  # 「Karpathy constant」

更新規則(簡化):

m ← β₁·m + (1-β₁)·g           # 梯度的指數平均
v ← β₂·v + (1-β₂)·g²          # 梯度平方的指數平均
m̂ = m / (1-β₁ᵗ)              # bias correction(早期 t 小時 m 偏向 0)
v̂ = v / (1-β₂ᵗ)
θ ← θ - lr · m̂ / (√v̂ + ε)    # 自適應 per-parameter lr

為什麼這套贏 SGD:

每個參數自動有自己的 effective lr(被 √v̂ 縮放)→ 不同層 / 不同維度自動平衡
bias correction 讓訓練最前 100 步不會「梯度很小走不動」
適合梯度量級劇烈變化的場景 = 大部分深度學習

1.3 AdamW(現在的標配)

optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=0.01)

跟 Adam 的差別 = 把 weight decay 從「混進梯度」變「直接縮參數」:

θ ← θ - lr · (m̂/(√v̂+ε) + λ·θ)   # Adam:weight decay 包在梯度裡
θ ← θ - lr · m̂/(√v̂+ε) - lr·λ·θ  # AdamW:weight decay 獨立縮參數

為什麼這個小改動有效:Adam 的 √v̂ 對 weight decay 也做了自適應縮放,結果 regularisation 強度跟著梯度量級走 → 不穩定。AdamW 把 weight decay 從那條 path 拿出來,regularisation 跟梯度估計脫鉤,訓練更穩、generalisation 更好。

1.4 哪個用哪個

場景	選	原因
純 CV(ResNet 等)	SGD with momentum	理論清楚 + 一致經驗:CNN 用 SGD final 結果常比 Adam 好
Transformer / NLP	AdamW	標配
大模型微調(7B+)	AdamW + 8-bit / LoRA	省顯存
想嚐鮮	Lion / SOAP / Sophia	新發表的,實證在某些 setup 上贏 AdamW 但驗證少

2. Learning Rate Scheduler:讓 lr 隨時間變

固定 lr 一路訓 → loss 一開始降很快,中後期震盪。Scheduler 解決這個。

2.1 三個常見排程

from torch.optim.lr_scheduler import (
    CosineAnnealingLR, OneCycleLR, LambdaLR
)

1. Linear warmup + cosine decay(LLM 預訓練 / 微調最常見):

lr ramps from 0 → peak_lr in N warmup_steps
then decays peak_lr → 0 along cosine curve over total_steps

Warmup 防止 epoch 0 lr 太大爆炸(weight 還在亂跳)
Cosine decay 後期慢慢收歛
HF Trainer 默認

2. OneCycleLR(fast.ai 推、CV 主流):

lr triangles: low → high → low across one training run
momentum inversely (high → low → high)

中期保持高 lr 探索 loss landscape
後期 lr 降 → 找到平坦 minima → generalisation 好

3. ReduceLROnPlateau(metric-driven):

scheduler = ReduceLROnPlateau(optimizer, mode='min', patience=3, factor=0.5)
# val loss 連 3 epoch 沒降 → lr ÷= 2

不需要預先決定 schedule shape,自動跟 val loss 走
適合資料量小、訓練曲線不可預測

2.2 哪個用哪個

場景	選
Transformer 預訓練 / 微調	linear warmup + cosine decay
CV from scratch(ImageNet 等)	OneCycleLR
小資料集 / 自己摸索	ReduceLROnPlateau
Multi-stage / discriminative lr(每層不同 lr)	自己寫 LambdaLR

# HF Trainer 一行搞定
from transformers import get_linear_schedule_with_warmup
scheduler = get_linear_schedule_with_warmup(
    optimizer, num_warmup_steps=100, num_training_steps=total_steps
)
# 訓練迴圈內每步呼叫
scheduler.step()

2.3 找最佳 lr 的兩個方法

方法 A:LR Range Test(Leslie Smith 的方法,fast.ai 用)

跑 100-200 step,lr 從 1e-7 線性升到 10
畫 loss vs lr 曲線
取「loss 下降最快、還沒爆炸」的點 → 那個是 peak_lr

from torch_lr_finder import LRFinder
lr_finder = LRFinder(model, optimizer, criterion, device='cuda')
lr_finder.range_test(train_loader, end_lr=10, num_iter=100)
lr_finder.plot()  # 看曲線

方法 B:lesson 14 的 ÷3 / ×3 經驗法

從 1e-3 起跳,看 50 step 內 loss 趨勢
不降 → ×3,震盪 → ÷3

3. Mixed Precision(必開,不開白丟 30-40% 效率)

3.1 為什麼能加速 + 省記憶體

float32 → float16 / bfloat16:

每元素 4 → 2 bytes,顯存省一半
Tensor Core(A100 / H100 / 4090) 對 16-bit matmul 有專用硬體加速,實測 1.5-3×

3.2 標準寫法(`torch.amp`)

from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()
for x, y in loader:
    optimizer.zero_grad()
    with autocast(dtype=torch.float16):     # forward 在 fp16 跑
        out = model(x)
        loss = criterion(out, y)
    scaler.scale(loss).backward()            # 放大 loss 防 underflow
    scaler.step(optimizer)                   # 自動 unscale + step
    scaler.update()                          # 動態調整 scale

重點:

參數本身仍是 fp32(model 沒搬到 fp16),只在 forward / backward 中間動 fp16
GradScaler 動態 scale loss 防止 fp16 表示範圍下溢(exp(-15) 就 0 了)
A100+ 推薦 dtype=torch.bfloat16(範圍跟 fp32 一樣,不需 GradScaler)

3.3 bfloat16 vs float16

	float16	bfloat16
範圍	±65,504	±3.4e+38(跟 fp32 同)
精度	較高(10-bit mantissa)	較低(7-bit)
Hardware	大部分 GPU	A100 / H100 / TPU / 4090
需要 GradScaler?	是	否

4. Gradient Accumulation:小 GPU 跑大 batch

batch_size = 64 跑 OOM?跑 batch_size=8 但累積 8 步 → 等效 batch_size = 64:

ACCUM = 8
for i, (x, y) in enumerate(loader):
    with autocast(dtype=torch.bfloat16):
        loss = model(x) / ACCUM            # 重點:除以 ACCUM 才等效平均
    loss.backward()                         # 梯度累加進 .grad

    if (i + 1) % ACCUM == 0:
        optimizer.step()
        scheduler.step()
        optimizer.zero_grad()

重點:loss / ACCUM 才等於「真實 batch_size = ACCUM × micro_batch」的平均梯度,不除 loss 會跟真實 batch 數值差 ACCUM 倍。

5. 完整訓練骨架(現代版)

🏁 里程碑 · 能照這個骨架對自己的資料跑通 = 你能做業界微調了

from torch.optim import AdamW
from torch.cuda.amp import autocast, GradScaler
from transformers import get_linear_schedule_with_warmup

model = MyModel().cuda()
optimizer = AdamW(model.parameters(), lr=3e-4, weight_decay=0.01)

total_steps = len(train_loader) * num_epochs // ACCUM
scheduler = get_linear_schedule_with_warmup(
    optimizer, num_warmup_steps=total_steps // 10, num_training_steps=total_steps
)
scaler = GradScaler()

for epoch in range(num_epochs):
    model.train()
    for i, batch in enumerate(train_loader):
        with autocast(dtype=torch.bfloat16):
            out = model(batch)
            loss = out.loss / ACCUM
        scaler.scale(loss).backward()

        if (i + 1) % ACCUM == 0:
            scaler.unscale_(optimizer)
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            scaler.step(optimizer)
            scaler.update()
            scheduler.step()
            optimizer.zero_grad()

    # eval
    model.eval()
    with torch.no_grad():
        for batch in val_loader:
            with autocast(dtype=torch.bfloat16):
                out = model(batch)
            # accumulate metrics ...

6. Gradient Clipping:防梯度爆掉

torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

把全部參數的梯度當一條向量,如果其 L2 norm > max_norm,整條按比例縮
LLM 必開(梯度偶爾會爆)
max_norm=1.0 是常見值,deep transformer 可試 0.5

7. 自我檢測

接下來

講義機制懂了,下一講 17 · Evaluation methodology — train/val/test 切錯就白訓,選錯 metric 就白做
想動手:回去看 code/02_tiny_gpt.py,換成 AdamW + warmup + bf16 + clip,看 loss 曲線跟原版差異

16 ・ 訓練迴圈深入:Optimizer / LR Scheduler / Mixed Precision

0. 從一行公式長出來的整個訓練

1. Optimizer:從 SGD 到 AdamW

1.1 SGD with momentum(最古典)

1.2 Adam(深度學習的工作馬)

1.3 AdamW(現在的標配)

1.4 哪個用哪個

2. Learning Rate Scheduler:讓 lr 隨時間變

2.1 三個常見排程

2.2 哪個用哪個

2.3 找最佳 lr 的兩個方法

3. Mixed Precision(必開,不開白丟 30-40% 效率)

3.1 為什麼能加速 + 省記憶體

3.2 標準寫法(torch.amp)

3.3 bfloat16 vs float16

4. Gradient Accumulation:小 GPU 跑大 batch

5. 完整訓練骨架(現代版)

6. Gradient Clipping:防梯度爆掉

7. 自我檢測

接下來

16 ・訓練迴圈深入:Optimizer / LR Scheduler / Mixed Precision

3.2 標準寫法(`torch.amp`)