Transformers documentation
Mamba
このモデルは2023年12月1日にHF papersで公開され、2024年3月5日にHugging Face Transformersに貢献されました。
Mamba
Mambaは、長文脈の処理におけるTransformersの計算効率の低さを回避するために設計された、選択型構造化状態空間モデル（SSM）です。完全にアテンション機構を排除したアーキテクチャであり、H3とゲート付きMLPブロック（Mambaブロック）の組み合わせで構成されています。Mambaの「コンテンツベースの推論」により、現在のトークンに応じて入力の特定の部分に焦点を当てることができます。また、Mambaは畳み込み演算の欠如を補うために、新しいハードウェア対応の並列アルゴリズムを採用しています。その結果、Mambaは高速な推論を実現し、非常に長いシーケンスにもスケーリングできます。
すべてのオリジナルの Mamba チェックポイントは、State Space Models 組織の下で確認できます。
このモデルは Molbap と AntonV によって提供されました。 Mamba を異なる言語タスクに適用する方法のさらなる例については、右側のサイドバーにある Mamba モデルをクリックしてください。
以下の例では、Pipeline、AutoModel、およびコマンドラインからテキストを生成する方法を示しています。
from transformers import pipeline
pipeline = pipeline(
task="text-generation",
model="state-spaces/mamba-130m-hf",
device=0
)
pipeline("Plants create energy through a process known as")from transformers import AutoModelForCausalLM, AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained("state-spaces/mamba-130m-hf")
model = AutoModelForCausalLM.from_pretrained("state-spaces/mamba-130m-hf", 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 torchao.quantization import Int4WeightOnlyConfig
from transformers import AutoModelForCausalLM, AutoTokenizer, TorchAoConfig
quantization_config = Int4WeightOnlyConfig(group_size=128)
quantization_config = TorchAoConfig(quant_type=quant_config)
tokenizer = AutoTokenizer.from_pretrained("state-spaces/mamba-2.8b-hf")
model = AutoModelForCausalLM.from_pretrained("state-spaces/mamba-2.8b-hf", 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))注記
現在の実装では、元のCUDAカーネルを使用しています。FlashAttentionに相当する実装は、mamba-ssmおよびcausal_conv1dリポジトリにホストされています。ハードウェアが対応している場合は、これらをインストールしてください。
Mambaスタックは、
Attentionレイヤーと同等のmixerレイヤーを積み重ねています。Mambaの主要なロジックはMambaMixerクラスにあります。以下の例では、PEFTを使用してMambaをファインチューニングする方法を示します。
from datasets import load_dataset from trl import SFTConfig, SFTTrainer from peft import LoraConfig model_id = "state-spaces/mamba-130m-hf" dataset = load_dataset("Abirate/english_quotes", split="train") training_args = SFTConfig(dataset_text_field="quote") 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()バッチ生成には左パディングを使用します。Padding sideを参照してください。
MambaConfig
クラス transformers.MambaConfig
< 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 = Nonevocab_size: int = 50280hidden_size: int = 768state_size: int = 16num_hidden_layers: int = 32layer_norm_epsilon: float = 1e-05pad_token_id: int | None = 0bos_token_id: int | None = 0eos_token_id: int | list[int] | None = 0expand: int = 2conv_kernel: int = 4use_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_scale: float = 1.0time_step_min: float = 0.001time_step_max: float = 0.1time_step_init_scheme: str = 'random'time_step_floor: float = 0.0001rescale_prenorm_residual: bool = Falseuse_cache: bool = Trueuse_mambapy: bool = Falseuse_associative_scan: bool = Truetie_word_embeddings: bool = True )
パラメータ
- vocab_size (
int、任意、デフォルトは50280) — モデルの語彙サイズ。input_idsで表現できるトークンの種類数を定義します。 - hidden_size (
int、任意、デフォルトは768) — 隠れ表現の次元数。 - state_size (
int、任意、デフォルトは16) — Mamba レイヤーにおける SSM 状態のサイズ（潜在状態の次元数）。 - num_hidden_layers (
int、任意、デフォルトは32) — Transformer デコーダーの隠れ層の数。 - layer_norm_epsilon (
float、任意、デフォルト値は 1e-05) — レイヤーノーマライゼーション層で使用する epsilon です。 - pad_token_id (
int、任意、デフォルトは0) — 語彙内でパディングに使用されるトークン ID。 - bos_token_id (
int、任意、デフォルトは0) — 語彙におけるストリーム開始に使用されるトークン ID。 - eos_token_id (
Union[int, list[int]]、任意、デフォルトは0) — 語彙内でストリーム終了に使用されるトークン ID。 - expand (
int、任意、デフォルトは 2) — 中間サイズを決定するために使用される拡張係数。 - conv_kernel (
int、任意、デフォルトは4) — 畳み込みカーネルのサイズです。 - 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_scale (
float、任意、デフォルトは1.0) — 離散化の前にタイムステップの差分に適用されるスケール。 - time_step_min (
float、任意、デフォルトは0.001) —dt_proj.biasを制限するために使用される最小のtime_step。 - time_step_max (
float、任意、デフォルトは0.1) —dt_proj.biasを制限するために使用されるtime_stepの最大値。 - time_step_init_scheme (
str、任意、デフォルトはrandom) — タイムステップのデルタの初期化スキームです。"random"または"uniform"を指定できます。 - time_step_floor (
float、任意、デフォルトは0.0001) — softplus 活性化後の離散タイムステップのデルタに許可される最小値です。 - rescale_prenorm_residual (
bool、任意、デフォルトはFalse) — 初期化時にout_projの重みをリスケールするかどうか。 - use_cache (
bool、任意、デフォルトはTrue) — モデルが最後のキー/バリューの注意を返すかどうか（すべてのモデルで利用されるわけではありません）。config.is_decoder=Trueの場合、またはモデルがデコーダー専用の生成モデルである場合にのみ関連します。 - use_mambapy (
bool、任意、デフォルトはFalse) — Mamba の CUDA ベースの公式実装が利用できない場合の、学習時のフォールバック戦略を決定します。Trueの場合、mamba.py 実装が使用されます。Falseの場合、単純で低速な実装が使用されます。メモリが限られている場合は、単純なバージョンへの切り替えを検討してください。 - use_associative_scan (
bool、任意、デフォルトはTrue) — 並列スキャンに、単純な逐次実装の代わりに PyTorch のtorch._higher_order_ops.associative_scanを使用するかどうかを指定します。associative scan はtorch.compileのトレース中のみ有効であり、torch >= 2.9.0 が必要です。両方のパスが数値的に同一の結果を生成することがテストで確認されています（test_associative_scan_matches_sequentialを参照）。Falseに設定すると、逐次ループにフォールバックします。 - tie_word_embeddings (
bool、任意、デフォルトはTrue) — モデルのtied_weights_keysマッピングに従って重み埋め込みを共有するかどうかを指定します。
これは MambaModel の設定を保存するための設定クラスです。指定された引数に基づいて Mamba モデルをインスタンス化し、モデルのアーキテクチャを定義するために使用されます。デフォルト値で設定をインスタンス化すると、state-spaces/mamba-2.8b と同様の設定が得られます。
設定オブジェクトはPreTrainedConfigを継承し、モデルの出力を制御するために使用できます。詳細については、PreTrainedConfigのドキュメントを参照してください。
例:
>>> from transformers import MambaConfig, MambaModel
>>> # Initializing a Mamba configuration
>>> configuration = MambaConfig()
>>> # Initializing a model (with random weights) from the configuration
>>> model = MambaModel(configuration)
>>> # Accessing the model configuration
>>> configuration = model.configMambaModel
クラス transformers.MambaModel
< source >( config: MambaConfig )
パラメータ
- config (MambaConfig) — モデルのパラメータをすべて含むモデル設定クラス。設定ファイルで初期化しても、モデルに関連する重みは読み込まれず、設定のみが読み込まれます。モデルの重みをロードするには、from_pretrained() メソッドを参照してください。
特定のヘッドを付加せず、生の隠れ状態を出力する素のMambaモデル。
このモデルは 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.LongTensor] = None**kwargs ) → MambaOutput or tuple(torch.Tensor)
パラメータ
- 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.LongTensorの形状(batch_size, sequence_length), 、任意) — パディングトークンのインデックスに対するアテンションを回避するためのマスク。[0, 1]:- マスクされていないトークンに対して1
- マスクされたトークンに対しては 0 です。
戻り値
MambaOutput または tuple(torch.Tensor)
MambaOutput、またはtorch.Tensorのタプル（return_dict=Falseが渡された場合、またはconfig.return_dict=Falseの場合）。設定（MambaConfig）と入力に応じて、さまざまな要素で構成されます。
MambaModelのforwardメソッドは、__call__特殊メソッドをオーバーライドします。
この関数内でフォワードパスの手順を定義する必要がありますが、その後、この関数ではなく
Moduleインスタンスを呼び出すべきです。前者は前処理と後処理の実行を適切に処理するのに対し、後者はそれらを黙って無視するためです。
last_hidden_state (
torch.Tensor、形状(batch_size, sequence_length, hidden_size)、任意) — モデルの最終層の出力における隠れ状態のシーケンス。cache_params (
~cache_utils.Cache、任意) — 最終タイムステップにおけるモデルの状態。次のinput_idsと forward メソッドで使用することで、古いinput_idsを提供することを回避できます。選択的スキャン後のState space modelの状態行列と、Convolutional statesの両方を含みます
hidden_states (
tuple[torch.Tensor]、optional、output_hidden_states=Trueが渡された場合またはconfig.output_hidden_states=Trueの場合に返されます) —torch.Tensorのタプル（モデルに埋め込み層がある場合は埋め込みの出力用 1 つ、各層の出力用 1 つずつ）で、形状は(batch_size, sequence_length, hidden_size)です。各層の出力におけるモデルのHidden-states、および任意の初期埋め込み出力。
MambaLMHeadModel
クラス transformers.MambaForCausalLM
< source >( config: MambaConfig )
パラメータ
- config (MambaConfig) — モデルのパラメータをすべて含むモデル設定クラス。設定ファイルで初期化しても、モデルに関連する重みは読み込まれず、設定のみが読み込まれます。モデルの重みをロードするには、from_pretrained() メソッドを参照してください。
入力埋め込みの重みと共有される線形層を上部に備えた言語モデリングヘッド付きのMAMBAモデルtransformer。
このモデルは PreTrainedModel を継承しています。すべてのモデルに実装されている汎用的なメソッド（ダウンロードや保存、入力埋め込みのリサイズ、ヘッドのプルーニングなど）については、スーパークラスのドキュメントを参照してください。
このモデルは、PyTorchのtorch.nn.Moduleのサブクラスでもあります。 通常のPyTorch Moduleとして使用し、一般的な使用方法や動作に関する詳細については、PyTorchのドキュメントを参照してください。
forward
< source >( input_ids: typing.Optional[torch.LongTensor] = Noneattention_mask: typing.Optional[torch.LongTensor] = Noneinputs_embeds: typing.Optional[torch.Tensor] = 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 = Nonelogits_to_keep: typing.Union[int, torch.Tensor] = 0**kwargs ) → MambaCausalLMOutput or tuple(torch.Tensor)
パラメータ
- input_ids (
torch.LongTensorの形状(batch_size, sequence_length), 、省略可) — 語彙における入力シーケンスのトークンインデックスです。デフォルトではパディングは無視されます。インデックスはAutoTokenizerを使用して取得できます。詳細については、PreTrainedTokenizer.encode()およびPreTrainedTokenizer.call()を参照してください。
- attention_mask (
torch.LongTensorの形状(batch_size, sequence_length), 、任意) — パディングトークンのインデックスに対するアテンションを回避するためのマスク。[0, 1]:- マスクされていないトークンに対して1
- マスクされたトークンに対しては 0 です。
- inputs_embeds (
torch.Tensor、形状(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 を高速に生成するために使用できます。 - logits_to_keep (
Union[int, torch.Tensor]、省略可、デフォルトは0) —intの場合、最後のlogits_to_keepトークンについて logits を計算します。0の場合、すべてのinput_idsについて logits を計算します（特殊ケース）。生成には最後のトークンの logits だけで十分であり、そのトークンのみについて計算することでメモリを節約できます。これは、長いシーケンスや大規模な語彙サイズの場合に特に顕著な効果があります。torch.Tensorの場合、シーケンス長次元で保持するインデックスに対応する 1 次元テンソルである必要があります。これは、パックされたテンソル形式（バッチとシーケンス長を単一の次元にまとめた形式）を使用する場合に有用です。
戻り値
MambaCausalLMOutput または tuple(torch.Tensor)
MambaCausalLMOutput、またはtorch.Tensorのタプル（return_dict=Falseが渡された場合、またはconfig.return_dict=Falseの場合）。構成要素は、設定（MambaConfig）と入力に応じて異なります。
MambaForCausalLM の forward メソッドは、__call__ という特殊メソッドをオーバーライドします。
この関数内でフォワードパスの手順を定義する必要がありますが、その後、この関数ではなく
Moduleインスタンスを呼び出すべきです。前者は前処理と後処理の実行を適切に処理するのに対し、後者はそれらを黙って無視するためです。
loss (
torch.Tensor、形状(1,)、任意、labelsが指定された場合に返されます) — 言語モデリングの損失(次のトークンの予測用)。logits (
torch.Tensor、形状(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.Tensor]、optional、output_hidden_states=Trueが渡された場合またはconfig.output_hidden_states=Trueの場合に返されます) —torch.Tensorのタプル（モデルに埋め込み層がある場合は埋め込みの出力用 1 つ、各層の出力用 1 つずつ）で、形状は(batch_size, sequence_length, hidden_size)です。各層の出力におけるモデルのHidden-states、および任意の初期埋め込み出力。