第 2 章环境准备与首次运行

文档摘要

第 2 章环境准备与首次运行

理论有了，这一章我们把环境搭好，并在五分钟内跑通第一次蒸馏训练。跑起来，是建立信心的第一步。

2.1 硬件要求

配置	是否可行	说明
有 CUDA 的 GPU（≥ 6GB 显存）	✅ 推荐	默认配置可流畅训练，约几分钟一个 epoch
仅 CPU	✅ 可行	训练较慢，建议用极小配置冒烟
无网络	⚠️ 受限	教师模型需首次下载，需提前联网缓存

教师是预训练标准 GPT2（约 124M 参数，权重文件约 548MB），学生默认约 10M 参数。两者同时驻留显存，对入门级 GPU 也算友好。

2.2 软件依赖

依赖	最低版本	用途
Python	3.9	项目使用类型注解与 dataclass 特性
PyTorch	2.0	深度学习框架
transformers	4.36	提供 GPT2 实现
tiktoken	0.5	GPT-2 分词
datasets	2.14	数据集下载（可选，缺失会自动降级）
numpy	1.23	数值处理
tqdm	4.65	进度条

2.3 安装步骤

步骤 1：创建虚拟环境（推荐）

隔离环境能避免污染全局包。


# 创建
python -m venv .venv

# 激活（Windows）
.venv\Scripts\activate

# 激活（Linux / macOS）
source .venv/bin/activate

步骤 2：安装依赖


# 方式一：使用项目依赖清单（推荐）
pip install -r requirements.txt

# 方式二：手动逐个安装
pip install torch transformers tiktoken datasets numpy tqdm

若你的机器有 NVIDIA GPU，建议先按 PyTorch 官网指引安装对应 CUDA 版本的 torch，再装其余依赖。

步骤 3：验证安装

逐项 import，确认无报错：


import torch, transformers, tiktoken, numpy, tqdm
print("torch:", torch.__version__)
print("transformers:", transformers.__version__)
print("CUDA 可用:", torch.cuda.is_available())

预期输出类似：


torch: 2.5.1
transformers: 4.47.1
CUDA 可用: True

2.4 教师模型预下载（强烈建议）

教师 GPT2 在首次运行时会从 HuggingFace Hub 下载（约 548MB）。如果你的训练机器网络不稳定，建议提前手动下载并缓存，避免训练中途因网络中断而失败。

最简单的预下载方式——让 transformers 把模型拉到本地缓存：


from transformers import GPT2LMHeadModel, GPT2TokenizerFast

# 这两行会触发下载并缓存到 ~/.cache/huggingface/
GPT2LMHeadModel.from_pretrained("gpt2")
GPT2TokenizerFast.from_pretrained("gpt2")
print("教师模型与分词器已缓存。")

下载完成后，后续运行会直接读本地缓存，不再联网。

网络受限环境：若 HuggingFace Hub 访问不畅，可设置镜像端点，例如配置环境变量 HF_ENDPOINT 指向可用镜像，再重试上述下载。

2.5 首次运行：极小配置冒烟测试

为了快速验证环境，我们用一个极小配置跑几个训练步。参数含义会在后续章节详解，这里照抄即可。


# 极小配置冒烟：block_size 缩到 32，只跑 3 步
python train.py --block-size 32 --batch-size 4 --max-iters 3 \
                --warmup-iters 1 --log-iter 1 --save-iter 2 --eval-iter 2

如果一切正常，你会看到类似输出：


[train] 使用设备: cuda
[dataset] ... tiny_shakespeare 下载/读取 ...
[dataset:train] ... 个样本
[dataset:eval] ... 个样本
[model] 加载教师模型: gpt2
[model] 教师参数量: 124.44M (已冻结)
[model] 学生配置: 2 层 / 4 头 / 256 维，参数量 10.12M
[train] 开始蒸馏: 共 3 步, T=2.0, alpha=0.5, cache=off, ...
蒸馏: 100%|████████| 3/3 ...
[train] step 1/3 | total=5.43 ce=4.89 kd=5.97 | lr=...
...
[train] 蒸馏训练完成。
[train] 最终学生模型已保存: checkpoints/student_final.pt

看到「损失分解日志」（total / ce / kd 三项）和「最终模型已保存」，说明整个训练流水线打通了。

2.6 跑一次完整的默认训练

确认冒烟测试通过后，可以跑默认配置（约 5000 步）。这一步会花几分钟到几十分钟，取决于硬件。


python train.py

想边训练边观察损失曲线？训练日志里每 log_iter 步会打印一次 total / ce / kd，把它们记下来画图即可。我们会在《第 7 章》讨论如何分析这些曲线。

2.7 常见安装问题速查

现象	可能原因	解决
`No module named 'tiktoken'`	tiktoken 未装	`pip install tiktoken`
教师下载报 `No space left on device`	磁盘空间不足	清理缓存或换盘，至少预留 1GB
教师下载超时/失败	网络问题	见 2.4 节镜像方案
训练卡在 CPU 且极慢	未检测到 CUDA	检查 torch 是否为 CUDA 版本
Windows 多进程报错	num_workers > 0	加 `--num-workers 0`

更多报错与排查，请参考后续《附录 C 常见问题排查》。

本章小结

推荐环境：Python 3.9 + 有 CUDA 的 GPU，CPU 也可用极小配置跑通。
核心依赖：torch、transformers、tiktoken、datasets、tqdm。
教师模型需首次联网下载约 548MB，建议提前缓存。
用极小配置做冒烟测试，确认全流程打通后再跑完整训练。

动手实验：跑通冒烟测试后，尝试把 --temperature 改成 1.0 和 8.0，观察 kd 损失值的变化，直观感受温度对软标签的影响。

下一站：环境就绪，我们开始深入代码。在《第 3 章配置体系》中，看看如何用一个数据类优雅地管理几十个超参数。

第 2 章 环境准备与首次运行

文档摘要