Last Updated on 2024-02-20 by Clay
問題描述
HuggingFace 之前曾經發表過文章表示現在的 LLM最好是依照 ChatML 格式去訓練，在一般情況下，會按照 system、user、assistant 的三種不同角色來進行生成，格式如下：
<|im_start|>system ...system prompt...<|im_end|> <|im_start|>user ...user message...<|im_end|> <|im_start|>assistant ...
通常在訓練時，我們會讓 <|im_start|>assistant\n 之前的 token 全都不參與 loss 計算，換言之即是讓模型只學習如何作為 assistant 進行回答，並在輸出特殊 eos_token <|im_end|> 後結束生成。
然而今天在我透過 trl 的 SFTTrainer() 進行 Mistral（實際上是 Mistral 的微調模型 OpenHermes）的微調後，我發現我的模型不再會生成 <|im_end|> 符號了！它會喋喋不休地一直往下生出看似有關聯的敘述。
像這樣的非預期行為，是完全不能應用在產品上的，所以我花費了一點時間仔細確認這是什麼原因造成的。
我在 GitHub - model produces multiple responses in one after training with conservation datasets using sft 上發現了跟我的問題十分相像的 issue，底下的討論分別提出了兩種建議：
- 確認資料集：有一種可能是資料集本身卻缺失了 eos_token 的結束符號，這會造成模型學不到在何時生成結束符號
- 嘗試設定
tokenizer.pad_token = tokenizer.unk_token：這是讓模型在訓練時所填充的特殊符號指定為 unk_token（unknown token）
我嘗試翻了一下 SFTTrainer() 的原始碼，發現第二種解決方法正是我所需要的。
解決方法
為什麼我的 Mistral 模型微調後不再生成 eos_token 的原因，我猜正是因為訓練時看過太多 eos_token 被填充在訓練資料的序列開頭，導致模型被強迫學習了忽略 eos_token。
這件事要從 SFTTraienr() 的原始碼說起：
if tokenizer is None:
tokenizer = AutoTokenizer.from_pretrained(model.config._name_or_path)
if getattr(tokenizer, "pad_token", None) is None:
tokenizer.pad_token = tokenizer.eos_tokentokenizer 是可以由我們主動輸入給 SFTTrainer() 的；然而在沒有傳入 tokenizer 的情況下，SFTTrainer() 會主動由我們傳入模型的設定位置，自動替我們建立 tokenizer。
但問題就出在於 SFTTrainer() 會主動替我們檢查 pad_token 是否存在、並且在不存在的情況下，自動使用 eos_token 同時作為 pad_token。
這個操作本身沒有問題，但是 SFTTrainer() 是預期我們在訓練時 padding 是從右側進行 padding 的。（這個操作的必要性是，在訓練時需要把一個 batch 的訓練資料通通 padding 成一樣長度才能批次訓練）
if tokenizer.padding_side is not None and tokenizer.padding_side != "right":
warnings.warn(
"You passed a tokenizer with `padding_side` not equal to `right` to the SFTTrainer. This might lead to some unexpected behaviour due to "
"overflow issues when training a model in half-precision. You might consider adding `tokenizer.padding_side = 'right'` to your code."
)
從右側 padding 本身也是沒有問題的，因為 SFTTrainer() 預期模型訓練時可能會看到的填充後資料是：
<|im_start|>assistant Today is a nice day!<|im_end|><|im_end|><|im_end|>...<|im_end|>
那模型自然會知道，生成結束後應該就是要輸出 eos_token。
然而 Mistral 模型的問題在於，它本來在接受訓練時的 padding 方向就是左邊而非右邊，並且這個設置仍然被 tokenizer 所儲存著。
tokenizer = AutoTokenizer.from_pretrained("teknium/OpenHermes-2.5-Mistral-7B/")
print(tokenizer.padding_side)
Output:
'left'
那麼這時候選擇使用 tokenizer.pad_token = tokenizer.eos_token 就不好了。模型實際上看到的訓練資料是：
<|im_end|><|im_end|>...<|im_end|><|im_start|>assistant Today is a nice day!<|im_end|>
在語言模型的建模中，起始位置、結束位置的特殊 token 是非常重要的。雖然我們預期模型在微調時也會學習到訓練資料序列尾端的 eos_token，但是填充特殊符號的意義在於讓模型學會省略那些填充符號，那在模型被強迫學會了忽略大量填充在左側的 eos_token 後，它可能沒辦法積極地學習識別與輸出尾端我們真正希望的 eos_token。
而這就是 Mistral 模型在透過 SFTTrainer() 微調後不輸出 eos_token 這一問題可能發生的原因。
所以我個人在目前幾次微調的測試後，會建議自己設定 tokenizer 並使用 tokenizer.pad_token = tokenizer.unk_token 這個方法。當然，如果 unk_token 對於你的任務來說很重要的話，那還是透過 add_special_tokens() 去添加真正的 pad_token 比較好，我這作法算是應急。
當然，另一種可行的辦法就是把 Mistral 的 tokenizer.padding_side = "right" 這樣進行設定。但我個人之前進行這樣的設定後，發現 Mistral 模型訓練時的 eval_loss 振幅非常大，到現在仍下意識地覺得還是按照 Mistral 當初微調時的初始設定來比較好。
那時候的 Loss 視覺化結果我甚至還有傳給同事看，讓他看看驚人的 eval_loss 振幅。
希望大家都可以不要踩到這個坑。
（2024/02/20 更新：根據同事的追蹤研究，後來發現確認就是 pad_token 不能跟 eos_token 一致，否則訓練資料中的 eos_token 會完全變成 -100 的標籤不參與訓練。總而言之就是設定 tokenizer.pad_token 跟 <unk> 即可。）
References
- Templates for Chat Models
- GitHub - model produces multiple responses in one after training with conservation datasets using sft
- https://github.com/huggingface/trl/blob/main/trl/trainer/sft_trainer.py