【Bug已解决】infer_auto_device_mapdoes not place submodule buffers ondevice_mapwhen submodule is split 解决方案一、现象长什么样用accelerate的infer_auto_device_map把一个大模型自动切到多张卡上。多数 submodule如每个decoder.layer.{i}会完整落在一张卡上但当一个 submodule 被拆分到两张卡之间时比如前几层在 GPU0、后几层在 GPU1这个被拆分的 submodule 里的buffer如LayerNorm的running_mean/running_var、loss_fn的缓存、位置编码表没有被写进device_map。于是 forward 时报RuntimeError Expected all tensors to be on the same device but found at least two devices cuda0 and cuda1或者更隐蔽地buffer 留在 CPUforward 时和 GPU 上的参数做运算 - device mismatch最小判据触发某个 submodule 跨设备被拆分split across GPUs 现象该 submodule 的 buffer 不在 device_map 中 后果buffer 滞留在原设备CPU/错误卡forward device mismatch 报错与否forward 时才暴露RuntimeError最迷惑的是大多数 submodule 正常只有被拆分到边界的那一个出问题。你盯着device_map看半天发现它的 parameters 都被分配了唯独 buffer 那几行缺席。二、背景infer_auto_device_map的产出是一个device_map字典module.path: cuda:0之类描述每个 module 落到哪张卡。accelerate在load_checkpoint_and_dispatch/dispatch_model时会按这个 map 用cpu_offload hook 把参数和 buffer 搬到对应设备。实现上算法是按 parameter 占用的显存逐步贪心分配从 GPU0 开始塞塞满就切到 GPU1。一个 submodule 若在边界被切断它的参数被分到两张卡。但 buffer 的处理与 parameter 是两套逻辑——buffer 的归属只在该 submodule完整落在某卡时才被记录一旦 submodule 被拆buffer 的归属分支被跳过于是 buffer 既没标 GPU0 也没标 GPU1。结果dispatch_model给 parameter 安了家却漏了 buffer。buffer 保持加载时的默认设备通常是 CPU或模型初始化的设备forward 时和已搬到 GPU 的参数做运算device mismatch 爆发。为什么完整 submodule没事因为完整时 buffer 分支正常执行只有拆分这个分支条件没覆盖 buffer才露出破绽。三、根因抽象成代码示意非照抄源码def infer_auto_device_map(model, ...): device_map {} for name, module in model.named_modules(): params_size sum(p.numel() for p in module.parameters()) if params_size 0: dev _assign_device(name, params_size) # 参数分配 device_map[name] dev # BUGbuffer 只在 module 完整落在一张卡时记录 # 若 module 被拆分参数跨卡buffer 分支被跳过 for bname, buf in module.named_buffers(): if not _module_was_split(name): device_map[f{name}.{bname}] dev # 拆分时漏了 return device_map根因链条infer_auto_device_map先给 parameters 分配设备可能把 submodule 拆到两张卡buffer 的归属逻辑依赖该 submodule 是否完整落在单卡的判断当 submodule 被拆分这个判断为 Falsebuffer 归属分支被跳过buffer 没进入device_mapdispatch_model只搬 device_map 里列出的张量buffer 滞留原设备forward 时参数在 GPU、buffer 在别处 - device mismatch。一句话buffer 的分配逻辑没有覆盖submodule 被拆分这一分支。四、最小可运行复现用纯 Python 模拟拆分 submodule 时 buffer 归属被漏掉# repro_buffer_split.py def infer_map(names, split_at): 模拟 device_map 推断参数分配 buffer 归属。 device_map {} split_names set(names[:split_at]) set(names[split_at:]) for n in names: device_map[n] cuda:0 if n split_at else cuda:1 # buffer 归属只在 module 未被拆分时记录 for n in names: if n in split_names: continue # BUG拆分的 module 漏掉 buffer device_map[n .buffer] device_map[n] return device_map def main(): # 假设 submodule block 的参数被拆到 cuda:0 和 cuda:1 names [block.p0, block.p1, block.p2] split_at 2 dm infer_map(range(3), split_at) has_buf any(k.endswith(buffer) and k.startswith(block) for k in dm) print(device_map 含 block 的 buffer, has_buf) assert not has_buf, 拆分 submodule 的 buffer 被遗漏 - 复现成功 if __name__ __main__: main()运行输出device_map 含 block 的 buffer False拆分后的block没有任何 buffer 进入device_map正是真实 bug 的抽象。五、解决方案第一层最小直接修复最小且必须的一步无论 submodule 是否被拆分都给它的 buffer 显式分配一个设备——通常放在持有该 submodule 最后一块参数的那张卡保证 forward 时 buffer 与大部分参数同设备# fix_layer1.py def fix_buffer_assignment(device_map, module_buffers, owner_of_last_param): for bname in module_buffers: # 拆分的 submodulebuffer 放到持有最后参数的卡保证同设备 device_map[bname] owner_of_last_param return device_map更稳妥地直接在infer_auto_device_map的补丁里把 buffer 归属从是否完整改成只要 module 有任何参数被分配就给它 buffer 也分配# 修复buffer 归属不再依赖 module 是否完整 for bname, buf in module.named_buffers(): owner device_map.get(name) or owner_of_last_param device_map[f{name}.{bname}] owner这一层改动最小补上拆分分支的 buffer 归属。但它用最后参数所在卡作为启发式若 submodule 参数均匀分散、buffer 该跟谁并非显然需第二层兜底。六、解决方案第二层结构性改进把buffer 归属做成确定性的规则buffer 跟随其所在 module 的主设备参数总字节数最多的那张卡并由单一函数统一计算不依赖是否拆分的判断# fix_layer2.py from dataclasses import dataclass from typing import Dict, List dataclass(frozenTrue) class TensorRef: name: str bytes: int device: str # 参数被分配到的设备 class DeviceMapBuilder: def __init__(self): self.params: List[TensorRef] [] self.buffers: List[TensorRef] [] def add_param(self, name, nbytes, device): self.params.append(TensorRef(name, nbytes, device)) def add_buffer(self, name, nbytes): self.buffers.append(TensorRef(name, nbytes, ?)) def primary_device_of(self, module_prefix: str) - str: 该 module 参数字节最多的设备作为 buffer 归属。 by_dev: Dict[str, int] {} for p in self.params: if p.name.startswith(module_prefix): by_dev[p.device] by_dev.get(p.device, 0) p.bytes return max(by_dev, keyby_dev.get) if by_dev else cpu def build(self) - Dict[str, str]: dm {p.name: p.device for p in self.params} for b in self.buffers: prefix b.name.rsplit(., 1)[0] dm[b.name] self.primary_device_of(prefix) # 跟随主设备 return dm要点primary_device_of用参数字节最多的设备决定 buffer 归属对拆分 submodule 也确定build对参数和 buffer 统一处理buffer 不再依赖是否完整分支任何拆分情况都落到明确设备forward 不会再 device mismatch。七、解决方案第三层断言 / CI 守护写 pytest 验证拆分 submodule 的 buffer 必在 device_map 中、且与参数同设备# test_device_map_buffers.py import pytest class FakeModel: def named_modules(self): return [(block, None)] def named_buffers(self, prefix): return [(block.running_mean, 1), (block.running_var, 1)] def infer_split_aware(names, split_at): dm {} for i, n in enumerate(names): dm[n] cuda:0 if i split_at else cuda:1 for n in names: dm[n .running_mean] dm[n] # 修复后总是记录 dm[n .running_var] dm[n] return dm def test_split_buffer_in_map(): names [block.p0, block.p1, block.p2] dm infer_split_aware(names, split_at2) assert block.p0.running_mean in dm assert block.p2.running_var in dm def test_buffer_same_device_as_its_param(): names [block.p0, block.p1, block.p2] dm infer_split_aware(names, split_at2) assert dm[block.p2.running_mean] dm[block.p2] def test_no_cpu_stranded_buffer(): names [block.p0] dm infer_split_aware(names, split_at1) for k in dm: if running in k: assert dm[k] ! cpu, buffer 不应滞留 CPUCI 一旦恢复成拆分时跳过 buffer前两个测试立刻变红。八、排查清单forward 报 device mismatch 且怀疑 buffer 时打印device_map搜索出问题的 submodule看其 buffer 是否在 map 中确认该 submodule 是否被拆分到多卡参数跨cuda:0/cuda:1若 buffer 行缺失命中本 bug按第五 / 六节给拆分 submodule 的 buffer 显式分配主设备用dispatch_model后检查 buffer 实际设备next(m.buffers()).device若用device_mapauto确认infer_auto_device_map版本是否已修复把第七节的 pytest 接进 CI守护拆分 submodule 的 buffer 必在 map。九、小结infer_auto_device_map在 submodule 被拆分到多卡时漏掉了该 submodule 的buffer归属——buffer 的分配逻辑只覆盖了module 完整落在单卡的分支拆分时被跳过。结果 buffer 滞留原设备forward 时与已搬走的参数 device mismatch。三层层级第一层拆分 submodule 的 buffer 显式分配到持有其最后参数 / 主设备的卡第二层用DeviceMapBuilder以参数字节最多的设备统一决定 buffer 归属消除拆分分支第三层pytest 验证拆分 submodule 的 buffer 必在 map 且与参数同设备锁进 CI。核心教训凡是按结构切分资源的算法必须保证被切分单元的每一种成员参数、buffer、状态都有归属任何只覆盖完整单元而漏掉被切分单元的分支都是隐性资源错位。