mamba

マンバ

Mamba: 選択的状態空間を使用した線形時間シーケンス モデリング
アルバート・グー*、トリ・ダオ*
論文: https://arxiv.org/abs/2312.00752

トランスフォーマーは SSM です: 一般化されたモデルと効率的なアルゴリズム
構造化された状態空間の二重性を通じて
トリ ダオ*、アルバート グ*
論文: https://arxiv.org/abs/2405.21060

について

Mamba は、言語モデリングなどの情報密度の高いデータに対して有望なパフォーマンスを示す新しい状態空間モデル アーキテクチャですが、以前の二次二次モデルでは Transformers には及ばません。これは、構造化状態空間モデルの進歩の流れに基づいており、FlashAttendant の精神に基づいた効率的なハードウェアを意識した設計と実装が行われています。

インストール

• [オプション] pip install causal-conv1d>=1.4.0 : Mamba ブロック内で使用される単純な因果 Conv1d レイヤーの効率的な実装。
• pip install mamba-ssm : コア Mamba パッケージ。
• pip install mamba-ssm[causal-conv1d] : コア Mamba パッケージと causal-conv1d をインストールします。
• pip install mamba-ssm[dev] : コア Mamba パッケージと開発依存関係をインストールします。

pip install .このリポジトリから。

pip PyTorch のバージョンについて不満を言う場合は、 --no-build-isolation pipに渡してみてください。

その他の要件:

• Linux
• NVIDIA GPU
• パイトーチ 1.12+
• CUDA 11.6+

AMD カードについては、以下の追加の前提条件を参照してください。

使用法

Mamba モデルとのいくつかのレベルのインターフェイスを公開します。

選択的 SSM

Mamba は選択的 SSM レイヤに基づいており、これがこの論文の焦点です (セクション 3; アルゴリズム 2)。

ソース: ops/selective_scan_interface.py。

マンバブロック

このリポジトリのメイン モジュールは、選択的 SSM をラップする Mamba アーキテクチャ ブロックです。

ソース: modules/mamba_simple.py。

使用法：

 import torch
from mamba_ssm import Mamba

batch , length , dim = 2 , 64 , 16
x = torch . randn ( batch , length , dim ). to ( "cuda" )
model = Mamba (
    # This module uses roughly 3 * expand * d_model^2 parameters
    d_model = dim , # Model dimension d_model
    d_state = 16 ,  # SSM state expansion factor
    d_conv = 4 ,    # Local convolution width
    expand = 2 ,    # Block expansion factor
). to ( "cuda" )
y = model ( x )
assert y . shape == x . shape

マンバ-2

Mamba-2 ブロックは modules/mamba2.py に実装されています。

より単純なバージョンは modules/mamba2_simple.py にあります。

使用法は Mamba(-1) と似ています。

 from mamba_ssm import Mamba2
model = Mamba2 (
    # This module uses roughly 3 * expand * d_model^2 parameters
    d_model = dim , # Model dimension d_model
    d_state = 64 ,  # SSM state expansion factor, typically 64 or 128
    d_conv = 4 ,    # Local convolution width
    expand = 2 ,    # Block expansion factor
). to ( "cuda" )
y = model ( x )
assert y . shape == x . shape 

SSD

「個別」SSM バージョンと「連続」SSM バージョン間の変換を備えた内部 SSD モジュールの最小バージョン (Mamba-2 論文のリスト 1) は modules/ssd_minimal.py にあります。

Mamba言語モデル

最後に、完全な言語モデルの例を示します。ディープ シーケンス モデル バックボーン (繰り返し Mamba ブロックを含む) + 言語モデル ヘッドです。

ソース: models/mixer_seq_simple.py。

これは、Mamba をエンドツーエンドのニューラル ネットワークに統合する方法の例です。この例は、以下の生成スクリプトで使用されます。

事前訓練されたモデル

事前トレーニングされたモデルは Hugging Face にアップロードされます: mamba-130m 、 mamba-370m 、 mamba-790m 、 mamba-1.4b 、 mamba-2.8b 、 mamba2-130m 、 mamba2-370m 、 mamba2-780m 、 mamba2-1.3b 、 mamba2-2.7b 、 transformerpp-2.7b 、 mamba2attn-2.7b (Pile 上の 300B トークンでトレーニング済み) およびmamba-2.8b-slimpj (SlimPajama データセット上の 600B トークンでトレーニング済み)。

モデルは以下の生成スクリプトによって自動ダウンロードされます。

これらのモデルは Pile でトレーニングされ、GPT-3 で記述され、多くのオープン ソース モデルが準拠している標準モデルの次元に従っています。

パラメータ	レイヤー	モデルが暗い。
130M	24	768
370M	48	1024
790M	48	1536年
1.4B	48	2048年
2.8B	64	2560

(Transformer の「レイヤー」(MHA ブロック + MLP ブロック) ごとに 2 つの Mamba ブロックが必要となるため、Mamba のレイヤー数は同様のサイズの Transformer のレイヤー数の 2 倍になります。)

注: これらは、ダウンストリームの変更 (命令の調整など) を一切行わずに、300B トークン専用にトレーニングされた基本モデルです。パフォーマンスは、同様のデータでトレーニングされた他のアーキテクチャと同等かそれ以上であることが期待されますが、大規模なモデルや微調整されたモデルには匹敵しません。

評価

モデルのゼロショット評価 (論文の表 3 に対応) を実行するには、lm-evaluation-harness ライブラリを使用します。

1. pip install lm-eval==0.4.2でlm-evaluation-harnessインストールします。
2. 次のコマンドで評価を実行します (詳細なドキュメントは lm-evaluation-harness リポジトリにあります)。

lm_eval --model mamba_ssm --model_args pretrained=state-spaces/mamba-130m --tasks lambada_openai,hellaswag,piqa,arc_easy,arc_challenge,winogrande,openbookqa --device cuda --batch_size 256
python evals/lm_harness_eval.py --model hf --model_args pretrained=EleutherAI/pythia-160m --tasks lambada_openai,hellaswag,piqa,arc_easy,arc_challenge,winogrande --device cuda --batch_size 64

ブログ投稿で報告されているmamba-2.8b-slimpjモデルで結果を再現するには、次のようにします。

lm_eval --model mamba_ssm --model_args pretrained=state-spaces/mamba-2.8b-slimpj --tasks boolq,piqa,hellaswag,winogrande,arc_easy,arc_challenge,openbookqa,race,truthfulqa_mc2 --device cuda --batch_size 256
lm_eval --model mamba_ssm --model_args pretrained=state-spaces/mamba-2.8b-slimpj --tasks mmlu --num_fewshot 5 --device cuda --batch_size 256

Mamba-2 モデルで評価を実行するには、モデル名を置き換えるだけです。

lm_eval --model mamba_ssm --model_args pretrained=state-spaces/mamba2-2.7b --tasks lambada_openai,hellaswag,piqa,arc_easy,arc_challenge,winogrande,openbookqa --device cuda --batch_size 256
lm_eval --model mamba_ssm --model_args pretrained=state-spaces/transformerpp-2.7b --tasks lambada_openai,hellaswag,piqa,arc_easy,arc_challenge,winogrande,openbookqa --device cuda --batch_size 256
lm_eval --model mamba_ssm --model_args pretrained=state-spaces/mamba2attn-2.7b --tasks lambada_openai,hellaswag,piqa,arc_easy,arc_challenge,winogrande,openbookqa --device cuda --batch_size 256

評価プロセスのノイズにより、各タスクの結果は報告された値と 0.1 ～ 0.3 異なる場合があることに注意してください。

推論

スクリプト benchmarks/benchmark_generation_mamba_simple.py

1. Hugging Face Hub からモデルを自動ロードします。
2. ユーザー指定のプロンプトの入力を生成します。
3. この世代の推論速度のベンチマークです。

その他の構成可能なオプションには、top-p (核サンプリング) 確率とソフトマックス温度が含まれます。

例

さまざまなサンプリング戦略を使用して生成レイテンシー (バッチ サイズ = 1 など) をテストするには、次のようにします。

python benchmarks/benchmark_generation_mamba_simple.py --model-name " state-spaces/mamba-2.8b " --prompt " My cat wrote all this CUDA code for a new language model and " --topp 0.9 --temperature 0.7 --repetition-penalty 1.2
python benchmarks/benchmark_generation_mamba_simple.py --model-name " EleutherAI/pythia-2.8b " --prompt " My cat wrote all this CUDA code for a new language model and " --topp 0.9 --temperature 0.7 --repetition-penalty 1.2
python benchmarks/benchmark_generation_mamba_simple.py --model-name " state-spaces/mamba-2.8b " --prompt " My cat wrote all this CUDA code for a new language model and " --minp 0.05 --topk 0 --temperature 0.7 --repetition-penalty 1.2

ランダムなプロンプト (大きなバッチサイズなど) を使用して生成スループットをテストするには:

python benchmarks/benchmark_generation_mamba_simple.py --model-name " state-spaces/mamba-2.8b " --batch 64
python benchmarks/benchmark_generation_mamba_simple.py --model-name " EleutherAI/pythia-2.8b " --batch 64

Mamba-2 では、モデル名を変更するだけです。

python benchmarks/benchmark_generation_mamba_simple.py --model-name " state-spaces/mamba2-2.7b " --prompt " My cat wrote all this CUDA code for a new language model and " --topp 0.9 --temperature 0.7 --repetition-penalty 1.2

トラブルシューティング

精度

私たちのモデルは、混合精度の PyTorch AMP を使用してトレーニングされました。 AMP はモデル パラメーターを float32 に保持し、必要に応じて半精度にキャストします。一方、DeepSpeed などの他のフレームワークは、パラメーターを float16 に保存し、必要に応じて (オプティマイザーの蓄積など) アップキャストを行います。

SSM は反復的なダイナミクスの影響を受けやすいため、メイン モデル パラメーターの精度を高める必要がある可能性があることがわかりました。不安定性が発生している場合は、最初のステップとして、fp32 にパラメータを保存するフレームワーク (AMP など) を試してください。

初期化

モデルの一部には、S4 モデルでの以前の作業から継承された初期化があります。たとえば、 $デルタ$パラメータは、その線形投影のバイアスを初期化することによってターゲット範囲を持ちます。ただし、一部のフレームワークには初期化後のフック ( nn.Linearモジュール内のすべてのバイアス項をゼロに設定するなど) がある場合があります。この場合、トレーニング フレームワークに固有のカスタム ロジックを追加する必要がある場合があります (たとえば、この行はトレーナーでは再初期化をオフにしますが、他のフレームワークでは何も行われません)。

AMD カードの追加の前提条件

ROCmのパッチ適用

ROCm 6.0 を使用している場合は、コンパイル中のエラーを回避するために次の手順を実行します。 ROCm 6.1 以降ではこれは必要ありません。

1. ROCm のインストール ディレクトリを見つけます。これは通常/opt/rocm/にありますが、インストールによって異なる場合があります。

2. パッチを適用します。権限の問題が発生した場合に備えて、 sudo使用して実行します。

    patch /opt/rocm/include/hip/amd_detail/amd_hip_bf16.h < rocm_patch/rocm6_0.patch 

引用

このコードベースを使用する場合、または私たちの仕事に価値があると思う場合は、Mamba を引用してください:

 @article{mamba,
  title={Mamba: Linear-Time Sequence Modeling with Selective State Spaces},
  author={Gu, Albert and Dao, Tri},
  journal={arXiv preprint arXiv:2312.00752},
  year={2023}
}

@inproceedings{mamba2,
  title={Transformers are {SSM}s: Generalized Models and Efficient Algorithms Through Structured State Space Duality},
  author={Dao, Tri and Gu, Albert},
  booktitle={International Conference on Machine Learning (ICML)},
  year={2024}
}