Transformers 文档

Mamba

Transformers

加入 Hugging Face 社区

并获得增强的文档体验

在模型、数据集和 Spaces 上进行协作

通过加速推理获得更快的示例

切换文档主题

开始使用

此模型于 2023 年 12 月 1 日在 HF 论文中发表，并于 2024 年 3 月 5 日贡献给 Hugging Face Transformers。

Mamba

Mamba 是一种选择性结构化状态空间模型（SSM），旨在解决 Transformer 在处理长序列时的计算效率低下问题。它是一个完全不含注意力（Attention-free）的架构，由 H3 和门控 MLP 模块（Mamba 模块）组合而成。Mamba 的“基于内容的推理”功能使其能够根据当前 token 聚焦于输入的特定部分。Mamba 还使用了一种新的硬件感知并行算法来弥补卷积运算的缺失。因此，Mamba 具有极快的推理速度，并且可以扩展到处理非常长的序列。

您可以在 State Space Models 组织下找到所有原始的 Mamba 检查点。

此模型由 Molbap 和 AntonV 贡献。点击右侧边栏中的 Mamba 模型，查看更多关于如何将 Mamba 应用于不同语言任务的示例。

下面的示例演示了如何使用 Pipeline、AutoModel 以及命令行来生成文本。

流水线

自动模型

量化通过以较低精度表示权重来减少大型模型的内存负担。有关更多可用量化后端，请参阅量化概述。

下面的示例使用 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 堆叠了 mixer 层，这与 Attention 层相当。您可以在 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

< 源代码 >

( transformers_version: str | None = None architectures: list[str] | None = None output_hidden_states: bool | None = False return_dict: bool | None = True dtype: typing.Union[str, ForwardRef('torch.dtype'), NoneType] = None chunk_size_feed_forward: int = 0 is_encoder_decoder: bool = False id2label: dict[int, str] | dict[str, str] | None = None label2id: dict[str, int] | dict[str, str] | None = None problem_type: typing.Optional[typing.Literal['regression', 'single_label_classification', 'multi_label_classification']] = None vocab_size: int = 50280 hidden_size: int = 768 state_size: int = 16 num_hidden_layers: int = 32 layer_norm_epsilon: float = 1e-05 pad_token_id: int | None = 0 bos_token_id: int | None = 0 eos_token_id: int | list[int] | None = 0 expand: int = 2 conv_kernel: int = 4 use_bias: bool = False use_conv_bias: bool = True hidden_act: str = 'silu' initializer_range: float = 0.1 residual_in_fp32: bool = True time_step_rank: str | int = 'auto' time_step_scale: float = 1.0 time_step_min: float = 0.001 time_step_max: float = 0.1 time_step_init_scheme: str = 'random' time_step_floor: float = 0.0001 rescale_prenorm_residual: bool = False use_cache: bool = True use_mambapy: bool = False use_associative_scan: bool = True tie_word_embeddings: bool = True )

参数

vocab_size (int, 可选, 默认为 50280) — 模型的词汇表大小。定义了 input_ids 可以表示的不同 token 的数量。
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) — 词汇表中用于填充（padding）的 token id。
bos_token_id (int, 可选, 默认为 0) — 词汇表中用于流开始（beginning-of-stream）的 token id。
eos_token_id (Union[int, list[int]], 可选, 默认为 0) — 词汇表中用于流结束（end-of-stream）的 token id。
expand (int, 可选, 默认为 2) — 用于确定中间尺寸的扩展因子。
conv_kernel (int, 可选, 默认为 4) — 卷积核的大小。
use_bias (bool, 可选, 默认为 False) — 是否在 mixer 块的 [“in_proj”, “out_proj”] 中使用偏置（bias）。
use_conv_bias (bool, 可选, 默认为 True) — 是否在 mixer 块的卷积层中使用偏置。
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) — 如果无法使用基于 CUDA 的官方 Mamba 实现，则确定训练期间的后备策略。如果为 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.config

MambaModel

class transformers.MambaModel

< 源代码 >

( config model_args: ~utils.generic.ModelArgs | None = None adapter_args: ~utils.generic.AdapterArgs | None = None lora_args: ~utils.generic.LoRAArgs | None = None tokenizer_args: ~utils.generic.TokenizerArgs | None = None dataset_args: ~utils.generic.DatasetArgs | None = None data_args: ~utils.generic.DataArgs | None = None training_args: ~utils.generic.TrainingArgs | None = None generation_args: ~utils.generic.GenerationArgs | None = None vision_tower_args: ~utils.generic.VisionTowerArgs | None = None qlora_args: ~utils.generic.QLoRAArgs | None = None vision_tower_template_args: ~utils.generic.VisionTowerTemplateArgs | None = None video_tower_args: ~utils.generic.VideoTowerArgs | None = None vision_config: ~utils.generic.VisionConfig | None = None video_config: ~utils.generic.VideoConfig | None = None load_dataset: bool | None = None load_data_collator: bool | None = None load_processor: bool | None = None load_lora_adapter: bool | None = None load_adapter: bool | None = None load_qlora_adapter: bool | None = None **kwargs: typing_extensions.Unpack[transformers.modeling_utils.PreTrainedModelKwargs] )

参数

config (MambaModel) — 包含模型所有参数的模型配置类。使用配置文件初始化不会加载与模型相关的权重，仅加载配置。请查看 from_pretrained() 方法以加载模型权重。

裸 Mamba 模型，输出原始隐藏状态，顶部没有任何特定的层（head）。

该模型继承自 PreTrainedModel。请查看超类文档以了解该库为所有模型实现的通用方法（例如下载或保存、调整输入嵌入大小、剪枝头部等）。

此模型也是一个 PyTorch torch.nn.Module 子类。像普通的 PyTorch Module 一样使用它，并参考 PyTorch 文档了解一般用法和行为的所有相关信息。

forward

< 源代码 >

( input_ids: torch.LongTensor | None = None inputs_embeds: torch.LongTensor | None = None cache_params: transformers.cache_utils.Cache | None = None use_cache: bool | None = None output_hidden_states: bool | None = None return_dict: bool | None = None attention_mask: torch.LongTensor | None = None **kwargs ) → MambaOutput 或 tuple(torch.FloatTensor)

参数

input_ids (torch.LongTensor，形状为 (batch_size, sequence_length)，可选) — 词汇表中输入序列标记的索引。默认情况下将忽略填充（Padding）。

索引可以使用 AutoTokenizer 获取。详细信息请参阅 PreTrainedTokenizer.encode() 和 PreTrainedTokenizer.call()。

什么是输入 ID？
inputs_embeds (torch.LongTensor，形状为 (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。
return_dict (bool, 可选) — 是否返回 ModelOutput 而不是普通元组。
attention_mask (torch.LongTensor，形状为 (batch_size, sequence_length)，可选) — 用于避免对填充标记索引执行注意力的掩码。掩码值在 [0, 1] 中选择：
- 1 表示未被掩码的标记，
- 0 表示被掩码的标记。
什么是注意力掩码？

MambaOutput 或 tuple(torch.FloatTensor)

根据配置（MambaConfig）和输入，MambaOutput 或 torch.FloatTensor 元组（如果传递了 return_dict=False 或当 config.return_dict=False 时）包含各种元素。

MambaModel 的前向传播方法，覆盖了 __call__ 特殊方法。

虽然 forward pass 的实现需要在此函数中定义，但你应该在之后调用 Module 实例而不是这个，因为前者负责运行预处理和后处理步骤，而后者会静默地忽略它们。

last_hidden_state (形状为 (batch_size, sequence_length, hidden_size) 的 torch.FloatTensor, 可选，默认为 None) — 模型最后一层输出的隐藏状态序列。
cache_params (~cache_utils.Cache, 可选, 默认值为 None) — 模型在最后一个时间步的状态。可以在带有下一个 input_ids 的 forward 方法中使用，以避免提供旧的 input_ids。

包括选择性扫描后的状态空间模型状态矩阵和卷积状态
hidden_states (tuple[torch.FloatTensor]，可选，在传入 output_hidden_states=True 或 config.output_hidden_states=True 时返回) — torch.FloatTensor 元组（一个用于嵌入层的输出（如果模型有嵌入层的话） + 一个用于每层输出），形状为 (batch_size, sequence_length, hidden_size)。

模型在每个层输出的隐藏状态以及可选的初始嵌入输出。

MambaLMHeadModel

class transformers.MambaForCausalLM

< 源代码 >

参数

config (MambaForCausalLM) — 包含模型所有参数的模型配置类。使用配置文件初始化不会加载与模型相关的权重，仅加载配置。请查看 from_pretrained() 方法以加载模型权重。

带有语言建模头部（线性层，权重与输入嵌入绑定）的 MAMBA 模型转换器。

该模型继承自 PreTrainedModel。请查看超类文档以了解该库为所有模型实现的通用方法（例如下载或保存、调整输入嵌入大小、剪枝头部等）。

此模型也是一个 PyTorch torch.nn.Module 子类。像普通的 PyTorch Module 一样使用它，并参考 PyTorch 文档了解一般用法和行为的所有相关信息。

forward

< 源代码 >

( input_ids: torch.LongTensor | None = None attention_mask: torch.LongTensor | None = None inputs_embeds: torch.FloatTensor | None = None cache_params: transformers.cache_utils.Cache | None = None labels: torch.LongTensor | None = None output_hidden_states: bool | None = None return_dict: bool | None = None use_cache: bool | None = None logits_to_keep: int | torch.Tensor = 0 **kwargs ) → MambaCausalLMOutput 或 tuple(torch.FloatTensor)

参数

input_ids (torch.LongTensor，形状为 (batch_size, sequence_length)，可选) — 词汇表中输入序列标记的索引。默认情况下将忽略填充。

索引可以使用 AutoTokenizer 获取。详细信息请参阅 PreTrainedTokenizer.encode() 和 PreTrainedTokenizer.call()。

什么是输入 ID？
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 (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。
return_dict (bool, 可选) — 是否返回一个 ModelOutput 而不是普通元组。
use_cache (bool, 可选) — 如果设置为 True，则返回 cache_params，可用于快速生成下一个 logits。
logits_to_keep (Union[int, torch.Tensor], 可选, 默认为 0) — 如果是 int，则计算最后 logits_to_keep 个 token 的 logits。如果为 0，则计算所有 input_ids 的 logits（特殊情况）。生成时通常只需要最后一个 token 的 logits，仅计算该 token 的 logits 可以节省内存，这对于长序列或大词表大小的情况非常有效。如果是 torch.Tensor，则必须是一维的，对应于要在序列长度维度中保留的索引。这在使用打包张量格式（批次和序列长度合并为单维度）时非常有用。

MambaCausalLMOutput 或 tuple(torch.FloatTensor)

一个 MambaCausalLMOutput 或一个 torch.FloatTensor 元组（如果传递了 return_dict=False 或者当 config.return_dict=False 时），根据配置 (MambaConfig) 和输入包含各种元素。

MambaForCausalLM 的 forward 方法，重写了 __call__ 特殊方法。

虽然 forward pass 的实现需要在此函数中定义，但你应该在之后调用 Module 实例而不是这个，因为前者负责运行预处理和后处理步骤，而后者会静默地忽略它们。

loss (torch.FloatTensor 形状为 (1,)，可选，当提供 labels 时返回) — 语言建模损失（用于下一个 token 预测）。
logits (形状为 (batch_size, sequence_length, config.vocab_size) 的 torch.FloatTensor) — 语言建模头部的预测分数（SoftMax 之前的每个词汇标记的分数）。
cache_params (~cache_utils.Cache, 可选, 默认值为 None) — 模型在最后一个时间步的状态。可以在带有下一个 input_ids 的 forward 方法中使用，以避免提供旧的 input_ids。

包括选择性扫描后的状态空间模型状态矩阵和卷积状态
hidden_states (tuple[torch.FloatTensor]，可选，在传入 output_hidden_states=True 或 config.output_hidden_states=True 时返回) — torch.FloatTensor 元组（一个用于嵌入层的输出（如果模型有嵌入层的话） + 一个用于每层输出），形状为 (batch_size, sequence_length, hidden_size)。

模型在每个层输出的隐藏状态以及可选的初始嵌入输出。

示例

在 GitHub 上更新

←MADLAD-400 Mamba2→