使用Hugging Face Trainer训练时，开启Checkpoint仍需手动保存模型吗？

问题描述 #

我在使用Hugging Face Trainer训练模型时已开启训练过程中的checkpoint保存机制，有两个疑问：

1. 执行trainer.train()后，是否仍需显式调用trainer.save_model()来确保最终/最优模型被保存？
2. 在已开启训练中checkpoint保存的前提下，训练结束后保存模型及可选tokenizer的标准方式是什么？

原有训练代码示例 #

# -- Training arguments and trainer instantiation
output_dir = Path(f'~/data/maf_data/results_{today}/').expanduser() if not debug else Path(f'~/data/maf_data/results/').expanduser()
print(f'{debug=} {output_dir=} \n {report_to=}')
training_args = TrainingArguments(
    output_dir=output_dir,  # 模型预测结果和checkpoint的输出目录
    # num_train_epochs = num_train_epochs, 
    max_steps=max_steps,  # TODO: hard to fix, see above
    per_device_train_batch_size=per_device_train_batch_size,
    gradient_accumulation_steps=gradient_accumulation_steps,  # 基于alpaca实现，支持累积梯度达到更大的有效batch size
    gradient_checkpointing = gradient_checkpointing,  # 根据硬件决定是否开启
    optim="paged_adamw_32bit",  # 使用32位优化器
    warmup_steps=500,  # 线性warmup步数
    warmup_ratio=0.03,  # 线性warmup比例
    weight_decay=0.00,  # 权重衰减系数
    learning_rate = 1e-5,  # 学习率
    max_grad_norm=1.0, # 梯度裁剪阈值
    lr_scheduler_type="cosine",  # 余弦学习率调度器
    logging_dir=Path('~/data/maf/logs').expanduser(),
    save_steps=2000,  # 每2000步保存一次checkpoint
    logging_steps=50,  # 每50步记录一次日志
    remove_unused_columns=False,  # 保留所有数据集列
    report_to=report_to,  # 日志报告平台
    fp16=False,  # 禁用fp16
    bf16=torch.cuda.get_device_capability(torch.cuda.current_device())[0] >= 8,  # 根据GPU能力自动开启bf16
    evaluation_strategy='steps',  # 按步数进行评估
    per_device_eval_batch_size=per_device_eval_batch_size,
    eval_accumulation_steps=eval_accumulation_steps,
    eval_steps=eval_steps,
)

print(f'{pretrained_model_name_or_path=}')
print(f'{train_dataset=}')
print(f'{eval_dataset=}')
trainer = Trainer(
    model=model,
    args=training_args,  
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    data_collator=lambda data: custom_collate_fn(data, tokenizer=tokenizer)
)

# 开始训练
cuda_visible_devices = os.environ.get('CUDA_VISIBLE_DEVICES')
if cuda_visible_devices is not None:
    print(f"CUDA_VISIBLE_DEVICES = {cuda_visible_devices}")
trainer.train()
trainer.save_model(output_dir=output_dir)  # 疑问：是否需要这一步？
print('Done!\a')

计划修改的保存代码 #

# 保存最终checkpoint 疑问：是否需要这一步？
final_ckpt_dir = output_dir / f'ckpt-{max_steps}'
final_ckpt_dir.mkdir(parents=True, exist_ok=True)
trainer.save_model(output_dir=final_ckpt_dir)
print('Done!\a')

解答 #

1. 是否需要显式调用trainer.save_model()？ #

分两种场景判断：

• 如果训练总步数max_steps是save_steps的整数倍，训练结束时最后一个checkpoint会自动保存，此时无需额外调用save_model()；
• 如果总步数不是save_steps的整数倍，训练结束时的最终模型状态不会被自动保存为checkpoint，必须显式调用trainer.save_model()才能保留最后一步的参数；
• 若开启了load_best_model_at_end=True和metric_for_best_model（最优模型自动加载），Trainer会在训练结束后加载最优模型，但不会自动将其保存到指定位置，此时也需要显式调用save_model()留存最优模型。

2. 训练结束后保存模型及tokenizer的标准方式 #

在已开启checkpoint保存的前提下，推荐以下流程：

• 保存最终/最优模型：调用trainer.save_model(save_directory)，该方法会自动保存模型权重、配置文件，若Trainer初始化时传入了tokenizer，也会同步保存分词器的基础配置；
• 单独保存tokenizer：为确保tokenizer的完整词汇表和配置被留存，显式调用tokenizer.save_pretrained(save_directory)，与模型保存到同一目录，方便后续统一加载；
• 保存训练参数：调用training_args.save_to_json(save_directory / "training_args.json")，留存训练参数配置，便于实验复现。

优化后的保存代码示例：

# 创建最终模型保存目录
final_ckpt_dir = output_dir / f'ckpt-{max_steps}'
final_ckpt_dir.mkdir(parents=True, exist_ok=True)

# 保存模型
trainer.save_model(final_ckpt_dir)

# 显式保存tokenizer
tokenizer.save_pretrained(final_ckpt_dir)

# 保存训练参数
training_args.save_to_json(final_ckpt_dir / "training_args.json")

print('Done!\a')

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