模型状态字典(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.