Bamba-9B: 推論効率の高いハイブリッド Mamba2 モデル

TL;DR

Bamba-9B は、IBM、Princeton、CMU、UIUC によってリリースされた推論効率の高いハイブリッド Mamba2 モデルです。vLLM において標準的な Transformer と比較して 2.5 倍のスループットと 2 倍のレイテンシ向上を示しており、transformers、vLLM、TRL、llama.cpp で即座に使用可能です。

動機

Transformer の推論は、コンテキスト長とともに増大する KV キャッシュのボトルネックによって制限されます。ハイブリッド Mamba2 アーキテクチャは KV キャッシュのサイズを一定に保つことで、このボトルネックに対処します。Bamba-9B は、完全にオープンなデータを使用して 7B-10B スケールでハイブリッド Mamba2 アプローチを検証し、コミュニティの実験を促進するために再現可能なチェックポイントを提供します。

transformers での使用

🤗 Transformers ライブラリで Bamba-9B を実行するには、AutoModelForCausalLM と AutoTokenizer でモデルとトークナイザーをロードし、generate を呼び出します。サンプルコード:

from transformers import AutoModelForCausalLM, AutoTokenizer
model = AutoModelForCausalLM.from_pretrained("ibm-fms/Bamba-9B
tokenizer = AutoTokenizer.from_pretrained("ibm-fms/Bamba-9B
txt = ["Mamba is a snake with following properties  
inputs = tokenizer(txt, return_tensors='pt', return_token_type_ids=False)
out = model.generate(**inputs, max_new_tokens=64)
print(tokenizer.batch_decode(out, skip_special_tokens=True)[0])

評価

SoTA Transformer モデルとの比較

Bamba-9B は HF OpenLLM v1 リーダーボードで平均 62.31 を記録し、Meta Llama 3.1 8B (63.51) をわずかに下回りますが、一部の指標では Olmo2 7B (66.17) や IBM Granite v3 8B (67.47) を上回っています。v2 リーダーボードでは、Bamba-9B は平均 10.91 で、Llama 3.1 8B (14.27) より低いですが、数学および MMLU タスクを除外すると同等になります。その場合、平均は Llama 3.1 8B の 44.68 に対し、Bamba-9B は 45.53 となります。

同様のトークン予算でトレーニングされた Transformer モデルとの比較

Bamba-9B (2.2T トークン) は平均 62.31 を記録し、2T トークンでトレーニングされた Olmo1.5 7B (55.8) を上回っています。2T トークンの Bamba-9B チェックポイント (59.11) を Llama2 7B (53.78) および IBM Granite 7B (52.07) と比較した場合でも、Bamba-9B は高い数値を維持しており、同等のデータ量であれば競争力があることを示しています。

他の Mamba/Mamba2 モデルとの比較

Bamba-9B は平均 62.31 で、一部の指標では NVIDIA Mamba2 Hybrid 8B* (58.78) や Zamba 7B (64.36) よりも高く、Falcon Mamba 7B (65.31) よりは低くなっています。表は、ハイブリッド Mamba2 モデルが、理論上最大 5 倍の推論効率を提供しながら、競争力のある結果を出せることを示しています。

推論効率

Bamba-9B は、NVIDIA H100 80GB GPU 上の vLLM を使用して、1K から 64K トークンのバッチサイズおよびシーケンス長にわたる測定において、Meta Llama 3.1 8B よりも最大 2.5 倍のスループット向上と 2 倍のレイテンシ低減を達成しています。演算強度分析では、デコード段階がメモリ制限（memory-bound）になった場合、最大 5 倍のスピードアップの可能性があることが予測されています。現在の vLLM の結果は、チャンク化された pre-fill のサポート不足、Transformer スタイルのメモリ割り当て、および H100 用に最適化されていない Mamba2 カーネルによって制限されています。

モデルアーキテクチャ

Bamba-9B は合計 32 層を使用しています：3 層のフルアテンション層と 29 層の Mamba2 層。MLP 拡張係数は 3.5、語彙サイズは 128k、RoPE エンベディング、および GQA (8 KV ヘッド、32 ヘッド) を備えています。NVIDIA のハイブリッド Mamba2 8B モデルと比較して、Bamba-9B はアテンション層を 4 層から 3 層に減らし、RoPE を追加しています。

データ

トレーニングには、第1フェーズで Dolma v1.7 を使用し、その後追加の 200B トークンとして FineWeb-edu と Cosmopedia を使用しました。すべてのデータは、Ray フレームワークを使用して内部の Red Hat OpenShift クラスター上でトークン化されました。第1フェーズのデータミックスはブログ投稿に可視化されています。

事前学習

事前学習は段階的に進められました：1.8B/100B トークンでのアブレーション、次に Dolma を使用した 3B/2T トークン、続いて 9B/2T トークンの実行、最後に FineWeb-edu と Cosmopedia を使用した 200B トークンのファインチューニングフェーズです。トレーニングハイパーパラメータ：cosine LR スケジュール、ピーク 3e-4、2000 ステップにわたる二次関数的ウォームアップ、減衰 0.033、終了 LR 1e-5、AdamW (β1=0.9, β2=0.95)、ウェイトデケイ 0.1、シーケンス長 4096、グローバルバッチサイズ 1.5M トークン、IBM Cloud Vela 上の 192 枚の A100 GPU を使用して約 2 か月間。デプロイエラーとハードウェア障害により、Autopilot システムによって検出された 3 回のジョブ中断が発生しました。

データローダー

リリースされたステートフルなデータローダーは、チェックポイントからの再開、自動スケーリング、オーバーヘッドなしのシャッフルストリーミング、ピア・ツー・ピア・トラフィックなしの非同期分散操作、動的なデータミキシングとオンザフライのトークン化を提供し、PyTorch ネイティブでモジュール化され、拡張可能です。数百のトレーニングジョブで実証され、Torch Titan と統合されています。

量子化

llm-compressor を使用した FMS Model Optimizer フレームワークを使用すると、Bamba-9B のチェックポイントは fp8 に量子化され、精度低下は無視できる程度でした：OpenLLM v1 の平均は 62.31 から 61.5 (-0.1) に、v2 の平均は 10.91 から 10.04 (-0.9) に低下しました。vLLM での fp8 推論の有効化は、Mamba2 層のカーネルアップデート待ちとなっています。

コンテキスト長の拡張

フルアテンション層に LongRoPE を適用することで、Bamba-9B のコンテキスト長を拡張できます。予備的な PhoneBook 検索テストでは、拡張されたモデルはチューニングなしで最大 16K トークンにおいてベースの Bamba-9B、Llama2-7B、Llama3-8B を上回り、Llama3.1-8B の性能に匹敵することを示しています。32K トークンでは Llama3.1-8B がリードしています。

まとめ

IBM、Princeton、CMU、UIUC による、2.2T のオープンなトークンでトレーニングされたハイブリッド Mamba2 モデルである Bamba-9B は、vLLM において Llama 3.1 8B に対して 2.5 倍のスループットと 2 倍のレイテンシ向上を実現し、transformers、vLLM、TRL、llama.cpp で即座に使用可能です。また、トレーニング、チューニング、拡張事前学習のレシピ、およびステートフルなデータローダーが付属しています。

今後の課題

追加データによる継続的な事前学習、コミュニティが提案するミックスによる SFT、Tulu-3、Orca-AgentInstruct、Daring-Anteater データセットを使用した教師ありファインチューニング、vLLM でのチャンク化された pre-fill と適切なメモリ割り当ての有効化、より高速な推論のための fp8 カーネルの追加、torch.compile と fp8 トレーニングの適用、およびコンテキスト長を 1M+ トークンまで拡張する計画が含まれています。

貢献者

データ収集およびキュレーション：AllenAI (Dolma) および Hugging Face (FineWeb-edu, Cosmopedia)。 データ前処理：IBM チームメンバー Tuan Hoang Trong, Syed Zawad, Jay Gala, Ryan Gordon (IBM Data Prep Kit を使用)。 モデルアーキテクチャ：Tri Dao (Princeton), Albert Gu (CMU), Linsong Chu (IBM), Davis Wertheimer (IBM), Minjia Zhang (UIUC), Mudhakar Srivatsa (IBM), Raghu Ganti (IBM)。 モデルトレーニング：IBM チームメンバー Linsong Chu, Divya Kumari, Davis Wertheimer, Raghu Ganti, Dakshi Agrawal。 モデルチューニング：IBM チームメンバー Sukriti Sharma, Anh Uong (TRL 経由)。 モデル推論：IBM およびコミュニティ貢献者 Fabian Lim, Antoni Viros i Martin, Adnan Hoque, Jamie Yang, Nelson Nimura Gonzalez, Joshua Rosenkranz, Nick Hill, Gabe Goodhart。 量子化：IBM チームメンバー Naigang Wang, Charlie Liu。 評価：Yotam Perlitz, Ofir Arviv, Michal Shmueli-Scheuer, Haoechen Shen, Minjia Zhang (UIUC) が率いる IBM 評価チーム。 リーダーシップへの謝辞：Priya Nagpurkar, David Cox, Sriram Raghavan, Aya Soffer, Ruchir Puri, Mukesh Khare。 コミュニティへの感謝：Pablo Montalvo-Leroux, Aritra Roy Gosthipaty, Vaibhav Srivastav (Hugging Face), Stas Bekman (Contextual AI), Tyler Michael Smith (Neural Magic)。 オープンソースへの貢献に対し、Meta PyTorch、AllenAI、Hugging Face にも感謝いたします。

付録：演算強度

付録では、アテンションおよび Bamba モデルの計算とメモリの式を導出し、Bamba-9B のデコード段階のメモリ上の利点が、長いシーケンス（>16K トークン）において Llama よりも最大 5 倍のスピードアップをもたらす可能性があることを示しています。現在の vLLM での 2.5 倍のスループットと 2 倍のレイテンシの測定値は、チャンク化された pre-fill、Transformer スタイルのメモリ割り当て、および H100 上での未最適化の Mamba2 カーネルによって制限されています。