作为一名深度学习初学者，最近在基于 Hugging Face Transformers 微调 BERT 模型做 QNLI 任务时，被Checkpoint 保存和TrainingArguments 配置这两个知识点卡了整整两天。从磁盘爆满、权重文件加载报错，到不知道如何控制 Checkpoint 数量，每一个问题都让我一度想放弃。好在最终逐一解决，特此整理成博客，希望能帮到同样踩坑的朋友。

一、核心背景：我在做什么？

本次任务是基于 GLUE 数据集的 QNLI（Question Natural Language Inference，问题自然语言推理）任务，用 Hugging Face 的run_glue.py脚本微调bert-base-cased模型。核心需求很简单：

1. 顺利完成模型微调，避免中途中断；
2. 控制 Checkpoint（模型快照）的保存数量，防止磁盘爆满；
3. 后续能正常加载 Checkpoint，用于后续的 TRAK 贡献度分析。

但实际操作中，光是 “Checkpoint 保存” 这一个环节，就暴露出我对TrainingArguments（训练参数配置类）的认知盲区。

二、先搞懂：TrainingArguments 是什么？为什么它很重要？

在解决问题前，必须先理清TrainingArguments的核心作用 —— 它是 Hugging Face Transformers 库中控制训练全流程的 “总开关”，几乎所有与训练相关的配置（如批次大小、学习率、Checkpoint 保存策略）都通过它定义。

1. TrainingArguments 的本质

TrainingArguments是一个数据类（dataclass） ，它将训练过程中需要的所有参数（从优化器设置到日志保存）封装成结构化对象，再传递给Trainer（训练器）实例。无需手动编写训练循环，只需配置好TrainingArguments，Trainer就能自动完成训练、验证、Checkpoint 保存等操作。

2. 常用核心参数（按功能分类）

我整理了本次任务中最常用的参数，按 “训练基础配置”“Checkpoint 控制”“日志与验证” 三类划分，新手直接套用即可：

3. TrainingArguments 的配置方式

TrainingArguments不支持在代码中硬编码（除非修改脚本），常用两种配置方式，新手推荐第二种：

方式 1：命令行参数（快速调试）

运行run_glue.py时，通过--参数名 参数值的格式传递，示例：

python run_glue.py \--model_name_or_path bert-base-cased \--task_name qnli \--output_dir /root/autodl-tmp/bert_qnli \--do_train \--do_eval \--per_device_train_batch_size 8 \--learning_rate 5e-5 \--num_train_epochs 3 \--save_strategy epoch \--save_total_limit 3 \--overwrite_output_dir

方式 2：JSON 配置文件（固定复用）

将所有参数写入 JSON 文件（如qnli_train_config.json），运行时直接指定文件，适合参数较多或多任务复用：

{"model_name_or_path": "bert-base-cased","task_name": "qnli","do_train": true,"do_eval": true,"max_seq_length": 128,"per_device_train_batch_size": 8,"learning_rate": 5e-5,"num_train_epochs": 3.0,"output_dir": "/root/autodl-tmp/bert_qnli_new","save_strategy": "epoch","save_total_limit": 3,"overwrite_output_dir": true,"logging_dir": "/root/autodl-tmp/bert_qnli_new/logs","logging_steps": 100
}

python run_glue.py qnli_train_config.json

三、我的踩坑实录：3 个经典问题与解决方案

接下来重点复盘我遇到的 3 个核心问题，每个问题都附 “报错现象→原因分析→解决步骤”，新手可直接对号入座。

问题 1：训练中途磁盘爆满，被迫中断

报错现象

训练到约 3 万步时，服务器提示 “磁盘空间不足”，查看output_dir发现有 10 多个 Checkpoint 文件夹，每个文件夹占用数百 MB，累计占用超过 20GB。

原因分析

默认情况下，TrainingArguments的save_strategy为"steps"（每 500 步保存一次），且save_total_limit未设置（不限制保存数量）。QNLI 任务 1 个 epoch 约 1.3 万步，3 个 epoch 会生成 6-8 个 Checkpoint，加上优化器状态文件（optimizer.pt），很容易撑爆磁盘。

解决步骤

1. 在TrainingArguments中添加save_total_limit: 3（最多保存 3 个 Checkpoint，超过自动删除最旧的）；
2. 选择合适的save_strategy：若追求稳定，用"epoch"（每轮保存一次，3 个 epoch 仅 3 个 Checkpoint）；若需中途恢复，用"steps"并设置较大的save_steps（如 5000 步）；
3. 可选添加save_only_model: true（只保存模型权重，不保存优化器状态，每个 Checkpoint 体积从 500MB 缩减到 300MB 左右）。

问题 2：加载 Checkpoint 时提示 “_pickle.UnpicklingError: invalid load key, '\xe0'”

报错现象

训练中断后，尝试加载已保存的 Checkpoint（路径/root/autodl-tmp/bert_qnli/checkpoint-31000），运行代码：

model.load_state_dict(torch.load(os.path.join(checkpoint, "model.safetensors"), map_location=DEVICE))

报错：_pickle.UnpicklingError: invalid load key, '\xe0'。

原因分析

• model.safetensors是Safetensors 格式的权重文件（更安全，但需专用方法加载）；
• torch.load()是 PyTorch 原生加载函数，更适合加载pytorch_model.bin（PyTorch 二进制格式），用它加载 Safetensors 格式会因 “格式不兼容” 报错。

解决步骤

有两种方案，按需选择：

方案 1：改用 Safetensors 专用加载函数（需先安装safetensors库）

pip install safetensors

from safetensors.torch import load_file
# 用load_file()替代torch.load()
model.load_state_dict(load_file(os.path.join(checkpoint, "model.safetensors"), device=DEVICE))

方案 2：让 Checkpoint 默认保存为pytorch_model.bin格式
在TrainingArguments中添加save_safetensors: false，后续生成的 Checkpoint 会默认保存为pytorch_model.bin，直接用torch.load()加载即可：

model.load_state_dict(torch.load(os.path.join(checkpoint, "pytorch_model.bin"), map_location=DEVICE))

问题 3：配置 TrainingArguments 后，Checkpoint 迟迟不生成

报错现象

设置save_strategy: "epoch"后，训练到 4000 步仍未生成任何 Checkpoint，怀疑配置未生效。

原因分析

save_strategy: "epoch"表示每轮训练结束后才保存 Checkpoint，而 QNLI 任务 1 个 epoch 约 1.3 万步（训练集约 10 万样本，batch_size=8时：100000÷8=12500 步）。4000 步仅完成第一个 epoch 的 1/3，未到保存时机，属于正常现象。

解决步骤

1. 若想快速验证配置是否生效，临时改用save_strategy: "steps"并设置较小的save_steps（如 2000 步），训练到 2000 步时会自动生成checkpoint-2000；
2. 若坚持按 epoch 保存，耐心等待第一个 epoch 结束（约 1.3 万步），日志会打印Saving model checkpoint to xxx，此时output_dir下会出现第一个 Checkpoint；
3. 查看日志确认配置：搜索save_strategy和save_total_limit，确认日志中显示的参数与 JSON 配置一致（避免 JSON 文件未被正确读取）。

四、总结：TrainingArguments 配置 “避坑指南”

经过这次踩坑，我总结出 3 条新手必看的 “避坑原则”，帮你少走弯路：

1. 优先用 JSON 配置文件：命令行参数容易遗漏，JSON 文件可固化配置，后续复用或修改时更清晰；
2. Checkpoint 配置 “三要素”：每次训练前必确认save_strategy（保存时机）、save_total_limit（保存数量）、output_dir（保存路径），这三个参数直接决定是否会出现磁盘爆满或 Checkpoint 丢失；
3. 加载 Checkpoint 前先看格式：先查看 Checkpoint 文件夹中的权重文件名（是model.safetensors还是pytorch_model.bin），再选择对应的加载函数，避免格式不兼容报错。

最后想说，深度学习中的 “环境配置” 和 “参数调试” 虽然繁琐，但每一次踩坑都是对知识点的深化。这次从 “完全不懂 TrainingArguments” 到 “能灵活控制 Checkpoint”，虽然花了两天时间，但后续再做其他 GLUE 任务（如 SST-2、MRPC）时，直接复用配置就能快速上手 —— 这大概就是踩坑的价值吧。

如果你也在做类似任务，欢迎在评论区交流更多踩坑经验～