[論文レビュー] Full Parameter Fine-tuning for Large Language Models with Limited Resources

Kai Lv, Yuqing Yang+4 more|arXiv (Cornell University)|Jun 16, 2023

Topic Modeling被引用数 14

ひとこと要約

LOw-Memory Optimization (LOMO) を紹介し、勾配計算とパラメータ更新を融合するオプティマイザにより、限られたハードウェア上で非常に大規模な言語モデルの全パラメータ微調整を可能にしてメモリ使用を大幅に削減する。65Bモデルまで8 RTX 3090 GPUs での実現性を示し、下流の性能も競争力がある。

ABSTRACT

Large Language Models (LLMs) have revolutionized Natural Language Processing (NLP) but demand massive GPU resources for training. Lowering the threshold for LLMs training would encourage greater participation from researchers, benefiting both academia and society. While existing approaches have focused on parameter-efficient fine-tuning, which tunes or adds a small number of parameters, few have addressed the challenge of tuning the full parameters of LLMs with limited resources. In this work, we propose a new optimizer, LOw-Memory Optimization (LOMO), which fuses the gradient computation and the parameter update in one step to reduce memory usage. By integrating LOMO with existing memory saving techniques, we reduce memory usage to 10.8% compared to the standard approach (DeepSpeed solution). Consequently, our approach enables the full parameter fine-tuning of a 65B model on a single machine with 8 RTX 3090, each with 24GB memory.Code and data are available at https://github.com/OpenLMLab/LOMO.

研究の動機と目的

• 限定的なGPU資源で大規模言語モデル（LLMs）の全パラメータ微調整を動機づけ、可能にする。
• トレーニング中のメモリ使用を削減しつつ、SGDの更新と同値性を保つ。
• 65Bパラメータモデルで実用的な実現性を示し、下流の性能を評価する。
• 新しいオプティマイザを用いた混合精度訓練を安定化する手法を探る。

提案手法

• 勾配計算とパラメータ更新を一歩で融合し、勾配テンソルを保存することを回避するLOw-Memory Optimization (LOMO)を提案。
• この文脈で全微調整のオプティマイザとしてSGDを有効な選択肢として提案し、滑らかな損失面と局所最適解が微調整タスクには許容されると主張。
• 活性化チェックポイント作成、混合精度などの一般的なメモリ節約技術とLOMOを統合し、動的損失スケーリングの安定性を保つために2パス方式を採用。
• 大規模で滑らかな損失地形に対するSGD下の暗黙バッチサイズと安定性について理論的議論を提供。
• 7Bから65BのLLaMAモデルでメモリ使用量、スループット、下流タスク性能（SuperGLUE）を経験的に評価し、AdamWとLoRAと比較。

実験結果

リサーチクエスチョン

• RQ1 SGD はメモリ制約下での LLMs の全パラメータ微調整に現実的なオプティマイザとなり得るか？
• RQ2 勾配計算とパラメータ更新を融合することでのメモリ節約（LOMO）および活性化チェックポイントと混合精度の組み合わせによる節約幅はどのくらいか？
• RQ3 LOMO を用いた7B–65Bモデルを家庭用GPU上で微調整する際のスループットとメモリのトレードオフは？
• RQ4 全パラメータ微調整を LOМO で行った場合の下流タスクのパフォーマンスは、ゼロショットおよび LoRA と比較してどうか？
• RQ5 LOMO は LoRA と互換・補完的に機能し、さらなる性能向上をもたらすか？

主な発見

AC	Params	Gradients	Optim States	Activations	Total Memory	Column? (inference)
AdamW	12.55	12.55	75.31	45.61	147.02	✓
SGD	12.55	12.55	25.10	45.61	96.81	✓
LOMO	12.55	0.24	0.00	45.61	59.40	✓

• LOMO は勾配メモリを最大勾配テンソルのサイズにまで圧縮し、総メモリ使用量を大幅に低減する。
• LOMO のメモリ使用量は、活性化チェックポイントと組み合わせた場合、AdamW の場合は 102.20 GB から LLaMA-7B で 14.58 GB に低下し、8× RTX 3090 GPU の単一マシンでの訓練を可能にする。
• LOMO は 8 台の RTX 3090 GPU で 65B モデルの訓練を可能にし、1GPU当たりのスループットは 4.93 トークン/秒（TGS）と報告されている。
• 下流の SuperGLUE タスクでは、65B パラメータの LOµMO は平均 89.9 の性能を達成し、ゼロショットを上回り、いくつかのデータセット（例：RTE、BoolQ、WSC、WIC、MultiRC、COPA など）で LoRA を概ね上回る。
• LoRA+LOMO はさらなる性能向上をもたらす可能性があり、全パラメータ微調整と低秩適応の補完的利点を示す。