使用device_map="auto"时A100 GPU触发CUDA断言错误的原因咨询

核心原因分析 #

出现这个CUDA设备端断言错误，主要和以下几个因素相关：

• 特殊Token语法错误：代码中tokenizer.convert_tokens_to_ids("<|eot_id|")少了闭合的>，正确写法应为tokenizer.convert_tokens_to_ids("<|eot_id|>")。单GPU/CPU模式下模型可能忽略该无效Token的异常，但多GPU自动分片时，模型内部的设备端断言会严格检查Token ID有效性，触发错误。
• 张量设备不匹配：使用device_map="auto"时，模型会被分片到多张GPU上，但input_ids.to(model.device)仅将输入张量移到模型的「主设备」（通常是cuda:0），生成过程中部分层在其他GPU运行，跨设备张量交互时因设备不一致触发断言。
• Transformers版本兼容性：旧版本Transformers库对Llama3多GPU自动分片的支持存在bug，尤其是特殊Token的跨设备处理逻辑不完善，导致生成阶段出现设备端错误。

针对性解决方案 #

1. 修复特殊Token语法
   修正terminators定义代码：

   terminators = [
       tokenizer.eos_token_id,
       tokenizer.convert_tokens_to_ids("<|eot_id|>")  # 补上闭合的>
   ]
   

2. 确保输入张量与模型层设备对齐
   可以手动指定生成时的设备，让输入与模型主设备保持一致：

   outputs = model.generate(
       input_ids,
       max_new_tokens=256,
       eos_token_id=terminators,
       do_sample=True,
       temperature=0.6,
       top_p=0.9,
       device=torch.device("cuda:0")
   )
   

3. 升级Transformers库
   执行命令更新到兼容Llama3多GPU的最新版本：

   pip install --upgrade transformers accelerate
   

4. 更换模型分片策略
   若自动分片仍有问题，可使用balanced策略让模型更均匀地分布在多张GPU上：

   model = AutoModelForCausalLM.from_pretrained(
       model_id,
       torch_dtype=torch.bfloat16,
       device_map="balanced",
   )
   

内容的提问来源于stack exchange，提问作者Stephen SUN