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` 实例而不是此函数，因为前者会处理运行前处理和后处理步骤，而后者会静默地忽略它们。