Transformers 文件
Mamba 2
並獲得增強的文件體驗
開始使用
Mamba 2
Mamba 2 基於狀態空間對偶 (SSD) 框架，該框架連線了結構化狀態空間模型 (SSM) 和注意力變體。它使用更高效的 SSD 演算法，比 Mamba 快 2-8 倍，並修改了架構以支援張量並行和分組值注意力 (GVA) 頭結構。
你可以在 State Space Models 組織下找到所有原始的 Mamba 2 checkpoint，但下面顯示的示例使用了 mistralai/Mamba-Codestral-7B-v0.1，因為原始 checkpoint 尚不支援 Hugging Face 實現。
點選右側邊欄中的 Mamba 模型，檢視更多如何將 Mamba 應用於不同語言任務的示例。
下面的示例演示瞭如何使用 Pipeline、AutoModel 以及從命令列生成文字。
hfoptions id=“usage”>
<hfoption id="Pipeline">import torch
from transformers import pipeline
pipeline = pipeline(
task="text-generation",
model="mistralai/Mamba-Codestral-7B-v0.1",
torch_dtype=torch.bfloat16,
device=0
)
pipeline("Plants create energy through a process known as")import torch
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", torch_dtype=torch.bfloat16, device_map="auto")
input_ids = tokenizer("Plants create energy through a process known as", return_tensors="pt").to("cuda")
output = model.generate(**input_ids)
print(tokenizer.decode(output[0], skip_special_tokens=True))echo -e "Plants create energy through a process known as" | transformers-cli run --task text-generation --model mistralai/Mamba-Codestral-7B-v0.1 --device 0量化透過以較低精度表示權重來減少大型模型的記憶體負擔。有關更多可用量化後端，請參閱量化概述。
下面的示例使用 torchao 僅將權重量子化為 4 位整數。
import torch
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", torch_dtype=torch.bfloat16, quantization_config=quantization_config, device_map="auto")
input_ids = tokenizer("Plants create energy through a process known as", return_tensors="pt").to("cuda")
output = model.generate(**input_ids)
print(tokenizer.decode(output[0], skip_special_tokens=True))說明
Codestral Mamba 有
groups=8，這類似於基於注意力的模型中的 kv 頭的數量。Codestral Mamba 有兩個不同的前向傳播過程，
torch_forward或cuda_kernels_forward，它們的結果預計會略有不同。此模型中沒有位置嵌入，但在兩個地方有一個
attention_mask和特定的邏輯來掩碼隱藏狀態，用於批處理生成（更多細節請參見此評論）。這（以及重新實現的 Mamba 2 核心的增加）導致批處理和快取生成之間存在輕微差異。SSM 演算法嚴重依賴張量收縮，它有等效的矩陣乘法，但操作順序略有不同。這使得在較小精度下的差異更大。
與填充標記相對應的隱藏狀態在兩個地方被關閉，並且主要使用左填充進行測試。右填充會向下傳播噪聲，不能保證產生滿意的結果。
tokenizer.padding_side = "left"確保你使用正確的填充側。以下示例演示瞭如何使用 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
class transformers.Mamba2Config
< 來源 >( num_heads = 128 head_dim = 64 vocab_size = 32768 hidden_size = 4096 state_size = 128 num_hidden_layers = 64 layer_norm_epsilon = 1e-05 pad_token_id = 1 bos_token_id = 0 eos_token_id = 2 expand = 2 conv_kernel = 4 n_groups = 8 use_bias = False use_conv_bias = True hidden_act = 'silu' initializer_range = 0.1 residual_in_fp32 = True time_step_rank = 'auto' time_step_min = 0.001 time_step_max = 0.1 time_step_floor = 0.0001 time_step_limit = (0.0, inf) rescale_prenorm_residual = False use_cache = True rms_norm = True chunk_size = 256 tie_word_embeddings = False **kwargs )
引數
- num_heads (
int，可選，預設為 128) — Mamba 2 演化矩陣的頭數。 - head_dim (
int，可選，預設為 64) — 每個頭的維度。 - vocab_size (
int，可選，預設為 32768) — MAMBA2 模型的詞彙表大小。定義了在呼叫 Mamba2Model 時，可以透過inputs_ids表示的不同標記的數量。 - hidden_size (
int，可選，預設為 4096) — 嵌入和隱藏狀態的維度。 - state_size (
int，可選，預設為 128) — 狀態空間潛變數的形狀。 - num_hidden_layers (
int，可選，預設為 64) — 模型中的隱藏層數量。 - layer_norm_epsilon (
float，可選，預設為 1e-05) — 層歸一化層中使用的 epsilon 值。 - pad_token_id (
int，可選，預設為 1) — 填充標記的 ID。 - bos_token_id (
int，可選，預設為 0) — 詞彙表中句子開頭標記的 ID。 - eos_token_id (
int，可選，預設為 2) — 詞彙表中句子結尾標記的 ID。 - expand (
int，可選，預設為 2) — 用於確定中間大小的擴充套件因子。 - conv_kernel (
int，可選，預設為 4) — 卷積核的大小。 - n_groups (
int，可選，預設為 8) — Mamba 2 演化矩陣的組數。 - use_bias (
bool，可選，預設為False) — 是否在混合器塊的 [“in_proj”, “out_proj”] 中使用偏置。 - use_conv_bias (
bool，可選，預設為True) — 是否在混合器塊的卷積層中使用偏置。 - hidden_act (
str，可選，預設為"silu") — 解碼器中的非線性啟用函式（函式或字串）。 - initializer_range (
float，可選，預設為 0.1) — 用於初始化所有權重矩陣的 truncated_normal_initializer 的標準差。 - residual_in_fp32 (
bool，可選，預設為True) — 殘差是否應為float32。如果設定為False，殘差將保持與模型其餘部分相同的dtype。 - time_step_rank (
Union[int,str]，可選，預設為"auto") — 離散化投影矩陣的秩。"auto"表示它將預設為math.ceil(self.hidden_size / 16)。 - 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) —dt_proj.bias層初始化的最小鉗位值。 - time_step_limit (
tuple，可選，預設為(0.0, inf)) — 接受的時間步值範圍。 - rescale_prenorm_residual (
bool，可選，預設為False) — 初始化時是否重新縮放out_proj權重。 - use_cache (
bool，可選，預設為True) — 是否應使用快取。 - rms_norm (
bool, 可選, 預設為True) — 是否使用 RMS 歸一化。 - chunk_size (
int, 可選, 預設為 256) — 構成序列的塊的大小。 - tie_word_embeddings (
bool, 可選, 預設為False) — 是否繫結詞嵌入。
這是一個用於儲存 Mamba2Model 配置的配置類。它根據指定的引數例項化一個 MAMBA2 模型，定義了模型的架構。使用預設值例項化配置將產生與 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
class transformers.Mamba2Model
< 來源 >( config )
引數
- config (Mamba2Model) — 包含模型所有引數的模型配置類。使用配置檔案進行初始化不會載入與模型相關的權重，只會載入配置。請檢視 from_pretrained() 方法來載入模型權重。
裸 Mamba2 模型，輸出原始的隱藏狀態，頂部沒有任何特定的頭。
該模型繼承自 PreTrainedModel。請檢視超類的文件，瞭解該庫為其所有模型實現的通用方法（例如下載或儲存、調整輸入嵌入大小、修剪頭等）。
該模型也是一個 PyTorch torch.nn.Module 的子類。可以像常規的 PyTorch 模組一樣使用它，並參考 PyTorch 文件瞭解所有與通用用法和行為相關的事項。
forward
< 來源 >( input_ids: typing.Optional[torch.LongTensor] = None inputs_embeds: typing.Optional[torch.LongTensor] = None cache_params: typing.Optional[transformers.models.mamba2.modeling_mamba2.Mamba2Cache] = None use_cache: typing.Optional[bool] = None output_hidden_states: typing.Optional[bool] = None return_dict: typing.Optional[bool] = None cache_position: typing.Optional[torch.LongTensor] = None attention_mask: typing.Optional[torch.Tensor] = None **kwargs ) → transformers.models.mamba2.modeling_mamba2.Mamba2Output 或 tuple(torch.FloatTensor)
引數
- input_ids (
torch.LongTensor，形狀為(batch_size, sequence_length)，可選) — 詞彙表中輸入序列標記的索引。預設情況下，填充將被忽略。可以使用 AutoTokenizer 獲取索引。有關詳細資訊，請參閱 PreTrainedTokenizer.encode() 和 PreTrainedTokenizer.call()。
- inputs_embeds (
torch.LongTensor，形狀為(batch_size, sequence_length, hidden_size)，可選) — 可選地，你可以不傳遞input_ids，而是直接傳遞嵌入表示。如果你想比模型內部的嵌入查詢矩陣更能控制如何將input_ids索引轉換為關聯向量，這會非常有用。 - cache_params (
Mamba2Cache, 可選) — 如果傳遞該引數，模型將在所有塊中使用先前的狀態（這將給出所提供input_ids的輸出，就好像模型將state_input_ids + input_ids作為上下文一樣）。 - use_cache (
bool, 可選) — 如果設定為True，將返回cache_params，可用於快速生成下一個 logits。 - output_hidden_states (
bool, 可選) — 是否返回所有層的隱藏狀態。有關更多詳細資訊，請參閱返回的張量下的hidden_states。 - return_dict (
bool, 可選) — 是返回一個 ModelOutput 而不是一個普通的元組。 - cache_position (
torch.LongTensor，形狀為(batch_size,)，可選) — 當前輸入在快取中的位置。這用於確保快取被正確更新。如果傳遞了cache_params，也應傳遞cache_position。 - attention_mask (
torch.Tensor，形狀為(batch_size, sequence_length)，可選) — 用於避免在填充標記索引上執行注意力的掩碼。掩碼值在[0, 1]中選擇：- 1 表示標記未被遮蓋，
- 0 表示標記被遮蓋。
返回
transformers.models.mamba2.modeling_mamba2.Mamba2Output 或 tuple(torch.FloatTensor)
一個 transformers.models.mamba2.modeling_mamba2.Mamba2Output 或一個 torch.FloatTensor 的元組（如果傳遞了 return_dict=False 或當 config.return_dict=False 時），根據配置（Mamba2Config）和輸入，包含各種元素。
-
last_hidden_state (形狀為
(batch_size, sequence_length, hidden_size)的torch.FloatTensor, 可選，預設為None) — 模型最後一層輸出的隱藏狀態序列。 -
cache_params (
~models.mamba2.modeling_mamba2.Mamba2Cache, 可選, 預設為None) — 模型在最後一個時間步的狀態。可以在前向方法中使用下一個input_ids，以避擴音供舊的input_ids。包括選擇性掃描後的狀態空間模型狀態矩陣和卷積狀態
-
hidden_states (
tuple[torch.FloatTensor], 可選, 當傳遞output_hidden_states=True或config.output_hidden_states=True時返回) —torch.FloatTensor的元組（如果模型有嵌入層，則一個用於嵌入層的輸出，另外每個層一個輸出），形狀為(batch_size, sequence_length, hidden_size)。模型在每個層輸出的隱藏狀態以及可選的初始嵌入輸出。
Mamba2Model 的 forward 方法，重寫了 __call__ 特殊方法。
儘管前向傳遞的邏輯需要在此函式內定義，但之後應該呼叫 `Module` 例項而不是這個函式，因為前者會處理前處理和後處理步驟，而後者會靜默地忽略它們。
Mamba2LMHeadModel
class transformers.Mamba2ForCausalLM
< 來源 >( config )
引數
- config (Mamba2ForCausalLM) — 包含模型所有引數的模型配置類。使用配置檔案進行初始化不會載入與模型相關的權重，只會載入配置。請檢視 from_pretrained() 方法來載入模型權重。
MAMBA2 模型，其頂部帶有一個語言建模頭（線性層，權重與輸入嵌入不繫結）。
該模型繼承自 PreTrainedModel。請檢視超類的文件，瞭解該庫為其所有模型實現的通用方法（例如下載或儲存、調整輸入嵌入大小、修剪頭等）。
該模型也是一個 PyTorch torch.nn.Module 的子類。可以像常規的 PyTorch 模組一樣使用它，並參考 PyTorch 文件瞭解所有與通用用法和行為相關的事項。
forward
< 來源 >( input_ids: typing.Optional[torch.LongTensor] = None inputs_embeds: typing.Optional[torch.FloatTensor] = None cache_params: typing.Optional[transformers.models.mamba2.modeling_mamba2.Mamba2Cache] = None labels: typing.Optional[torch.LongTensor] = None output_hidden_states: typing.Optional[bool] = None return_dict: typing.Optional[bool] = None use_cache: typing.Optional[bool] = None cache_position: typing.Optional[torch.Tensor] = None attention_mask: typing.Optional[torch.Tensor] = None **kwargs ) → transformers.models.mamba2.modeling_mamba2.Mamba2CausalLMOutput 或 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 (
Mamba2Cache, 可選) — 如果傳遞該引數，模型將在所有塊中使用先前的狀態（這將給出所提供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。 - return_dict (
bool, 可選) — 是返回一個 ModelOutput 而不是一個普通的元組。 - use_cache (
bool, 可選) — 如果設定為True，將返回cache_params，可用於快速生成下一個 logits。 - cache_position (
torch.LongTensor，形狀為(batch_size,)，可選) — 當前輸入在快取中的位置。這用於確保快取被正確更新。如果傳遞了cache_params，也應傳遞cache_position。 - attention_mask (
torch.Tensor，形狀為(batch_size, sequence_length)，可選) — 用於避免在填充標記索引上執行注意力的掩碼。掩碼值在[0, 1]中選擇：- 1 表示標記未被遮蓋，
- 0 表示標記被遮蓋。
返回
transformers.models.mamba2.modeling_mamba2.Mamba2CausalLMOutput 或 tuple(torch.FloatTensor)
一個 transformers.models.mamba2.modeling_mamba2.Mamba2CausalLMOutput 或一個 torch.FloatTensor 的元組（如果傳遞了 return_dict=False 或當 config.return_dict=False 時），根據配置（Mamba2Config）和輸入，包含各種元素。
-
loss (
torch.FloatTensor形狀為(1,)，可選，當提供labels時返回) — 語言建模損失（用於下一個 token 預測）。 -
logits (形狀為
(batch_size, sequence_length, config.vocab_size)的torch.FloatTensor) — 語言建模頭部的預測分數（SoftMax 之前的每個詞彙標記的分數）。 -
cache_params (
~models.mamba2.modeling_mamba2.Mamba2Cache, 可選, 預設為None) — 模型在最後一個時間步的狀態。可以在前向方法中使用下一個input_ids，以避擴音供舊的input_ids。包括選擇性掃描後的狀態空間模型狀態矩陣和卷積狀態
-
hidden_states (
tuple[torch.FloatTensor], 可選, 當傳遞output_hidden_states=True或config.output_hidden_states=True時返回) —torch.FloatTensor的元組（如果模型有嵌入層，則一個用於嵌入層的輸出，另外每個層一個輸出），形狀為(batch_size, sequence_length, hidden_size)。模型在每個層輸出的隱藏狀態以及可選的初始嵌入輸出。
Mamba2ForCausalLM 的 forward 方法，重寫了 __call__ 特殊方法。
儘管前向傳遞的邏輯需要在此函式內定義，但之後應該呼叫 `Module` 例項而不是這個函式，因為前者會處理前處理和後處理步驟，而後者會靜默地忽略它們。