TACOの論文がarXivで公開されました。LLMのファインチューニングにおける厄介なメモリボトルネックのひとつ、オプティマイザーの状態に取り組んでいます。
多くの人はモデルのパラメーターに注目しますが、オプティマイザーの状態(モメンタムバッファーやAdamの分散推定値)は、パラメーターのメモリ使用量の2〜3倍に達することがあります。7Bモデルでは、かなり深刻な問題です。
TACOのアプローチは、オプティマイザーの状態を三値量子化することです。float32やbfloat16の値をそのまま保存する代わりに、{-1, 0, +1}の表現に圧縮します。論文によると、オプティマイザーのメモリ使用量を大幅に削減しながら、収束品質を維持できるとのことです。
重要なポイントは次のとおりです。
- 同じハードウェアで、より大きなモデルをファインチューニングできる
- 学習中のメモリ帯域への負荷を軽減できる
- オプティマイザー状態の同期コストが下がり、マルチGPU環境の効率が向上する可能性がある
LoRA、QLoRA、あるいはフルファインチューニングを行っていて、OOMエラーに悩まされているなら、この手法で必要な余裕を確保できるかもしれません。本当に有効かどうかは、さまざまなモデルアーキテクチャやタスクで効果が持続するかにかかっています。
多くの人はモデルのパラメーターに注目しますが、オプティマイザーの状態(モメンタムバッファーやAdamの分散推定値)は、パラメーターのメモリ使用量の2〜3倍に達することがあります。7Bモデルでは、かなり深刻な問題です。
TACOのアプローチは、オプティマイザーの状態を三値量子化することです。float32やbfloat16の値をそのまま保存する代わりに、{-1, 0, +1}の表現に圧縮します。論文によると、オプティマイザーのメモリ使用量を大幅に削減しながら、収束品質を維持できるとのことです。
重要なポイントは次のとおりです。
- 同じハードウェアで、より大きなモデルをファインチューニングできる
- 学習中のメモリ帯域への負荷を軽減できる
- オプティマイザー状態の同期コストが下がり、マルチGPU環境の効率が向上する可能性がある
LoRA、QLoRA、あるいはフルファインチューニングを行っていて、OOMエラーに悩まされているなら、この手法で必要な余裕を確保できるかもしれません。本当に有効かどうかは、さまざまなモデルアーキテクチャやタスクで効果が持続するかにかかっています。