Zamba 7B V1 Phase1

🚀 Zambaのモデルカード

Zamba-7B-v1-phase1は、状態空間モデルであるMambaとTransformerのハイブリッドモデルです。6ブロックごとに共有のTransformer層を持つMambaバックボーンを使用しています。Zambaは次のトークン予測を使用して訓練され、Mistral v0.1トークナイザーを使用しています。このアーキテクチャは、小規模での一連のアブレーション実験の結果に基づいています。Zamba-7B-v1-phase-1は、オープンウェブデータセットから取得した1Tトークンのテキストとコードデータで事前学習されています。Zamba-v1とは異なり、このモデルはウェブデータセットのみでの純粋な事前学習後のチェックポイントを表しています。主にアニーリングプロセスの効果を調べるための比較ツールとしての使用を想定しています。

注意: 現在のHuggingfaceでのZambaの実装は、内部実装よりも低速です。Huggingfaceチームと協力してこの問題を解決するために取り組んでいます。

Zambaの訓練について説明する技術レポートはこちらで入手できます。

🚀 クイックスタート

前提条件

Zambaをダウンロードするには、Zyphraのtransformersのフォークをクローンします。

1. git clone https://github.com/Zyphra/transformers_zamba
2. cd transformers_zamba
3. リポジトリをインストールします: pip install -e .

CUDAデバイスで最適化されたMamba実装を実行するには、mamba-ssmとcausal-conv1dをインストールする必要があります。

pip install mamba-ssm causal-conv1d>=1.2.0

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

CPUで実行する場合は、AutoModelForCausalLM.from_pretrainedを使用してモデルをロードする際にuse_mamba_kernels=Falseを指定してください。

推論

from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

tokenizer = AutoTokenizer.from_pretrained("Zyphra/Zamba-7B-v1-phase1")
model = AutoModelForCausalLM.from_pretrained("Zyphra/Zamba-7B-v1-phase1", device_map="auto", 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]))

別のチェックポイントをロードするには、例えばイテレーション2500の場合は次のようにします。

model = AutoModelForCausalLM.from_pretrained("Zyphra/Zamba-7B-v1-phase1", device_map="auto", torch_dtype=torch.bfloat16, revision="iter2500")

デフォルトのイテレーションは、完全に訓練されたフェーズ1のモデルで、イテレーション462070に対応しています。これは、ランダムな初期化から始めてモデルを訓練した際のイテレーション数です。訓練の詳細についてはarXiv:2405.16712を参照してください。

📚 ドキュメント

モデルの詳細

Zambaは独自のハイブリッドSSMアーキテクチャを利用しています。このアーキテクチャは、Mamba層のバックボーンに共有のアテンション層が散在した構造になっています。このアテンションは共有重みを持っており、モデルのパラメータコストを最小化しています。元のモデル埋め込みをこのアテンションブロックの入力に連結することで、性能が向上することがわかっています。これはおそらく、深さにわたって情報がよりよく維持されるためです。

性能

この規模では、Zambaは既存のオープンモデル（オープンデータセットと訓練詳細を持つ）よりも大幅に優れた性能を発揮することがわかっています。ただし、7B規模の主要なオープン重みモデルよりはわずかに劣っています。この差の大部分は、MMLUと推論評価に由来しています。しかし、Zambaはこれらのモデルよりも大幅に少ないトークンで訓練されており、訓練トークンあたりの性能の点で最もサンプル効率の高いモデルです。

SSMアーキテクチャのため、Zambaは推論において非常に効率的で、同等の7Bおよび8Bモデルよりも大幅に優れた推論遅延と生成のメモリコストを実現しています。これは、大幅に縮小されたKVキャッシュによるものです。

引用

もしあなたの研究でZambaが役立った場合は、次のように引用してください。

@article{glorioso2024zamba,
  title={Zamba: A Compact 7B SSM Hybrid Model},
  author={Glorioso, Paolo and Anthony, Quentin and Tokpanov, Yury and Whittington, James and Pilault, Jonathan and Ibrahim, Adam and Millidge, Beren},
  journal={arXiv preprint arXiv:2405.16712},
  year={2024}
}

注意事項

Zambaは事前学習されたベースモデルであり、モデレーションメカニズムを持っていません。また、このモデルはチャット用にファインチューニングされていないため、良好なチャット性能を期待することはできません。

📄 ライセンス

このモデルはApache-2.0ライセンスの下で提供されています。