算法工程师的工具箱：从一个想法到一次能跑的实验

01 · Python 使用层

结论：训练代码里反复出现的那一小撮语法各是一个 Python 协议，PyTorch 建在上面；生成器让峰值内存与文件大小无关；GIL 让 CPU 密集的线程不并行，进程池才有用。

两个能测出来的数字

• 内存：19.4 MB 的 JSONL 读成 list 峰值 102.9 MB（5.3 倍）；生成器流水线 11.5 MB——与文件多大无关
• 并行：CPU 密集的预处理，串行 0.46 s、8 线程 0.46 s（GIL）、8 进程 0.15 s（3.1×）
• 进程池的规矩：传给 Pool.map 的函数要在模块顶层、全局变量在子进程里是拷贝——DataLoader worker、datasets.map(num_proc=8)、DDP 的每个 rank 都是同一套规矩
• 读 traceback：从下往上、跳过 _call_impl、找自己文件的最后一帧

02 · 数据科学三剑客

结论：所有「沿哪个维度」是一个概念——轴；广播太宽容所以形状错误常不报错；与参考实现对数值到浮点精度，是验证手写算子的标准。

一个数字不算结论：多 seed 与置信区间

• 评测表带上 ci95 = 1.96 √(p(1−p)/n)：80 道题 ±10.9，5 个点的差分辨不出
• loss 曲线：对数 x 轴、多 seed 阴影带；不重叠才算差别

手写 attention 与 PyTorch 对拍

scores = np.einsum("btd,bsd->bts", q, k) / np.sqrt(d)      # [B, T, S]
scores = np.where(mask, scores, -np.inf)                   # 因果 mask
p = np.exp(scores - scores.max(-1, keepdims=True))         # 减最大值再 exp
p = p / p.sum(-1, keepdims=True)
out = np.einsum("bts,bsd->btd", p, v)                      # [B, T, D]

• 广播三条规则；(3,) + (3, 1) 悄悄变成 (3, 3)——关键处 assert 形状
• 与 F.scaled_dot_product_attention 最大误差 2.65 × 10⁻⁷：float32 约 1e-6、float64 约 1e-12 才叫「对了」

03 · PyTorch 使用层（上）：五个对象与二十行

结论：使用层只有五个对象；Autograd 做三件事（记图、累加、no_grad）；二十行训练循环里每一行对应一个概念，Trainer 是它加外壳。

%%{init: {"flowchart": {"wrappingWidth": 120}}}%%
flowchart LR
    DS["<b>Dataset</b><br/>__getitem__ 取一条"] --> DL["<b>DataLoader</b><br/>组 batch、打乱、预取"]
    DL -- "batch" --> M["<b>nn.Module</b><br/>forward：参数 × 输入"]
    M -- "logits" --> L["<b>loss</b><br/>cross_entropy(logits, labels)"]
    L -- "backward()" --> AG["<b>Autograd</b><br/>沿图反向，梯度累加到 .grad"]
    AG -- ".grad" --> OPT["<b>Optimizer</b><br/>step() 更新、zero_grad() 清零"]
    OPT -- "新参数" --> M
    classDef data fill:#eefaf0,stroke:#4d9a5c,color:#222
    classDef model fill:#fff7e0,stroke:#c98a00,stroke-width:2px,color:#222
    classDef grad fill:#eef6ff,stroke:#5b8fd6,color:#222
    class DS,DL data
    class M,L model
    class AG,OPT grad

二十行里最容易踩的三行

for step, (x, y) in enumerate(loader):
    logits = model(x)                                   # [B, T, V]
    loss = F.cross_entropy(logits.view(-1, V).float(),  # loss 用 fp32 算
                           y.view(-1), ignore_index=-100)
    loss.backward()                                     # 梯度是【累加】的
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    opt.step(); sched.step()
    opt.zero_grad(set_to_none=True)                     # 所以要清零

• 不 zero_grad 不是 bug 就是梯度累积（小 batch 模拟大 batch）
• 评测不加 no_grad：中间量等一个不会来的 backward，显存涨到 OOM
• 首步 loss 应 ≈ ln V：84 万参数、V = 128 → 5.07 ✓；一分钟 PPL 128 → 9.6

学习率调度：warmup + cosine

• loss.item() 为什么慢：CPU 异步发 kernel，.item() 要等 GPU 队列清空——别每步都取
• CPU 上开 bf16 autocast 慢 30 倍：精度的收益来自硬件

04 · PyTorch 使用层（下）：显存的账

结论：每个可训练参数 16 字节；激活是账外的一块、与参数量无关；「全量还是 LoRA」的第一道约束是有几张卡。

• 激活另算：B = 1、T = 4096 约 16.5 GiB；gradient checkpointing 多 30% 计算换 16.5 → 1 GiB
• FSDP 8 卡切状态：每卡 16 GB 放得下全量

混合精度下一步训练的数据流

flowchart LR
    W32["fp32 主权重 W<br/>4 字节 / 参数，常驻"] -- "① autocast 转 bf16" --> W16["bf16 权重副本<br/>2 字节"]
    W16 -- "② 矩阵乘（Tensor Core）" --> ACT["bf16 激活"]
    ACT -- "③ reduction 升回 fp32" --> LOSS["fp32 loss"]
    LOSS -- "④ backward" --> G["梯度"]
    G -- "⑤" --> OPT["AdamW 状态 m、v<br/>4 + 4 字节，fp32 常驻"]
    OPT -- "⑥ W ← W − lr·m̂/(√v̂+ε)，在 fp32 上加" --> W32

• autocast 不把参数变成 bf16：减半的是矩阵乘的中间结果，训练状态仍是 16 字节 / 参数
• OOM 按爆的时机归因：加载后 → 权重；第一步 backward 后 → 梯度 + 状态；加了 LoRA 还 OOM → 看激活

05 · Hugging Face 生态：六个库与六行 LoRA SFT

结论：六个库各管一段，六行组装；背后的每件事都在二十行里有位置；从 compute_loss 往下追源码是学后训练最快的路。

%%{init: {"flowchart": {"wrappingWidth": 200}}}%%
flowchart LR
    HUB["`**Hub**
config.json · safetensors · tokenizer.json`"]
    DS["`**datasets**
load · map · streaming`"]
    TOK["`**tokenizers**
BPE 编码（Rust）`"]
    TF["`**transformers**
AutoModel · generate · Trainer`"]
    PEFT["`**peft**
LoraConfig · get_peft_model`"]
    TRL["`**trl**
SFTTrainer · DPOTrainer · GRPOTrainer`"]
    ACC["`**accelerate**
launch · DDP / FSDP`"]
    HUB --> DS & TF
    DS --> TOK --> TRL
    TF --> PEFT --> TRL
    TRL --> ACC
    classDef lib fill:#eef4fb,stroke:#5b8dc9,color:#222
    classDef hub fill:#fff7e0,stroke:#c98a00,stroke-width:2px,color:#222
    class DS,TOK,TF,PEFT,TRL,ACC lib
    class HUB hub

Hub 上的三个文件

• config.json 里的数字就是结构：Qwen2.5-0.5B 24 层、hidden 896、词表 151936、tie_word_embeddings
• tokenizer.json + chat template：对话怎么变成 token、哪些 token 算 loss

六行 LoRA SFT，与它背后的六件事

ds    = load_dataset("...", split="train")                  # datasets
tok   = AutoTokenizer.from_pretrained(name)                 # chat template → __getitem__
model = AutoModelForCausalLM.from_pretrained(name, dtype=torch.bfloat16)
model = get_peft_model(model, LoraConfig(r=16, lora_alpha=32, ...))  # ΔW = (α/r)·BA
trainer = SFTTrainer(model, train_dataset=ds, args=SFTConfig(...))   # loss mask = ignore_index=-100
trainer.train()                                             # 二十行加外壳

• Qwen2.5-0.5B：494M 参数、可训练 8.80M（1.78%）、训练状态 141 MB——CPU 都能微调
• 一条对话 85% 的 token 被 mask；20 步 loss 5.3 → 1.7，答案对了却不会停——结束符要进 loss

06 · GPU 直觉与实验管理

结论：两个上限（算力、带宽）之比是 ridge；decode 强度 ≈ 1 远低于 295，是 memory-bound，所以 batch 大才快；显存四块；七项记录齐了才谈复现。

用两个数字解释快慢

• 量化对推理有效、对训练用处不大，根本原因就是 decode 在 ridge 左边、训练在右边

三个月后能复现：七项记录

• 代码版本（commit）· 数据版本（revision）· 配置（asdict(cfg) → config.json）· 环境（pip freeze）· seed · 硬件 · 日志与 checkpoint
• 同 seed 两次完全一致要 use_deterministic_algorithms；换个 seed 差 0.14——20 步训练里 seed 的影响比很多「方法改进」都大
• profiler 表怎么读：慢先归到 GEMM、attention、copy_、launch 开销或 GPU 空转之一