Transformers documentation
Mamba 2
このモデルは2024年5月31日にHF papersで公開され、2024年8月6日にHugging Face Transformersに貢献されました。
Mamba 2
Mamba 2は、構造化状態空間モデル（SSM）とアテンションのバリアントを結びつける状態空間対偶性（SSD）フレームワークに基づいています。Mambaよりも2〜8倍高速なより効率的なSSDアルゴリズムを使用し、テンソル並列処理を可能にするようアーキテクチャを変更するとともに、グループ化値アテンション（GVA）ヘッド構造を導入しています。
すべてのオリジナルの Mamba 2 チェックポイントは State Space Models 組織配下で確認できますが、以下の例では、オリジナルのチェックポイントにはまだ Hugging Face での実装がサポートされていないため、mistralai/Mamba-Codestral-7B-v0.1 を使用しています。
Mamba 2をベースとした他のアーキテクチャには、Bamba、FalconH1、Zamba2があります。
このモデルは ArthurZ によって提供されました。 Mamba を異なる言語タスクに適用する方法のより多くの例については、右側のサイドバーにある Mamba モデルをクリックしてください。
以下の例では、Pipeline、AutoModel、およびコマンドラインからテキストを生成する方法を示しています。
from transformers import pipeline
pipeline = pipeline(
task="text-generation",
model="mistralai/Mamba-Codestral-7B-v0.1",
device=0
)
pipeline("Plants create energy through a process known as")from transformers import AutoModelForCausalLM, AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained("mistralai/Mamba-Codestral-7B-v0.1")
model = AutoModelForCausalLM.from_pretrained("mistralai/Mamba-Codestral-7B-v0.1", device_map="auto")
input_ids = tokenizer("Plants create energy through a process known as", return_tensors="pt").to(model.device)
output = model.generate(**input_ids)
print(tokenizer.decode(output[0], skip_special_tokens=True))Quantization は、重みを低精度で表現することで、大規模モデルのメモリ負荷を軽減します。利用可能な Quantization バックエンドの詳細については、Quantization の概要を参照してください。
以下の例では、torchaoを使用して、重みを4ビット整数にのみ量子化しています。
from transformers import AutoModelForCausalLM, AutoTokenizer, TorchAoConfig
quantization_config = TorchAoConfig("int4_weight_only", group_size=128)
tokenizer = AutoTokenizer.from_pretrained("mistralai/Mamba-Codestral-7B-v0.1")
model = AutoModelForCausalLM.from_pretrained("mistralai/Mamba-Codestral-7B-v0.1", quantization_config=quantization_config, device_map="auto")
input_ids = tokenizer("Plants create energy through a process known as", return_tensors="pt").to(model.device)
output = model.generate(**input_ids)
print(tokenizer.decode(output[0], skip_special_tokens=True))注記
Codestral Mamba は
groups=8を持ち、これは attention ベースのモデルにおける kv heads の数に相当します。Codestral Mambaには2つの異なるフォワードパス、
torch_forwardまたはcuda_kernels_forwardがあり、その結果はわずかに異なることが期待されます。このモデルには位置埋め込みはありませんが、
attention_maskが存在し、バッチ生成の場合には2つの場所で隠れ状態をマスクするための特定のロジックが実装されています（詳細はこのコメントを参照）。これと、再実装された Mamba 2 カーネルの追加により、バッチ生成とキャッシュ付き生成の間にわずかな差異が生じます。SSMアルゴリズムはテンソル収縮に大きく依存しており、これにはmatmulの同等物がありますが、演算の順序がわずかに異なります。そのため、精度が低い場合、その差がより大きくなります。
Padding は左側 padding のみでテストされています。右側 padding はキャッシュに pad 位置を更新するため、信頼性の低い結果を返します。バッチ生成を行う場合は、
tokenizer.padding_side="left"を設定してください。Padding sideを参照してください。以下の例では、PEFTを使用してMamba 2をファインチューニングする方法を示しています。
from datasets import load_dataset
from peft import LoraConfig
from trl import SFTConfig, SFTTrainer
model_id = "mistralai/Mamba-Codestral-7B-v0.1"
dataset = load_dataset("Abirate/english_quotes", split="train")
training_args = SFTConfig(dataset_text_field="quote", gradient_checkpointing=True, per_device_train_batch_size=4)
lora_config = LoraConfig(target_modules=["x_proj", "embeddings", "in_proj", "out_proj"])
trainer = SFTTrainer(
model=model_id,
args=training_args,
train_dataset=dataset,
peft_config=lora_config,
)
trainer.train()Mamba2Config
クラス transformers.Mamba2Config
< source >( transformers_version: str | None = Nonearchitectures: list[str] | None = Noneoutput_hidden_states: bool | None = Falsereturn_dict: bool | None = Truedtype: str | torch.dtype | None = Nonechunk_size_feed_forward: int = 0is_encoder_decoder: bool = Falseid2label: dict[int, str] | dict[str, str] | None = Nonelabel2id: dict[str, int] | dict[str, str] | None = Noneproblem_type: Literal['regression', 'single_label_classification', 'multi_label_classification'] | None = Nonenum_heads: int = 128head_dim: int = 64vocab_size: int = 32768hidden_size: int = 4096state_size: int = 128num_hidden_layers: int = 64layer_norm_epsilon: float = 1e-05pad_token_id: int | None = 1bos_token_id: int | None = 0eos_token_id: int | list[int] | None = 2expand: int = 2conv_kernel: int = 4n_groups: int = 8use_bias: bool = Falseuse_conv_bias: bool = Truehidden_act: str = 'silu'initializer_range: float = 0.1residual_in_fp32: bool = Truetime_step_rank: str | int = 'auto'time_step_min: float = 0.001time_step_max: float = 0.1time_step_floor: float = 0.0001time_step_limit: list[float] | tuple[float, ...] = (0.0, inf)rescale_prenorm_residual: bool = Falseuse_cache: bool = Truechunk_size: int = 256tie_word_embeddings: bool = False )
パラメータ
- num_heads (
int、任意、デフォルトは128) — Transformerデコーダーの各アテンション層におけるアテンションヘッドの数。 - head_dim (
int、任意、デフォルトは64) — アテンションヘッドの次元です。None の場合、hidden_size // num_attention_heads がデフォルトとして使用されます。 - vocab_size (
int、任意、デフォルトは32768) — モデルの語彙サイズです。input_idsで表現できるトークンの種類数を定義します。 - hidden_size (
int、任意、デフォルトは4096) — 隠れ表現の次元数。 - state_size (
int、任意、デフォルトは128) — Mamba レイヤーにおける SSM 状態のサイズ（潜在状態の次元数）。 - num_hidden_layers (
int、任意、デフォルトは64) — Transformer デコーダーの隠れ層の数。 - layer_norm_epsilon (
float、任意、デフォルト値は 1e-05) — Layer Normalization 層で使用する epsilon です。 - pad_token_id (
int、任意、デフォルトは1) — 語彙内でパディングに使用されるトークン ID。 - bos_token_id (
int、任意、デフォルトは0) — 語彙におけるストリーム開始に使用されるトークン ID。 - eos_token_id (
Union[int, list[int]]、任意、デフォルトは2) — 語彙内でストリーム終了に使用されるトークン ID。 - expand (
int、任意、デフォルトは 2) — 中間サイズを決定するために使用される拡張係数。 - conv_kernel (
int、任意、デフォルトは4) — 畳み込みカーネルのサイズです。 - n_groups (
int、任意、デフォルトは 8) — mamba 2 の進化行列のグループ数です。 - use_bias (
bool、任意、デフォルトはFalse) — mixerブロックの["in_proj", "out_proj"]でバイアスを使用するかどうか - use_conv_bias (
bool、任意、デフォルトはTrue) — ミキサーブロックの畳み込み層でバイアスを使用するかどうか。 - hidden_act (
str、任意、デフォルトはsilu) — デコーダーで使用する非線形活性化関数（関数または文字列）。例えば、"gelu"、"relu"、"silu"など。 - initializer_range (
float、任意、デフォルトは0.1) — すべての重み行列の初期化に使用される truncated_normal_initializer の標準偏差です。 - residual_in_fp32 (
bool、省略可能、デフォルトはTrue) — 残差をfloat32で保持するかどうかを指定します。Falseに設定すると、残差はモデルの他の部分と同じdtypeが維持されます。 - time_step_rank (
Union[str, int]、任意、デフォルトはauto) — デルタ（タイムステップ）投影のランクです。"auto"を指定すると自動的に設定されます。 - time_step_min (
float、任意、デフォルトは0.001) —dt_proj.biasを制限するために使用される最小のtime_step。 - time_step_max (
float、任意、デフォルトは0.1) —dt_proj.biasを制限するために使用されるtime_stepの最大値。 - time_step_floor (
float、任意、デフォルトは0.0001) — softplus 活性化後の離散タイムステップのデルタに許可される最小値です。 - time_step_limit (
Union[list[float], tuple[float, ...]]、省略可、デフォルトは(0.0, inf)) — クランプに使用するタイムステップ値の許容範囲。 - rescale_prenorm_residual (
bool、任意、デフォルトはFalse) — 初期化時にout_projの重みをリスケールするかどうか。 - use_cache (
bool、任意、デフォルトはTrue) — モデルが最後のキー/バリューの注意を返すかどうか（すべてのモデルで利用されるわけではありません）。config.is_decoder=Trueの場合、またはモデルがデコーダー専用の生成モデルである場合にのみ関連します。 - chunk_size (
int、任意、デフォルトは 256) — シーケンスを構成するチャンクのサイズ。 - tie_word_embeddings (
bool、任意、デフォルトはFalse) — モデルのtied_weights_keysマッピングに従って重み埋め込みを共有するかどうかを指定します。
これはMamba2Modelの設定を保存するための設定クラスです。指定された引数に基づいてMamba2モデルをインスタンス化し、モデルのアーキテクチャを定義するために使用されます。デフォルト値で設定をインスタンス化すると、state-spaces/mamba2-2.8bと類似した設定が得られます。
設定オブジェクトはPreTrainedConfigを継承し、モデルの出力を制御するために使用できます。詳細については、PreTrainedConfigのドキュメントを参照してください。
例:
>>> from transformers import Mamba2Config, Mamba2Model
>>> # Initializing a Mamba2 configuration
>>> configuration = Mamba2Config()
>>> # Initializing a model (with random weights) from the configuration
>>> model = Mamba2Model(configuration)
>>> # Accessing the model configuration
>>> configuration = model.configMamba2Model
クラス transformers.Mamba2Model
< source >( config: Mamba2Config )
パラメータ
- config (Mamba2Config) — モデルのすべてのパラメータを含むモデル設定クラスです。設定ファイルで初期化しても、モデルに関連する重みは読み込まれず、設定のみが読み込まれます。モデルの重みをロードするには、from_pretrained() メソッドを参照してください。
特定のヘッドを付加せず、生の隠れ状態を出力する素のMamba2モデル。
このモデルは PreTrainedModel を継承しています。すべてのモデルに実装されている汎用的なメソッド（ダウンロードや保存、入力埋め込みのリサイズ、ヘッドのプルーニングなど）については、スーパークラスのドキュメントを参照してください。
このモデルは、PyTorchのtorch.nn.Moduleのサブクラスでもあります。 通常のPyTorch Moduleとして使用し、一般的な使用方法や動作に関する詳細については、PyTorchのドキュメントを参照してください。
forward
< source >( input_ids: typing.Optional[torch.LongTensor] = Noneinputs_embeds: typing.Optional[torch.LongTensor] = Nonecache_params: transformers.cache_utils.Cache | None = Noneuse_cache: bool | None = Noneoutput_hidden_states: bool | None = Nonereturn_dict: bool | None = Noneattention_mask: typing.Optional[torch.Tensor] = None**kwargs ) → Mamba2Output or tuple(torch.FloatTensor)
パラメータ
- input_ids (
torch.LongTensorの形状(batch_size, sequence_length), 、省略可) — 語彙における入力シーケンスのトークンインデックスです。デフォルトではパディングは無視されます。インデックスはAutoTokenizerを使用して取得できます。詳細については、PreTrainedTokenizer.encode()およびPreTrainedTokenizer.call()を参照してください。
- inputs_embeds (
torch.LongTensor、shape(batch_size, sequence_length, hidden_size)、任意) —input_idsを渡す代わりに、埋め込み表現を直接渡すこともできます。これは、モデル内部の埋め込みルックアップ行列よりも、input_idsのインデックスを対応するベクトルに変換する方法をより細かく制御したい場合に有用です。 - cache_params (
Cache、任意) — 指定された場合、モデルはすべてのブロックで前回の状態を使用します（input_idsに対して、モデルがstate_input_ids + input_idsをコンテキストとして追加したかのような出力が得られます）。 - use_cache (
bool、省略可) —Trueに設定すると、cache_paramsが返され、次の logits を高速に生成するために使用できます。 - output_hidden_states (
bool, 省略可) — すべての層の hidden states を返すかどうかを指定します。詳細は、返されるテンソルのhidden_statesを参照してください。 - return_dict (
bool, 任意) — プレーンなタプルではなく ModelOutput を返すかどうか。 - attention_mask (
torch.Tensorの形状(batch_size, sequence_length), 、任意) — パディングトークンのインデックスに対するアテンションを回避するためのマスク。[0, 1]:- マスクされていないトークンに対して1
- マスクされたトークンに対しては 0 です。
戻り値
Mamba2Output または tuple(torch.FloatTensor)
Mamba2Output、またはtorch.FloatTensorのタプル（return_dict=Falseが渡された場合、またはconfig.return_dict=Falseの場合）。設定（Mamba2Config）や入力に応じて、さまざまな要素で構成されます。
Mamba2Modelのforwardメソッドは、__call__特殊メソッドをオーバーライドします。
この関数内でフォワードパスの手順を定義する必要がありますが、その後、この関数ではなく
Moduleインスタンスを呼び出すべきです。前者は前処理と後処理の実行を適切に処理するのに対し、後者はそれらを黙って無視するためです。
last_hidden_state (
torch.FloatTensor、shape(batch_size, sequence_length, hidden_size)、optional) — モデルの最終層の出力における隠れ状態のシーケンス。cache_params (
~cache_utils.Cache、任意) — 最終タイムステップにおけるモデルの状態。次のinput_idsと forward メソッドで使用することで、古いinput_idsを提供することを回避できます。選択的スキャン後のState space modelの状態行列と、Convolutional statesの両方を含みます
hidden_states (
tuple[torch.FloatTensor]、optional、output_hidden_states=Trueが渡された場合またはconfig.output_hidden_states=Trueの場合に返される) —torch.FloatTensorのタプル（埋め込み層を持つモデルの場合、埋め込みの出力用が 1 つ、各層の出力用が 1 つずつ）で、形状は(batch_size, sequence_length, hidden_size)です。各層の出力におけるモデルのHidden-states、および任意の初期埋め込み出力。
Mamba2LMHeadModel
クラス transformers.Mamba2ForCausalLM
< source >( config: Mamba2Config )
パラメータ
- config (Mamba2Config) — モデルのすべてのパラメータを含むモデル設定クラスです。設定ファイルで初期化しても、モデルに関連する重みは読み込まれず、設定のみが読み込まれます。モデルの重みをロードするには、from_pretrained() メソッドを参照してください。
MAMBA2 モデルの Transformer で、上部に言語モデリングヘッド（入力埋め込みの重みと共有されない線形層）を備えています。
このモデルは PreTrainedModel を継承しています。すべてのモデルに実装されている汎用的なメソッド（ダウンロードや保存、入力埋め込みのリサイズ、ヘッドのプルーニングなど）については、スーパークラスのドキュメントを参照してください。
このモデルは、PyTorchのtorch.nn.Moduleのサブクラスでもあります。 通常のPyTorch Moduleとして使用し、一般的な使用方法や動作に関する詳細については、PyTorchのドキュメントを参照してください。
forward
< source >( input_ids: typing.Optional[torch.LongTensor] = Noneinputs_embeds: typing.Optional[torch.FloatTensor] = Nonecache_params: transformers.cache_utils.Cache | None = Nonelabels: typing.Optional[torch.LongTensor] = Noneoutput_hidden_states: bool | None = Nonereturn_dict: bool | None = Noneuse_cache: bool | None = Noneattention_mask: typing.Optional[torch.Tensor] = Nonelogits_to_keep: typing.Union[int, torch.Tensor] = 0**kwargs ) → Mamba2CausalLMOutput or tuple(torch.FloatTensor)
パラメータ
- input_ids (
torch.LongTensorの形状(batch_size, sequence_length), 、省略可) — 語彙における入力シーケンスのトークンインデックスです。デフォルトではパディングは無視されます。インデックスはAutoTokenizerを使用して取得できます。詳細については、PreTrainedTokenizer.encode()およびPreTrainedTokenizer.call()を参照してください。
- inputs_embeds (
torch.FloatTensor、形状(batch_size, sequence_length, hidden_size)、任意) —input_idsを渡す代わりに、埋め込み表現を直接渡すこともできます。これは、モデル内部の埋め込みルックアップ行列よりも、input_idsのインデックスを対応するベクトルに変換する方法をより細かく制御したい場合に有用です。 - cache_params (
Cache、任意) — 指定された場合、モデルはすべてのブロックで前回の状態を使用します（input_idsに対して、モデルがstate_input_ids + input_idsをコンテキストとして追加したかのような出力が得られます）。 - labels (
torch.LongTensor、形状(batch_size, sequence_length)、任意) — 言語モデリング用のラベルです。モデル内部ではラベルがシフトされるため、labels = input_idsと設定できます。インデックスは[-100, 0, ..., config.vocab_size]から選択されます。-100に設定されたすべてのラベルは無視（マスク）され、損失は[0, ..., config.vocab_size]内のラベルに対してのみ計算されます。 - output_hidden_states (
bool, 省略可) — すべての層の hidden states を返すかどうかを指定します。詳細は、返されるテンソルのhidden_statesを参照してください。 - return_dict (
bool, 任意) — プレーンなタプルではなく ModelOutput を返すかどうか。 - use_cache (
bool、省略可) —Trueに設定すると、cache_paramsが返され、次の logits を高速に生成するために使用できます。 - attention_mask (
torch.Tensorの形状(batch_size, sequence_length), 、任意) — パディングトークンのインデックスに対するアテンションを回避するためのマスク。[0, 1]:- マスクされていないトークンに対して1
- マスクされたトークンに対しては 0 です。
- logits_to_keep (
Union[int, torch.Tensor]、省略可、デフォルトは0) —intの場合、最後のlogits_to_keepトークンについて logits を計算します。0の場合、すべてのinput_idsについて logits を計算します（特殊ケース）。生成には最後のトークンの logits だけで十分であり、そのトークンのみについて計算することでメモリを節約できます。これは、長いシーケンスや大規模な語彙サイズの場合に特に顕著な効果があります。torch.Tensorの場合、シーケンス長次元で保持するインデックスに対応する 1 次元テンソルである必要があります。これは、パックされたテンソル形式（バッチとシーケンス長を単一の次元にまとめた形式）を使用する場合に有用です。
戻り値
Mamba2CausalLMOutput または tuple(torch.FloatTensor)
Mamba2CausalLMOutput、またはtorch.FloatTensorのタプル（return_dict=Falseが渡された場合、またはconfig.return_dict=Falseの場合）。設定（Mamba2Config）や入力に応じて、さまざまな要素で構成されます。
Mamba2ForCausalLM の forward メソッドは、__call__ という特殊メソッドをオーバーライドします。
この関数内でフォワードパスの手順を定義する必要がありますが、その後、この関数ではなく
Moduleインスタンスを呼び出すべきです。前者は前処理と後処理の実行を適切に処理するのに対し、後者はそれらを黙って無視するためです。
loss (
torch.FloatTensor、shape(1,)、optional、labelsが指定された場合に返されます) — 言語モデリングの損失（次のトークン予測用）。logits (
torch.FloatTensor、shape(batch_size, sequence_length, config.vocab_size)) — 言語モデリングヘッドの予測スコア（SoftMax 適用前の各語彙トークンのスコア）。cache_params (
~cache_utils.Cache、任意) — 最終タイムステップにおけるモデルの状態。次のinput_idsと forward メソッドで使用することで、古いinput_idsを提供することを回避できます。選択的スキャン後のState space modelの状態行列と、Convolutional statesの両方を含みます
hidden_states (
tuple[torch.FloatTensor]、optional、output_hidden_states=Trueが渡された場合またはconfig.output_hidden_states=Trueの場合に返される) —torch.FloatTensorのタプル（埋め込み層を持つモデルの場合、埋め込みの出力用が 1 つ、各層の出力用が 1 つずつ）で、形状は(batch_size, sequence_length, hidden_size)です。各層の出力におけるモデルのHidden-states、および任意の初期埋め込み出力。