メモリ効率の良い学習

🤗 PEFTはファインチューニングをパラメータ効率の高いものにするものの、メモリ効率を自動的に高めるものではありません。この概要では、トレーニング時のメモリ使用量を削減するためのヒントと、詳細なガイドへのリンクをまとめました。

メモリ使用量を削減するには、常にベースモデルのサイズを小さくする、バッチサイズを小さくする、シーケンス長を短くするといった基本的な選択肢を検討してください。

トレーニングメモリ概要

学習メモリ分布を分析し、潜在的な対策について検討できるようにしましょう。ここでは、Adamオプティマイザで学習された大規模言語モデルを例に説明します。フルファインチューニングを行う場合、以下の要素がメモリを占有します。

1. ベースモデルのパラメータ: メモリ消費量は選択したdtypeに大きく依存します。パラメータあたりのビット数（float16の場合は16ビット）が少ないほど、消費するメモリは少なくなります。float16（パラメータあたり16ビット/2バイト）の1Bモデルは、おおよそ1e9 × 2 byte = 1.863 GiBのメモリを消費します。
2. 学習可能なベースモデルのパラメータすべて ×3（勾配用に1倍、Adamオプティマイザの状態用に2倍）、したがってメモリは5.59GB
3. レイヤー間の中間活性化値のメモリは予測が困難ですが、主に使用される計算のdtypeとシーケンス長/バッチサイズに依存します。

より小さいベースモデルを使用するか、より小さい計算用データ型を使用すると、すべての項目が削減されます。一方、より短いシーケンスやより小さいバッチサイズを使用すると、主に勾配とアクティベーションのメモリに影響します。PEFT手法を採用すると、学習可能なパラメータ数が減少するため、勾配とオプティマイザの状態の両方が大幅に削減され、多くのメモリを節約できます。

適切な手法の選択

すべてのPEFT手法が同じように構築されているわけではなく、メモリ効率を重視して構築しやすい定式化もあります。メモリに制約がある場合は、PEFT手法比較スイートを確認し、最大アクセラレータメモリ使用量で絞り込むことをお勧めします。平均的なアクセラレータメモリ使用量は手法間でほぼ同等である場合が多いですが、アクティベーションやシーケンス長に対するスケーリングは手法によって異なり、メモリスパイクが発生しやすい手法とそうではない手法があります。

言語モデリングヘッドや埋め込みレイヤーのような大規模なレイヤーを対象として、特定のトークンのファインチューニングを行う場合は、学習可能なトークンの使用を検討してください。

量子化

量子化は、ベースモデルのメモリ消費量を削減する最も有効な方法の一つであり、採用する量子化手法によっては、アクティベーションのメモリも削減されます。PEFT手法はパラメータ総数のごく一部しか占めないため、PEFTはデフォルトでベースモデルよりも高い精度を使用します。これにより、アダプタが量子化手法による品質劣化の一部を緩和する効果も期待できます。PEFT量子化ガイドをご覧ください。

コンパイル

学習するモデルは、行列積、和、代入などの操作で構成されており、各操作は新しい結果を生成し、その結果としてメモリを消費します。これらの中間結果が不要な場合は、操作を融合（フュージョン）してメモリを節約できます。これはtorch.compileが実行できる多くの最適化の一つに過ぎませんので、PEFT torch.compile ガイドをご覧ください。

Gradient Checkpointing

レイヤー間でn番目の勾配のみを保存し、残りはオンザフライで計算することで、メモリと計算量をトレードオフできます。詳細については、Transformersの勾配チェックポイントのドキュメントを参照してください。

DiffusersやTransformersを使用しない場合、使用している学習フレームワークに応じて、独自の勾配チェックポイントロジックを実装する必要がある場合があります。

チャンク化されたNLL損失

言語モデル（または分類タスク）の学習において、NLLLossを使用することは非常に一般的です。batch × sequence × vocabularyサイズの行列を割り当てます。特にシーケンスや語彙が長い場合、このコストは急速に高くなります。

TRL を使用する際は、Liger kernel integration を使うか、Chunked NLLLoss を使うかのいずれかを選択できます。後者はシーケンスをサイズ 256 のチャンクに分割することで、最大メモリ消費量を一定に保ちます。

デフォルトのチャンクサイズがあなたの設定に最適でない場合は、チャンクサイズを調整する方法に関する詳細については、元のTRL PRを参照してください。