模型加载优化：从FP32到INT8的显存博弈
模型加载这块，说白了就是怎么把训练好的参数塞进显存里。我见过太多团队，模型写得很漂亮，结果一加载就OOM，连训练都跑不起来。今天咱们聊聊三种主流存储格式，以及怎么聪明地加载模型。
1. 模型参数存储格式：FP32 / FP16 / INT8
先看个直观对比。我习惯把模型参数想象成「数字的精度和体积」——精度越高，占地方越大。
| 格式 | 位宽 | 显存占用（相对） | 适用场景 |
|---|---|---|---|
| FP32 | 32位 | 1x | 训练初期、高精度需求 |
| FP16 | 16位 | 0.5x | 混合精度训练、推理 |
| INT8 | 8位 | 0.25x | 推理部署、量化模型 |
FP32是「老大哥」，精度最高，但显存消耗也最狠。一个70B的模型，光参数就要140GB显存——这还没算梯度和优化器状态。FP16能省一半，但训练时容易溢出。INT8更极端，省是省，但精度损失你得掂量掂量。
我的经验：训练阶段我一般用FP32或混合精度（FP16+FP32）。推理阶段才考虑INT8。别为了省显存盲目降精度，模型训崩了更亏。
2. torch.load内存映射：懒加载的艺术
torch.load默认会把整个模型加载到内存里。你想想看，一个70B的模型，140GB参数全塞进内存，再往显存里搬——内存先爆了。
这时候就需要mmap（内存映射）。它不一次性加载全部数据，而是按需读取。就像你看一本厚书，不需要把整本书背下来，看到哪页翻哪页。
# 传统加载方式（容易OOM）
model = torch.load('model_fp32.pt')
# 内存映射加载（推荐）
model = torch.load('model_fp32.pt', map_location='cpu', mmap=True)
嗯，这里要注意：mmap=True只在PyTorch 2.0+版本支持。我遇到过有人用老版本，加了参数直接报错——升级PyTorch吧。
小技巧：配合map_location='cpu'使用，先把模型映射到CPU内存，再按需搬到显存。这样显存压力小很多。
3. 模型分片加载策略：大模型的「化整为零」
当模型大到单卡装不下时，分片加载就是救命稻草。说白了，就是把模型切成几块，每块单独加载到不同GPU上。
我个人习惯用accelerate库来做分片。它封装了底层细节，你只需要指定分片策略就行。
from accelerate import init_empty_weights, load_checkpoint_and_dispatch
# 初始化空模型
with init_empty_weights():
model = MyLargeModel()
# 分片加载到多GPU
model = load_checkpoint_and_dispatch(
model,
checkpoint='model_sharded/',
device_map='auto',
no_split_module_classes=['TransformerLayer']
)
这里device_map='auto'会自动分配各层到不同GPU。no_split_module_classes指定哪些模块不能拆分——比如TransformerLayer内部有残差连接，拆开反而麻烦。
避坑指南：我曾经把注意力层拆到两张卡上，结果通信开销比计算还大，训练速度反而慢了。分片不是越细越好，要平衡计算和通信。
4. 实战建议：三步走策略
我一般按这个顺序优化模型加载：
- 先看模型大小：计算参数总量，估算显存需求。比如7B模型用FP32需要28GB，单卡A100（80GB）能装下。
- 再选加载方式：能单卡装下就用
mmap加载；装不下就分片。 - 最后调精度：如果显存还有余量，保持FP32；紧张就降FP16或INT8。
举个例子，我最近训一个13B的LLaMA模型：
- FP32需要52GB显存，单卡A100装不下
- 改用FP16，26GB，单卡刚好
- 配合
mmap加载，内存占用从52GB降到2GB
你看，一个技巧就能省下50GB内存。这就是优化的魅力。
5. 知识体系总览
下面这张图总结了模型加载优化的核心逻辑。我画的时候特意把「精度-显存-速度」的三角关系标出来了——你优化时得权衡这三者。
这张图把三个核心模块串起来了。从左到右：先选存储格式，再定加载方式，最后决定分片策略。底部的三角关系时刻提醒你——没有免费的午餐，省显存往往要牺牲精度或速度。
总结一下：模型加载优化的本质是「用时间换空间」或「用精度换空间」。mmap用读取时间换内存空间，INT8用精度换显存空间。具体怎么选，看你手头有什么硬件，模型有多大。
好了，今天就聊到这儿。记住，别一上来就上FP32全量加载——先算算显存，再选策略。我当年就是吃了这个亏，一个70B模型加载了半小时，最后OOM了。从那以后，我每次加载模型前都先做这三步评估。
想要更多资料，联系可可