模型狀態字典(state_dict)與張量序列化機制
Agent 輔助工程檢查
前置要求
完成第 54 天 Transformer 考核,掌握 PyTorch Module 核心原理。
學習目標
深入掌握 model.state_dict() 的底層字典映射結構、可學習參數與持久化緩衝區的區別,以及跨設備序列化(CPU/CUDA)映射機制。
操作指引
學習來源:在本活動現有的閱讀與練習時段內,參考 CS336 Assignment 1: Basics (Spring 2025 tagged release) 的 tests/test_serialization.py: test_checkpointing; model, optimizer and iteration state,用自己的話解釋一處與本地任務的聯繫。保留現有練習與驗收檢查;若外部資料無法存取,使用課程自寫的解釋。1. 遍歷 model.state_dict(),觀察鍵名與子模組樹形結構的對應關係(如 'blocks.0.attn.c_attn.weight')。
2. 理解 torch.save 底層依賴 Python pickle 機制打包張量元數據與原始存儲區。
3. 對比保存完整 Python 模型對象與保存純 state_dict 的差異,領會官方標準推薦 state_dict 的解耦與版本穩定性優勢。
4. 實踐 map_location='cpu' 選項,消除在無 GPU 或不同卡號設備上加載權重時的 CUDA 報錯風險。
自主實作與記錄
實例化一個小型的 2 層 GPT 模型,將其 state_dict 持久化至 'ckpt_model.pt'。重新創建一個新的隨機初始化模型,使用 map_location='cpu' 恢復權重,並斷言兩者的每一個參數張量完全逐元素相等。
保存前可列印 keys()。恢復時調用 model.load_state_dict(torch.load('ckpt_model.pt', map_location='cpu'))。
import torch
state = model.state_dict()
torch.save(state, 'ckpt_model.pt')
loaded = torch.load('ckpt_model.pt', map_location='cpu')
model2.load_state_dict(loaded)
for k, v in model.state_dict().items():
assert torch.equal(v, model2.state_dict()[k])
能準確提取與序列化 state_dict,並在新建模型上達成 100% 精度參數還原。
權威來源資源與權利證明
CS336 Assignment 1: Basics (Spring 2025 tagged release)
對應章節 / 時間戳: tests/test_serialization.py: test_checkpointing; model, optimizer and iteration state
驗證版本與日期: spring2025 tag; README, tests and MIT LICENSE checked 2026-11-01
Lead assignment-driven source for pretraining component practice and tests. nanoGPT remains a compact historical code-reading companion, not the current recommended training curriculum.
Fluent Python, PyTorch, and linear algebra
Local self-contained pretraining scripts run on standard local CPUs or budget GPUs without Stanford cluster dependency.
對比候選資源與選型分析 (2)
Simpler code layout, useful complementary reference for minimal PyTorch Transformer.
High-level API focus rather than building from raw tensor fundamentals.
PyTorch Tutorial: Saving and Loading Models
對應章節 / 時間戳: What is a state_dict?
驗證版本與日期: PyTorch 2.x Documentation (Verified: 2026-10-30)
Official documentation detailing state_dict serialization, optimizer state saving, and safe loading patterns.
Python, PyTorch tensors and nn.Module basics.
All checkpoint serialization semantics, dictionary schema keys, and safe recovery steps are self-contained in lesson code.
對比候選資源與選型分析 (2)
TensorFlow trackable checkpoint architecture; PyTorch's explicit dictionary is our target implementation standard.
Safetensors provides safe zero-copy tensor storage without pickle, compared directly when discussing pickle execution risks.
nanoGPT: train.py pretraining loop
對應章節 / 時間戳: train.py: checkpoint dict {'model', 'optimizer', 'model_args', 'iter_num', 'best_val_loss', 'config'} written with torch.save to out_dir/ckpt.pt; the init_from == 'resume' branch reloads it
驗證版本與日期: Git master commit 71c3905 (Verified: 2026-10-30)
Compact historical code reference for the named implementation functions; CS336 spring2025 provides the lead component exercises. Do not infer top-p, atomic persistence or cloud budget enforcement from nanoGPT.
Transformer architecture from Days 46-48 and training loop from Days 49-51.
Complete checkpoint saving patterns and resuming loops are directly implemented in lesson exercises.
對比候選資源與選型分析 (2)
High-complexity multi-GPU distributed model-parallel checkpointing; nanoGPT provides accessible single-GPU clarity.
High-level framework callback abstraction; nanoGPT keeps explicit control over state dict dictionaries.