# Mamba

[Mamba](https://huggingface.co/papers/2312.00752) 是一种选择性结构化状态空间模型（SSMs），旨在解决 Transformers 在处理长序列时的计算低效问题。 它是一种完全无注意力机制的架构，由 H3 和门控 MLP 块（Mamba 块）组合而成。Mamba 的“基于内容的推理”使其能够根据当前 token 聚焦于输入中的特定部分。Mamba 还采用了一种新的硬件感知并行算法，以弥补卷积操作的缺失。因此，Mamba 具有快速的推理速度，并且可以扩展到非常长的序列。

您可以在 [State Space Models](https://huggingface.co/state-spaces) 组织下找到所有原始的 Mamba 检查点。

> [!TIP]
> 该模型由 [Molbap](https://huggingface.co/Molbap) 和 [AntonV](https://huggingface.co/AntonV) 贡献。点击右侧边栏中的 Mamba 模型，以查看更多将 Mamba 应用于不同语言任务的示例。

以下示例演示了如何使用 [Pipeline](/docs/transformers/main/zh/main_classes/pipelines#transformers.Pipeline)、[AutoModel](/docs/transformers/main/zh/model_doc/auto#transformers.AutoModel) 以及命令行来生成文本。

```python
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")
```

```python
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))
```

量化通过以较低精度表示权重来降低大模型的内存负担。有关更多可用的量化后端，请参阅 [量化](../quantization/overview) 概述。

下面的示例使用 [torchao](../quantization/torchao) 仅将权重量化为 4 位整数。

```python
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](https://github.com/state-spaces/mamba) 和 [causal_conv1d](https://github.com/Dao-AILab/causal-conv1d) 仓库中。如果您的硬件支持，请务必安装它们！
- Mamba 堆叠了与 `Attention` 层等效的 `mixer` 层。您可以在 `MambaMixer` 类中找到 Mamba 的主要逻辑。
- 以下示例演示了如何使用 [PEFT](https://huggingface.co/docs/peft) 对 Mamba 进行微调。

  ```py
  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()
   ```

- 在批量生成时使用左侧填充。请参阅 [填充侧](../llm_tutorial#padding-side)。

## MambaConfig[[transformers.MambaConfig]]

#### transformers.MambaConfig[[transformers.MambaConfig]]

```python
transformers.MambaConfig(transformers_version: str | None = None, architectures: list[str] | None = None, output_hidden_states: bool | None = False, return_dict: bool | None = True, dtype: str | torch.dtype | None = 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: Literal['regression', 'single_label_classification', 'multi_label_classification'] | None = 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)
```

[来源](https://github.com/huggingface/transformers/blob/main/src/transformers/models/mamba/configuration_mamba.py#L26)

**参数：**

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`)：增量（时间步）投影的秩。可以设置为 `"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 激活后，离散时间步 delta 允许的最小值。

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](https://huggingface.co/state-spaces/mamba-2.8b) 类似的配置。

配置对象继承自 [PreTrainedConfig](/docs/transformers/main/zh/main_classes/configuration#transformers.PreTrainedConfig)，可用于控制模型输出。有关更多信息，请参阅 [PreTrainedConfig](/docs/transformers/main/zh/main_classes/configuration#transformers.PreTrainedConfig) 的文档。

示例：

```python
>>> 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[[transformers.MambaModel]]

#### transformers.MambaModel[[transformers.MambaModel]]

```python
transformers.MambaModel(config: MambaConfig)
```

[来源](https://github.com/huggingface/transformers/blob/main/src/transformers/models/mamba/modeling_mamba.py#L614)

**参数：**

config ([MambaConfig](/docs/transformers/main/zh/model_doc/mamba#transformers.MambaConfig))：模型配置类，包含模型的所有参数。通过配置文件进行初始化时，不会加载与模型关联的权重，仅加载配置。请查看 [from_pretrained()](/docs/transformers/main/zh/main_classes/model#transformers.PreTrainedModel.from_pretrained) 方法以加载模型权重。

输出原始隐藏状态、顶部未添加任何特定头部的裸 Mamba 模型。

该模型继承自 [PreTrainedModel](/docs/transformers/main/zh/main_classes/model#transformers.PreTrainedModel)。请查阅父类文档，了解该库为所有模型实现的通用方法（例如下载或保存、调整输入嵌入大小、剪枝头等）。

该模型也是 PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) 的子类。请将其作为常规的 PyTorch Module 使用，有关通用用法和行为的所有事项，请参阅 PyTorch 文档。

#### forward[[transformers.MambaModel.forward]]

```python
forward(input_ids: typing.Optional[torch.LongTensor] = None, inputs_embeds: typing.Optional[torch.LongTensor] = 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: typing.Optional[torch.LongTensor] = None, **kwargs)
```

[来源](https://github.com/huggingface/transformers/blob/main/src/transformers/models/mamba/modeling_mamba.py#L639)

**参数：**

input_ids (`torch.LongTensor`，形状为 `(batch_size, sequence_length)`，*可选*)：输入序列中 token 在词汇表中的索引。默认情况下，填充（padding）将被忽略。可以使用 [AutoTokenizer](/docs/transformers/main/zh/model_doc/auto#transformers.AutoTokenizer) 获取索引。有关详细信息，请参阅 [PreTrainedTokenizer.encode()](/docs/transformers/main/zh/internal/tokenization_utils#transformers.PreTrainedTokenizerBase.encode) 和 [PreTrainedTokenizer.__call__()](/docs/transformers/main/zh/internal/tokenization_utils#transformers.PreTrainedTokenizerBase.__call__)。[什么是输入 ID？](../glossary#input-ids)

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](/docs/transformers/main/zh/main_classes/output#transformers.utils.ModelOutput) 而不是普通元组。

attention_mask (`torch.LongTensor`，形状为 `(batch_size, sequence_length)`，*可选*)：用于避免对填充（padding）标记索引执行注意力机制的掩码。掩码值选自 `[0, 1]`： - 对于**未掩码**的标记，值为 1； - 对于**已掩码**的标记，值为 0。 [什么是注意力掩码？](../glossary#attention-mask)

**返回：** `MambaOutput` 或 `tuple(torch.Tensor)`

`MambaOutput` 或 `torch.Tensor` 元组（如果传入 `return_dict=False` 或 `config.return_dict=False`），其包含的元素取决于配置（[MambaConfig](/docs/transformers/main/zh/model_doc/mamba#transformers.MambaConfig)）和输入。

[MambaModel](/docs/transformers/main/zh/model_doc/mamba#transformers.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[[transformers.MambaForCausalLM]]

#### transformers.MambaForCausalLM[[transformers.MambaForCausalLM]]

```python
transformers.MambaForCausalLM(config: MambaConfig)
```

[来源](https://github.com/huggingface/transformers/blob/main/src/transformers/models/mamba/modeling_mamba.py#L709)

**参数：**

config ([MambaConfig](/docs/transformers/main/zh/model_doc/mamba#transformers.MambaConfig))：模型配置类，包含模型的所有参数。通过配置文件进行初始化时，不会加载与模型关联的权重，仅加载配置。请查看 [from_pretrained()](/docs/transformers/main/zh/main_classes/model#transformers.PreTrainedModel.from_pretrained) 方法以加载模型权重。

带有语言建模头（线性层，其权重与输入嵌入共享）的 MAMBA 模型 Transformer。

该模型继承自 [PreTrainedModel](/docs/transformers/main/zh/main_classes/model#transformers.PreTrainedModel)。请查阅父类文档，了解该库为所有模型实现的通用方法（例如下载或保存、调整输入嵌入大小、剪枝头等）。

该模型也是 PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) 的子类。请将其作为常规的 PyTorch Module 使用，有关通用用法和行为的所有事项，请参阅 PyTorch 文档。

#### forward[[transformers.MambaForCausalLM.forward]]

```python
forward(input_ids: typing.Optional[torch.LongTensor] = None, attention_mask: typing.Optional[torch.LongTensor] = None, inputs_embeds: typing.Optional[torch.Tensor] = None, cache_params: transformers.cache_utils.Cache | None = None, labels: typing.Optional[torch.LongTensor] = None, output_hidden_states: bool | None = None, return_dict: bool | None = None, use_cache: bool | None = None, logits_to_keep: typing.Union[int, torch.Tensor] = 0, **kwargs)
```

[来源](https://github.com/huggingface/transformers/blob/main/src/transformers/models/mamba/modeling_mamba.py#L750)

**参数：**

input_ids (`torch.LongTensor`，形状为 `(batch_size, sequence_length)`，*可选*)：输入序列中 token 在词汇表中的索引。默认情况下，填充（padding）将被忽略。可以使用 [AutoTokenizer](/docs/transformers/main/zh/model_doc/auto#transformers.AutoTokenizer) 获取索引。有关详细信息，请参阅 [PreTrainedTokenizer.encode()](/docs/transformers/main/zh/internal/tokenization_utils#transformers.PreTrainedTokenizerBase.encode) 和 [PreTrainedTokenizer.__call__()](/docs/transformers/main/zh/internal/tokenization_utils#transformers.PreTrainedTokenizerBase.__call__)。[什么是输入 ID？](../glossary#input-ids)

attention_mask (`torch.LongTensor`，形状为 `(batch_size, sequence_length)`，*可选*)：用于避免对填充（padding）标记索引执行注意力机制的掩码。掩码值选自 `[0, 1]`： - 对于**未掩码**的标记，值为 1； - 对于**已掩码**的标记，值为 0。 [什么是注意力掩码？](../glossary#attention-mask)

inputs_embeds (`torch.Tensor`，形状为 `(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](/docs/transformers/main/zh/main_classes/output#transformers.utils.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](/docs/transformers/main/zh/model_doc/mamba#transformers.MambaConfig)）和输入。

[MambaForCausalLM](/docs/transformers/main/zh/model_doc/mamba#transformers.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)`。

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

示例：

```python
```