Transformers 文档
OLMoE
并获得增强的文档体验
开始使用
OLMoE
概述
OLMoE 模型在 OLMoE: Open Mixture-of-Experts Language Models 这篇论文中被提出，作者是 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 是一系列开放语言模型，使用稀疏的混合专家模型，旨在促进语言模型的科学研究。我们发布了所有代码、检查点、日志以及训练这些模型的详细信息。
论文摘要如下：
我们介绍 OLMoE，这是一个完全开放、最先进的语言模型，利用稀疏的混合专家模型 (MoE)。OLMoE-1B-7B 具有 70 亿 (B) 参数，但每个输入 token 仅使用 10 亿。我们在 5 万亿 token 上对其进行预训练，并进一步调整以创建 OLMoE-1B-7B-Instruct。我们的模型优于所有具有相似活跃参数的可用模型，甚至超越了更大的模型，如 Llama2-13B-Chat 和 DeepSeekMoE-16B。我们展示了关于 MoE 训练的各种实验，分析了我们模型中的路由，显示出高度专业化，并开源了我们工作的所有方面：模型权重、训练数据、代码和日志。
此模型由Muennighoff贡献。 原始代码可以在这里找到。
OlmoeConfig
class transformers.OlmoeConfig
< source >( 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 norm_topk_prob = False **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）的 key_value 头部的数量。 如果num_key_value_heads=num_attention_heads，模型将使用多头注意力（Multi Head Attention, MHA）；如果num_key_value_heads=1，模型将使用多查询注意力（Multi Query Attention, MQA）；否则使用 GQA。 当将多头检查点转换为 GQA 检查点时，每个组 key 和 value 头部应通过对该组内所有原始头部进行平均池化来构建。 有关更多详细信息，请查看这篇论文。 如果未指定，则默认为num_attention_heads。 - hidden_act (
str或function, 可选，默认为"silu") — 解码器中的非线性激活函数（函数或字符串）。 - max_position_embeddings (
int, 可选，默认为 4096) — 此模型可能使用的最大序列长度。 - initializer_range (
float, 可选，默认为 0.02) — 用于初始化所有权重矩阵的 truncated_normal_initializer 的标准差。 - rms_norm_eps (
float, 可选，默认为 1e-05) — rms 归一化层使用的 epsilon 值。 - use_cache (
bool, 可选，默认为True) — 模型是否应返回最后的 key/values 注意力（并非所有模型都使用）。 仅当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 嵌入的缩放配置的字典。 目前支持两种缩放策略：线性（linear）和动态（dynamic）。 它们的缩放因子必须是大于 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) — 在自注意力期间，是否在 query、key、value 和输出投影层中使用偏置。 - attention_dropout (
float, 可选，默认为 0.0) — 注意力概率的 dropout 比率。 - clip_qkv (
float, 可选) — 如果不为None，则 query、key 和 value 注意力状态的元素将被裁剪，使其绝对值不超过此值。 - num_experts_per_tok (
int, 可选，默认为 8) — 每个 token 选择的专家数量。 - num_experts (
int, 可选，默认为 64) — 路由专家的数量。 - output_router_logits (
bool, 可选，默认为False) — 模型是否应返回路由器 logits。 启用此选项还将允许模型输出辅助损失，包括负载均衡损失和路由器 z-loss。 - 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() 方法以加载模型权重。
- config — OlmoeConfig
裸 Olmoe 模型输出原始隐藏状态，顶部没有任何特定的头部。此模型继承自 PreTrainedModel。查看超类文档以了解库为其所有模型实现的通用方法（例如下载或保存、调整输入嵌入大小、剪枝头部等）。
此模型也是 PyTorch torch.nn.Module 子类。将其用作常规 PyTorch 模块，并参考 PyTorch 文档以了解与一般用法和行为相关的所有事项。
Transformer 解码器由 config.num_hidden_layers 层组成。每一层都是一个 OlmoeDecoderLayer
forward
< 源代码 >( input_ids: LongTensor = None attention_mask: typing.Optional[torch.Tensor] = None position_ids: typing.Optional[torch.LongTensor] = None past_key_values: typing.Union[typing.List[torch.FloatTensor], transformers.cache_utils.Cache, 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 )
参数
- input_ids (`torch.LongTensor`，形状为 `(batch_size, sequence_length)`) — 词汇表中输入序列 tokens 的索引。如果您提供填充，默认情况下填充将被忽略。
索引可以使用 AutoTokenizer 获得。请参阅 PreTrainedTokenizer.encode() 和 PreTrainedTokenizer.call() 以了解详情。
- attention_mask (`torch.Tensor`，形状为 `(batch_size, sequence_length)`，*可选*) — 用于避免对填充 token 索引执行 attention 的掩码。掩码值在 `[0, 1]` 中选择：
- 1 表示 未被掩盖 的 tokens，
- 0 表示 被掩盖 的 tokens。
索引可以使用 AutoTokenizer 获得。请参阅 PreTrainedTokenizer.encode() 和 PreTrainedTokenizer.call() 以了解详情。
如果使用
past_key_values，则可以选择仅输入最后的input_ids（请参阅past_key_values）。如果您想更改填充行为，您应该阅读
modeling_opt._prepare_decoder_attention_mask并根据您的需求进行修改。有关默认策略的更多信息，请参阅 论文 中的图 1。- 1 表示 head 未被掩盖，
- 0 表示 head 被掩盖。
- position_ids (`torch.LongTensor`，形状为 `(batch_size, sequence_length)`，*可选*) — 每个输入序列 tokens 在位置嵌入中的位置索引。在范围 `[0, config.n_positions - 1]` 中选择。
- past_key_values (`Cache` 或 `tuple(tuple(torch.FloatTensor))`，*可选*) — 预先计算的隐藏状态（self-attention 块和 cross-attention 块中的键和值），可用于加速顺序解码。这通常包含由模型在先前解码阶段返回的 `past_key_values`，当 `use_cache=True` 或 `config.use_cache=True` 时。
允许两种格式：
- Cache 实例；
- 长度为 `config.n_layers` 的 `tuple(torch.FloatTensor)` 元组，每个元组有两个形状为 `(batch_size, num_heads, sequence_length, embed_size_per_head)` 的 tensors。这也称为旧版缓存格式。
模型将输出与输入相同的缓存格式。如果未传递
past_key_values，则将返回旧版缓存格式。如果使用
past_key_values，用户可以选择仅输入最后的input_ids（那些没有将其过去的键值状态提供给此模型的输入 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`，*可选*) — 是否返回所有 attention 层的 attention tensors。有关更多详细信息，请参阅返回 tensors 下的 `attentions`。
- output_hidden_states (`bool`，*可选*) — 是否返回所有层的隐藏状态。有关更多详细信息，请参阅返回 tensors 下的 `hidden_states`。
- output_router_logits (`bool`，*可选*) — 是否返回所有 routers 的 logits。它们对于计算 router 损失很有用，并且不应在推理期间返回。
- return_dict (`bool`，*可选*) — 是否返回 ModelOutput 而不是普通元组。
- cache_position (`torch.LongTensor`，形状为 `(sequence_length)`，*可选*) — 描述输入序列 tokens 在序列中位置的索引。与 `position_ids` 相反，此 tensor 不受填充的影响。它用于在正确的位置更新缓存并推断完整序列长度。
OlmoeModel 的 forward 方法覆盖了 `__call__` 特殊方法。
尽管 forward 传递的配方需要在该函数内定义，但之后应该调用 `Module` 实例而不是此函数，因为前者负责运行预处理和后处理步骤，而后者会默默地忽略它们。
OlmoeForCausalLM
forward
< 源代码 >( input_ids: LongTensor = None attention_mask: typing.Optional[torch.Tensor] = None position_ids: typing.Optional[torch.LongTensor] = None past_key_values: typing.Optional[typing.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 or tuple(torch.FloatTensor)
参数
- input_ids (
torch.LongTensor，形状为(batch_size, sequence_length)) — 词汇表中输入序列 tokens 的索引。如果您提供填充，默认情况下填充将被忽略。索引可以使用 AutoTokenizer 获得。 请参阅 PreTrainedTokenizer.encode() 和 PreTrainedTokenizer.call() 以了解详细信息。
- attention_mask (
torch.Tensor，形状为(batch_size, sequence_length)，可选) — 用于避免在填充 token 索引上执行 attention 的掩码。掩码值在[0, 1]中选择：- 1 表示 未被掩码 的 tokens，
- 0 表示 被掩码 的 tokens。
索引可以使用 AutoTokenizer 获得。 请参阅 PreTrainedTokenizer.encode() 和 PreTrainedTokenizer.call() 以了解详细信息。
如果使用了
past_key_values，则可以选择仅输入最后的input_ids（请参阅past_key_values）。如果您想更改填充行为，您应该阅读
modeling_opt._prepare_decoder_attention_mask并根据您的需要进行修改。 有关默认策略的更多信息，请参阅 论文 中的图 1。- 1 表示 head 未被掩码，
- 0 表示 head 被掩码。
- position_ids (
torch.LongTensor，形状为(batch_size, sequence_length)，可选) — 每个输入序列 token 在位置嵌入中的位置索引。在范围[0, config.n_positions - 1]中选择。 - past_key_values (
Cache或tuple(tuple(torch.FloatTensor))，可选) — 预先计算的隐藏状态（self-attention 块和 cross-attention 块中的 key 和 values），可用于加速顺序解码。 这通常包括在解码的先前阶段由模型返回的past_key_values，当use_cache=True或config.use_cache=True时。允许两种格式：
- Cache 实例；
- 长度为
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（那些没有将其过去 key value 状态提供给此模型的 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，可选) — 是否返回所有 attention 层的 attention 张量。 有关更多详细信息，请参阅返回张量下的attentions。 - output_hidden_states (
bool，可选) — 是否返回所有层的隐藏状态。 有关更多详细信息，请参阅返回张量下的hidden_states。 - output_router_logits (
bool，可选) — 是否返回所有路由器的 logits。 它们对于计算路由器损失很有用，并且不应在推理期间返回。 - return_dict (
bool，可选) — 是否返回 ModelOutput 而不是纯元组。 - cache_position (
torch.LongTensor，形状为(sequence_length)，可选) — 索引，描述输入序列 tokens 在序列中的位置。 与position_ids相反，此张量不受填充的影响。 它用于在正确的位置更新缓存并推断完整序列长度。 - labels (
torch.LongTensor，形状为(batch_size, sequence_length)，可选) — 用于计算掩码语言建模损失的标签。 索引应在[0, ..., config.vocab_size]或 -100 中（请参阅input_ids文档字符串）。 索引设置为-100的 tokens 将被忽略（掩码），损失仅针对标签在[0, ..., config.vocab_size]中的 tokens 计算。 - logits_to_keep (
int或torch.Tensor，可选) — 如果是int，则计算最后logits_to_keep个 tokens 的 logits。 如果为0，则计算所有input_ids的 logits（特殊情况）。 只需要最后一个 token logits 用于生成，并且仅针对该 token 计算它们可以节省内存，这对于长序列或大词汇量大小而言变得非常重要。 如果是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 (
torch.FloatTensor，形状为(batch_size, sequence_length, config.vocab_size)) — 语言建模 head 的预测分数（SoftMax 之前每个词汇表 token 的分数）。 -
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时返回) — 形状为(batch_size, sequence_length, num_experts)的torch.FloatTensor元组（每层一个）。原始路由器 logits（后 softmax），由 MoE 路由器计算，这些项用于计算专家混合模型的辅助损失。
-
past_key_values (
tuple(tuple(torch.FloatTensor))，可选，当传递use_cache=True或当config.use_cache=True时返回) — 长度为config.n_layers的tuple(torch.FloatTensor)元组，其中每个元组有 2 个形状为(batch_size, num_heads, sequence_length, embed_size_per_head)的张量）包含预先计算的隐藏状态（self-attention 块中的 key 和 values），可用于加速顺序解码（请参阅
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时返回) — 形状为(batch_size, num_heads, sequence_length, sequence_length)的torch.FloatTensor元组（每层一个）。attention softmax 之后的 Attention 权重，用于计算 self-attention heads 中的加权平均值。
OlmoeForCausalLM forward 方法，覆盖了 __call__ 特殊方法。
尽管 forward 传递的配方需要在该函数内定义，但之后应该调用 `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'