使用google/gemma-2b-it模型时HuggingFace Trainer报错求助

针对google/gemma-2b-it模型在HuggingFace Trainer中出现的Tried to track the number of tokens seen...报错，有两种直接的解决方式：

方式一：给模型添加main_input_name属性 #

Gemma模型的核心输入是input_ids，加载模型后直接设置该属性即可：

from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained("google/gemma-2b-it")
tokenizer = AutoTokenizer.from_pretrained("google/gemma-2b-it")

# 添加main_input_name属性
model.main_input_name = "input_ids"

如果是使用PEFT做微调，需要给底层的基础模型设置该属性：

from peft import PeftModel

peft_model = PeftModel.from_pretrained(base_model, "your_peft_adapter_path")
peft_model.base_model.main_input_name = "input_ids"

方式二：关闭token计数功能 #

在初始化TrainerArguments时，禁用token总数跟踪，这样就不会触发输入检查：

from transformers import TrainerArguments

training_args = TrainerArguments(
    output_dir="./output",
    # 其他训练参数...
    track_total_tokens=False  # 关闭token计数
)

内容的提问来源于stack exchange，提问作者Rylan Schaeffer