如何解决Transformers中google/gemma-2-2b的Index put dtype不匹配错误？

这个错误的核心是评估阶段张量 dtype 不匹配：模型计算用了BFloat16，但目标张量是Float32。Gemma系列对 dtype 一致性的要求比GPT-2、Llama-3更严格，以下是针对性的解决和调试步骤：

1. 强制统一模型与训练配置的 dtype #

初始化模型时明确指定 dtype，避免自动推断带来的不一致：

from transformers import AutoModelForCausalLM
import torch

model = AutoModelForCausalLM.from_pretrained(
    "google/gemma-2-2b",
    torch_dtype=torch.bfloat16,  # 硬件支持的话优先用bf16，否则换torch.float16
    device_map="auto"
)

同时在TrainingArguments中对齐混合精度设置，不要同时开启fp16和bf16：

training_args = TrainingArguments(
    # 其他参数...
    bf16=True,  # 和模型dtype一致
    fp16=False,
    torch_compile=False,  # 暂时关闭编译，避免隐式dtype转换
)

2. 检查评估数据集的预处理 dtype #

确保评估数据的张量 dtype 和模型完全一致，避免预处理时被强制转换：
在数据预处理函数中显式设置张量类型：

def preprocess_function(examples):
    outputs = tokenizer(
        examples["text"],
        truncation=True,
        padding="max_length",
        max_length=512
    )
    # 强制对齐模型dtype
    for key in outputs:
        outputs[key] = torch.tensor(outputs[key], dtype=torch.bfloat16)
    return outputs

如果使用DataCollatorForLanguageModeling，确保它不会修改数据类型：

from transformers import DataCollatorForLanguageModeling

data_collator = DataCollatorForLanguageModeling(
    tokenizer=tokenizer,
    mlm=False,
    return_tensors="pt"  # 保持PyTorch张量类型，避免自动转换
)

3. 修正评估阶段的 dtype 对齐逻辑 #

自定义评估函数，强制将logits和标签的 dtype 统一：

def compute_metrics(eval_pred):
    logits, labels = eval_pred
    # 强制转换为模型使用的dtype
    logits = logits.to(torch.bfloat16)
    labels = labels.to(torch.bfloat16)
    # 后续评估指标计算...
    return {"perplexity": perplexity}

另外，评估前手动将模型切换到评估模式并统一 dtype：

# 在Trainer初始化后添加
model.eval()
model = model.to(torch.bfloat16)

4. 排查硬件与环境问题 #

• 禁止跨设备评估：确保训练和评估都在同一设备（比如GPU）运行，避免CPU/GPU之间的 dtype 转换异常：

  training_args = TrainingArguments(
      # 其他参数...
      eval_on_cpu=False,
  )
  

• 升级依赖库：Gemma-2是较新模型，旧版本Transformers可能存在兼容性问题，执行以下命令升级：

  pip install --upgrade transformers accelerate torch
  

5. 定位问题根源的调试技巧 #

在关键步骤打印张量 dtype，找到不匹配的具体环节：

# 预处理后打印数据dtype
processed_data = preprocess_function(eval_dataset[0])
print("Input IDs dtype:", processed_data["input_ids"].dtype)

# 模型前向传播后打印logits dtype
with torch.no_grad():
    inputs = {k: v.unsqueeze(0).to("cuda") for k, v in processed_data.items()}
    outputs = model(**inputs)
    print("Logits dtype:", outputs.logits.dtype)

内容的提问来源于stack exchange，提问作者Charlie Parker