Transformers 文档
Mamba 2
并获得增强的文档体验
开始使用
该模型于 2024-05-31 发布，并于 2024-08-06 添加到 Hugging Face Transformers。
Mamba 2
Mamba 2 基于状态空间对偶（SSD）框架，该框架连接了结构化状态空间模型（SSM）和注意力变体。它使用更高效的 SSD 算法，比 Mamba 快 2-8 倍，并修改了架构以支持张量并行和分组值注意力（GVA）头结构。
您可以在 State Space Models 组织下找到所有原始 Mamba 2 检查点，但下面显示的示例使用了 mistralai/Mamba-Codestral-7B-v0.1，因为 Hugging Face 尚未支持原始检查点的实现。
其他基于 Mamba 2 的架构包括 Bamba、FalconH1 和 Zamba2。
此模型由 ArthurZ 贡献。点击右侧边栏的 Mamba 模型，可了解更多将 Mamba 应用于不同语言任务的示例。
以下示例展示了如何使用 Pipeline、AutoModel 和命令行生成文本。
import torch
from transformers import pipeline
pipeline = pipeline(
task="text-generation",
model="mistralai/Mamba-Codestral-7B-v0.1",
dtype=torch.bfloat16,
device=0
)
pipeline("Plants create energy through a process known as")量化通过以较低精度表示权重来减少大型模型的内存负担。有关更多可用量化后端，请参阅量化概述。
下面的示例使用 torchao 将权重量化为 4 位整数。
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, TorchAoConfig
quantization_config = TorchAoConfig("int4_weight_only", group_size=128)
tokenizer = AutoTokenizer.from_pretrained("mistralai/Mamba-Codestral-7B-v0.1")
model = AutoModelForCausalLM.from_pretrained("mistralai/Mamba-Codestral-7B-v0.1", 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(model.device)
output = model.generate(**input_ids)
print(tokenizer.decode(output[0], skip_special_tokens=True))注意事项
Codestral Mamba 具有
groups=8，这与基于注意力的模型中的 kv 头数量相似。Codestral Mamba 有两种不同的前向传播方式：
torch_forward或cuda_kernels_forward，它们的结果预计会略有不同。此模型中没有位置嵌入，但有一个
attention_mask和特定的逻辑，在批量生成（batched generation）的情况下会掩盖两个位置的隐藏状态（有关更多详细信息，请参阅 此评论）。这（以及重写的 Mamba 2 内核的添加）导致了批量生成和缓存生成之间存在微小的差异。SSM 算法在很大程度上依赖于张量收缩（tensor contractions），它具有与矩阵乘法（matmul）等效的操作，但运算顺序略有不同。这使得在较低精度下差异更大。
与填充标记（padding tokens）对应的隐藏状态在两个地方被关闭，并且主要使用左填充（left-padding）进行测试。右填充（right-padding）会沿线传播噪声，并且不能保证产生令人满意的结果。
tokenizer.padding_side = "left"可确保您使用正确的填充侧。下面的示例演示了如何使用 PEFT 微调 Mamba 2。
from datasets import load_dataset
from peft import LoraConfig
from trl import SFTConfig, SFTTrainer
model_id = "mistralai/Mamba-Codestral-7B-v0.1"
dataset = load_dataset("Abirate/english_quotes", split="train")
training_args = SFTConfig(dataset_text_field="quote", gradient_checkpointing=True, per_device_train_batch_size=4)
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()Mamba2Config
class transformers.Mamba2Config
< source >( num_heads = 128 head_dim = 64 vocab_size = 32768 hidden_size = 4096 state_size = 128 num_hidden_layers = 64 layer_norm_epsilon = 1e-05 pad_token_id = 1 bos_token_id = 0 eos_token_id = 2 expand = 2 conv_kernel = 4 n_groups = 8 use_bias = False use_conv_bias = True hidden_act = 'silu' initializer_range = 0.1 residual_in_fp32 = True time_step_rank = 'auto' time_step_min = 0.001 time_step_max = 0.1 time_step_floor = 0.0001 time_step_limit = (0.0, inf) rescale_prenorm_residual = False use_cache = True rms_norm = True chunk_size = 256 tie_word_embeddings = False **kwargs )
参数
- num_heads (
int, optional, defaults to 128) — Mamba 2 的演化矩阵的头数。 - head_dim (
int, optional, defaults to 64) — 每个头的维度。 - vocab_size (
int, optional, defaults to 32768) — MAMBA2 模型词汇表大小。定义调用 Mamba2Model 时传递的inputs_ids可以表示的不同标记的数量。 - hidden_size (
int, optional, defaults to 4096) — 嵌入和隐藏状态的维度。 - state_size (
int, optional, defaults to 128) — 状态空间潜在向量的形状。 - num_hidden_layers (
int, optional, defaults to 64) — 模型中的隐藏层数量。 - layer_norm_epsilon (
float, optional, defaults to 1e-05) — 层归一化层中使用的 epsilon。 - pad_token_id (
int, optional, defaults to 1) — 填充标记 ID。 - bos_token_id (
int, optional, defaults to 0) — 词汇表中开始标记的 ID。 - eos_token_id (
int, optional, defaults to 2) — 词汇表中结束标记的 ID。 - expand (
int, optional, defaults to 2) — 用于确定中间大小的扩展因子。 - conv_kernel (
int, optional, defaults to 4) — 卷积核的大小。 - n_groups (
int, optional, defaults to 8) — Mamba 2 演化矩阵的组数。 - use_bias (
bool, optional, defaults toFalse) — 是否在 mixer 块的 [“in_proj”, “out_proj”] 中使用偏置。 - use_conv_bias (
bool, optional, defaults toTrue) — 是否在 mixer 块的卷积层中使用偏置。 - hidden_act (
str, optional, defaults to"silu") — 解码器中的非线性激活函数（函数或字符串）。 - initializer_range (
float, optional, defaults to 0.1) — 用于初始化所有权重矩阵的 truncated_normal_initializer 的标准差。 - residual_in_fp32 (
bool, optional, defaults toTrue) — 残差是否应为float32。如果设置为False，残差将保持与模型其余部分相同的dtype。 - time_step_rank (
Union[int,str], optional, defaults to"auto") — 离散化投影矩阵的秩。"auto"表示它将默认为math.ceil(self.hidden_size / 16)。 - time_step_min (
float, optional, defaults to 0.001) — 用于限制dt_proj.bias的最小time_step。 - time_step_max (
float, optional, defaults to 0.1) — 用于限制dt_proj.bias的最大time_step。 - time_step_floor (
float, optional, defaults to 0.0001) —dt_proj.bias层初始化的最小钳位值。 - time_step_limit (
tuple, optional, defaults to(0.0, inf)) — 可接受的时间步长值范围。 - rescale_prenorm_residual (
bool, optional, defaults toFalse) — 在初始化时是否重缩放out_proj权重。 - use_cache (
bool, optional, defaults toTrue) — 是否使用缓存。 - rms_norm (
bool, optional, defaults toTrue) — 是否使用 RMS norm。 - chunk_size (
int, optional, defaults to 256) — 构成序列的块的大小。 - tie_word_embeddings (
bool, optional, defaults toFalse) — 是否绑定词嵌入。
这是用于存储 Mamba2Model 配置的类。它用于根据指定的参数实例化 MAMBA2 模型，定义模型架构。使用默认值实例化一个配置将得到一个与 MAMBA2 state-spaces/mamba2-2.8b 架构类似的配置。
配置对象继承自 PreTrainedConfig，可用于控制模型输出。有关更多信息，请阅读 PreTrainedConfig 的文档。
示例
>>> from transformers import Mamba2Config, Mamba2Model
>>> # Initializing a Mamba2 configuration
>>> configuration = Mamba2Config()
>>> # Initializing a model (with random weights) from the configuration
>>> model = Mamba2Model(configuration)
>>> # Accessing the model configuration
>>> configuration = model.configMamba2Model
class transformers.Mamba2Model
< source >( 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 (Mamba2Model) — 模型配置类，包含模型的所有参数。使用配置文件初始化不会加载模型的权重，只加载配置。请查看 from_pretrained() 方法以加载模型权重。
Mamba2 模型，输出原始的 hidden states，没有任何额外的头部。
此模型继承自 PreTrainedModel。查看其父类文档，了解库为所有模型实现的通用方法（例如下载或保存、调整输入嵌入大小、修剪头等）。
此模型也是一个 PyTorch torch.nn.Module 子类。像普通的 PyTorch Module 一样使用它，并参考 PyTorch 文档了解一般用法和行为的所有相关信息。
forward
< source >( input_ids: torch.LongTensor | None = None inputs_embeds: torch.LongTensor | None = None cache_params: transformers.models.mamba2.modeling_mamba2.Mamba2Cache | None = None use_cache: bool | None = None output_hidden_states: bool | None = None return_dict: bool | None = None cache_position: torch.LongTensor | None = None attention_mask: torch.Tensor | None = None **kwargs ) → transformers.models.mamba2.modeling_mamba2.Mamba2Output 或 tuple(torch.FloatTensor)
参数
- input_ids (
torch.LongTensorof shape(batch_size, sequence_length), optional) — 词汇表中输入序列 token 的索引。默认情况下会忽略填充。可以使用 AutoTokenizer 获取索引。有关详细信息，请参阅 PreTrainedTokenizer.encode() 和 PreTrainedTokenizer.call()。
- inputs_embeds (
torch.LongTensorof shape(batch_size, sequence_length, hidden_size), optional) — 可选地，您可以直接传入嵌入表示，而不是传入input_ids。如果您希望比模型内部的嵌入查找矩阵对input_ids索引到关联向量的转换有更多控制，这会很有用。 - cache_params (
Mamba2Cache, optional) — 如果传入，模型将使用所有块的先前状态（这将提供与模型添加state_input_ids + input_ids作为上下文时相同的输出）。 - use_cache (
bool, optional) — 如果设置为True，则返回cache_params，可用于快速生成下一个 logits。 - output_hidden_states (
bool, optional) — 是否返回所有层的隐藏状态。有关更多详细信息，请参阅返回张量下的hidden_states。 - return_dict (
bool, optional) — 是否返回ModelOutput而不是普通元组。 - cache_position (
torch.LongTensorof shape(batch_size,), optional) — 当前输入在缓存中的位置。这用于确保缓存被正确更新。如果传入了cache_params，则也应传入cache_position。 - attention_mask (
torch.Tensorof shape(batch_size, sequence_length), optional) — 用于避免对 padding token 索引执行 attention 的掩码。掩码值选择在[0, 1]：- 1 表示 未掩码 的 token，
- 0 表示 已掩码 的 token。
返回
transformers.models.mamba2.modeling_mamba2.Mamba2Output 或 tuple(torch.FloatTensor)
一个 transformers.models.mamba2.modeling_mamba2.Mamba2Output 或一个 torch.FloatTensor 的元组（如果传入了 return_dict=False 或当 config.return_dict=False 时），根据配置（Mamba2Config）和输入的不同而包含各种元素。
-
last_hidden_state (
torch.FloatTensor | None.last_hidden_state, shape(batch_size, sequence_length, hidden_size), 默认为None) — 模型最后一层的输出的隐藏状态序列。 -
cache_params (
~models.mamba2.modeling_mamba2.Mamba2Cache | None.cache_params, defaults toNone) — 上一个时间步的模型状态。可以在下一个input_ids的 forward 方法中使用，以避免提供旧的input_ids。包括选择性扫描后的状态空间模型状态矩阵和卷积状态
-
hidden_states (
tuple[torch.FloatTensor] | None.hidden_states, 当传递output_hidden_states=True或当config.output_hidden_states=True时返回) —torch.FloatTensor的元组（一个用于嵌入的输出，如果模型有嵌入层，+ 每个层的输出），形状为(batch_size, sequence_length, hidden_size)。模型在每个层输出的隐藏状态以及可选的初始嵌入输出。
Mamba2Model 的 forward 方法，重写了 __call__ 特殊方法。
虽然 forward pass 的实现需要在此函数中定义，但你应该在之后调用
Module实例而不是这个，因为前者负责运行预处理和后处理步骤，而后者会静默地忽略它们。
Mamba2LMHeadModel
class transformers.Mamba2ForCausalLM
< source >( 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 (Mamba2ForCausalLM) — 模型配置类，包含模型的所有参数。使用配置文件初始化不会加载模型的权重，只加载配置。请查看 from_pretrained() 方法以加载模型权重。
MAMBA2 模型 transformer，顶部带有一个语言模型头（线性层，权重不与输入嵌入绑定）。
此模型继承自 PreTrainedModel。查看其父类文档，了解库为所有模型实现的通用方法（例如下载或保存、调整输入嵌入大小、修剪头等）。
此模型也是一个 PyTorch torch.nn.Module 子类。像普通的 PyTorch Module 一样使用它，并参考 PyTorch 文档了解一般用法和行为的所有相关信息。
forward
< source >( input_ids: torch.LongTensor | None = None inputs_embeds: torch.FloatTensor | None = None cache_params: transformers.models.mamba2.modeling_mamba2.Mamba2Cache | None = None labels: torch.LongTensor | None = None output_hidden_states: bool | None = None return_dict: bool | None = None use_cache: bool | None = None cache_position: torch.Tensor | None = None attention_mask: torch.Tensor | None = None logits_to_keep: int | torch.Tensor = 0 **kwargs ) → transformers.models.mamba2.modeling_mamba2.Mamba2CausalLMOutput 或 tuple(torch.FloatTensor)
参数
- input_ids (
torch.LongTensorof shape(batch_size, sequence_length), optional) — 词汇表中输入序列 token 的索引。默认情况下会忽略填充。可以使用 AutoTokenizer 获取索引。有关详细信息，请参阅 PreTrainedTokenizer.encode() 和 PreTrainedTokenizer.call()。
- inputs_embeds (
torch.FloatTensorof shape(batch_size, sequence_length, hidden_size), optional) — 可选地，您可以直接传入嵌入表示，而不是传入input_ids。如果您希望比模型内部的嵌入查找矩阵对input_ids索引到关联向量的转换有更多控制，这会很有用。 - cache_params (
Mamba2Cache, optional) — 如果传入，模型将使用所有块的先前状态（这将提供与模型添加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，可用于快速生成下一个 logit。 - cache_position (
torch.LongTensor, 形状为(batch_size,), 可选) — 当前输入在缓存中的位置。这用于确保缓存被正确更新。如果传递了cache_params，也应传递cache_position。 - attention_mask (
torch.Tensor, 形状为(batch_size, sequence_length), 可选) — 用于避免在填充标记索引上执行注意力的掩码。掩码值选择范围为[0, 1]：- 1 表示未屏蔽的标记，
- 0 表示被屏蔽的标记。
- logits_to_keep (
Union[int, torch.Tensor], 可选, 默认为0) — 如果是int，则计算最后logits_to_keep个标记的 logit。如果为0，则计算所有input_ids的 logit（特殊情况）。仅最后一个标记的 logit 对于生成是必需的，并且仅为该标记计算它们可以节省内存，这对于长序列或大词汇量来说非常可观。如果为torch.Tensor，则必须是 1D 且对应于序列长度维度中要保留的索引。这在使用打包张量格式（批量和序列长度的单维）时很有用。
返回
transformers.models.mamba2.modeling_mamba2.Mamba2CausalLMOutput 或 tuple(torch.FloatTensor)
根据配置（Mamba2Config）和输入，一个 transformers.models.mamba2.modeling_mamba2.Mamba2CausalLMOutput 或一个 torch.FloatTensor 的元组（如果传递了 return_dict=False 或当 config.return_dict=False 时），包含各种元素。
-
loss (
torch.FloatTensor形状为(1,)，可选，当提供labels时返回) — 语言建模损失（用于下一个 token 预测）。 -
logits (形状为
(batch_size, sequence_length, config.vocab_size)的torch.FloatTensor) — 语言建模头部的预测分数（SoftMax 之前的每个词汇标记的分数）。 -
cache_params (
~models.mamba2.modeling_mamba2.Mamba2Cache | None.cache_params, defaults toNone) — 上一个时间步的模型状态。可以在下一个input_ids的 forward 方法中使用，以避免提供旧的input_ids。包括选择性扫描后的状态空间模型状态矩阵和卷积状态
-
hidden_states (
tuple[torch.FloatTensor] | None.hidden_states, 当传递output_hidden_states=True或当config.output_hidden_states=True时返回) —torch.FloatTensor的元组（一个用于嵌入的输出，如果模型有嵌入层，+ 每个层的输出），形状为(batch_size, sequence_length, hidden_size)。模型在每个层输出的隐藏状态以及可选的初始嵌入输出。
transformers.Mamba2ForCausalLM 的 forward 方法，覆盖了 __call__ 特殊方法。
虽然 forward pass 的实现需要在此函数中定义，但你应该在之后调用
Module实例而不是这个，因为前者负责运行预处理和后处理步骤，而后者会静默地忽略它们。