Transformers 文档

Bamba

Hugging Face's logo
加入 Hugging Face 社区

并获得增强的文档体验

开始使用

PyTorch FlashAttention SDPA

Bamba

Bamba 是一个拥有 90 亿参数的仅解码器语言模型,构建于 Mamba-2 架构之上。它分两个阶段进行预训练——首先在 Dolma v1.7 数据集的 2T 词元上进行训练,然后在来自 FineWebCosmopedia 的额外 200B 词元上进行训练。

你可以在 Bamba 合集中找到所有原始的 Bamba checkpoints。

该模型由 ani300fabianlim 贡献。

点击右侧边栏中的 Bamba 模型,查看更多如何将 Bamba 应用于不同文本生成任务的示例。

下面的示例演示了如何使用 PipelineAutoModel 以及从命令行生成文本。

流水线
自动模型
Transformers CLI
import torch
from transformers import pipeline

pipeline = pipeline(
    task="text-generation",
    model="ibm-ai-platform/Bamba-9B-v2",
    torch_dtype=torch.bfloat16,
    device=0
)
pipeline("Plants create energy through a process known as")

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

以下示例使用torchao仅将权重量化为int4。

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, TorchAoConfig

quantization_config = TorchAoConfig("int4_weight_only", group_size=128)
tokenizer = AutoTokenizer.from_pretrained("ibm-ai-platform/Bamba-9B-v2")
model = AutoModelForCausalLM.from_pretrained(
   "ibm-ai-platform/Bamba-9B-v2",
   quantization_config=quantization_config,
   device_map="auto",
   attn_implementation="sdpa"
)

inputs = tokenizer("Plants create energy through a process known as", return_tensors="pt").to("cuda")
output = model.generate(**inputs)
print(tokenizer.decode(output[0], skip_special_tokens=True))

注意事项

  • Bamba 支持无填充训练,这种方式会将不同的训练样本连接起来,同时仍将输入作为单独的批次进行处理。如果样本长度不同,它可以通过避免填充词元带来的不必要的计算和内存开销,将推理速度显著提升 约 2 倍(具体取决于模型和数据分布)并减少内存使用。

    无填充训练需要 `flash-attn`、`mamba-ssm` 和 `causal-conv1d` 包,并且除了 `input_ids` 和 `labels` 之外,还必须向模型传递以下参数。

    • position_ids: torch.LongTensor: 每个序列中每个词元的位置索引。
    • seq_idx: torch.IntTensor: 批次中每个序列的索引。
    • 每个 `FlashAttentionKwargs`
      • cu_seq_lens_q: torch.LongTensor: 所有查询的累积序列长度。
      • cu_seq_lens_k: torch.LongTensor: 所有键的累积序列长度。
      • max_length_q: int: 批次中最长查询的长度。
      • max_length_k: int: 批次中最长键的长度。

    不应提供 `attention_mask` 输入。DataCollatorWithFlattening 通过使用 `return_seq_idx=True` 和 `return_flash_attn_kwargs=True` 来以编程方式生成上述附加参数集。有关更多信息,请参阅 通过 Flash Attention 封装提升 Hugging Face 训练效率博客文章。

    from transformers import DataCollatorWithFlattening
    
    # Example of using padding-free training
    data_collator = DataCollatorWithFlattening(
        tokenizer=tokenizer,
        return_seq_idx=True,
        return_flash_attn_kwargs=True
    )

BambaConfig

class transformers.BambaConfig

< >

( vocab_size = 128000 tie_word_embeddings = False hidden_size = 4096 intermediate_size = 14336 num_hidden_layers = 32 num_attention_heads = 32 num_key_value_heads = 8 hidden_act = 'silu' initializer_range = 0.02 rms_norm_eps = 1e-05 use_cache = True num_logits_to_keep = 1 pad_token_id = 0 bos_token_id = 1 eos_token_id = 2 max_position_embeddings = 262144 attention_dropout = 0.0 attn_layer_indices = None mamba_n_heads = 128 mamba_d_head = 'auto' mamba_n_groups = 1 mamba_d_state = 256 mamba_d_conv = 4 mamba_expand = 2 mamba_chunk_size = 256 mamba_conv_bias = True mamba_proj_bias = False z_loss_coefficient = 0.0 **kwargs )

参数

  • vocab_size (int, 可选, 默认为 128000) — Bamba 模型的词汇表大小。定义了在调用 BambaModel 时传入的 inputs_ids 可以表示的不同词元的数量。
  • tie_word_embeddings (bool, 可选, 默认为 False) — 是否将模型的输入和输出词嵌入绑定。请注意,这仅在模型具有输出词嵌入层时才相关。
  • hidden_size (int, 可选, 默认为 4096) — 隐藏表示的维度。
  • intermediate_size (int, 可选, 默认为 14336) — MLP 表示的维度。
  • num_hidden_layers (int, 可选, 默认为 32) — Transformer 编码器中的隐藏层数量。
  • num_attention_heads (int, 可选, 默认为 32) — Transformer 编码器中每个注意力层的注意力头数量。
  • num_key_value_heads (int, 可选, 默认为 8) — 这是用于实现分组查询注意力 (Grouped Query Attention) 的键值头数量。如果 num_key_value_heads=num_attention_heads,模型将使用多头注意力 (MHA);如果 num_key_value_heads=1,模型将使用多查询注意力 (MQA);否则使用 GQA。将多头检查点转换为 GQA 检查点时,每个组的键和值头应通过对该组内所有原始头进行均值池化来构建。更多细节,请查看这篇论文。如果未指定,将默认为 8
  • hidden_act (strfunction, 可选, 默认为 "silu") — 解码器中的非线性激活函数(函数或字符串)。
  • initializer_range (float, 可选, 默认为 0.02) — 用于初始化所有权重矩阵的 truncated_normal_initializer 的标准差。
  • rms_norm_eps (float, 可选, 默认为 1e-05) — RMS 归一化层使用的 epsilon。
  • use_cache (bool, 可选, 默认为 True) — 模型是否应返回最后的键/值注意力(并非所有模型都使用)。仅在 config.is_decoder=True 时相关。
  • num_logits_to_keep (intNone, 可选, 默认为 1) — 在生成过程中要计算的提示 logits 的数量。如果为 None,将计算所有 logits。如果为整数值,则只计算最后的 num_logits_to_keep 个 logits。默认为 1,因为生成只需要最后一个提示词元的 logits。对于长序列,整个序列的 logits 可能会占用大量内存,因此设置 num_logits_to_keep=1 将显著减少内存占用。
  • pad_token_id (int, 可选, 默认为 0) — 填充词元的 ID。
  • bos_token_id (int, 可选, 默认为 1) — “序列开始”词元的 ID。
  • eos_token_id (int, 可选, 默认为 2) — “序列结束”词元的 ID。
  • max_position_embeddings (int, 可选, 默认为 262144) — 模型的最大缓存序列长度
  • attention_dropout (float, 可选, 默认为 0.0) — 注意力概率的 dropout 比例。
  • attn_layer_indices (list, 可选) — 指定将具有完全注意力的层的索引。必须包含最多为 num_hidden_layers 的值。
  • mamba_n_heads (int, 可选, 默认为 128) — v2 实现中使用的 mamba 头的数量。
  • mamba_d_head (int, 可选, 默认为 "auto") — 头嵌入维度大小
  • mamba_n_groups (int, 可选, 默认为 1) — v2 实现中使用的 mamba 组的数量。
  • mamba_d_state (int, 可选, 默认为 256) — mamba 状态空间潜变量的维度
  • mamba_d_conv (int, 可选, 默认为 4) — mamba 卷积核的大小
  • mamba_expand (int, 可选, 默认为 2) — 用于确定 mamba 中间大小的扩展因子(相对于 hidden_size)
  • mamba_chunk_size (int, 可选, 默认为 256) — 在进行预填充/训练时,用于拆分序列的块大小。
  • mamba_conv_bias (bool, 可选, 默认为 True) — 指示是否在 mamba 混合器块的卷积层中使用偏置的标志。
  • mamba_proj_bias (bool, 可选, 默认为 False) — 指示是否在 mamba 混合器块的输入和输出投影([“in_proj”, “out_proj”])中使用偏置的标志。
  • z_loss_coefficient (float, 可选, 默认为 0.0) — 用于在训练期间控制 logit 增长的辅助 z 损失的系数。

这是用于存储 BambaModel 配置的配置类。它根据指定的参数实例化一个 BambaModel 模型,定义模型架构。使用 ibm-fms/Bamba-9.8b-2.2T-hf 的默认值实例化配置。

BambaModel 是一个混合了 mamba2 架构和 SwiGLU 的模型。其检查点由 IBM、普林斯顿大学和伊利诺伊大学香槟分校联合训练。

配置对象继承自 PretrainedConfig,可用于控制模型输出。请阅读 PretrainedConfig 的文档以获取更多信息。

BambaModel

class transformers.BambaModel

< >

( config: BambaConfig )

参数

  • config (BambaConfig) — 包含模型所有参数的模型配置类。使用配置文件进行初始化不会加载与模型相关的权重,只会加载配置。请查阅 from_pretrained() 方法来加载模型权重。

基础的 Bamba 模型,输出原始的隐藏状态,顶部没有任何特定的头(head)。

该模型继承自 PreTrainedModel。请查阅超类文档以了解该库为所有模型实现的通用方法(例如下载或保存、调整输入嵌入大小、修剪注意力头等)。

该模型也是 PyTorch torch.nn.Module 的子类。可以像常规 PyTorch 模块一样使用它,并参考 PyTorch 文档了解所有与通用用法和行为相关的事项。

forward

< >

( input_ids: typing.Optional[torch.LongTensor] = None attention_mask: typing.Optional[torch.Tensor] = None position_ids: typing.Optional[torch.LongTensor] = None past_key_values: typing.Optional[transformers.models.bamba.modeling_bamba.HybridMambaAttentionDynamicCache] = None inputs_embeds: typing.Optional[torch.FloatTensor] = None use_cache: typing.Optional[bool] = None output_attentions: typing.Optional[bool] = None output_hidden_states: typing.Optional[bool] = None cache_position: typing.Optional[torch.LongTensor] = None **kwargs: typing_extensions.Unpack[transformers.models.bamba.modeling_bamba.BambaFlashAttentionKwargs] ) transformers.modeling_outputs.BaseModelOutputWithPasttuple(torch.FloatTensor)

参数

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

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

    什么是输入 ID?

  • attention_mask (torch.Tensor,形状为 (batch_size, sequence_length)可选) — 用于避免对填充标记索引执行注意力的掩码。掩码值选自 [0, 1]

    • 1 表示标记未被屏蔽
    • 0 表示标记被屏蔽

    什么是注意力掩码?

  • position_ids (torch.LongTensor,形状为 (batch_size, sequence_length)可选) — 位置嵌入中每个输入序列标记的位置索引。选自范围 [0, config.n_positions - 1]

    什么是位置 ID?

  • past_key_values (~models.bamba.modeling_bamba.HybridMambaAttentionDynamicCache, 可选) — 预先计算的隐藏状态(自注意力块和交叉注意力块中的键和值),可用于加速序列解码。这通常包括模型在解码的前一阶段返回的 `past_key_values`,当 `use_cache=True` 或 `config.use_cache=True` 时。

    允许两种格式:

    • 一个 Cache 实例,请参阅我们的 kv 缓存指南
    • 一个长度为 `config.n_layers` 的 `tuple(torch.FloatTensor)` 元组,每个元组包含 2 个形状为 `(batch_size, num_heads, sequence_length, embed_size_per_head)` 的张量。这也称为旧版缓存格式。

    模型将输出与输入相同的缓存格式。如果没有传递 `past_key_values`,将返回旧版缓存格式。

    如果使用 `past_key_values`,用户可以选择只输入最后一个 `input_ids`(即那些没有为其提供过去键值状态的 `input_ids`),形状为 `(batch_size, 1)`,而不是所有形状为 `(batch_size, sequence_length)` 的 `input_ids`。

  • inputs_embeds (torch.FloatTensor,形状为 (batch_size, sequence_length, hidden_size)可选) — 可选地,你可以选择直接传递嵌入表示,而不是传递 `input_ids`。如果你想比模型内部的嵌入查找矩阵更好地控制如何将 `input_ids` 索引转换为相关向量,这会很有用。
  • use_cache (bool, 可选) — 如果设置为 True,将返回 past_key_values 键值状态,并可用于加速解码(请参阅 past_key_values)。
  • output_attentions (bool, 可选) — 是否返回所有注意力层的注意力张量。有关更多详细信息,请参阅返回张量下的 `attentions`。
  • output_hidden_states (bool, 可选) — 是否返回所有层的隐藏状态。有关更多详细信息,请参阅返回张量下的 `hidden_states`。
  • cache_position (torch.LongTensor,形状为 (sequence_length)可选) — 描绘输入序列标记在序列中位置的索引。与 `position_ids` 相反,此张量不受填充影响。它用于在正确的位置更新缓存并推断完整的序列长度。

返回

transformers.modeling_outputs.BaseModelOutputWithPasttuple(torch.FloatTensor)

一个 transformers.modeling_outputs.BaseModelOutputWithPast 或一个 `torch.FloatTensor` 的元组(如果传递了 `return_dict=False` 或当 `config.return_dict=False` 时),根据配置(BambaConfig)和输入,包含各种元素。

  • last_hidden_state (torch.FloatTensor, 形状为 (batch_size, sequence_length, hidden_size)) — 模型最后一层输出的隐藏状态序列。

    如果使用了 past_key_values,则只输出形状为 (batch_size, 1, hidden_size) 的序列的最后一个隐藏状态。

  • past_key_values (Cache, 可选, 当传递了 `use_cache=True` 或 `config.use_cache=True` 时返回) — 这是一个 Cache 实例。有关更多详细信息,请参阅我们的 kv 缓存指南

    包含预先计算的隐藏状态(自注意力块中的键和值,以及当 `config.is_encoder_decoder=True` 时,交叉注意力块中的键和值),可用于(请参阅 `past_key_values` 输入)加速序列解码。

  • hidden_states (tuple(torch.FloatTensor), 可选, 当传递了 `output_hidden_states=True` 或 `config.output_hidden_states=True` 时返回) — `torch.FloatTensor` 的元组(如果模型有嵌入层,则第一个是嵌入层的输出,然后是每个层的输出),形状为 `(batch_size, sequence_length, hidden_size)`。

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

  • attentions (tuple(torch.FloatTensor), 可选, 当传递了 `output_attentions=True` 或 `config.output_attentions=True` 时返回) — `torch.FloatTensor` 的元组(每层一个),形状为 `(batch_size, num_heads, sequence_length, sequence_length)`。

    注意力 softmax 后的注意力权重,用于计算自注意力头中的加权平均值。

BambaModel 的 forward 方法,重写了 `__call__` 特殊方法。

虽然前向传播的流程需要在此函数内定义,但之后应调用 `Module` 实例而不是此函数,因为前者会处理运行前处理和后处理步骤,而后者会静默地忽略它们。

BambaForCausalLM

class transformers.BambaForCausalLM

< >

( config )

参数

  • config (BambaForCausalLM) — 包含模型所有参数的模型配置类。使用配置文件进行初始化不会加载与模型相关的权重,只会加载配置。请查阅 from_pretrained() 方法来加载模型权重。

用于因果语言建模的 Bamba 模型。

该模型继承自 PreTrainedModel。请查阅超类文档以了解该库为所有模型实现的通用方法(例如下载或保存、调整输入嵌入大小、修剪注意力头等)。

该模型也是 PyTorch torch.nn.Module 的子类。可以像常规 PyTorch 模块一样使用它,并参考 PyTorch 文档了解所有与通用用法和行为相关的事项。

forward

< >

( input_ids: typing.Optional[torch.LongTensor] = None attention_mask: typing.Optional[torch.Tensor] = None position_ids: typing.Optional[torch.LongTensor] = None past_key_values: typing.Optional[transformers.models.bamba.modeling_bamba.HybridMambaAttentionDynamicCache] = None inputs_embeds: typing.Optional[torch.FloatTensor] = None labels: typing.Optional[torch.LongTensor] = None use_cache: typing.Optional[bool] = None output_attentions: typing.Optional[bool] = None output_hidden_states: typing.Optional[bool] = None cache_position: typing.Optional[torch.LongTensor] = None logits_to_keep: typing.Union[int, torch.Tensor] = 0 **kwargs ) transformers.modeling_outputs.CausalLMOutputWithPasttuple(torch.FloatTensor)

参数

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

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

    什么是输入 ID?

  • attention_mask (torch.Tensor,形状为 (batch_size, sequence_length)可选) — 用于避免对填充标记索引执行注意力的掩码。掩码值选自 [0, 1]

    • 1 表示标记未被屏蔽
    • 0 表示标记被屏蔽

    什么是注意力掩码?

  • position_ids (torch.LongTensor,形状为 (batch_size, sequence_length)可选) — 位置嵌入中每个输入序列标记的位置索引。选自范围 [0, config.n_positions - 1]

    什么是位置 ID?

  • past_key_values (~models.bamba.modeling_bamba.HybridMambaAttentionDynamicCache, 可选) — 预先计算的隐藏状态(自注意力块和交叉注意力块中的键和值),可用于加速序列解码。这通常包括模型在解码的前一阶段返回的 `past_key_values`,当 `use_cache=True` 或 `config.use_cache=True` 时。

    允许两种格式:

    • 一个 Cache 实例,请参阅我们的 kv 缓存指南
    • 一个长度为 `config.n_layers` 的 `tuple(torch.FloatTensor)` 元组,每个元组包含 2 个形状为 `(batch_size, num_heads, sequence_length, embed_size_per_head)` 的张量。这也称为旧版缓存格式。

    模型将输出与输入相同的缓存格式。如果没有传递 `past_key_values`,将返回旧版缓存格式。

    如果使用 `past_key_values`,用户可以选择只输入最后一个 `input_ids`(即那些没有为其提供过去键值状态的 `input_ids`),形状为 `(batch_size, 1)`,而不是所有形状为 `(batch_size, sequence_length)` 的 `input_ids`。

  • inputs_embeds (torch.FloatTensor,形状为 (batch_size, sequence_length, hidden_size)可选) — 可选地,你可以选择直接传递嵌入表示,而不是传递 `input_ids`。如果你想比模型内部的嵌入查找矩阵更好地控制如何将 `input_ids` 索引转换为相关向量,这会很有用。
  • labels (torch.LongTensor,形状为 (batch_size, sequence_length)可选) — 用于计算掩码语言建模损失的标签。索引应在 [0, ..., config.vocab_size] 或 -100 之间(请参阅 `input_ids` 文档字符串)。索引设置为 `-100` 的标记将被忽略(屏蔽),损失仅对标签在 [0, ..., config.vocab_size] 中的标记计算。
  • use_cache (bool, 可选) — 如果设置为 True,将返回 past_key_values 键值状态,并可用于加速解码(请参阅 past_key_values)。
  • output_attentions (bool, 可选) — 是否返回所有注意力层的注意力张量。有关更多详细信息,请参阅返回张量下的 `attentions`。
  • output_hidden_states (bool, 可选) — 是否返回所有层的隐藏状态。有关更多详细信息,请参阅返回张量下的 `hidden_states`。
  • cache_position (torch.LongTensor,形状为 (sequence_length)可选) — 描绘输入序列标记在序列中位置的索引。与 `position_ids` 相反,此张量不受填充影响。它用于在正确的位置更新缓存并推断完整的序列长度。
  • logits_to_keep (Union[int, torch.Tensor], 默认为 0) — 如果是 `int`,则计算最后 `logits_to_keep` 个标记的 logits。如果是 `0`,则计算所有 `input_ids` 的 logits(特殊情况)。生成时只需要最后一个标记的 logits,仅为该标记计算它们可以节省内存,这对于长序列或大词汇表来说非常重要。如果是 `torch.Tensor`,则必须是 1D,对应于要在序列长度维度中保留的索引。这在使用打包张量格式(批处理和序列长度的单个维度)时很有用。

返回

transformers.modeling_outputs.CausalLMOutputWithPasttuple(torch.FloatTensor)

一个 transformers.modeling_outputs.CausalLMOutputWithPast 或一个 `torch.FloatTensor` 的元组(如果传递了 `return_dict=False` 或当 `config.return_dict=False` 时),根据配置(BambaConfig)和输入,包含各种元素。

  • loss (torch.FloatTensor 形状为 (1,)可选,当提供 labels 时返回) — 语言建模损失(用于下一个 token 预测)。

  • logits (形状为 (batch_size, sequence_length, config.vocab_size)torch.FloatTensor) — 语言建模头部的预测分数(SoftMax 之前的每个词汇标记的分数)。

  • past_key_values (Cache, 可选, 当传递了 `use_cache=True` 或 `config.use_cache=True` 时返回) — 这是一个 Cache 实例。有关更多详细信息,请参阅我们的 kv 缓存指南

    包含预计算的隐藏状态(自注意力块中的键和值),可用于(参见 past_key_values 输入)加速顺序解码。

  • hidden_states (tuple(torch.FloatTensor), 可选, 当传递了 `output_hidden_states=True` 或 `config.output_hidden_states=True` 时返回) — `torch.FloatTensor` 的元组(如果模型有嵌入层,则第一个是嵌入层的输出,然后是每个层的输出),形状为 `(batch_size, sequence_length, hidden_size)`。

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

  • attentions (tuple(torch.FloatTensor), 可选, 当传递了 `output_attentions=True` 或 `config.output_attentions=True` 时返回) — `torch.FloatTensor` 的元组(每层一个),形状为 `(batch_size, num_heads, sequence_length, sequence_length)`。

    注意力 softmax 后的注意力权重,用于计算自注意力头中的加权平均值。

BambaForCausalLM 的 forward 方法,重写了 `__call__` 特殊方法。

虽然前向传播的流程需要在此函数内定义,但之后应调用 `Module` 实例而不是此函数,因为前者会处理运行前处理和后处理步骤,而后者会静默地忽略它们。

示例

>>> from transformers import AutoTokenizer, BambaForCausalLM

>>> model = BambaForCausalLM.from_pretrained("...")
>>> tokenizer = AutoTokenizer.from_pretrained("...")

>>> prompt = "Hey, are you conscious? Can you talk to me?"
>>> inputs = tokenizer(prompt, return_tensors="pt")

>>> # Generate
>>> generate_ids = model.generate(inputs.input_ids, max_length=30)
>>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
"Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
< > 在 GitHub 上更新