0.2 环境与 GPU 准备
0.2 环境与 GPU 准备
本节摘要：本节解决开跑前的最后一个问题：算力从哪来、要多少。先给结论：SFT 是 nanochat 流水线里最"便宜"的阶段——124M 量级模型的 SFT 用一块 24GB 消费级显卡就能完整复现，云上租几小时即可。随后给出三套方案的对比表（本地 GPU / 云按租 / Modal 云函数，后者来自 AI Engineering Academy 的社区方案），一份按"参数量 × 精度 × 优化器状态"估算显存的速算法，以及全书统一的运行环境约定与信息时效声明。
学习目标
阅读完本节，你应当能够：
- 在本地 GPU、云租用、Modal 三方案中选出适合自己的组合。
- 用速算公式估算 SFT 所需显存，并知道三个最有效的省显存开关。
- 按全书记载的环境约定（Python/PyTorch/CUDA 版本与目录结构）准备好工作区。
一、三套算力方案对比
| 维度 | 本地 GPU | 云租用（按小时） | Modal（云函数） |
|---|---|---|---|
| 典型硬件 | RTX 3090/4090（24GB） | A100/H100 单卡或 8 卡 | 按需调度 GPU（社区方案常用 A100/H100） |
| 成本模型 | 一次性购机 + 电费 | 约 1~4 美元/小时（单卡，随型号浮动） | 按秒计费，脚本即基础设施 |
| 适合人群 | 有闲置卡、迭代频繁 | 想复现完整流水线一次跑通 | 会写 Python、想免运维 |
| 主要缺点 | 显存/卡数受限；预训练跑不动 | 环境配置、数据上传费时 | 依赖第三方平台，调试链路长 |
| 本书用法 | 第 2~4 章 SFT 实战足够 | 复现完整四阶段时租用 | 备选：社区有 nanochat 的 Modal 完整方案 |
💡 社区方案：AI Engineering Academy 发布过 nanochat 的 Modal 云训练方案（脚本化的 GPU 云函数，免运维跑通全流程），对应本表第三列；细节以其仓库/文档为准。国内读者也可换成任意"按小时租 GPU"的平台，思路完全一致：上传权重与数据 → 跑训练脚本 → 取回产物。
二、显存估算：SFT 要多少
速算公式（混合精度 + Adam 优化器，最常见配置）：
显存 ≈ 参数量 × 2（权重 fp16/bf16） + 参数量 × 2（梯度） + 参数量 × 8（Adam 的一阶/二阶动量，fp32 存两份 + 主权重一份，共 12 份的保守计法取 8） + 激活值（随 batch × 序列长度增长，梯度检查点可大幅压缩）
对 nanochat 的 124M 模型做 SFT（示意估算，实测因框架与序列长度而异）：
| 项 | 估算 |
|---|---|
| 权重 + 梯度 + Adam 状态 | 124M × (2+2+8) bytes ≈ 1.5GB |
| 激活值（batch 8 × 4096 tokens，开梯度检查点） | 约 2~6GB |
| 总计 | 约 4~8GB，24GB 卡非常宽裕 |
同公式换到 1.5B 模型：约 18GB 静态部分 + 激活值——24GB 卡需要开梯度检查点并减小 batch。先小后大：所有实验先在 124M 上跑通再放大，是本书的默认纪律。
三个最有效的省显存开关（第 3 章实战脚本会用到）：
- 混合精度（bf16）训练；
- 梯度检查点（gradient checkpointing）——用约 30% 重算换大幅激活值下降；
- 小 batch + 梯度累积（accumulate steps）凑有效 batch。
三、软件环境
全书代码统一按此环境书写（版本号为主流稳定版即可，不要求完全一致）：
# setup_env.sh —— 全书运行环境（写法示意） conda create -n nanochat python=3.11 -y conda activate nanochat pip install torch --index-url https://download.pytorch.org/whl/cu121 pip install transformers datasets numpy matplotlib # 部署与语音阶段（第 6~7 章）另需 llama.cpp 与 whisper 环境 nvidia-smi # 确认驱动与卡被识别
工作区目录约定（后续章节的代码都以此为根）：
workspace/ ├── base/ # 基座模型权重（前半程的产物，或社区复现版） ├── data_sft/ # 第 2 章产出的指令数据（jsonl） ├── runs/sft_124m/ # 第 3 章训练输出（checkpoint、日志） └── chat_model/ # 训练后的对话模型，供第 4~7 章使用
⚠️ Windows 用户建议在 WSL2 里跑训练（部分训练依赖的编译工具链在原生 Windows 上安装成本高）；推理与数据处理在原生 Windows 下没有问题。
四、全书约定与声明
- 代码来源标注：每个代码块首行注释标注来源文件名；nanochat 原生做法注明出处（官方仓库/发布帖），本书改写的教学版标注「写法示意，以官方仓库 README 为准」。
- 数字属性：官方自报（约 4 小时、约 100 美元跑通四阶段）、社区实测（预训练 loss 10.4→2.32、SFT loss ~1.3 起步，来自公开 walkthrough）、示意估算（本节显存表）三类数字在正文中分别标注，不混用。
- 信息时效：nanochat 仓库高速演进（当前 master 已转向 GPT-2 speedrun 路线、SFT 数据换用 SmolTalk），本书以 2025 年 10 月发布帖（Discussion #1）的四阶段流程为叙事主轴；跑代码前先对照官方 README 的最新脚本名。
- 算力兜底：没有 GPU 也能读完全书并完成第 2 章全部数据实验（CPU 足够）；只有第 3 章训练与第 6~7 章部署必须 GPU/CPU 推理环境。
本节要点回顾
- 三方案：本地 24GB 卡足够做 124M 的 SFT；复现完整流水线租云；免运维选 Modal（社区方案）。
- 显存速算：参数量 × 12 bytes 是静态底盘；bf16、梯度检查点、梯度累积是三大开关。
- 全书代码按统一工作区与环境书写；数字标注属性；nanochat 以官方 README 为最终依据。
环境就绪，地图先行——第 1 章把 nanochat 四阶段流水线完整铺开，标出本书覆盖的后半程。