MaskLLM

MaskLLM

NVIDIA Research · コーディング · その他

MaskLLMは、N:Mスパースのマスクを学習して大規模言語モデルを小さくするNVIDIA Researchのオープンソースなプルーニング手法だ。推論コストを抑えつつ、精度を元モデルの近くに保つ。

MaskLLM の画面プレビュー

MaskLLM について

What Is MaskLLM

MaskLLMは大規模言語モデルを圧縮するための研究手法だ。モデルに不要なパラメータを取り除くプルーニングツールの一種で、半構造化スパース(N:Mスパースとも呼ばれる)を使う。考え方は単純で、M個の連続した重みのブロックごとにN個だけ残す。2:4のパターンなら4つの重みのうち2つを保つので、ハードウェアが動かすデータ量は半分になる。

MaskLLMの違いは、どの重みを残すかの決め方にある。従来の手法は固定のルールに頼っていた。小さな重みを捨てる、統計で推定する、といった具合だ。MaskLLMはマスクの選択そのものをモデルが学ぶ対象として扱う。N:Mパターンを確率分布としてモデル化し、Gumbel Softmaxというサンプリング手法で訓練するので、大規模なデータセット上でより良いマスクを見つけられる。マスクを学ぶ間、重みは凍結したままだ。この点が、特定タスクに対して2:4スパースを「ロスレス」にする理由である。

実用的な魅力は数字にある。NVIDIAによれば、密なLLaMA-2 7BモデルはWikitextでperplexity 5.12を記録する。従来のプルーニング手法はこれを10以上に押し上げた。MaskLLMは重みを凍結したままマスクだけを学び、6.72から6.78まで下げる。プロジェクトはNVIDIAライセンスの下でオープンソース化されており、研究者もエンジニアも自分で動かせる。ただし本物のGPUハードウェアと機械学習の知識は要る。

Getting Started

  1. Pythonパッケージを入れ、計算済みのマスクをHugging Faceから取得するスクリプトを実行する
  2. LLaMA-2 7BやLLaMA-3 8Bなど目的のモデルをPyTorchで読み込む
  3. 取得したマスクを適用し、付属のperplexityスクリプトで精度を測る
  4. 自分で訓練する場合はMegatron-LM環境を用意し、データを整えて自領域のマスクを新たに訓練する

AIツール情報

MaskLLMの料金、対応プラットフォーム、性能を簡単に確認できます。

無料プランはい
有料プランFree (open-source)
プラットフォームPython, PyTorch, Hugging Face Transformers, Megatron-LM, NVIDIA GPUs
開発元NVIDIA Research
カテゴリコーディング · その他
リリース日Sep 2024
最終更新Dec 2024
サイト訪問数N/A
サイト世界ランキングN/A
API提供状況いいえ

こんな人におすすめ

このツールが最も力を発揮するユーザー、タスク、シーン。

ユーザー

  • 機械学習の研究者
  • モデルを配備するMLエンジニア
  • 大学院生と研究者
  • フレームワークやツールの開発者

タスク

  • 大規模言語モデルのプルーニング
  • 一発圧縮
  • 領域特化のスパース
  • スパースの転移
  • プルーニング手法の比較

シーン

  • 少ないGPUで7Bモデルを動かすチーム
  • 論文の数字を再現する必要がある研究
  • 全体は良いがニッチ領域で崩れるモデル
  • モデルファミリー向けの圧縮パイプライン構築
  • ライセンスやプライバシーの制約を検討

主な機能

学習可能なN:Mマスク

MaskLLMは各重みを手作業で採点する代わりに、マスクそのものを学ぶ。有効なN:Mパターンの集合を分布として扱い、訓練中にそこからサンプリングするので、モデルはデータセットに最適なマスクへ収束していく。この学習ステップが、マスクを一度選んで終わる一発方式との根本的な違いだ。

Gumbel Softmaxサンプリング

この手法は離散サンプリングの技法を借り、マスクの選択をエンドツーエンドで訓練可能にする。Gumbel Softmaxは、本来なら白黒はっきりした硬い選択に勾配を通す。だからこそ、N:Mマスク上の分布を推定ではなく実データで最適化できる。

Hugging Faceモデル向けの計算済みマスク

プロジェクトには既成のマスクが用意され、訓練なしで試せる。LLaMA-2 7B、LLaMA-3 8B、LLaMA-3.1 8B向けのマスクが公開され、それぞれC4のサブセットで訓練されている。スクリプトがHugging Faceから取得し、評価用に公式モデルへ適用する。

半構造化2:4スパース

MaskLLMが狙うのはN:Mスパース、多くの場合は4つの連続した重みに2つの非ゼロ値を残す2:4パターンだ。この形式が重要なのは、ハードウェア加速に対応するからである。小さくなったモデルは場所を取らないだけでなく、適切なGPUなら高速にも動く。

疎パターンのマスク転移

マスクを学ぶ同じ確率モデルが、マスクをある領域から別の領域へ運ぶ。このマスク転移は、マスク分布を一度訓練すれば、新しいタスクやデータセットへ移すときの手間がゼロから始めるより少ない。著者はこれを転移性と呼ぶ。モデルを複数の場面で鋭く保ちたいときに役立つ。

Megatron-LMの訓練パイプライン

本格的な訓練には、LLaMAモデルを複数GPUに分散するMegatron-LMのワークフローが付属する。リポジトリには、チェックポイントの変換、データのトークン化、SparseGPTやWanda、マグニチュードプルーニングによる事前マスクの生成、マスクの訓練、公開前のチェックポイント削減までをこなすスクリプトがある。

メリットとデメリット

メリット

  • 固定の重要度ルールに頼らずデータからマスクを学ぶため、同じスパース率で精度が上がる
  • コード、重み、計算済みマスクをオープンライセンスで公開し、結果を再現できる
  • 転移学習に対応し、1つのマスク分布を複数のタスクと領域で使い回せる
  • Hugging FaceとMegatron-LMを通じ、LLaMA-2やLLaMA-3、LLaMA-3.1など主要モデルで動く
  • 重みを凍結するため、学んだ2:4スパースは調整したタスクに対してロスレスになる

デメリット

  • 新しいマスクの訓練にはマルチGPU環境が要り、リポジトリは本格実行で1GPUあたり約40GBのメモリを示す
  • 研究用の成果物であり製品ではないため、サポート窓口もホスティングサービスもAPIもない
  • 既成マスクは少数のLLaMAモデル向けで、他のアーキテクチャは自分で訓練する必要がある
  • 効果が出るのは2:4加速に対応したハードウェアで、旧型やNVIDIA以外のGPUでは恩恵が限られる

よくある質問

MaskLLMは、大規模言語モデルにN:Mスパースを加える学習型プルーニング手法だ。各重みを採点する代わりに、マスクを確率分布としてモデル化し、残す重みを学ぶ。NVIDIA ResearchがNeurIPS 2024で発表し、MaskLLMのコードはGitHubで公開されている。