Appleが大規模モデル圧縮技術を革新、大規模モデルを携帯電話に搭載可能に

大規模言語モデル (LLM)、特に生成事前トレーニング済みトランスフォーマー (GPT) モデルは、多くの複雑な言語タスクで優れたパフォーマンスを示しています。この画期的な進歩により、ユーザーのプライバシーを保護するために、これらの LLM をモバイル デバイス上でローカルに実行できる可能性が高まります。ただし、小さな LLM でもこれらのデバイスで実行するには大きすぎます。

たとえば、小型の LLaMA には 7B のパラメータがあり、FP16 バージョンのサイズは 14GB ですが、モバイル デバイスには 18GB の DRAM しかありません。したがって、スパース化、量子化、重みクラスタリングなどのトレーニング時の最適化を通じて LLM を圧縮することは、デバイス上の LLM 展開の重要なステップです。しかし、LLM のトレーニング時間の最適化は、モデルのサイズと計算リソースのオーバーヘッドのために非常にコストがかかります。重み付けクラスタリング用の SOTA アルゴリズムの 1 つである DKM は、すべての重みとすべての可能なクラスタリング オプション間の相互作用を分析する必要があるため、トレーニング時間と可変重みクラスタリングのためにコンピューティング リソースに対する需要が高くなります。

したがって、GTPQ や AWQ などの既存の LLM 圧縮技術の多くは、トレーニング後の最適化に依存しています。この論文では、研究者らは、トレーニング時間加重クラスタリングを実現するためのメモリ最適化手法と、DKM におけるその応用、つまり eDKM を提案しています。

この論文で使用されている技術には、デバイス間のテンソル オーケストレーション、重み行列の一意性とシャーディングが含まれます。研究者らは、eDKM を使用して LLaMA 7B モデルを微調整し、重み係数ごとに 3 ビットに圧縮することで、デコーダー スタックのメモリ フットプリントを約 130 倍削減し、既存の 3 ビット圧縮技術を上回るパフォーマンスを実現しました。

DKMメモリ効率の向上

図 1 に示すように、プルーニング、量子化、正規化はすべて一般的な重み最適化手法です。これらの方法は、元の重み W を最適化して、推論の待ち時間、精度、またはモデル サイズを最適化する重みを取得します。これらの技術の中で、本論文の研究者は主に重み付けクラスタリング、特に重み付けクラスタリングアルゴリズム DKM に焦点を当てています。

重みクラスタリングは、重みマトリックスがルックアップ テーブルに圧縮され、低精度のインデックスのリストがルックアップ テーブルに格納され、最新の推論アクセラレータで処理できる非線形重み離散化です。 DKM は、重み (W で示される) と中心点 (C で示される) 間の相互作用を分析することによって微分可能な重みクラスタリングを実行し、圧縮率と精度の間でトレードオフを行います。

したがって、LLM 圧縮に DKM を使用すると、高品質の結果が得られます。しかし、DKM 計算中に生成されるアテンション マップは大きく、フォワード/バックワード パスのメモリ複雑度は O(|W||C|) (つまり、図 1 のマトリックス) であり、これは LLM 圧縮にとって特に困難です。たとえば、LLaMA 7B モデルでは、4 ビット重みクラスタリングのアテンション マップを計算するためだけに、少なくとも 224 GB のメモリが必要です。

図 1: 重量最適化システムの概要。 DKM では、微分可能な重み付きクラスター化された注意グラフがシステムによって内部的に作成されます。

したがって、研究者はこのような大きなメモリ要件を処理するために CPU メモリを使用する必要があります。つまり、最初に情報を CPU メモリに保存し、必要に応じてそれを GPU にコピーし直す必要があります。ただし、これにより GPU と CPU の間で大量のトラフィックが生成され (トレーニングの速度が低下します)、膨大な CPU メモリ容量が必要になります。つまり、CPU と GPU 間のトランザクション数を減らし、トランザクションあたりのトラフィックを最小限に抑えることが重要です。これらの課題に対処するために、研究者は PyTorch に 2 つの新しいメモリ最適化手法を導入しました。

• デバイス間でのテンソルのオーケストレーション: デバイス間でコピーされたテンソルを追跡し、冗長なコピーを回避することで、メモリ使用量を削減し、トレーニングを高速化します。
• 重みの一意性とシャーディング: 16 ビットの重みには 216 個の一意の値しかないという事実を利用して、アテンション マップの表現を削減し (図 1 を参照)、さらに複数の学習モデルに分割します。

クロスデバイス テンソル オーケストレーション

PyTorch は、実際のデータ レイアウトにリンクされたデータ ストアと、テンソルの形状や型などを格納するメタデータを使用してテンソルを表します。このテンソル アーキテクチャにより、PyTorch はデータ ストレージを可能な限り再利用し、メモリ使用量を効果的に削減できます。ただし、テンソルが別のデバイス (GPU から CPU など) に移動されると、データ ストレージを再利用することができなくなり、新しいテンソルを作成する必要があります。

表 1 は、PyTorch デバイス間で移動するときのテンソルのメモリ使用量を示しています。行 0 に割り当てられたテンソル x0 は、GPU 上で 4MB を消費します。 1 行目でビューが変更されると、基礎となるデータ ストアを再利用できるため (つまり、x0 と x1 は実質的に同じ)、追加の GPU メモリは必要ありません。しかし、2行目と3行目のようにx0とx1をCPUに移動すると、y0とy1はCPU上で同じデータストレージを共有できますが、CPUメモリの消費量は8MBとなり、CPUメモリの冗長性が発生し、GPUからCPUへのトラフィックが増加します。

表 1: LLM の微調整では、CPU メモリを使用して GPU のメモリ使用量をオフロードする必要がある場合があります。デバイス間のテンソル管理が不足すると、デバイス間での冗長なコピーが発生し (特に計算グラフが複雑な場合)、LLM のトレーニング時間の最適化に特に悪影響を及ぼします。たとえば、x0 と x1 は異なるビューを持つ同じテンソルですが、CPU にコピーされると、結果のテンソル y0 と y1 はデータ ストレージを共有しませんが、GPU では x0 と x1 はデータ ストレージを共有します。

この非効率性を解決するために、研究者は図 2 (b) にオーケストレーション層を配置しました。ここで、黒は実際のデータ ストレージとメタデータを表し、灰色はメタデータのみを表します。図 2 (a) は表 1 の例を示しており、x1 は x0 とデータ レイアウトを共有していますが、y0 と y1 は CPU 上で重複したデータ ストレージを持っています。図 2 (b) に示すように、研究者はオーケストレーション層を挿入することでこの冗長性を回避し、GPU から CPU へのトラフィックを削減しました。研究者らは、PyTorch の save-tensor-hook を使用してこのようなスワッピング スキームを実装し、同じデータ ストレージがコピーされているかどうかを確認しました。

ただし、このような方式を使用して、ターゲット デバイスに同じテンソルが存在するかどうかを確認するのはコストがかかります。図 2 (b) の例では、x1 を CPU にコピーする代わりに、y0 への参照を返して、x1 と y0 の間でビュー操作を実行します。

図 2: 表 1 の状況にクロスデバイス テンソル オーケストレーションを適用すると、CPU 側での重複を回避でき、メモリとトラフィックを節約できます。

計算グラフをナビゲートすると余分な計算サイクルが追加されますが、これは不要なコピーを保存することで相殺できます。元の DKM 実装の計算グラフですべての適格なケースを検出するには、4 ホップ以内の検索で十分であることがわかりました。

重みの一意性とシャーディング

ほとんどの LLM トレーニングでは、重みは通常 16 ビット ストレージ (BF16 や FP16 など) を使用して保存されます。つまり、LLM には数十億のパラメーターがありますが、ビット幅のため、一意の係数は 216 個しかありません。これにより、図 3 に示すように、重みと中心の間の注目マップを大幅に圧縮する機会が提供されます。

図3: 重みの一意性とシャーディング

実験結果

LLMの精度

この論文では、eDKM を、RTN、SmoothQuant、GPTQ、AWQ、LLM-QAT などの他の量子化ベースの圧縮方式と比較します。 eDKM の場合、研究者らは埋め込み層で 8 ビット圧縮も実行しました。最終的に、以下の結論が導き出されました。

• eDKM により、3 ビット圧縮 LLaMA 7B モデルは他のすべての 3 ビット圧縮方式よりも優れたパフォーマンスを発揮します。
• eDKM は、3 ビット構成と 4 ビット構成の両方で ARC-e ベンチマークで最高の精度を実現します。
• 4 ビット圧縮モデルを使用した PIQA および MMLU ベンチマークでは、eDKM は非常に競争力のあるパフォーマンスを実現します。

アブレーション実験

アブレーション実験では、研究者らは LLaMA 7B デコーダー スタックのアテンション レイヤーを例にとり、3 ビット圧縮によるメモリ使用量と前後方向の速度のトレードオフを測定しました。クロスデバイス テンソル オーケストレーションだけでも、実行時のオーバーヘッドを最小限に抑えながらメモリ使用量を 2.9 倍削減できます。一方、シャーディング モジュールと一意性モジュールでは、それぞれ 23.5 倍と 16.4 倍の節約が実現します。すべての技術を組み合わせると、eDKM は約 130 倍の節約を実現できます。これらの手順では追加の計算と通信のオーバーヘッドが必要になりますが、GPU と CPU 間のトラフィックが大幅に削減されるため、実行時のオーバーヘッドはごくわずかです。

詳細は原文をご参照ください。