最近、LLMトレーニング中のメモリを大幅に削減できる技術が登場したので、紹介したいと思います。この技術は「Adam-mini」と呼ばれ、従来のAdamの改良版です。
ディープラーニングにおいて、Adamは最も一般的に使われている最適化アルゴリズムであり、LLMのトレーニングではほぼ標準となっています。最適化アルゴリズムは主にモデルのパラメータを更新・計算し、lossを最小化するために使用されます。
Adamは計算速度が速く、性能も優れていますが、メモリ使用量が比較的多いという課題があります。例えば、7Bのモデルをトレーニングする場合、optimizerの状態（m、vというパラメータ）を保存するために、少なくとも約50GB以上のメモリが必要になります。
最近公開されたAdam-miniは、トレーニング中のメモリ使用量を半分に削減できると謳っています。基本的には、transformerの設定において、トレーニング時に各パラメータブロックに異なる学習率を設定する必要があります。しかし、Adamは精度を求めるために、パラメータブロック単位ではなく、各パラメータに異なる学習率を割り当てています。（通常、パラメータは数十億個ありますが、パラメータブロックは約数百個程度です）。
なので、Adam-miniの開発チームは学習率をパラメータレベルまで精密に設定する必要はなく、パラメータブロックレベルで十分ではないかという仮説を立て、検証しました。その結果、パラメータブロックレベルでも精度はそれほど損なうことなく、処理速度は大幅に向上させることができることが分かりました。
⇩はllama3-7Bモデルのトレーニングを行った結果です。メモリ使用量が約50%削減され、スループットも50%向上したとのことです。
私も試してみましたが、なぜかうまくいきませんでした。トレーニング後のモデル（チェックポイント）のサイズは確かに大幅小さくなったものの、トレーニング中のメモリ使用量は変化しませんでした。メモリ使用量が変わらなかったため、当然ながらスループットも向上しませんでした。
この問題をgitHubのissueで質問しようとしたところ、既に複数の人が同じような状況を報告しています。現時点では開発チームからまだ回答はないですが、回答があり次第、もう一度試してみたいと思います。
興味ある方：