使用TensorFlow训练GPT2模型时遇Checkpoint相关ValueError求助

我微调模型时用GPT2Tokenizer和GPT2LMHeadModel替换了教程中的BERT，训练阶段触发以下错误：

ValueError: `Checkpoint` was expecting model to be a trackable object (an object derived from `Trackable`), got GPT2LMHeadModel(
  (transformer): GPT2Model(
    (wte): Embedding(50257, 768)
    (wpe): Embedding(1024, 768)
    (drop): Dropout(p=0.1, inplace=False)
    (h): ModuleList(
      (0-11): 12 x GPT2Block(
        (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)
        (attn): GPT2Attention(
          (c_attn): Conv1D()
          (c_proj): Conv1D()
          (attn_dropout): Dropout(p=0.1, inplace=False)
          (resid_dropout): Dropout(p=0.1, inplace=False)
        )
        (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)
        (mlp): GPT2MLP(
          (c_fc): Conv1D()
          (c_proj): Conv1D()
          (act): NewGELUActivation()
          (dropout): Dropout(p=0.1, inplace=False)
        )
      )
    )
    (ln_f): LayerNorm((768,), eps=1e-05, elementwise_affine=True)
  )
  (lm_head): Linear(in_features=768, out_features=50257, bias=False)
). If you believe this object should be trackable (i.e. it is part of the TensorFlow Python API and manages state), please open an issue.

原训练代码：

training_args = TFTrainingArguments(
    output_dir='./results',          
    num_train_epochs=2,          
    per_device_train_batch_size=8, 
    per_device_eval_batch_size=16
    warmup_steps=500,               
    weight_decay=0.01,              
    logging_dir='./logs',           
    logging_steps=10,
)


with training_args.strategy.scope():
#     model = AutoModelForCausalLM.from_pretrained("gpt2")
    model = GPT2LMHeadModel.from_pretrained("gpt2")

trainer = TFTrainer(
    model=model,                       
    args=training_args,                 
    train_dataset=train_dataset,       
    eval_dataset=test_dataset       
)

trainer.train()

错误原因 #

你加载的是PyTorch版本的GPT2LMHeadModel，但搭配了TensorFlow专属的TFTrainingArguments和TFTrainer，两者不兼容。PyTorch模型不属于TensorFlow的Trackable对象体系，因此checkpoint模块无法识别。

解决方案 #

1. 加载TensorFlow版本的GPT2模型 #

将模型加载代码替换为TensorFlow兼容的版本，使用TFGPT2LMHeadModel或TFAutoModelForCausalLM。

2. 修正后的完整代码 #

注意原代码中per_device_eval_batch_size=16后缺少逗号，属于语法错误，需一并修正：

from transformers import TFTrainingArguments, TFTrainer, TFGPT2LMHeadModel

training_args = TFTrainingArguments(
    output_dir='./results',          
    num_train_epochs=2,          
    per_device_train_batch_size=8, 
    per_device_eval_batch_size=16,  # 补上缺失的逗号
    warmup_steps=500,               
    weight_decay=0.01,              
    logging_dir='./logs',           
    logging_steps=10,
)


with training_args.strategy.scope():
    model = TFGPT2LMHeadModel.from_pretrained("gpt2")  # 改用TensorFlow版本模型

trainer = TFTrainer(
    model=model,                       
    args=training_args,                 
    train_dataset=train_dataset,       
    eval_dataset=test_dataset       
)

trainer.train()

额外注意事项 #

• 确保你的transformers库是稳定版，版本过低可能存在TF/PyTorch模型兼容问题
• 训练数据集需适配TensorFlow格式，若之前用的是PyTorch Dataset，需转换为TF Dataset格式

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