Transformers 文件
Mamba
並獲得增強的文件體驗
開始使用
Mamba
Mamba 是一種選擇性結構化狀態空間模型（SSM），旨在解決 Transformer 在處理長序列時計算效率低下的問題。它是一個完全無注意力機制的架構，由 H3 和門控 MLP 塊（Mamba 塊）組合而成。Mamba 的“基於內容的推理”使其能夠根據當前 token 專注於輸入的特定部分。Mamba 還使用一種新的硬體感知並行演算法來彌補卷積運算的不足。因此，Mamba 具有快速的推理能力，並且可以擴充套件到非常長的序列。
您可以在 State Space Models 組織下找到所有原始的 Mamba 模型權重。
點選右側邊欄中的 Mamba 模型，檢視更多關於如何將 Mamba 應用於不同語言任務的示例。
下面的示例演示瞭如何使用 Pipeline、AutoModel 以及從命令列生成文字。
import torch
from transformers import pipeline
pipeline = pipeline(
task="text-generation",
model="state-spaces/mamba-130m-hf",
torch_dtype=torch.float16,
device=0
)
pipeline("Plants create energy through a process known as")量化透過以較低精度表示權重來減少大型模型的記憶體負擔。有關更多可用量化後端，請參閱量化概述。
下面的示例使用 torchao 僅將權重化為 4 位整數。
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, TorchAoConfig
from torchao.quantization import Int4WeightOnlyConfig
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", 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))注意
當前的實現使用了原始的 CUDA 核心。等效的 FlashAttention 實現託管在 mamba-ssm 和 causal_conv1d 倉庫中。如果你的硬體支援，請確保安裝它們！
Mamba 堆疊了等效於 `Attention` 層的 `mixer` 層。你可以在 `MambaMixer` 類中找到 Mamba 的主要邏輯。
下面的示例演示瞭如何使用 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()
MambaConfig
class transformers.MambaConfig
< 源 >( vocab_size = 50280 hidden_size = 768 state_size = 16 num_hidden_layers = 32 layer_norm_epsilon = 1e-05 pad_token_id = 0 bos_token_id = 0 eos_token_id = 0 expand = 2 conv_kernel = 4 use_bias = False use_conv_bias = True hidden_act = 'silu' initializer_range = 0.1 residual_in_fp32 = True time_step_rank = 'auto' time_step_scale = 1.0 time_step_min = 0.001 time_step_max = 0.1 time_step_init_scheme = 'random' time_step_floor = 0.0001 rescale_prenorm_residual = False use_cache = True use_mambapy = False **kwargs )
引數
- vocab_size (
int, 可選, 預設為 50280) — MAMBA 模型的詞彙表大小。定義了在呼叫 MambaModel 時傳入的inputs_ids可以表示的不同詞元數量。 - hidden_size (
int, 可選, 預設為 768) — 嵌入和隱藏狀態的維度。 - state_size (
int, 可選, 預設為 16) — 狀態空間潛變數的形狀。 - num_hidden_layers (
int, 可選, 預設為 32) — 模型中的隱藏層數量。 - layer_norm_epsilon (
float, 可選, 預設為 1e-05) — 在層歸一化層中使用的 epsilon 值。 - pad_token_id (
int, 可選, 預設為 0) — 填充詞元的 ID。 - bos_token_id (
int, 可選, 預設為 0) — 詞彙表中句子開頭詞元的 ID。 - eos_token_id (
int, 可選, 預設為 0) — 詞彙表中句子結尾詞元的 ID。 - expand (
int, 可選, 預設為 2) — 用於確定中間大小的擴充套件因子。 - conv_kernel (
int, 可選, 預設為 4) — 卷積核的大小。 - 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_scale (
float, 可選, 預設為 1.0) — 用於縮放dt_proj.bias的比例。 - 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 (
float, 可選, 預設為"random") — 用於初始化dt_proj.weight的方案。應為["random","uniform"]之一 - time_step_floor (
float, 可選, 預設為 0.0001) —dt_proj.bias層初始化的最小截斷值。 - rescale_prenorm_residual (
bool, 可選, 預設為False) — 初始化時是否重新縮放out_proj的權重。 - use_cache (
bool, 可選, 預設為True) — 是否應使用快取。 - use_mambapy (
bool, 可選, 預設為False) — 如果基於 CUDA 的 Mamba 官方實現不可用，確定訓練期間的回退策略。如果為True，則使用 mamba.py 實現。如果為False，則使用樸素且較慢的實現。如果記憶體有限，可以考慮切換到樸素版本。
這是用於儲存 MambaModel 配置的配置類。它用於根據指定的引數例項化一個 MAMBA 模型，定義模型架構。使用預設值例項化配置將產生與 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
class transformers.MambaModel
< 源 >( config )
引數
- config (MambaModel) — 包含模型所有引數的模型配置類。使用配置檔案進行初始化不會載入與模型相關的權重，只會載入配置。請檢視 from_pretrained() 方法來載入模型權重。
基礎的 Mamba 模型，輸出原始的隱藏狀態，沒有附加任何特定的頭部。
該模型繼承自 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.cache_utils.MambaCache] = 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.LongTensor] = None ) → transformers.models.mamba.modeling_mamba.MambaOutput 或 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 (
MambaCache, 可選) — 如果傳遞此引數，模型將在所有塊中使用先前的狀態（這將為提供的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，形狀為(sequence_length)，可選) — 描述輸入序列標記在序列中位置的索引。與position_ids相反，此張量不受填充影響。它用於在正確的位置更新快取並推斷完整的序列長度。 - attention_mask (
torch.LongTensor，形狀為(batch_size, sequence_length)，可選) — 用於避免對填充標記索引執行注意力操作的掩碼。掩碼值選自[0, 1]：- 1 表示標記未被遮蓋，
- 0 表示標記被遮蓋。
返回
transformers.models.mamba.modeling_mamba.MambaOutput 或 tuple(torch.FloatTensor)
一個 transformers.models.mamba.modeling_mamba.MambaOutput 或一個 `torch.FloatTensor` 的元組（如果傳遞了 `return_dict=False` 或當 `config.return_dict=False` 時），根據配置（MambaConfig）和輸入包含各種元素。
-
last_hidden_state (形狀為
(batch_size, sequence_length, hidden_size)的torch.FloatTensor, 可選，預設為None) — 模型最後一層輸出的隱藏狀態序列。 -
cache_params (`~cache_utils.MambaCache`, *可選*，預設為 `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)。模型在每個層輸出的隱藏狀態以及可選的初始嵌入輸出。
MambaModel 的前向方法，覆蓋了 `__call__` 特殊方法。
儘管前向傳遞的配方需要在此函式內定義，但之後應呼叫 `Module` 例項而不是此函式，因為前者會處理執行前處理和後處理步驟，而後者會靜默地忽略它們。
MambaLMHeadModel
class transformers.MambaForCausalLM
< 原始碼 >( config )
引數
- config (MambaForCausalLM) — 包含模型所有引數的模型配置類。用配置檔案初始化不會載入與模型關聯的權重，只會載入配置。請檢視 from_pretrained() 方法以載入模型權重。
MAMBA 模型，其頂部帶有一個語言建模頭（線性層，權重與輸入嵌入繫結）。
該模型繼承自 PreTrainedModel。請檢視超類的文件，瞭解該庫為其所有模型實現的通用方法（例如下載或儲存、調整輸入嵌入的大小、修剪頭部等）。
該模型也是 PyTorch 的 torch.nn.Module 子類。可以像常規的 PyTorch 模組一樣使用它，並參考 PyTorch 文件瞭解所有與常規用法和行為相關的事項。
forward
< 原始碼 >( input_ids: typing.Optional[torch.LongTensor] = None attention_mask: typing.Optional[torch.LongTensor] = None inputs_embeds: typing.Optional[torch.FloatTensor] = None cache_params: typing.Optional[transformers.cache_utils.MambaCache] = 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 **kwargs ) → transformers.models.mamba.modeling_mamba.MambaCausalLMOutput 或 tuple(torch.FloatTensor)
引數
- 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.FloatTensor，形狀為(batch_size, sequence_length, hidden_size)，可選) — 可選地，你可以不傳遞input_ids，而是直接傳遞嵌入表示。如果你想比模型內部的嵌入查詢矩陣更好地控制如何將input_ids索引轉換為關聯向量，這很有用。 - cache_params (
MambaCache, 可選) — 如果傳遞此引數，模型將在所有塊中使用先前的狀態（這將為提供的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.Tensor，形狀為(sequence_length)，可選) — 描述輸入序列標記在序列中位置的索引。與position_ids相反，此張量不受填充影響。它用於在正確的位置更新快取並推斷完整的序列長度。
返回
transformers.models.mamba.modeling_mamba.MambaCausalLMOutput 或 tuple(torch.FloatTensor)
一個 transformers.models.mamba.modeling_mamba.MambaCausalLMOutput 或一個 `torch.FloatTensor` 的元組（如果傳遞了 `return_dict=False` 或當 `config.return_dict=False` 時），根據配置（MambaConfig）和輸入包含各種元素。
-
loss (
torch.FloatTensor形狀為(1,)，可選，當提供labels時返回) — 語言建模損失（用於下一個 token 預測）。 -
logits (形狀為
(batch_size, sequence_length, config.vocab_size)的torch.FloatTensor) — 語言建模頭部的預測分數（SoftMax 之前的每個詞彙標記的分數）。 -
cache_params (`~cache_utils.MambaCache`, *可選*，預設為 `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)。模型在每個層輸出的隱藏狀態以及可選的初始嵌入輸出。
MambaForCausalLM 的前向方法，覆蓋了 `__call__` 特殊方法。
儘管前向傳遞的配方需要在此函式內定義，但之後應呼叫 `Module` 例項而不是此函式，因為前者會處理執行前處理和後處理步驟，而後者會靜默地忽略它們。