Transformers documentation
Mamba
该模型于 2023-12-01 在 HF papers 上发布，并于 2024-03-05 贡献至 Hugging Face Transformers。
Mamba
Mamba 是一种选择性结构化状态空间模型（SSMs），旨在解决 Transformer 在处理长序列时的计算效率问题。它采用完全无注意力的架构，由 H3 和门控 MLP 块（即 Mamba 块）组合而成。Mamba 的“基于内容的推理”机制使其能够根据当前 token 聚焦于输入中的特定部分。此外，Mamba 还采用了一种新的硬件感知并行算法，以弥补缺乏卷积操作的不足。因此，Mamba 具有快速的推理速度，并且能够扩展到非常长的序列。
您可以在 State Space Models 组织下找到所有原始的 Mamba 检查点。
该模型由 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))量化通过以较低精度表示权重来降低大模型的内存负担。有关更多可用的量化后端，请参阅 量化 概述。
下面的示例使用 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层。您可以在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
< 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可以表示的不同 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) — 词表中用于填充的 token id。 - bos_token_id (
int，可选，默认为0) — 词汇表中用于表示流起始的 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”] 中使用偏置 - 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) — Delta（时间步）投影的秩。可以设置为"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来执行并行扫描，而不是使用朴素的顺序实现。关联扫描仅在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
class 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), ，可选) — 输入序列中各 token 在词汇表中的索引。默认情况下，填充（padding）将被忽略。可以使用 AutoTokenizer 获取索引。有关详细信息，请参阅 PreTrainedTokenizer.encode() 和 PreTrainedTokenizer.call()。
- inputs_embeds（形状为
(batch_size, sequence_length, hidden_size)的torch.LongTensor，可选） — 可选地，您可以选择直接传入嵌入表示，而不是传入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), ，可选) — 用于避免对填充（padding）标记索引执行注意力机制的掩码。掩码值取自 [0, 1]：1 表示未屏蔽的标记，0 表示被屏蔽的标记。什么是注意力掩码？[0, 1]:- 1 表示未被屏蔽的 token，
- 对于被掩码的 token，值为 0。
返回值
MambaOutput 或 tuple(torch.Tensor)
MambaOutput 或 torch.Tensor 元组（当传入 return_dict=False 或 config.return_dict=False 时），其包含的元素取决于配置（MambaConfig）和输入。
MambaModel 的 forward 方法重写了 __call__ 特殊方法。
尽管前向传播的实现需要在此函数中定义，但之后应调用
Module实例，而不是直接调用此函数，因为前者会负责执行预处理和后处理步骤，而后者则会静默地忽略这些步骤。
last_hidden_state（形状为
(batch_size, sequence_length, hidden_size)的torch.Tensor，可选）— 模型最后一层输出的隐藏状态序列。cache_params (
~cache_utils.Cache，可选) — 模型在上一个时间步的状态。可与下一个input_ids一起用于前向方法，从而避免提供旧的input_ids。包含选择性扫描之后的状态空间模型状态矩阵，以及卷积状态
hidden_states (
tuple[torch.Tensor]，可选，当传入output_hidden_states=True或config.output_hidden_states=True时返回) — 一个torch.Tensor元组（如果模型具有嵌入层，则包含嵌入层输出的一个张量，以及每个层输出的一个张量），形状为(batch_size, sequence_length, hidden_size)。模型在每一层输出处的隐藏状态，以及可选的初始嵌入输出。
MambaLMHeadModel
class 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), ，可选) — 输入序列中各 token 在词汇表中的索引。默认情况下，填充（padding）将被忽略。可以使用 AutoTokenizer 获取索引。有关详细信息，请参阅 PreTrainedTokenizer.encode() 和 PreTrainedTokenizer.call()。
- attention_mask (
torch.LongTensor形状为(batch_size, sequence_length), ，可选) — 用于避免对填充（padding）标记索引执行注意力机制的掩码。掩码值取自 [0, 1]：1 表示未屏蔽的标记，0 表示被屏蔽的标记。什么是注意力掩码？[0, 1]:- 1 表示未被屏蔽的 token，
- 对于被掩码的 token，值为 0。
- inputs_embeds（形状为
(batch_size, sequence_length, hidden_size)的torch.Tensor，可选） — 可选地，您可以选择直接传入嵌入表示，而不是传入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 进行计算可以节省内存，对于长序列或大词汇量而言，这种节省相当显著。 如果为torch.Tensor，则必须是一维张量，对应于序列长度维度中需要保留的索引。 这在使用打包张量格式（将批次和序列长度合并为单一维度）时非常有用。
返回值
MambaCausalLMOutput 或 tuple(torch.Tensor)
MambaCausalLMOutput 或 torch.Tensor 元组（如果传入 return_dict=False 或 config.return_dict=False），其包含的元素取决于配置（MambaConfig）和输入。
MambaForCausalLM 的 forward 方法重写了 __call__ 特殊方法。
尽管前向传播的实现需要在此函数中定义，但之后应调用
Module实例，而不是直接调用此函数，因为前者会负责执行预处理和后处理步骤，而后者则会静默地忽略这些步骤。
loss (
torch.Tensor，形状为(1,)，可选，当提供labels时返回) — 语言建模损失（用于下一个 token 预测）。logits（形状为
(batch_size, sequence_length, config.vocab_size)的torch.Tensor）— 语言建模头的预测分数（SoftMax 之前每个词汇表 token 的分数）。cache_params (
~cache_utils.Cache，可选) — 模型在上一个时间步的状态。可与下一个input_ids一起用于前向方法，从而避免提供旧的input_ids。包含选择性扫描之后的状态空间模型状态矩阵，以及卷积状态
hidden_states (
tuple[torch.Tensor]，可选，当传入output_hidden_states=True或config.output_hidden_states=True时返回) — 一个torch.Tensor元组（如果模型具有嵌入层，则包含嵌入层输出的一个张量，以及每个层输出的一个张量），形状为(batch_size, sequence_length, hidden_size)。模型在每一层输出处的隐藏状态，以及可选的初始嵌入输出。