モデル状態辞書(state_dict)とテンソル直列化の基礎
Agent支援エンジニアリングチェック
前提条件
第54日Transformer評価課題の完了およびPyTorchモジュールの基礎知識。
学習目標
model.state_dict() の内部構造、パラメータと登録バッファの違い、torch.save/torch.load における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 は内部で pickle を使用し、テンソルメタデータとストレージポインタを記録します。
3. モデルオブジェクト全体の保存と state_dict 保存を比較し、コードの結合度低減とアーキテクチャ安定性の観点から公式が state_dict を推奨する理由を把握します。
4. GPUで保存した重みをロードする際の map_location='cpu' の必須性を確認します。
演習と予測記録
小さな2層GPTモデルを作成し、その state_dict を 'ckpt_model.pt' に保存してください。新しいモデルインスタンスに map_location='cpu' を指定してロードし、torch.equal() で全テンソルが完全一致することを確認してください。
ロード前に model.state_dict().keys() を確認してください。torch.load('ckpt_model.pt', map_location='cpu') を使用します。
import torch
# 状態辞書の直列化と安全な再読み込み
state = model.state_dict()
torch.save(state, 'ckpt_model.pt')
new_model.load_state_dict(torch.load('ckpt_model.pt', map_location='cpu'))
assert all(torch.equal(p1, p2) for p1, p2 in zip(model.parameters(), new_model.parameters()))
状態辞書を正確にシリアライズし、新規インスタンスへ完全一致でロードできる。
公式ソースリソースと権利証跡
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.