【Bug已解决】GPT2 cannot be used with device_mapauto; Report found at least two devices 解决方案一、现象长什么样想用accelerate的device_mapauto把 GPT2 自动切到多张卡或 CPUGPU 混合做大模型推理from transformers import GPT2LMHeadModel, AutoTokenizer model GPT2LMHeadModel.from_pretrained( gpt2, device_mapauto, torch_dtypeauto, ) tok AutoTokenizer.from_pretrained(gpt2) out model.generate(tok(Hello, return_tensorspt).input_ids)报错ValueError: Youre trying to load a model that was saved with a different device mapping. ... Found at least two devices: cuda:0 and cpu (or cuda:1)或者更常见的变体RuntimeError: GPT2LMHeadModel: weight wte.weight is on cuda:0 but lm_head.weight expects cpu — found at least two devices最迷惑的是用device_mapauto本意是「让 accelerate 帮我自动放」结果它反而因为「两个模块共享同一份权重却在两个设备」而拒绝。单卡model.cuda()完全正常一上device_map就炸。二、背景GPT2 的wte词嵌入和lm_head语言模型头是**权重共享tied**的lm_head.weight直接复用wte.weight不单独存参数。accelerate的device_mapauto做自动切分时会逐个子模块决定放哪张卡。问题在于它默认把wte和lm_head当成两个独立模块分别规划设备。如果自动规划把wte放cuda:0、把lm_head放cpu或cuda:1就出现「同一份逻辑权重被要求同时在两个设备」的矛盾——生成时lm_head要拿wte的权重但权重在另一张卡于是报「found at least two devices」。正确行为应该是accelerate 知道lm_head和wte是 tied把它们强制放在同一设备。但 GPT2 早期的实现里lm_head没有被显式标注为 tied或_tied_weights_keys没登记accelerate 无从得知就各放各的触发矛盾。三、根因根因一句话GPT2 的lm_head与wte是共享权重但device_mapauto的自动切分把它们当成独立模块分到不同设备导致「同一权重需跨设备」的矛盾触发 found-at-least-two-devices 错误。三点展开tied 未登记lm_head没在_tied_weights_keys/ tied 标记里登记accelerate 不知道它和wte是同一份。自动切分各放各的device_mapauto把wte和lm_head分别规划到cuda:0/cpu或不同卡产生设备矛盾。缺乏 co-locate 兜底当检测到 tied 权重被分到不同设备时没有「强制合并到同一设备」的兜底直接抛出ValueError/RuntimeError。不是卡不够是「权重共享关系没告诉切分器」。四、最小可运行复现不依赖真实大模型模拟「tied 权重被分到两个设备」from dataclasses import dataclass from typing import List dataclass class FakeModule: name: str device: str class FakeAutoMap: def __init__(self, modules): self.modules modules def validate_tied(self, tied_pairs): # 检查 tied 的两个模块是否同设备 by_name {m.name: m.device for m in self.modules} errors [] for a, b in tied_pairs: if by_name[a] ! by_name[b]: errors.append((a, by_name[a], b, by_name[b])) return errors # 模拟 device_mapauto 把 wte 和 lm_head 分到不同设备 modules [FakeModule(wte, cuda:0), FakeModule(lm_head, cpu)] automap FakeAutoMap(modules) errs automap.validate_tied([(wte, lm_head)]) print(tied 冲突:, errs) # [(wte,cuda:0,lm_head,cpu)] - 报错跑出来wte在cuda:0、lm_head在cputied 冲突非空 → 触发 found-at-least-two-devices。这就是精确复现。五、解决方案第一层最小直接修复最小修复让device_mapauto知道lm_head与wte是 tied强制同设备。两条路路线 A推荐确保模型声明 tied让 accelerate 自动 co-locate。from transformers import GPT2LMHeadModel # 关键显式登记 tied 权重accelerate 会把它们放同一设备 model GPT2LMHeadModel.from_pretrained( gpt2, device_mapauto, torch_dtypeauto, # 若模型未自动登记可手动 ) # 手动确保 tied部分版本需要 if model.lm_head.weight is not model.wte.weight: model.lm_head.weight model.wte.weight # 复用而非独立参数 out model.generate(...)路线 B稳妥如果仍冲突用显式device_map把共享权重的模块放同一设备或干脆整模型放一张卡# 整模型一张卡彻底避免跨设备 model GPT2LMHeadModel.from_pretrained(gpt2, device_map{: 0}) # 或显式指定wte 和 lm_head 必须在同一设备 model GPT2LMHeadModel.from_pretrained( gpt2, device_map{wte: 0, lm_head: 0, h: 0, ln_f: 0}, )要点tied 权重必须物理复用同一 Parameter 对象lm_head.weight wte.weightaccelerate 才不会各放各的。若自动切分仍冲突用显式device_map字典把共享模块锁在同一设备。纯多卡且必须切分时确保所有引用同一权重的模块都在同一device_map条目下。这一步单独就让device_mapauto不再报 two-devices。六、解决方案第二层结构性改进第一层是「在加载处补一行 tied」。但 GPT2 及所有 tied 权重的模型OPT、BLOOM、GPT-Neo 等都面临同样问题。更稳的做法把「tied 权重如何登记、如何强制同设备」收敛成单一策略对象。from dataclasses import dataclass, field from typing import Dict, List, Tuple dataclass class DeviceMapTieResolver: 解决 tied 权重在 device_map 下被分到多设备的单一策略。 # tied 权重对共享同一物理参数的模块路径 tied_pairs: List[Tuple[str, str]] field(default_factorylist) def register(self, a: str, b: str): self.tied_pairs.append((a, b)) def co_located_map(self, auto_map: Dict[str, str]) - Dict[str, str]: 把每个 tied 对的成员强制放到同一设备取第一个出现的设备。 resolved dict(auto_map) for a, b in self.tied_pairs: dev_a resolved.get(a) dev_b resolved.get(b) if dev_a is not None and dev_b is not None and dev_a ! dev_b: # 强制 b 跟随 a 的设备 resolved[b] dev_a elif dev_a is None and dev_b is not None: resolved[a] dev_b elif dev_b is None and dev_a is not None: resolved[b] dev_a return resolved def validate(self, final_map: Dict[str, str]) - List[str]: errors [] for a, b in self.tied_pairs: if final_map.get(a) ! final_map.get(b): errors.append(ftied 冲突: {a}{final_map.get(a)} vs {b}{final_map.get(b)}) return errors # 用法针对 GPT2 登记 wte-lm_head resolver DeviceMapTieResolver() resolver.register(wte, lm_head) auto {wte: cuda:0, lm_head: cpu, h: cuda:0} final resolver.co_located_map(auto) print(修正后映射:, final) # lm_head 被拉回 cuda:0 print(冲突:, resolver.validate(final)) # []结构收益单一策略所有 tied 对集中登记device_map自动 co-locate。可校验validate在加载前抓出任何 tied 跨设备避免运行时报错。可复用OPT/BLOOM/GPT-Neo 等只需追加register即可。七、解决方案第三层断言 / CI 守护写 pytest 守三条(1) tied 对最终同设备(2) 冲突被validate抓出(3) 修正后无 two-devices。import pytest from your_lib import DeviceMapTieResolver pytest.fixture def resolver(): r DeviceMapTieResolver() r.register(wte, lm_head) return r def test_co_located(resolver): auto {wte: cuda:0, lm_head: cpu, h: cuda:0} final resolver.co_located_map(auto) assert final[wte] final[lm_head] cuda:0 def test_validate_catches_conflict(resolver): bad {wte: cuda:0, lm_head: cuda:1} errs resolver.validate(bad) assert len(errs) 1, 应抓出 tied 跨设备冲突 def test_no_conflict_after_fix(resolver): auto {wte: cuda:0, lm_head: cpu} final resolver.co_located_map(auto) assert resolver.validate(final) [], 修正后不应有冲突 def test_multiple_tied_pairs(): r DeviceMapTieResolver() r.register(wte, lm_head) r.register(shared1, shared2) auto {wte: cuda:0, lm_head: cpu, shared1: cuda:1, shared2: cpu} final r.co_located_map(auto) assert final[wte] final[lm_head] assert final[shared1] final[shared2]CI 常驻跑这四条后任何「tied 权重又被分到多设备」的回归都会立刻爆红。八、排查清单GPT2 / tied 模型上device_map报 two-devices 时按顺序查先确认报错含found at least two devices或wte.weight ... lm_head.weight—— 是的话定位 tied 跨设备。检查模型是否有 tied 权重model.lm_head.weight is model.wte.weight应为True。若为False手动model.lm_head.weight model.wte.weight复用同一对象。确认_tied_weights_keys含lm_head.weight或对应路径让 accelerate 识别 tied。若自动切分仍冲突用显式device_map字典把 tied 模块锁同一设备。多卡时确认所有「共享同一物理参数」的模块都在同一device_map条目。升级 transformers/accelerate 后重跑一次device_mapauto加载冒烟断言 tied 同设备。九、小结GPT2 上device_mapauto报 found-at-least-two-devices根子是lm_head与wte是 tied 共享权重却没被登记进 tied 关系accelerate 的自动切分把它们分到不同设备造成「同一权重需跨设备」矛盾。修复三层次第一层确保lm_head.weight物理复用wte.weight或用显式device_map锁同设备第二层用DeviceMapTieResolverdataclass 把 tied 对登记、自动 co-locate 并校验第三层用 pytest 守「tied 同设备」「冲突被抓」「修正后无 two-devices」。工程启示凡是带 tied 权重的模型走device_map切分都必须先把共享关系告诉切分器登记_tied_weights_keys/ 物理复用 Parameter否则自动切分必然把共享权重拆到多设备而拒绝加载。这是多卡/CPU-offload 推理最高频的坑之一。