【Bug已解决】GPT2 cannot be used with device_map=‘auto‘； Report “found at least two devices“ 解决方案
【Bug已解决】GPT2 cannot be used with device_map='auto'; Report "found at least two devices" 解决方案
一、现象长什么样
想用accelerate的device_map="auto"把 GPT2 自动切到多张卡（或 CPU+GPU 混合）做大模型推理：
from transformers import GPT2LMHeadModel, AutoTokenizer model = GPT2LMHeadModel.from_pretrained( "gpt2", device_map="auto", torch_dtype="auto", ) tok = AutoTokenizer.from_pretrained("gpt2") out = model.generate(tok("Hello", return_tensors="pt").input_ids)报错：
ValueError: You're 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_map="auto"本意是「让 accelerate 帮我自动放」，结果它反而因为「两个模块共享同一份权重却在两个设备」而拒绝。单卡model.cuda()完全正常，一上device_map就炸。
二、背景
GPT2 的wte（词嵌入）和lm_head（语言模型头）是**权重共享（tied）**的：lm_head.weight直接复用wte.weight，不单独存参数。
accelerate的device_map="auto"做自动切分时，会逐个子模块决定放哪张卡。问题在于：它默认把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_map="auto"的自动切分把它们当成独立模块分到不同设备，导致「同一权重需跨设备」的矛盾，触发 found-at-least-two-devices 错误。
三点展开：
- tied 未登记：
lm_head没在_tied_weights_keys/ tied 标记里登记，accelerate 不知道它和wte是同一份。 - 自动切分各放各的：
device_map="auto"把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_map="auto" 把 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在cpu，tied 冲突非空 → 触发 found-at-least-two-devices。这就是精确复现。
五、解决方案（第一层：最小直接修复）
最小修复：让device_map="auto"知道lm_head与wte是 tied，强制同设备。两条路：
路线 A（推荐）：确保模型声明 tied，让 accelerate 自动 co-locate。
from transformers import GPT2LMHeadModel # 关键：显式登记 tied 权重，accelerate 会把它们放同一设备 model = GPT2LMHeadModel.from_pretrained( "gpt2", device_map="auto", torch_dtype="auto", # 若模型未自动登记，可手动： ) # 手动确保 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.weight），accelerate 才不会各放各的。 - 若自动切分仍冲突，用显式
device_map字典把共享模块锁在同一设备。 - 纯多卡且必须切分时，确保所有引用同一权重的模块都在同一
device_map条目下。
这一步单独就让device_map="auto"不再报 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_factory=list) 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(f"tied 冲突: {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_map="auto"加载冒烟，断言 tied 同设备。
九、小结
GPT2 上device_map="auto"报 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 推理最高频的坑之一。