Zamba2 2.7B

🚀 Zamba2-2.7Bモデルカード

Zamba2-2.7Bは、状態空間とTransformerモジュールから構成される混合モデルです。このモデルは、大まかにZambaアーキテクチャに従っており、このアーキテクチャはMambaのメイン部分と共有Transformerモジュールが交互に配置されています（詳細はモデルの詳細の図を参照）。Zamba1と比較すると、Zamba2-2.7Bには3つの主要な改良点があります：

1. Mamba1モジュールをMamba2モジュールに置き換えました。
2. 単一の共有注意力モジュールを使用する代わりに、2つの共有注意力モジュールを採用し、これらはABABパターンでネットワーク全体にわたって使用されます。
3. 各共有MLPモジュールにLoRA投影器を適用しました。これにより、ネットワークは深さ方向で共有層を呼び出す際にMLPを特化させることができます。LoRAを使用することで、総パラメータ数をわずかに増やすだけで、深さ方向の特化を実現できます。

Zamba2-2.7BはMistral v0.1トークナイザーを使用し、オープンなネットワークデータセット（Zydaを含む）から取得した3Tのテキストとコードデータで事前学習されています。その後、第2段階では、100Bの高品質トークンの混合データでアニーリング処理が行われました。

⚠️ 重要提示

これはZamba2-2.7BのHuggingFace上の一時的な実装です。すべてのHuggingFaceモデルとの相互作用を目的としたフレームワークやツールと完全に互換性があるとは限りません。

Zamba2-2.7Bの独立したPyTorch実装はこちらで見ることができます。

🚀 クイックスタート

前提条件

Zamba2-2.7Bを使用するには、ソースコードからtransformersをインストールする必要があります。

1. git clone https://github.com/huggingface/transformers.git
2. cd transformers && pip install .

Mamba2カーネルを実行するために必要な依存関係をインストールするには、ソースコードからmamba-ssm（PyTorchとの互換性の問題があるため）とcausal-conv1dをインストールする必要があります。

1. git clone https://github.com/state-spaces/mamba.git
2. cd mamba && git checkout v2.1.0 && pip install .
3. pip install causal-conv1d

最適化されたMamba2カーネルを使用せずにモデルを実行することもできますが、推奨されません。これは、遅延が大幅に増加し、メモリ使用量が増えるためです。

推論

from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

tokenizer = AutoTokenizer.from_pretrained("Zyphra/Zamba2-2.7B")
model = AutoModelForCausalLM.from_pretrained("Zyphra/Zamba2-2.7B", device_map="cuda", torch_dtype=torch.bfloat16)

input_text = "What factors contributed to the fall of the Roman Empire?"
input_ids = tokenizer(input_text, return_tensors="pt").to("cuda")

outputs = model.generate(**input_ids, max_new_tokens=100)
print(tokenizer.decode(outputs[0]))

✨ 主な機能

• アーキテクチャの革新：混合状態空間とTransformerアーキテクチャを採用し、Mambaモジュールと共有注意力モジュールを組み合わせて性能を向上させます。
• パラメータの最適化：共有注意力重みとLoRA投影器を使用することで、パラメータ数を抑えながら深さ方向の特化を実現します。
• 高性能：30億未満のパラメータを持つモデルの中でトップクラスの性能を発揮し、より大規模なモデルとも競争力があります。
• 低遅延と低メモリ使用量：独自の混合SSMアーキテクチャにより、極低い推論遅延、高速な生成速度、および少ないメモリ使用量を実現します。

📦 インストール

transformersのインストール

git clone https://github.com/huggingface/transformers.git
cd transformers && pip install .

Mamba2カーネルの依存関係のインストール

git clone https://github.com/state-spaces/mamba.git
cd mamba && git checkout v2.1.0 && pip install .
pip install causal-conv1d

💻 使用例

基本的な使用法

from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

tokenizer = AutoTokenizer.from_pretrained("Zyphra/Zamba2-2.7B")
model = AutoModelForCausalLM.from_pretrained("Zyphra/Zamba2-2.7B", device_map="cuda", torch_dtype=torch.bfloat16)

input_text = "What factors contributed to the fall of the Roman Empire?"
input_ids = tokenizer(input_text, return_tensors="pt").to("cuda")

outputs = model.generate(**input_ids, max_new_tokens=100)
print(tokenizer.decode(outputs[0]))

📚 ドキュメント

モデルの詳細

Zamba2-2.7Bは、最初のZamba混合SSM-注意力アーキテクチャを利用し、拡張しています。核心となるZambaアーキテクチャは、Mamba層のメイン部分と1つまたは複数の共有注意力層が交互に配置されています（Zamba1では1つの共有注意力層、Zamba2では2つ）。この注意力は共有重みを持ち、モデルのパラメータコストを最小化します。元のモデル埋め込みをこの注意力モジュールの入力に接続することで性能が向上することがわかっており、これは深さ方向での情報の維持が良好になるためと考えられます。Zamba2アーキテクチャはまた、共有MLPにLoRA投影行列を適用して、各モジュールで追加の表現力を得るとともに、各共有モジュールが独自の位置に応じて微調整できるようにし、追加のパラメータコストを抑えています。

性能表現

Zamba2-2.7Bは、30億未満のパラメータを持つモデルの中でトップクラスの最先端性能を達成しており、はるかに大規模なモデルとも競争力があります。さらに、独自の混合SSMアーキテクチャにより、極低い推論遅延と高速な生成を実現し、同等のTransformerベースのモデルと比較してメモリ使用量が大幅に少なくなっています。

Zamba2-2.7Bの高性能と少ない推論計算およびメモリ使用量は、デバイス上でのアプリケーションに最適な汎用モデルとしての魅力を高めています。

最初のトークン時間 (TTFT)	出力生成

🔧 技術詳細

Zamba2-2.7Bの技術的な実装詳細には以下のものが含まれます：

• アーキテクチャ設計：混合状態空間とTransformerアーキテクチャを採用し、Mambaモジュールと共有注意力モジュールを組み合わせて性能と効率を向上させます。
• パラメータ最適化：共有注意力重みとLoRA投影器を使用することで、パラメータ数を抑えながら深さ方向の特化を実現します。
• データ処理：Mistral v0.1トークナイザーを使用し、3Tのテキストとコードデータで事前学習し、100Bの高品質トークンの混合データでアニーリング処理を行います。

📄 ライセンス

このプロジェクトはApache-2.0ライセンスの下で提供されています。