Transformers 文档

OLMoE

Hugging Face's logo
加入 Hugging Face 社区

并获得增强的文档体验

开始使用

该模型于 2024-09-03 发布,并于 2024-09-03 添加到 Hugging Face Transformers。

PyTorch FlashAttention SDPA

OLMoE

OLMoE 是一个稀疏的专家混合 (MoE) 语言模型,拥有 7B 参数,但每个输入 token 仅使用 1B 参数。它具有与密集模型相似的推理成本,但训练速度快约 3 倍。OLMoE 在每层使用 64 个小型专家进行细粒度路由,并采用无丢弃 token 的路由算法。

您可以在 OLMoE 集合下找到所有原始 OLMoE 检查点。

此模型由 Muennighoff 贡献。

单击右侧边栏中的 OLMoE 模型,了解更多关于如何将 OLMoE 应用于不同语言任务的示例。

下面的示例演示了如何使用 PipelineAutoModel 类生成文本。

<hfoptions id="usage"> <hfoption id="Pipeline">
import torch
from transformers import pipeline

pipe = pipeline(
    task="text-generation",
    model="allenai/OLMoE-1B-7B-0125",
    dtype=torch.float16,
    device=0,
)

result = pipe("Dionysus is the god of")
print(result)
</hfoption> <hfoption id="AutoModel">
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from accelerate import Accelerator

device = Accelerator().device

model = AutoModelForCausalLM.from_pretrained("allenai/OLMoE-1B-7B-0924", attn_implementation="sdpa", dtype="auto", device_map="auto").to(device)
tokenizer = AutoTokenizer.from_pretrained("allenai/OLMoE-1B-7B-0924")

inputs = tokenizer("Bitcoin is", return_tensors="pt")
inputs = {k: v.to(device) for k, v in inputs.items()}
output = model.generate(**inputs, max_length=64)
print(tokenizer.decode(output[0]))

量化

量化通过使用较低精度的权重来降低大型模型的内存负担。有关更多可用的量化后端,请参阅 量化概述。下面的示例使用 bitsandbytes 将权重量化为 4 位。

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from accelerate import Accelerator

device = Accelerator().device

quantization_config = BitsAndBytesConfig(
   load_in_4bit=True,
   bnb_4bit_compute_dtype=torch.float16,
   bnb_4bit_use_double_quant=True,
   bnb_4bit_quant_type="nf4"
)

model = AutoModelForCausalLM.from_pretrained("allenai/OLMoE-1B-7B-0924", attn_implementation="sdpa", dtype="auto", device_map="auto", quantization_config=quantization_config).to(device)
tokenizer = AutoTokenizer.from_pretrained("allenai/OLMoE-1B-7B-0924")

inputs = tokenizer("Bitcoin is", return_tensors="pt")
inputs = {k: v.to(device) for k, v in inputs.items()}
output = model.generate(**inputs, max_length=64)
print(tokenizer.decode(output[0]))

OlmoeConfig

class transformers.OlmoeConfig

< >

( vocab_size: int | None = 50304 hidden_size: int | None = 2048 intermediate_size: int | None = 2048 num_hidden_layers: int | None = 16 num_attention_heads: int | None = 16 num_key_value_heads: int | None = None hidden_act: str | None = 'silu' max_position_embeddings: int | None = 4096 initializer_range: float | None = 0.02 rms_norm_eps: int | None = 1e-05 use_cache: bool | None = True pad_token_id: int | None = 1 bos_token_id: int | None = None eos_token_id: int | None = 50279 tie_word_embeddings: int | None = False rope_parameters: transformers.modeling_rope_utils.RopeParameters | dict[str, transformers.modeling_rope_utils.RopeParameters] | None = None attention_bias: bool | None = False attention_dropout: float | None = 0.0 clip_qkv: bool | None = None num_experts_per_tok: int | None = 8 num_experts: int | None = 64 output_router_logits: bool | None = False router_aux_loss_coef: float | None = 0.01 norm_topk_prob: bool | None = False **kwargs )

参数

  • vocab_size (int, optional, defaults to 50304) — OLMoE 模型词汇表大小。定义了在调用 OlmoeModel 时传入的 inputs_ids 可以表示的不同 token 的数量。
  • hidden_size (int, optional, defaults to 2048) — 隐藏表示的维度。
  • intermediate_size (int, optional, defaults to 2048) — MLP 表示的维度。
  • num_hidden_layers (int, optional, defaults to 16) — Transformer 解码器中的隐藏层数。
  • num_attention_heads (int, optional, defaults to 16) — Transformer 解码器中每个注意力层的注意力头数。
  • num_key_value_heads (int, optional) — 这是实现分组查询注意力 (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 or function, optional, defaults to "silu") — 解码器中的非线性激活函数(函数或字符串)。
  • max_position_embeddings (int, optional, defaults to 4096) — 该模型可能使用的最大序列长度。
  • initializer_range (float, optional, defaults to 0.02) — 用于初始化所有权重矩阵的 truncated_normal_initializer 的标准差。
  • rms_norm_eps (float, optional, defaults to 1e-05) — rms normalization 层使用的 epsilon。
  • use_cache (bool, optional, defaults to True) — 模型是否应返回最后一个 key/values 注意力(并非所有模型都使用)。仅当 config.is_decoder=True 时相关。
  • pad_token_id (int, optional, defaults to 1) — 填充 token ID。
  • bos_token_id (int, optional) — 流开始 token ID。
  • eos_token_id (int, optional, defaults to 50279) — 流结束 token ID。
  • tie_word_embeddings (bool, optional, defaults to False) — 是否绑定词嵌入。
  • rope_parameters (RopeParameters, optional) — 包含 RoPE 嵌入配置参数的字典。该字典应包含 rope_theta 的值,并且如果希望使用更长的 max_position_embeddings 进行 RoPE,还可以包含用于缩放的参数。
  • attention_bias (bool, defaults to False, optional, defaults to False) — 在自注意力期间是否在查询、键、值和输出投影层中使用偏置。
  • attention_dropout (float, optional, defaults to 0.0) — 注意力概率的 dropout 比率。
  • clip_qkv (float, optional) — 如果不为 None,则查询、键和值注意力状态的元素将被裁剪,使其绝对值不超过此值。
  • num_experts_per_tok (int, optional, defaults to 8) — 选择的专家数量。
  • num_experts (int, optional, defaults to 64) — 专家数量。
  • output_router_logits (bool, optional, defaults to False) — 是否返回路由器的 logits。启用此选项还将允许模型输出辅助损失,包括负载均衡损失和路由器 z-loss。
  • router_aux_loss_coef (float, optional, defaults to 0.01) — 用于总损失的辅助损失因子。
  • norm_topk_prob (bool, optional, defaults to 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.config

OlmoeModel

class transformers.OlmoeModel

< >

( config: OlmoeConfig )

参数

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

不带任何特定头部的 OLMoE 裸模型,输出原始隐藏状态。

此模型继承自 PreTrainedModel。查看其父类文档,了解库为所有模型实现的通用方法(例如下载或保存、调整输入嵌入大小、修剪头等)。

此模型也是一个 PyTorch torch.nn.Module 子类。像普通的 PyTorch Module 一样使用它,并参考 PyTorch 文档了解一般用法和行为的所有相关信息。

forward

< >

( input_ids: torch.LongTensor | None = None attention_mask: torch.Tensor | None = None position_ids: torch.LongTensor | None = None past_key_values: transformers.cache_utils.Cache | None = None inputs_embeds: torch.FloatTensor | None = None use_cache: bool | None = None cache_position: torch.LongTensor | None = None **kwargs: typing_extensions.Unpack[transformers.utils.generic.TransformersKwargs] ) transformers.modeling_outputs.MoeModelOutputWithPasttuple(torch.FloatTensor)

参数

  • input_ids (torch.LongTensor of shape (batch_size, sequence_length), optional) — 词汇表中输入序列 token 的索引。默认情况下会忽略填充。

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

    什么是输入 ID?

  • attention_mask (torch.Tensor of shape (batch_size, sequence_length), optional) — 用于避免对填充 token 索引执行 attention 的掩码。掩码值选择在 [0, 1] 中:

    • 1 表示未被掩码的 token,
    • 0 表示被掩码的 token。

    什么是 attention mask?

  • position_ids (torch.LongTensor of shape (batch_size, sequence_length), optional) — 每个输入序列 token 在位置嵌入中的位置索引。选择在范围 [0, config.n_positions - 1] 内。

    什么是位置 ID?

  • past_key_values (~cache_utils.Cache, optional) — 可以用于加速顺序解码的预计算隐藏状态(自注意力块和交叉注意力块中的键和值)。这通常包括在 use_cache=Trueconfig.use_cache=True 时,模型在先前解码阶段返回的 past_key_values

    只有 Cache 实例被允许作为输入,请参阅我们的 kv cache 指南。如果不传递 past_key_values,默认将初始化 DynamicCache

    模型将输出与输入相同的缓存格式。

    如果使用了 past_key_values,用户应只输入未处理的 input_ids(那些没有将它们的 past key value 状态传递给此模型的),形状为 (batch_size, unprocessed_length),而不是所有 input_ids 的形状 (batch_size, sequence_length)

  • inputs_embeds (torch.FloatTensor of shape (batch_size, sequence_length, hidden_size), optional) — 可选地,您可以直接传递嵌入表示,而不是传递 input_ids。如果您想比模型内部的嵌入查找矩阵更好地控制如何将 input_ids 索引转换为关联向量,这将非常有用。
  • use_cache (bool, optional) — 如果设置为 True,则会返回 past_key_values 键值状态,并可用于加速解码(请参阅 past_key_values)。
  • cache_position (torch.LongTensor of shape (sequence_length), optional) — 描述输入序列 token 在序列中位置的索引。与 position_ids 不同,此张量不受填充影响。它用于在正确位置更新缓存并推断完整的序列长度。

返回

transformers.modeling_outputs.MoeModelOutputWithPasttuple(torch.FloatTensor)

根据配置 (OlmoeConfig) 和输入,包含各种元素的 transformers.modeling_outputs.MoeModelOutputWithPasttorch.FloatTensor 元组(如果传递了 return_dict=Falseconfig.return_dict=False)。

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

  • past_key_values (Cache, optional, 当传递 use_cache=True 或当 config.use_cache=True 时返回) — 它是 Cache 实例。更多详情,请参阅我们的 kv cache 指南

    Contains pre-computed hidden-states (key and values in the self-attention blocks and optionally if config.is_encoder_decoder=True in the cross-attention blocks) that can be used (see past_key_values input) to speed up sequential decoding.

  • hidden_states (tuple(torch.FloatTensor), optional, 当传递 output_hidden_states=True 或当 config.output_hidden_states=True 时返回) — torch.FloatTensor 的元组(一个用于嵌入层的输出,如果模型有嵌入层;+一个用于每个层的输出),形状为 (batch_size, sequence_length, hidden_size)

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

  • attentions (tuple(torch.FloatTensor), optional, 当传递 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=Trueconfig.add_router_probs=True 时,或 config.output_router_probs=True 时返回) — 形状为 (batch_size, sequence_length, num_experts)torch.FloatTensor 元组(每一层一个)。

    由 MoE 路由器计算的原始路由器对数(softmax 后),这些术语用于计算专家混合模型的辅助损失。

覆盖 __call__ 特殊方法的 OlmoeModel forward 方法。

虽然 forward pass 的实现需要在此函数中定义,但你应该在之后调用 Module 实例而不是这个,因为前者负责运行预处理和后处理步骤,而后者会静默地忽略它们。

OlmoeForCausalLM

class transformers.OlmoeForCausalLM

< >

( 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 (OlmoeForCausalLM) — 模型的配置类,包含模型的所有参数。使用配置文件初始化不会加载与模型相关的权重,只加载配置。请查看 from_pretrained() 方法来加载模型权重。

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

此模型继承自 PreTrainedModel。查看其父类文档,了解库为所有模型实现的通用方法(例如下载或保存、调整输入嵌入大小、修剪头等)。

此模型也是一个 PyTorch torch.nn.Module 子类。像普通的 PyTorch Module 一样使用它,并参考 PyTorch 文档了解一般用法和行为的所有相关信息。

forward

< >

( input_ids: torch.LongTensor | None = None attention_mask: torch.Tensor | None = None position_ids: torch.LongTensor | None = None past_key_values: transformers.cache_utils.Cache | None = None inputs_embeds: torch.FloatTensor | None = None labels: torch.LongTensor | None = None use_cache: bool | None = None output_router_logits: bool | None = None cache_position: torch.LongTensor | None = None logits_to_keep: int | torch.Tensor = 0 **kwargs: typing_extensions.Unpack[transformers.utils.generic.TransformersKwargs] ) transformers.modeling_outputs.MoeCausalLMOutputWithPasttuple(torch.FloatTensor)

参数

  • input_ids (torch.LongTensor of shape (batch_size, sequence_length), optional) — 词汇表中输入序列 token 的索引。默认情况下会忽略填充。

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

    什么是输入 ID?

  • attention_mask (torch.Tensor of shape (batch_size, sequence_length), optional) — Mask to avoid performing attention on padding token indices. Mask values selected in [0, 1]:

    • 1 for tokens that are not masked,
    • 0 for tokens that are masked.

    What are attention masks?

  • position_ids (torch.LongTensor of shape (batch_size, sequence_length), optional) — Indices of positions of each input sequence tokens in the position embeddings. Selected in the range [0, config.n_positions - 1].

    What are position IDs?

  • past_key_values (~cache_utils.Cache, optional) — Pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention blocks) that can be used to speed up sequential decoding. This typically consists in the past_key_values returned by the model at a previous stage of decoding, when use_cache=True or config.use_cache=True.

    Only Cache instance is allowed as input, see our kv cache guide. If no past_key_values are passed, DynamicCache will be initialized by default.

    The model will output the same cache format that is fed as input.

    If past_key_values are used, the user is expected to input only unprocessed input_ids (those that don’t have their past key value states given to this model) of shape (batch_size, unprocessed_length) instead of all input_ids of shape (batch_size, sequence_length).

  • inputs_embeds (torch.FloatTensor of shape (batch_size, sequence_length, hidden_size), optional) — Optionally, instead of passing input_ids you can choose to directly pass an embedded representation. This is useful if you want more control over how to convert input_ids indices into associated vectors than the model’s internal embedding lookup matrix.
  • labels (torch.LongTensor of shape (batch_size, sequence_length), optional) — Labels for computing the masked language modeling loss. Indices should either be in [0, ..., config.vocab_size] or -100 (see input_ids docstring). Tokens with indices set to -100 are ignored (masked), the loss is only computed for the tokens with labels in [0, ..., config.vocab_size].
  • use_cache (bool, optional) — If set to True, past_key_values key value states are returned and can be used to speed up decoding (see past_key_values).
  • output_router_logits (bool, optional) — Whether or not to return the logits of all the routers. They are useful for computing the router loss, and should not be returned during inference.
  • cache_position (torch.LongTensor of shape (sequence_length), optional) — Indices depicting the position of the input sequence tokens in the sequence. Contrarily to position_ids, this tensor is not affected by padding. It is used to update the cache in the correct position and to infer the complete sequence length.
  • logits_to_keep (Union[int, torch.Tensor], optional, defaults to 0) — If an int, compute logits for the last logits_to_keep tokens. If 0, calculate logits for all input_ids (special case). Only last token logits are needed for generation, and calculating them only for that token can save memory, which becomes pretty significant for long sequences or large vocabulary size. If a torch.Tensor, must be 1D corresponding to the indices to keep in the sequence length dimension. This is useful when using packed tensor format (single dimension for batch and sequence length).

返回

transformers.modeling_outputs.MoeCausalLMOutputWithPasttuple(torch.FloatTensor)

A transformers.modeling_outputs.MoeCausalLMOutputWithPast or a tuple of torch.FloatTensor (if return_dict=False is passed or when config.return_dict=False) comprising various elements depending on the configuration (MoeCausalLMOutputWithPast) and inputs.

  • 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=Trueconfig.add_router_probs=True 时,或 config.output_router_probs=True 时返回) — 形状为 (batch_size, sequence_length, num_experts)torch.FloatTensor 元组(每一层一个)。

    由 MoE 路由器计算的原始路由器对数(softmax 后),这些术语用于计算专家混合模型的辅助损失。

  • past_key_values (Cache, optional, 当传递 use_cache=True 或当 config.use_cache=True 时返回) — 它是 Cache 实例。更多详情,请参阅我们的 kv cache 指南

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

  • hidden_states (tuple(torch.FloatTensor), optional, 当传递 output_hidden_states=True 或当 config.output_hidden_states=True 时返回) — torch.FloatTensor 的元组(一个用于嵌入层的输出,如果模型有嵌入层;+一个用于每个层的输出),形状为 (batch_size, sequence_length, hidden_size)

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

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

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

The OlmoeForCausalLM forward method, overrides the __call__ special method.

虽然 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'
在 GitHub 上更新

© . This site is unofficial and not affiliated with Hugging Face, Inc.