Transformers 文档
OLMoE
并获得增强的文档体验
开始使用
OLMoE
概述
OLMoE 模型由 Niklas Muennighoff, Luca Soldaini, Dirk Groeneveld, Kyle Lo, Jacob Morrison, Sewon Min, Weijia Shi, Pete Walsh, Oyvind Tafjord, Nathan Lambert, Yuling Gu, Shane Arora, Akshita Bhagia, Dustin Schwenk, David Wadden, Alexander Wettig, Binyuan Hui, Tim Dettmers, Douwe Kiela, Ali Farhadi, Noah A. Smith, Pang Wei Koh, Amanpreet Singh, Hannaneh Hajishirzi 在 OLMoE: Open Mixture-of-Experts Language Models 中提出。
OLMoE 是一系列使用稀疏 **专**家 **混**合 **开**放 **语**言模型的模型（**O**pen **L**anguage **Mo**dels using sparse **M**ixture-**o**f-**E**xperts），旨在推动语言模型科学的发展。我们发布了训练这些模型所涉及的所有代码、检查点、日志和详细信息。
论文摘要如下：
我们引入了 OLMoE，一个完全开放的、最先进的语言模型，利用了稀疏专家混合（MoE）。OLMoE-1B-7B 拥有 70 亿（B）参数，但每个输入 token 仅使用 1B 参数。我们对其进行了 5 万亿 token 的预训练，并进一步调整以创建 OLMoE-1B-7B-Instruct。我们的模型在所有具有相似活跃参数的模型中表现优异，甚至超越了像 Llama2-13B-Chat 和 DeepSeekMoE-16B 这样更大的模型。我们展示了各种 MoE 训练实验，分析了模型中的路由，显示出高度专业化，并开源了我们工作的所有方面：模型权重、训练数据、代码和日志。
该模型由 Muennighoff 贡献。原始代码可以在这里找到。
OlmoeConfig
class transformers.OlmoeConfig
< 源 >( vocab_size = 50304 hidden_size = 2048 intermediate_size = 2048 num_hidden_layers = 16 num_attention_heads = 16 num_key_value_heads = None hidden_act = 'silu' max_position_embeddings = 4096 initializer_range = 0.02 rms_norm_eps = 1e-05 use_cache = True pad_token_id = 1 bos_token_id = None eos_token_id = 50279 tie_word_embeddings = False rope_theta = 10000.0 rope_scaling = None attention_bias = False attention_dropout = 0.0 clip_qkv = None num_experts_per_tok = 8 num_experts = 64 output_router_logits = False router_aux_loss_coef = 0.01 **kwargs )
参数
- vocab_size (
int, 可选, 默认为 50304) — OLMoE 模型的词汇表大小。定义了在调用 OlmoeModel 时inputs_ids可以表示的不同 token 的数量。 - hidden_size (
int, 可选, 默认为 2048) — 隐藏表示的维度。 - intermediate_size (
int, 可选, 默认为 2048) — MLP 表示的维度。 - num_hidden_layers (
int, 可选, 默认为 16) — Transformer 解码器中隐藏层的数量。 - num_attention_heads (
int, 可选, 默认为 16) — Transformer 解码器中每个注意力层的注意力头数量。 - num_key_value_heads (
int, 可选) — 用于实现分组查询注意力（Grouped Query Attention）的键值头数量。如果num_key_value_heads=num_attention_heads，模型将使用多头注意力（MHA）；如果num_key_value_heads=1，模型将使用多查询注意力（MQA），否则使用 GQA。将多头检查点转换为 GQA 检查点时，每个分组键和值头应通过对该组内的所有原始头进行均值池化来构建。有关这些缩放策略如何表现的更多详细信息，请查看 此论文。如果未指定，将默认为num_attention_heads。 - hidden_act (
str或function, 可选, 默认为"silu") — 解码器中的非线性激活函数（函数或字符串）。 - max_position_embeddings (
int, 可选, 默认为 4096) — 此模型可能使用的最大序列长度。 - initializer_range (
float, 可选, 默认为 0.02) — 用于初始化所有权重矩阵的截断正态初始化器的标准差。 - rms_norm_eps (
float, 可选, 默认为 1e-05) — rms 归一化层使用的 epsilon 值。 - use_cache (
bool, 可选, 默认为True) — 模型是否应返回最后一个键/值注意力（并非所有模型都使用）。仅在config.is_decoder=True时相关。 - pad_token_id (
int, 可选, 默认为 1) — 填充 token ID。 - bos_token_id (
int, 可选) — 流开始 token ID。 - eos_token_id (
int, 可选, 默认为 50279) — 流结束 token ID。 - tie_word_embeddings (
bool, 可选, 默认为False) — 是否绑定词嵌入。 - rope_theta (
float, 可选, 默认为 10000.0) — RoPE 嵌入的基期。 - rope_scaling (
Dict, 可选) — 包含 RoPE 嵌入缩放配置的字典。目前支持两种缩放策略：线性和动态。它们的缩放因子必须是大于 1 的浮点数。预期格式为{"type": 策略名称, "factor": 缩放因子}。使用此标志时，请勿将max_position_embeddings更新为预期的新最大值。有关这些缩放策略如何表现的更多信息，请参阅以下帖子：https://www.reddit.com/r/LocalLLaMA/comments/14mrgpr/dynamically_scaled_rope_further_increases/。这是一个实验性功能，未来版本可能会有重大 API 更改。 - attention_bias (
bool, 默认为False, 可选, 默认为False) — 在自注意力过程中，查询、键、值和输出投影层是否使用偏置。 - attention_dropout (
float, 可选, 默认为 0.0) — 注意力概率的 dropout 比率。 - clip_qkv (
float, 可选) — 如果不是None，查询、键和值注意力状态的元素将被裁剪，使其绝对值不超过此值。 - num_experts_per_tok (
int, 可选, 默认为 8) — 每个 token 选择的专家数量。 - num_experts (
int, 可选, 默认为 64) — 路由专家数量。 - output_router_logits (
bool, 可选, 默认为False) — 模型是否应返回路由器 logits。启用此选项还将允许模型输出辅助损失，包括负载平衡损失和路由器 z 损失。 - router_aux_loss_coef (
float, 可选, 默认为 0.01) — 总损失的辅助损失因子。 - norm_topk_prob (
bool, 可选, 默认为False) — 是否归一化 topk 概率。
这是用于存储 OlmoeModel 配置的配置类。它用于根据指定的参数实例化 OLMoE 模型，定义模型架构。使用默认值实例化配置将生成与 allenai/OLMoE-1B-7B-0924 类似的配置。
配置对象继承自 PretrainedConfig，可用于控制模型输出。有关更多信息，请阅读 PretrainedConfig 的文档。
>>> from transformers import OlmoeModel, OlmoeConfig
>>> # Initializing a OLMoE 7B A1B style configuration
>>> configuration = OlmoeConfig()
>>> # Initializing a model from the OLMoE 7B A1B style configuration
>>> model = OlmoeModel(configuration)
>>> # Accessing the model configuration
>>> configuration = model.configOlmoeModel
class transformers.OlmoeModel
< 源 >( config: OlmoeConfig )
参数
- config (OlmoeConfig) — 包含模型所有参数的模型配置类。用配置文件初始化并不会加载与模型相关的权重，只加载配置。请查看 from_pretrained() 方法以加载模型权重。
裸 Olmoe 模型，输出原始隐藏状态，顶部没有任何特定头部。
此模型继承自 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.Union[transformers.cache_utils.Cache, list[torch.FloatTensor], NoneType] = 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 output_router_logits: typing.Optional[bool] = None return_dict: typing.Optional[bool] = None cache_position: typing.Optional[torch.LongTensor] = None ) → transformers.modeling_outputs.MoeModelOutputWithPast 或 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。
- position_ids (
torch.LongTensor，形状为(batch_size, sequence_length)，可选) — 每个输入序列标记在位置嵌入中的位置索引。选择范围为[0, config.n_positions - 1]。 - past_key_values (
Union[~cache_utils.Cache, list[torch.FloatTensor], NoneType]) — 预先计算的隐藏状态（自注意力块和交叉注意力块中的键和值），可用于加速顺序解码。这通常由模型在解码的先前阶段返回的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，用户可以选择只输入形状为(batch_size, 1)的最后一个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不同，此张量不受填充影响。它用于在正确位置更新缓存并推断完整的序列长度。
返回
transformers.modeling_outputs.MoeModelOutputWithPast 或 tuple(torch.FloatTensor)
一个 transformers.modeling_outputs.MoeModelOutputWithPast 或一个 torch.FloatTensor 元组（如果传入 return_dict=False 或 config.return_dict=False），根据配置 (OlmoeConfig) 和输入包含各种元素。
-
last_hidden_state (
torch.FloatTensor, 形状为(batch_size, sequence_length, 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 后的注意力权重，用于计算自注意力头中的加权平均值。
-
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 路由器计算的原始路由器对数（softmax 后），这些术语用于计算专家混合模型的辅助损失。
OlmoeModel 的 forward 方法，覆盖了 __call__ 特殊方法。
尽管 forward pass 的配方需要在此函数中定义，但此后应调用 Module 实例而不是此函数，因为前者负责运行预处理和后处理步骤，而后者则默默地忽略它们。
OlmoeForCausalLM
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[list[torch.FloatTensor]] = 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 output_router_logits: typing.Optional[bool] = None return_dict: typing.Optional[bool] = None cache_position: typing.Optional[torch.LongTensor] = None logits_to_keep: typing.Union[int, torch.Tensor] = 0 **loss_kwargs ) → 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。
- position_ids (
torch.LongTensor，形状为(batch_size, sequence_length)，可选) — 每个输入序列标记在位置嵌入中的位置索引。选择范围为[0, config.n_positions - 1]。 - past_key_values (
list[torch.FloatTensor]，可选) — 预先计算的隐藏状态（自注意力块和交叉注意力块中的键和值），可用于加速顺序解码。这通常由模型在解码的先前阶段返回的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，用户可以选择只输入形状为(batch_size, 1)的最后一个input_ids（那些没有将其过去的键值状态提供给此模型的），而不是所有形状为(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。 - output_router_logits (
bool，可选) — 是否返回所有路由器的 logits。它们对于计算路由器损失很有用，不应在推理期间返回。 - return_dict (
bool，可选) — 是否返回 ModelOutput 而不是普通的元组。 - 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.MoeCausalLMOutputWithPast 或 tuple(torch.FloatTensor)
一个 transformers.modeling_outputs.MoeCausalLMOutputWithPast 或一个 torch.FloatTensor 元组（如果传入 return_dict=False 或 config.return_dict=False），根据配置 (OlmoeConfig) 和输入包含各种元素。
-
loss (
torch.FloatTensor形状为(1,)，可选，当提供labels时返回) — 语言建模损失（用于下一个 token 预测）。 -
logits (形状为
(batch_size, sequence_length, config.vocab_size)的torch.FloatTensor) — 语言建模头部的预测分数（SoftMax 之前的每个词汇标记的分数）。 -
aux_loss (
torch.FloatTensor，可选，当提供labels时返回) — 稀疏模块的辅助损失。 -
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 路由器计算的原始路由器对数（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 后的注意力权重，用于计算自注意力头中的加权平均值。
OlmoeForCausalLM 的 forward 方法，覆盖了 __call__ 特殊方法。
尽管 forward pass 的配方需要在此函数中定义，但此后应调用 Module 实例而不是此函数，因为前者负责运行预处理和后处理步骤，而后者则默默地忽略它们。
示例
>>> from transformers import AutoTokenizer, OlmoeForCausalLM
>>> model = OlmoeForCausalLM.from_pretrained("allenai/OLMoE-1B-7B-0924")
>>> tokenizer = AutoTokenizer.from_pretrained("allenai/OLMoE-1B-7B-0924")
>>> 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 sure if you’re conscious of this, but I’m'