Jamba
概述
Jamba 是一种最先进的混合 SSM-Transformer LLM。它是第一个生产规模的 Mamba 实现，它打开了有趣的研究和应用机会。虽然这种初步的实验表明令人鼓舞的成果，但我们预计随着未来的优化和探索，这些成果将得到进一步增强。
有关此模型的完整详细信息，请阅读 发布博客文章。
模型详细信息
Jamba 是一个经过预训练的专家混合 (MoE) 生成文本模型，具有 120 亿个活动参数，所有专家总共拥有 520 亿个参数。它支持 256K 上下文长度，并且可以在单个 80GB GPU 上容纳高达 140K 个标记。
如下图所示，Jamba 的架构采用块和层方法，使 Jamba 能够成功地将 Transformer 和 Mamba 架构整合在一起。每个 Jamba 块包含一个注意力层或一个 Mamba 层，之后是一个多层感知器 (MLP)，在所有 8 层中平均 1 层 Transformer 层。
使用
先决条件
Jamba 要求您使用 transformers 版本 4.39.0 或更高版本
pip install transformers>=4.39.0
为了运行优化的 Mamba 实现，您首先需要安装 mamba-ssm 和 causal-conv1d
pip install mamba-ssm causal-conv1d>=1.2.0
您还必须将模型置于 CUDA 设备上。
您可以运行模型而不使用优化的 Mamba 内核，但不推荐，因为这将导致明显更低的延迟。为了实现这一点，您需要在加载模型时指定 use_mamba_kernels=False。
运行模型
from transformers import AutoModelForCausalLM, AutoTokenizer
model = AutoModelForCausalLM.from_pretrained("ai21labs/Jamba-v0.1")
tokenizer = AutoTokenizer.from_pretrained("ai21labs/Jamba-v0.1")
input_ids = tokenizer("In the recent Super Bowl LVIII,", return_tensors='pt').to(model.device)["input_ids"]
outputs = model.generate(input_ids, max_new_tokens=216)
print(tokenizer.batch_decode(outputs))
# ["<|startoftext|>In the recent Super Bowl LVIII, the Kansas City Chiefs emerged victorious, defeating the San Francisco 49ers in a thrilling overtime showdown. The game was a nail-biter, with both teams showcasing their skills and determination.\n\nThe Chiefs, led by their star quarterback Patrick Mahomes, displayed their offensive prowess, while the 49ers, led by their strong defense, put up a tough fight. The game went into overtime, with the Chiefs ultimately securing the win with a touchdown.\n\nThe victory marked the Chiefs' second Super Bowl win in four years, solidifying their status as one of the top teams in the NFL. The game was a testament to the skill and talent of both teams, and a thrilling end to the NFL season.\n\nThe Super Bowl is not just about the game itself, but also about the halftime show and the commercials. This year's halftime show featured a star-studded lineup, including Usher, Alicia Keys, and Lil Jon. The show was a spectacle of music and dance, with the performers delivering an energetic and entertaining performance.\n"]以半精度加载模型
发布的检查点以 BF16 格式保存。为了以 BF16/FP16 格式将其加载到 RAM 中，您需要指定 torch_dtype
from transformers import AutoModelForCausalLM
import torch
model = AutoModelForCausalLM.from_pretrained("ai21labs/Jamba-v0.1", torch_dtype=torch.bfloat16)
# you can also use torch_dtype=torch.float16使用半精度时，您可以启用 FlashAttention2 的注意力块实现。要使用它，您还需要将模型置于 CUDA 设备上。由于在这种精度下，模型太大而无法容纳在单个 80GB GPU 上，因此您还需要使用 accelerate 对其进行并行化
from transformers import AutoModelForCausalLM
import torch
model = AutoModelForCausalLM.from_pretrained("ai21labs/Jamba-v0.1",
torch_dtype=torch.bfloat16,
attn_implementation="flash_attention_2",
device_map="auto")以 8 位加载模型
使用 8 位精度，可以在单个 80GB GPU 上容纳高达 140K 的序列长度。您可以使用 bitsandbytes 轻松将模型量化为 8 位。为了不降低模型质量，我们建议将 Mamba 块排除在量化之外。
from transformers import AutoModelForCausalLM, BitsAndBytesConfig
quantization_config = BitsAndBytesConfig(load_in_8bit=True, llm_int8_skip_modules=["mamba"])
model = AutoModelForCausalLM.from_pretrained(
"ai21labs/Jamba-v0.1", torch_dtype=torch.bfloat16, attn_implementation="flash_attention_2", quantization_config=quantization_config
)JambaConfig
class transformers.JambaConfig
< 源代码 >( vocab_size = 65536 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-06 use_cache = True num_logits_to_keep = 1 output_router_logits = False router_aux_loss_coef = 0.001 pad_token_id = 0 bos_token_id = 1 eos_token_id = 2 sliding_window = None max_position_embeddings = 262144 attention_dropout = 0.0 num_experts_per_tok = 2 num_experts = 16 expert_layer_period = 2 expert_layer_offset = 1 attn_layer_period = 8 attn_layer_offset = 4 use_mamba_kernels = True mamba_d_state = 16 mamba_d_conv = 4 mamba_expand = 2 mamba_dt_rank = 'auto' mamba_conv_bias = True mamba_proj_bias = False **kwargs )
参数
- vocab_size (
int, 可选, 默认值为 65536) — Jamba 模型的词汇量。定义了调用 JambaModel 时，通过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) — 用于实现分组查询注意力的键值头数量。如果num_key_value_heads=num_attention_heads，模型将使用多头注意力 (MHA)，如果num_key_value_heads=1，模型将使用多查询注意力 (MQA)，否则使用 GQA。当将多头检查点转换为 GQA 检查点时，每个组键和值头应通过对该组中所有原始头进行平均池化来构建。有关更多详细信息，请查看 这篇论文。如果没有指定，将默认为8。 - hidden_act (
str或function, 可选, 默认为"silu") — 解码器中的非线性激活函数（函数或字符串）。 - initializer_range (
float, 可选, 默认为 0.02) — 初始化所有权重矩阵的截断正态初始化的标准差。 - rms_norm_eps (
float, 可选, 默认为 1e-06) — RMS 规范化层使用的 epsilon。 - use_cache (
bool, 可选, 默认为True) — 模型是否应返回最后键值注意力（并非所有模型都使用）。仅在config.is_decoder=True时相关。 - num_logits_to_keep (
int或None, 可选, 默认为 1) — 在生成期间要计算的提示 logits 数量。如果为None，将计算所有 logits。如果为整数，则只计算最后num_logits_to_keep个 logits。默认值为 1，因为生成只需要最后提示令牌的 logits。对于长序列，整个序列的 logits 可能占用大量内存，因此将num_logits_to_keep设置为 1 将显着减少内存占用。 - output_router_logits (
bool, 可选, 默认为False) — 模型是否应返回路由 logits。启用此功能还将允许模型输出辅助损失。有关更多详细信息，请查看 此处。 - router_aux_loss_coef (
float, 可选, 默认为 0.001) — 针对总损失的辅助损失因子。 - pad_token_id (
int, 可选, 默认为 0) — 填充令牌的 id。 - bos_token_id (
int, 可选, 默认为 1) — “序列开始”令牌的 id。 - 最大位置嵌入 (
int, 可选，默认为 262144) — 此值没有实际影响。此模型用于的最大序列长度。它可以用于更长的序列，但性能可能会下降。 - 注意力丢弃率 (
float, 可选，默认为 0.0) — 注意力概率的丢弃率。 - 每个词元的专家数量 (
int, 可选，默认为 2) — 每个词元根节点的专家数量，也可以解释为top-p路由参数 - 专家数量 (
int, 可选，默认为 16) — 每个稀疏 MLP 层的专家数量。 - 专家层周期 (
int, 可选，默认为 2) — 每隔这么多层，我们将会有一个专家层 - 专家层偏移量 (
int, 可选，默认为 1) — 包含专家 mlp 层的第一个层索引 - 注意力层周期 (
int, 可选，默认为 8) — 每隔这么多层，我们将会有一个普通注意力层 - 注意力层偏移量 (
int, 可选，默认为 4) — 包含普通注意力 mlp 层的第一个层索引 - 使用 Mamba 内核 (
bool, 可选，默认为True) — 指示是否使用快速 Mamba 内核的标志。只有在安装了mamba-ssm和causal-conv1d且 Mamba 模块在 CUDA 设备上运行时才可以使用。如果为True且内核不可用，则会引发 ValueError - mamba_expand (
int, optional, defaults to 2) — 用于确定 Mamba 中间大小的扩展因子（相对于 hidden_size） - mamba_dt_rank (
Union[int,str], optional, defaults to"auto") — Mamba 离散化投影矩阵的秩。"auto"表示它将默认为math.ceil(self.hidden_size / 16) - mamba_conv_bias (
bool, optional, defaults toTrue) — 指示是否在 Mamba 混合器块的卷积层中使用偏差的标志。 - mamba_proj_bias (
bool, optional, defaults toFalse) — 指示是否在 Mamba 混合器块的输入和输出投影（[“in_proj”, “out_proj”]）中使用偏差的标志。
这是一个用于存储 JambaModel 配置的配置类。它用于根据指定的参数实例化 Jamba 模型，定义模型架构。使用默认值实例化配置将生成与 Jamba-v0.1 模型类似的配置。
配置对象继承自 PretrainedConfig，可用于控制模型输出。阅读 PretrainedConfig 文档以获取更多信息。
JambaModel
class transformers.JambaModel
< source >( config: JambaConfig )
参数
- config (JambaConfig) — 模型配置类，包含模型的所有参数。使用配置文件初始化不会加载与模型关联的权重，只会加载配置。查看 from_pretrained() 方法以加载模型权重。 config — JambaConfig
不带任何特定头部，输出原始隐藏状态的 Jamba 模型。该模型继承自 PreTrainedModel。查看超类文档以了解库为所有模型实现的通用方法（如下载或保存、调整输入嵌入大小、裁剪头部等）。
该模型也是一个 PyTorch torch.nn.Module 子类。将其用作常规 PyTorch 模块，并参考 PyTorch 文档以获取有关一般用法和行为的所有事项。
由 config.num_hidden_layers 层组成的 Transformer 解码器。每一层都是一个 JambaDecoderLayer
forward
< source >( input_ids: LongTensor = None attention_mask: Optional = None position_ids: Optional = None past_key_values: Optional = None inputs_embeds: Optional = None use_cache: Optional = None output_attentions: Optional = None output_hidden_states: Optional = None output_router_logits: Optional = None return_dict: Optional = None cache_position: Optional = None )
参数
- input_ids (
torch.LongTensor形状为(batch_size, sequence_length)) — 词汇表中输入序列标记的索引。如果您提供填充，则默认情况下填充将被忽略。可以使用 AutoTokenizer 获取索引。有关详细信息，请参见 PreTrainedTokenizer.encode() 和 PreTrainedTokenizer.call()。
- attention_mask (
torch.Tensor形状为(batch_size, sequence_length), 可选) — 掩码以避免对填充标记索引执行注意力。[0, 1]中选择掩码值：- 1 表示 未掩码 的标记，
- 0 表示 掩码 的标记。
可以使用 AutoTokenizer 获取索引。有关详细信息，请参见 PreTrainedTokenizer.encode() 和 PreTrainedTokenizer.call()。
如果使用了
past_key_values，则可以选择仅输入最后一个input_ids（见past_key_values）。如果您想更改填充行为，您应该阅读
modeling_opt._prepare_decoder_attention_mask并修改以满足您的需求。有关默认策略的更多信息，请参见 论文 中的图 1。- 1 表示头部 未掩码，
- 0 表示头部 掩码。
- position_ids (
torch.LongTensor形状为(batch_size, sequence_length), 可选) — 每个输入序列标记在位置嵌入中的位置索引。在[0, config.n_positions - 1]范围内选择。 - past_key_values (
HybridMambaAttentionDynamicCache, 可选, 当传递use_cache=True或config.use_cache=True时返回) — 包含预先计算的隐藏状态的 HybridMambaAttentionDynamicCache 对象（自注意力块中的键和值，以及 mamba 块中的卷积和 ssm 状态），可用于（见past_key_values输入）加速顺序解码。键和值缓存张量形状为(batch_size, num_heads, seq_len, head_dim)。卷积和 ssm 状态张量形状分别为(batch_size, d_inner, d_conv)和(batch_size, d_inner, d_state)。有关更多详细信息，请参见HybridMambaAttentionDynamicCache类。如果使用了
past_key_values，则用户可以选择仅输入最后一个input_ids（那些没有提供其过去键值状态的模型）的形状为(batch_size, 1)，而不是所有input_ids的形状为(batch_size, sequence_length)。 - 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。 - output_router_logits (
bool, 可选) — 是否返回所有路由器的 logits。它们对于计算路由器损失很有用，在推理期间不应该返回。 - return_dict (
bool, 可选) — 是否返回一个 ModelOutput 而不是一个简单的元组。 - cache_position (
torch.LongTensor形状为(sequence_length), 可选) — 指示输入序列标记在序列中的位置的索引。与position_ids相反，此张量不受填充影响。它用于在正确的位置更新缓存并推断完整的序列长度。
The JambaModel forward method, overrides the __call__ special method.
Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the pre and post processing steps while the latter silently ignores them.
JambaForCausalLM
forward
< 源代码 >( input_ids: LongTensor = None attention_mask: Optional = None position_ids: Optional = None past_key_values: Optional = None inputs_embeds: Optional = None labels: Optional = None use_cache: Optional = None output_attentions: Optional = None output_hidden_states: Optional = None output_router_logits: Optional = None return_dict: Optional = None cache_position: Optional = None num_logits_to_keep: Optional = None ) → transformers.modeling_outputs.MoeCausalLMOutputWithPast 或 tuple(torch.FloatTensor)
参数
- input_ids (
torch.LongTensor形状为(batch_size, sequence_length)) — 输入序列标记在词汇表中的索引。如果您提供填充，则默认情况下将忽略填充。可以使用 AutoTokenizer 获取索引。有关详细信息，请参阅 PreTrainedTokenizer.encode() 和 PreTrainedTokenizer.call() 。
- attention_mask (
torch.Tensor形状为(batch_size, sequence_length), 可选) — 用于避免对填充令牌索引执行注意力的掩码。在[0, 1]中选择的掩码值：- 1 代表未屏蔽的令牌，
- 0 代表屏蔽的令牌。
可以使用 AutoTokenizer 获取索引。有关详细信息，请参阅 PreTrainedTokenizer.encode() 和 PreTrainedTokenizer.call()。
如果使用
past_key_values，可以选择仅输入最后一个input_ids（请参阅past_key_values）。如果您想更改填充行为，则应阅读
modeling_opt._prepare_decoder_attention_mask并根据您的需要进行修改。有关默认策略的更多信息，请参阅论文中的图 1。- 1 表示头部未屏蔽，
- 0 表示头部屏蔽。
- position_ids (
torch.LongTensor形状为(batch_size, sequence_length), 可选) — 每个输入序列令牌在位置嵌入中的位置索引。在[0, config.n_positions - 1]范围内选择。 - past_key_values (
HybridMambaAttentionDynamicCache, 可选，在传递use_cache=True或config.use_cache=True时返回) — 包含预计算隐藏状态 (自注意力块中的键和值，以及 mamba 块中的卷积和 ssm 状态) 的 HybridMambaAttentionDynamicCache 对象，可用于 (请参阅past_key_values输入) 加速顺序解码。键和值缓存张量形状为(batch_size, num_heads, seq_len, head_dim)。卷积和 ssm 状态张量形状分别为(batch_size, d_inner, d_conv)和(batch_size, d_inner, d_state)。有关更多详细信息，请参阅HybridMambaAttentionDynamicCache类。如果使用
past_key_values，用户可以选择仅输入最后一个input_ids（那些没有提供给该模型的过去键值状态的输入）形状为(batch_size, 1)，而不是所有input_ids形状为(batch_size, sequence_length)。 - 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。 - output_router_logits (
bool, 可选) — 是否返回所有路由器的 logits。它们对于计算路由器损失很有用，在推理过程中不应该返回。 - return_dict (
bool, 可选) — 是否返回 ModelOutput 而不是普通元组。 - cache_position (
torch.LongTensor形状为(sequence_length)，可选) — 指示输入序列标记在序列中的位置的索引。与position_ids相反，此张量不受填充的影响。它用于在正确的位置更新缓存并推断完整序列长度。参数 — labels (
torch.LongTensor形状为(batch_size, sequence_length)，可选): 用于计算掩码语言建模损失的标签。索引应在[0, ..., config.vocab_size]或 -100 中（参见input_ids文档字符串）。索引设置为-100的标记将被忽略（掩码），损失仅针对标签在[0, ..., config.vocab_size]中的标记计算。num_logits_to_keep (
int或None，可选): 为最后num_logits_to_keep个标记计算 logits。如果为None，则为所有input_ids计算 logits。生成仅需要最后一个标记的 logits，并且仅针对该标记计算它们可以节省内存，这对于长序列来说非常重要。
返回值
transformers.modeling_outputs.MoeCausalLMOutputWithPast 或 tuple(torch.FloatTensor)
一个 transformers.modeling_outputs.MoeCausalLMOutputWithPast 或 torch.FloatTensor 的元组（如果传递了 return_dict=False 或 config.return_dict=False），包含根据配置 (JambaConfig) 和输入的不同元素。
-
loss (
torch.FloatTensor形状为(1,)，可选，在提供labels时返回) — 语言建模损失（用于下一个标记预测）。 -
logits (
torch.FloatTensor形状为(batch_size, sequence_length, config.vocab_size)) — 语言建模头的预测分数（每个词汇标记在 SoftMax 之前的分数）。 -
aux_loss (
torch.FloatTensor，可选，在提供labels时返回) — 稀疏模块的 aux_loss。 -
router_logits (
tuple(torch.FloatTensor)，可选，在传递output_router_probs=True和config.add_router_probs=True或config.output_router_probs=True时返回) —torch.FloatTensor的元组（每层一个），形状为(batch_size, sequence_length, num_experts)。由 MoE 路由器计算的原始路由器 logits（softmax 后），这些项用于计算专家混合模型的辅助损失。
-
past_key_values (
tuple(tuple(torch.FloatTensor))，可选，在传递use_cache=True或config.use_cache=True时返回) —tuple(torch.FloatTensor)的元组，长度为config.n_layers，每个元组包含 2 个形状为(batch_size, num_heads, sequence_length, embed_size_per_head)的张量。包含预先计算的隐藏状态（自注意力块中的键和值），可用于（参见
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 后的注意力权重，用于计算自注意力头中的加权平均值。
JambaForCausalLM 正向方法，覆盖 __call__ 特殊方法。
Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the pre and post processing steps while the latter silently ignores them.
示例
>>> from transformers import AutoTokenizer, JambaForCausalLM
>>> model = JambaForCausalLM.from_pretrained("ai21labs/Jamba-v0.1")
>>> tokenizer = AutoTokenizer.from_pretrained("ai21labs/Jamba-v0.1")
>>> 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."JambaForSequenceClassification
class transformers.JambaForSequenceClassification
< source >( config )
参数
- config (JambaConfig) — 模型配置类，包含模型的所有参数。使用配置文件初始化不会加载与模型关联的权重，仅加载配置。查看 from_pretrained() 方法来加载模型权重。
顶部带有序列分类头的 Jamba 模型（线性层）。
JambaForSequenceClassification 使用最后一个标记来进行分类，就像其他因果模型（例如 GPT-2）一样。
由于它对最后一个标记进行分类，因此需要知道最后一个标记的位置。如果在配置中定义了 pad_token_id，它会找到每行中不是填充标记的最后一个标记。如果没有定义 pad_token_id，它会简单地取批次中每行的最后一个值。由于在传递了 inputs_embeds 而不是 input_ids 时它无法猜测填充标记，因此它会执行相同的操作（取批次中每行的最后一个值）。
此模型继承自 PreTrainedModel。查看超类文档，了解库为其所有模型实现的通用方法（例如下载或保存、调整输入嵌入的大小、修剪头等）。
该模型也是一个 PyTorch torch.nn.Module 子类。将其用作常规 PyTorch 模块，并参考 PyTorch 文档以获取有关一般用法和行为的所有事项。
forward
< source >( input_ids: Optional = None attention_mask: Optional = None position_ids: Optional = None past_key_values: Union = None inputs_embeds: Optional = None labels: Optional = None use_cache: Optional = None output_attentions: Optional = None output_hidden_states: Optional = None return_dict: Optional = None )
参数
- input_ids (
torch.LongTensor形状为(batch_size, sequence_length)) — 词汇表中输入序列标记的索引。如果您提供填充，默认情况下将忽略填充。索引可以使用 AutoTokenizer 获得。有关详细信息，请参阅 PreTrainedTokenizer.encode() 和 PreTrainedTokenizer.call()。
- attention_mask (
torch.Tensor形状为(batch_size, sequence_length)，可选) — 掩码以避免对填充标记索引执行注意力。在[0, 1]中选择掩码值：- 1 表示未被掩盖的标记，
- 0 表示被掩盖的标记。
索引可以使用 AutoTokenizer 获得。有关详细信息，请参阅 PreTrainedTokenizer.encode() 和 PreTrainedTokenizer.call()。
如果使用
past_key_values，可以选择仅输入最后一个input_ids（见past_key_values）。如果您想更改填充行为，您应该阅读
modeling_opt._prepare_decoder_attention_mask并根据您的需要进行修改。有关默认策略的更多信息，请参阅 论文 中的图表 1。- 1 表示头部未被掩盖，
- 0 表示头部被掩盖。
- position_ids (
torch.LongTensor形状为(batch_size, sequence_length)，可选) — 位置嵌入中每个输入序列标记的位置索引。在[0, config.n_positions - 1]范围内选择。 - past_key_values (
HybridMambaAttentionDynamicCache，可选，当传递use_cache=True或config.use_cache=True时返回) — 一个 HybridMambaAttentionDynamicCache 对象，包含预先计算的隐藏状态（自注意力块中的键和值，以及 Mamba 块中的卷积和 SSM 状态），这些状态可用于（见past_key_values输入）加速顺序解码。键和值缓存张量具有形状(batch_size, num_heads, seq_len, head_dim)。卷积和 SSM 状态张量分别具有形状(batch_size, d_inner, d_conv)和(batch_size, d_inner, d_state)。有关详细信息，请参阅HybridMambaAttentionDynamicCache类。如果使用
past_key_values，用户可以选择仅输入形状为(batch_size, 1)的最后一个input_ids（那些没有将其过去键值状态提供给此模型的input_ids），而不是形状为(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。 - output_router_logits (
bool，可选) — 是否返回所有路由器的 logits。它们对于计算路由器损失很有用，在推理期间不应返回。 - return_dict (
bool, 可选) — 是否返回一个 ModelOutput 而不是一个简单的元组。 - cache_position (
torch.LongTensor形状为(sequence_length), 可选) — 指示输入序列令牌在序列中的位置的索引。与position_ids相反，此张量不受填充的影响。它用于在正确的位置更新缓存并推断完整的序列长度。 - labels (
torch.LongTensor形状为(batch_size,), 可选) — 用于计算序列分类/回归损失的标签。索引应在[0, ..., config.num_labels - 1]中。如果config.num_labels == 1，则计算回归损失（均方损失），如果config.num_labels > 1，则计算分类损失（交叉熵）。
The JambaForSequenceClassification forward method, overrides the __call__ special method.
Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the pre and post processing steps while the latter silently ignores them.