Adam-mini:使用更少的學習率獲得更多Adam-mini: Use Fewer Learning Rates To Gain More
我們提出了Adam-mini,這是一種優化器,其在記憶體佔用上比AdamW表現相當或更好,並且記憶體佔用量減少了45%至50%。Adam-mini通過降低Adam中學習率資源(即1/v)來減少記憶體使用量。我們發現,在v中,超過90%的這些學習率可以被安全地移除,方法是:(1)根據我們提出的海森矩陣結構原則,仔細將參數劃分為塊;(2)為每個參數塊分配一個單一但良好的學習率。我們進一步發現,對於每個這些參數塊,存在一個高質量的單一學習率,可以勝過Adam,前提是有足夠的資源來搜索它。然後,我們提供了一種成本效益的方法來找到良好的學習率,並提出Adam-mini。在實驗中,我們驗證了Adam-mini在從125M到7B的各種語言模型上進行預訓練、監督微調和RLHF時的表現與AdamW相當或更好。Adam-mini的減少記憶體佔用量還減輕了GPU和CPU之間的通信開銷,從而提高了吞吐量。例如,當在2倍A800-80GB GPU上對Llama2-7B進行預訓練時,Adam-mini的吞吐量比AdamW高出49.6%,節省了33%的預訓練時間。