调用GPT2的trainer.train()触发torch.embedding索引越界错误求助

问题根源 #

你手动设置的GPT2Config里vocab_size=10000，但使用的GPT2Tokenizer.from_pretrained('gpt2')的词汇表大小是50257。模型的embedding层权重只初始化了10000个词的向量，但tokenizer生成的部分token id会超过10000，导致torch.embedding时索引越界。

解决方案 #

有两种可行的修复方式，根据需求选择：

方式一：使用预训练GPT2模型的默认配置（推荐） #

直接加载预训练的GPT2模型，不用手动创建config，模型和tokenizer的词汇表大小会自动匹配：

from transformers import GPT2LMHeadModel, GPT2Tokenizer, TextDataset, DataCollatorForLanguageModeling, Trainer, TrainingArguments

# 直接加载预训练模型和对应的tokenizer
tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
model = GPT2LMHeadModel.from_pretrained('gpt2')

train_data = TextDataset(tokenizer=tokenizer, file_path='train.txt', block_size=256)

training_args = TrainingArguments(
    output_dir='./models',
    overwrite_output_dir=True,
    num_train_epochs=1,
    per_device_train_batch_size=4,
    save_steps=1000,
    save_total_limit=2,
    prediction_loss_only=True,
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_data,
    data_collator=DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False),
)
trainer.train()

方式二：自定义模型配置并匹配tokenizer词汇表大小 #

如果确实需要自定义模型参数（层数、隐藏层维度等），需确保vocab_size和tokenizer的实际词汇表大小一致：

from transformers import GPT2Config, GPT2LMHeadModel, GPT2Tokenizer, TextDataset, DataCollatorForLanguageModeling, Trainer, TrainingArguments

tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
# 获取tokenizer的实际词汇表大小
vocab_size = tokenizer.vocab_size

# 自定义config时使用该vocab_size
config = GPT2Config(
    vocab_size=vocab_size, 
    n_positions=256, 
    n_ctx=256, 
    n_embd=512, 
    n_layer=12, 
    n_head=8
)

model = GPT2LMHeadModel(config=config)

train_data = TextDataset(tokenizer=tokenizer, file_path='train.txt', block_size=256)

training_args = TrainingArguments(
    output_dir='./models',
    overwrite_output_dir=True,
    num_train_epochs=1,
    per_device_train_batch_size=4,
    save_steps=1000,
    save_total_limit=2,
    prediction_loss_only=True,
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_data,
    data_collator=DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False),
)
trainer.train()

注意：Trainer初始化时无需重复设置prediction_loss_only=True，TrainingArguments中已配置该参数。

内容的提问来源于stack exchange，提问作者Shubhankar Kumar