隨著語言模型的發展
語言模型越來越大，它們的詞彙量也隨之增加。這使得在訓練過程中，LLM（大型語言模型）的記憶使用量主要集中在一個層面上：損失計算中的交叉熵。交叉熵會建立一個對應矩陣，這個矩陣包含每一對輸入標記和詞彙項的數據。對於小型模型來說，這會消耗比其他部分多得多的記憶體。
我們的解決方案：Cut Cross-Entropy (CCE)
我們提出了一種方法，稱為Cut Cross-Entropy（CCE），這種方法在計算交叉熵損失時，不需要將所有標記的logits（對數幾率）存入全局記憶體。相反，CCE只計算正確標記的logit，並即時計算所有logits的對數和指數。
如何實現CCE
我們實現了一個自訂的內核，這個內核在快閃記憶體中執行矩陣乘法和對詞彙的對數和指數計算，這樣使得交叉熵計算的全局記憶體消耗變得微不足道。
顯著的效果
以Gemma 2（2B）模型為例，CCE將損失計算的記憶體使用量從24 GB減少到1 MB，並將分類器頭的總訓練記憶體消耗從28 GB減少到1 GB。
提高CCE的效率
為了提高CCE的處理速度，我們利用softmax的固有稀疏性，提出跳過對梯度計算中貢獻微不足道的元素（即低於數值精度的部分）。實驗證明，記憶體消耗的顯著減少是在不影響訓練速度或收斂性的情況下實現的。
本文由 AI 台灣 運用 AI 技術編撰，內容僅供參考，請自行核實相關資訊。
歡迎加入我們的 AI TAIWAN 台灣人工智慧中心 FB 社團，
隨時掌握最新 AI 動態與實用資訊！
