Vertex AI中GPT-2微调进程停滞，请求技术排查

我按教程编写的GPT-2训练代码在Google Colab能正常运行，但迁移到Vertex AI Workbench后，执行trainer.train()仅弹出以下AdamW弃用警告就停滞了：

/opt/conda/lib/python3.7/site-packages/transformers/optimization.py:395: FutureWarning: This implementation of AdamW is deprecated and will be removed in a future version. Use the PyTorch implementation torch.optim.AdamW instead, or set `no_deprecation_warning=True` to disable this warning
  FutureWarning,

已确认GPU配额、结算账户、相关API均配置完成，nvidia-smi显示GPU正在被使用，但训练进程毫无进展。当前使用Tesla P100-PCIE-16GB GPU、GPT-2 small模型，数据集仅1000行，理论上训练应该很快完成。

分词器代码 #

def tokenize_function(examples):
        return base_tokenizer(examples['Prompt_Final'], padding=True)
    
# Split in train and test
df_train, df_val = train_test_split(df, train_size = 0.8)

# load the dataset from the dataframes
train_dataset = Dataset.from_pandas(df_train[['Prompt_Final']])
val_dataset = Dataset.from_pandas(df_val[['Prompt_Final']])

# tokenize the training and validation
tokenized_train_dataset = train_dataset.map(
    tokenize_function,
    batched=True,
    num_proc=1
)

tokenized_val_dataset = val_dataset.map(
    tokenize_function,
    batched=True,
    num_proc=1
)

模型初始化代码 #

bos = '<|startoftext|>'
eos = '<|endoftext|>'
body = '<|body|>'

special_tokens_dict = {'eos_token': eos, 'bos_token': bos, 'pad_token': '<pad>',
                       'sep_token': body} 

# the new tokens are added to the tokenizer
num_added_toks = base_tokenizer.add_special_tokens(special_tokens_dict)

# model configuration
config = AutoConfig.from_pretrained('gpt2', 
                                    bos_token_id=base_tokenizer.bos_token_id,
                                    eos_token_id=base_tokenizer.eos_token_id,
                                    pad_token_id=base_tokenizer.pad_token_id,
                                    sep_token_id=base_tokenizer.sep_token_id,
                                    output_hidden_states=False)

# we load the pre-trained model with custom settings
base_model = GPT2LMHeadModel.from_pretrained('gpt2', config=config)

# model embeding resizing
base_model.resize_token_embeddings(len(base_tokenizer))

# make sure its using the gpu
base_model = base_model.to(device)

训练参数与Trainer代码 #

model_articles_path = r'Model/Model_Path'

training_args = TrainingArguments(
    output_dir=model_articles_path,  # output directory
    num_train_epochs=1,              # total num of training epochs
    per_device_train_batch_size=10,   # batch size per device during training
    per_device_eval_batch_size=10,    # batch size for evaluation
    warmup_steps=200,                # number of warmup steps for learning rate scheduler
    weight_decay=0.01,               # strength of weight decay
    logging_dir=model_articles_path, # directory for storing logs
    prediction_loss_only=True,
    evaluation_strategy= "steps",
    save_steps=10,
    gradient_accumulation_steps=1,
    # gradient_checkpointing=True,
    eval_accumulation_steps=1,
    fp16=True
)

data_collator = DataCollatorForLanguageModeling(
        tokenizer=base_tokenizer,
        mlm=False
    )

trainer = Trainer(
    model=base_model,                      # the instantiated Transformers model to be trained
    args=training_args,                    # training arguments, defined above
    data_collator=data_collator,
    train_dataset=tokenized_train_dataset, # training dataset
    eval_dataset=tokenized_val_dataset     # validation dataset
)

排查建议 #

• 强制输出训练日志：在TrainingArguments中添加logging_steps=1，每步训练都会打印日志，确认进程是否真的在运行还是卡住。
• 关闭混合精度训练：把fp16=True改为fp16=False，排除Vertex AI环境对FP16支持的潜在问题。
• 验证数据集格式：打印tokenized_train_dataset[0]，确认生成的input_ids、attention_mask等字段是否符合要求，避免数据集格式错误导致训练无法推进。
• 调整训练参数：将per_device_train_batch_size调小（比如改为2），同时把evaluation_strategy设为epoch、save_steps设为一个较大的值（比如等于训练总步数），减少中间保存和评估的开销。
• 替换AdamW优化器：使用PyTorch原生AdamW替代transformers的旧实现，手动指定优化器：

  from torch.optim import AdamW
  optimizer = AdamW(base_model.parameters(), lr=5e-5)
  trainer = Trainer(
      model=base_model,
      args=training_args,
      data_collator=data_collator,
      train_dataset=tokenized_train_dataset,
      eval_dataset=tokenized_val_dataset,
      optimizers=(optimizer, None)
  )
  

• 对齐依赖版本：对比Colab和Vertex AI Workbench中transformers、datasets、torch的版本，确保两边版本一致，版本差异可能导致隐性兼容性问题。

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