作為一名深度學習初學者，最近在基于 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）時，直接復用配置就能快速上手 —— 這大概就是踩坑的價值吧。

如果你也在做類似任務，歡迎在評論區交流更多踩坑經驗～