📚 课程目录第4章模型加载与初始化

4. 模型加载与初始化

模型加载，说白了就是把训练好的权重从硬盘搬到内存里，再放到 GPU 上。这一步看似简单，但坑不少。我见过不少同学在加载时卡住，要么 OOM，要么设备不匹配，要么加载完发现参数对不上。今天咱们就把这块彻底捋清楚。

4.1 PyTorch 模型加载：state_dict 与 torch.load

PyTorch 保存模型有两种常见方式。我个人习惯只保存 state_dict，也就是参数的字典。为什么？因为这样更灵活，模型结构可以单独定义，权重只是“填充物”。

核心区别：torch.save(model.state_dict(), "model.pt") 只存参数；torch.save(model, "model.pt") 连结构一起存。后者容易出兼容性问题，我建议你别用。

import torch
import torch.nn as nn

# 定义一个简单的模型
class MyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = nn.Linear(768, 1024)

    def forward(self, x):
        return self.fc(x)

model = MyModel()

# 保存 state_dict
torch.save(model.state_dict(), "my_model.pt")

# 加载时先实例化，再 load
model_loaded = MyModel()
state_dict = torch.load("my_model.pt", map_location="cpu")
model_loaded.load_state_dict(state_dict)

这里有个关键点：map_location。我曾经在项目里忘了加这个参数，结果模型在 GPU 0 上保存，加载时却跑到了 GPU 1，导致后续张量设备不一致，报错找了好久。所以，加载时明确指定 map_location 是个好习惯。

避坑指南：load_state_dict 要求 key 完全匹配。如果你改了模型结构（比如层名变了），加载会报 missing keys 或 unexpected keys。我曾经因为改了个 fc 为 classifier，结果加载失败，排查了半天。解决办法是用 strict=False 加载，然后手动处理不匹配的 key。

4.2 Hugging Face Transformers：from_pretrained 的魔法

实际做大模型推理，我们很少手写 torch.load。更多是用 Hugging Face 的 from_pretrained。这个接口封装了下载、缓存、加载、设备映射等一系列操作。

from transformers import AutoModelForCausalLM, AutoTokenizer

model_name = "meta-llama/Llama-2-7b-hf"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.float16,      # 半精度加载，省显存
    device_map="auto",              # 自动分配设备
    low_cpu_mem_usage=True          # 减少 CPU 内存占用
)

你想想看，device_map="auto" 到底做了什么？它会根据你的 GPU 显存大小，自动把模型的各层分配到不同设备上。如果显存不够，它会把部分层放到 CPU。这就是模型并行的雏形。

我的经验：low_cpu_mem_usage=True 这个参数一定要加。不加的话，加载 7B 模型可能会吃掉 30GB 以上的 CPU 内存，因为 PyTorch 会先加载 float32 版本再转换。加了之后，直接加载 float16，CPU 内存占用能降一半。

4.3 模型并行基础：单卡放不下怎么办？

大模型动辄几十 GB，单张 GPU 根本放不下。这时候就需要模型并行。说白了，就是把模型切成几块，每块放在不同的 GPU 上。

我画了一张图，帮你理解模型并行的核心逻辑：

上图展示的是流水线并行（Pipeline Parallelism）。输入先经过 GPU 0，算完把中间结果传给 GPU 1，依次类推。每张卡只存一部分参数，这样就能跑更大的模型。

关键点：模型并行不是免费的。卡间通信有开销，尤其是当模型层数很多时，传输中间激活会占用大量带宽。所以实际部署时，我们会尽量让每张卡的计算量均衡，减少通信次数。

4.4 设备映射：CPU/GPU 的协同

设备映射，就是决定模型的哪部分放在哪个设备上。Hugging Face 的 device_map 参数帮我们做了这件事。但如果你想手动控制，可以用 PyTorch 的 to() 方法。

# 手动设备映射示例
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16)

# 把 embedding 层放到 GPU 0
model.model.embed_tokens = model.model.embed_tokens.to("cuda:0")

# 把前 16 层放到 GPU 0
for i in range(16):
    model.model.layers[i] = model.model.layers[i].to("cuda:0")

# 把后 16 层放到 GPU 1
for i in range(16, 32):
    model.model.layers[i] = model.model.layers[i].to("cuda:1")

# 把 lm_head 放到 GPU 1
model.lm_head = model.lm_head.to("cuda:1")

手动映射的好处是精细控制，但容易出错。我曾经手写映射时，漏掉了某个 layer norm，结果推理时设备不匹配，报错信息还特别隐晦。所以，除非你有特殊需求，否则我建议直接用 device_map="auto"。

注意：device_map="auto" 不是万能的。它基于显存大小做贪心分配，但不会考虑通信拓扑。比如，如果你有 4 张卡通过 NVLink 连接，而另外 4 张卡通过 PCIe 连接，自动分配可能把相邻层放到跨 PCIe 的卡上，导致通信变慢。这时候，手动映射反而更好。

4.5 加载后的初始化检查

模型加载完，别急着推理。先做几个检查，能省去后面很多调试时间。

检查项	方法	常见问题
参数设备	`next(model.parameters()).device`	参数还在 CPU 上，推理极慢
参数 dtype	`next(model.parameters()).dtype`	误加载为 float32，显存翻倍
参数数量	`sum(p.numel() for p in model.parameters())`	加载了错误的 checkpoint
前向测试	用 dummy input 跑一次 `model()`	shape 不匹配、NaN 输出

我的习惯：加载完模型后，我会先打印模型结构的前 3 层和后 3 层，确认层名和参数 shape 跟预期一致。这一步虽然简单，但能避免很多“模型加载了但推理结果不对”的坑。

4.6 总结：加载流程一览

最后，我把整个模型加载与初始化的流程总结成一张图，方便你对照操作：

嗯，到这里，模型加载与初始化的核心内容就讲完了。记住：加载不是简单的 torch.load，而是一个涉及设备、精度、并行策略的系统工程。每一步都踩过坑，才能做到心里有数。