跳到主要内容

数据截至 (上游 commit fd01e35c83d8)

05 · 大模型 offload 与 device_map

这一章讲什么： 一张 24GB 的卡怎么跑 70B 模型？答案不是并行，而是「时间换空间」：参数平时躺在 CPU 内存甚至磁盘上，forward 执行到某一层时才把该层权重搬上 GPU，算完立刻卸掉。这套机制就是 transformers 里 device_map="auto" 的底层。

1. 它要解决的小问题​

模型权重大于单卡显存时，朴素方案全部失效：from_pretrained 先在 CPU 建完整模型再 .to(cuda)，CPU 内存先爆一次；就算加载成功，单卡也放不下。你需要：加载时永不实体化完整模型；运行时只有正在算的那一层占显存。

2. 直觉​

两个关键解耦：

  • 形状与数据解耦：PyTorch 的 meta 设备可以创建「有形状、有 dtype、不占存储」的张量。先在 meta 上把模型结构完整建出来（零内存），形状信息已足够算出「每层多大」。
  • 驻留与执行解耦：给每个 block 的 forward 包一层钩子——pre_forward 把本层权重从 CPU/磁盘物化到执行设备，post_forward 立刻把权重扔回 meta。显存占用峰值 ≈ 常驻层 + 最大单层。

device_map 就是一张「模块名 → 设备」的字典，值可以是 0、1、"cpu"、"disk"。整个推理过程变成：按图施工，逐层搬运。

3. 图示​

加载期与运行期两条流水线：

加载期:
init_empty_weights      infer_auto_device_map      load_checkpoint_in_model     dispatch_model
 nn.Linear 等构造时       按 max_memory 装箱:         逐 shard 读盘, 按 map       给每个 block 挂
 参数落在 meta 设备       每层分配到 GPU/CPU/disk     直接落到目标设备            AlignDevicesHook
 (零内存)                 (GPU 预留最大层空间)        (CPU/disk 层保持 meta)      (执行时搬运)

运行期 (某 offloaded block 的一次 forward):
pre_forward: weights_map[name] ──▶ set_module_tensor_to_device(execution_device) ──▶ 原 forward ──▶ post_forward: 权重重新放回 meta
                              （输入张量同时被 send_to_device 到执行设备）

4. 原理演示（示意代码）​

整套机制可以压缩成三十行：

@contextmanager
def init_empty_weights():
    old = nn.Module.register_parameter
    nn.Module.register_parameter = lambda m, n, p: old(m, n, p.to("meta") if p is not None else p)
    try:
        yield
    finally:
        nn.Module.register_parameter = old

class AlignDevicesHook:
    def pre_forward(self, module, *args, **kwargs):
        for name, _ in module.named_parameters():
            set_module_tensor_to_device(module, name, self.execution_device,
                                        value=self.weights_map[name])   # 物化
        return send_to_device(args, self.execution_device), send_to_device(kwargs, self.execution_device)

    def post_forward(self, module, output):
        for name, _ in module.named_parameters():
            set_module_tensor_to_device(module, name, "meta")           # 卸回 meta
        return output

# 挂钩子 = 换掉模块的 forward
def add_hook_to_module(module, hook):
    module._old_forward, module._hf_hook = module.forward, hook
    def new_forward(module, *args, **kwargs):
        args, kwargs = hook.pre_forward(module, *args, **kwargs)
        return hook.post_forward(module, module._old_forward(*args, **kwargs))
    module.forward = functools.partial(new_forward, module)

5. 真实实现​

5.1 空壳初始化：init_empty_weights / init_on_device​

init_empty_weights（src/accelerate/big_modeling.py:62-96）就是 init_on_device(torch.device("meta")) 的别名。后者（src/accelerate/big_modeling.py:98-177）在 context 内做两类猴子补丁：

  • 替换 nn.Module.register_parameter（src/accelerate/big_modeling.py:131-141）：先走原注册，再把参数 .to(device)；保留 requires_grad 与 HF 的 _is_hf_initialized 标记；
  • include_buffers=True 时同时替换 register_buffer，并给 torch.empty/zeros/ones/full 打补丁强制 device（src/accelerate/big_modeling.py:157-163）。

context 退出时在 finally 里全部还原（src/accelerate/big_modeling.py:170-177）。补丁是全局的，所以嵌套第三方建模代码（如 from_config）同样生效。

5.2 装箱：get_balanced_memory 与 infer_auto_device_map​

get_balanced_memory（src/accelerate/utils/modeling.py:931）先算每张卡的预算：单卡时打 9 折留缓冲（src/accelerate/utils/modeling.py:1001-1008），多卡时按卡间均衡目标压低预算，low_zero=True（balanced_low_0）额外压低 GPU 0 给 generate 留地方。

infer_auto_device_map（src/accelerate/utils/modeling.py:1304）是主循环装箱：

  • 初始化时算出各模块尺寸、不可拆分模块（no_split_module_classes，如带残差连接的 block）、tied 参数组；
  • 每处理一个模块前，主设备的可用预算要减去「剩余层中最大那层」的尺寸（src/accelerate/utils/modeling.py:1419-1423）——这个预留保证任何 offload 到 CPU 的层都能随时搬回 GPU 执行，是整套 offload 的可行性不变量；
  • 模块连同它的 tied 参数一起试放当前设备（module_size_with_ties <= current_max_size，src/accelerate/utils/modeling.py:1429），放不下就换下一台设备，顺序是 GPU 0 → GPU 1 → … → CPU → disk；
  • 全部放完后 clean_result 把同设备的连续模块合并成前缀，得到紧凑的 device_map。

5.3 落盘与派发：load_checkpoint_and_dispatch → dispatch_model​

load_checkpoint_and_dispatch（src/accelerate/big_modeling.py:520-660）串起全流程：字符串型 device_map（auto/balanced/balanced_low_0/sequential）先换算成真实 map（src/accelerate/big_modeling.py:618-633），再 load_checkpoint_in_model 逐 shard 加载——加载时用 set_module_tensor_to_device（src/accelerate/utils/modeling.py:227）把权重直接写进目标设备，offload 的层保持 meta，最后调 dispatch_model。

dispatch_model（src/accelerate/big_modeling.py:315-518）：

  1. check_device_map 校验 map 覆盖完整（src/accelerate/utils/modeling.py:1612）；
  2. 只有一台设备且无 disk 时不挂钩子，直接 model.to(device)（src/accelerate/big_modeling.py:501-517）；
  3. 否则算出每个模块的执行设备（cpu/disk 上的模块执行设备 = main_device，src/accelerate/big_modeling.py:400-403），构造 weights_map（磁盘 offload 时是 OffloadedWeightsLoader 懒加载映射，src/accelerate/big_modeling.py:407-413）；
  4. 登记 tied 参数的 data_ptr（src/accelerate/big_modeling.py:418-432），供 hook 复用已物化的共享权重；
  5. attach_align_device_hook_on_blocks 递归挂钩（src/accelerate/big_modeling.py:434-442）；
  6. 给 model.to / model.cuda 等加警告包装——被 dispatch 的模型不该再整体搬设备，且还有 meta 参数时直接 raise（src/accelerate/big_modeling.py:458-476）。

5.4 钩子本体：AlignDevicesHook​

add_hook_to_module（src/accelerate/hooks.py:147-203）实现「换 forward」：原 forward 存 _old_forward，新 forward 按 pre_forward → _old_forward → post_forward 串联（src/accelerate/hooks.py:186-194），GraphModule 走类级补丁特例（src/accelerate/hooks.py:196-201）。

AlignDevicesHook（src/accelerate/hooks.py:242）的三个生命期：

  • init_hook（src/accelerate/hooks.py:298-343）：非 offload 模块直接把参数搬到执行设备；offload 模块先把真权重收进 weights_map（或挂懒加载），再把模块参数全部设为 meta——注意 meta 设备上 data_ptr() 恒为 0，tied 记账此时禁用（src/accelerate/hooks.py:300-301）；
  • pre_forward（src/accelerate/hooks.py:359-399）：offload 模块从 weights_map[name] 取值物化到 execution_device；tied 权重若已在本设备物化过则直接复用指针（src/accelerate/hooks.py:377-386）；int8 量化权重顺带捎上 SCB 统计量（_maybe_get_fp16_statistics，src/accelerate/hooks.py:345-356）；最后把输入 send_to_device；
  • post_forward（src/accelerate/hooks.py:402-430）：参数立刻放回 meta，清理本轮登记的 tied 指针让 GC 能回收；io_same_device=True 时把输出送回输入原设备——于是跨设备边界对用户完全透明。

挂接策略由 attach_align_device_hook_on_blocks（src/accelerate/hooks.py:586-718）递归决定：offload 的 block 挂「物化-卸载」型 hook 并给子模块挂执行设备 hook；GPU 常驻的 block 只挂执行对齐 hook；根模块挂 io_same_device 的总入口 hook。每个设备边界恰好一对 pre/post，不会重复搬运。

6. 坑​

  1. 这是推理机制，不是训练机制。offload 的层参数在 meta 上，梯度、optimizer state 无从谈起；prepare 对 device_map 模型直接 raise（第 1 章）。单进程下做 naive pipeline 微调也仅限「全部层都在同一台设备」的退化形态。
  2. 速度按搬运量线性劣化。每次 forward 都有 CPU→GPU 的全量权重拷贝（disk offload 还多一次读盘），没有任何预取/重叠——它换的是「能跑」，不是「跑得快」。
  3. init_empty_weights 出来的模型不能直接 .to()。参数在 meta 上没有数据，to 搬不动欠条；必须走 load_checkpoint_and_dispatch，且自定义 device_map 时 dispatch_model 才生效。
  4. tied weights 的特例链。共享权重靠 data_ptr 记账复用；但 disk offload 时 safetensors 不保证每次 get_tensor 返回同指针，只能退化到运行时按名追踪（src/accelerate/hooks.py:320-331 注释），这也是为什么 tied 参数多的模型 offload 开销更高。
  5. forward 被改写的外部性。_hf_hook/_old_forward 挂在模块上，torch.compile 需要用 @_compiler_disable 包住 hook 的 pre/post（src/accelerate/hooks.py:358、src/accelerate/hooks.py:401）避免 graph break 扩大化；自己再包一层 forward 时要想清楚顺序。
  6. disk offload 必须有 offload_dir。缺目录且 map 含 "disk" 时 dispatch_model 直接 raise（src/accelerate/big_modeling.py:388-392）；磁盘格式走 safetensors + index.json，与 transformers 的 shard 格式同源。