Transformers 文档
OLMoE
并获得增强的文档体验
开始使用
该模型于 2024-09-03 发布,并于 2024-09-03 添加到 Hugging Face Transformers。
OLMoE
OLMoE 是一个稀疏的专家混合 (MoE) 语言模型,拥有 7B 参数,但每个输入 token 仅使用 1B 参数。它具有与密集模型相似的推理成本,但训练速度快约 3 倍。OLMoE 在每层使用 64 个小型专家进行细粒度路由,并采用无丢弃 token 的路由算法。
您可以在 OLMoE 集合下找到所有原始 OLMoE 检查点。
此模型由 Muennighoff 贡献。
单击右侧边栏中的 OLMoE 模型,了解更多关于如何将 OLMoE 应用于不同语言任务的示例。
下面的示例演示了如何使用 Pipeline 或 AutoModel 类生成文本。
<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)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
< source >( 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 (
strorfunction, 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 toTrue) — 模型是否应返回最后一个 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 toFalse) — 是否绑定词嵌入。 - rope_parameters (
RopeParameters, optional) — 包含 RoPE 嵌入配置参数的字典。该字典应包含rope_theta的值,并且如果希望使用更长的max_position_embeddings进行 RoPE,还可以包含用于缩放的参数。 - attention_bias (
bool, defaults toFalse, optional, defaults toFalse) — 在自注意力期间是否在查询、键、值和输出投影层中使用偏置。 - 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 toFalse) — 是否返回路由器的 logits。启用此选项还将允许模型输出辅助损失,包括负载均衡损失和路由器 z-loss。 - router_aux_loss_coef (
float, optional, defaults to 0.01) — 用于总损失的辅助损失因子。 - norm_topk_prob (
bool, optional, defaults toFalse) — 是否对 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
< source >( config: OlmoeConfig )
参数
- config (OlmoeConfig) — 模型的配置类,包含模型的所有参数。使用配置文件初始化不会加载与模型相关的权重,只加载配置。请查看 from_pretrained() 方法来加载模型权重。
不带任何特定头部的 OLMoE 裸模型,输出原始隐藏状态。
此模型继承自 PreTrainedModel。查看其父类文档,了解库为所有模型实现的通用方法(例如下载或保存、调整输入嵌入大小、修剪头等)。
此模型也是一个 PyTorch torch.nn.Module 子类。像普通的 PyTorch Module 一样使用它,并参考 PyTorch 文档了解一般用法和行为的所有相关信息。
forward
< source >( 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.MoeModelOutputWithPast 或 tuple(torch.FloatTensor)
参数
- input_ids (
torch.LongTensorof shape(batch_size, sequence_length), optional) — 词汇表中输入序列 token 的索引。默认情况下会忽略填充。可以使用 AutoTokenizer 获取索引。有关详细信息,请参阅 PreTrainedTokenizer.encode() 和 PreTrainedTokenizer.call()。
- attention_mask (
torch.Tensorof shape(batch_size, sequence_length), optional) — 用于避免对填充 token 索引执行 attention 的掩码。掩码值选择在[0, 1]中:- 1 表示未被掩码的 token,
- 0 表示被掩码的 token。
- position_ids (
torch.LongTensorof shape(batch_size, sequence_length), optional) — 每个输入序列 token 在位置嵌入中的位置索引。选择在范围[0, config.n_positions - 1]内。 - past_key_values (
~cache_utils.Cache, optional) — 可以用于加速顺序解码的预计算隐藏状态(自注意力块和交叉注意力块中的键和值)。这通常包括在use_cache=True或config.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.FloatTensorof 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.LongTensorof shape(sequence_length), optional) — 描述输入序列 token 在序列中位置的索引。与position_ids不同,此张量不受填充影响。它用于在正确位置更新缓存并推断完整的序列长度。
返回
transformers.modeling_outputs.MoeModelOutputWithPast 或 tuple(torch.FloatTensor)
根据配置 (OlmoeConfig) 和输入,包含各种元素的 transformers.modeling_outputs.MoeModelOutputWithPast 或 torch.FloatTensor 元组(如果传递了 return_dict=False 或 config.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=Truein the cross-attention blocks) that can be used (seepast_key_valuesinput) 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=True且config.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
< source >( 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
< source >( 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.MoeCausalLMOutputWithPast 或 tuple(torch.FloatTensor)
参数
- input_ids (
torch.LongTensorof shape(batch_size, sequence_length), optional) — 词汇表中输入序列 token 的索引。默认情况下会忽略填充。可以使用 AutoTokenizer 获取索引。有关详细信息,请参阅 PreTrainedTokenizer.encode() 和 PreTrainedTokenizer.call()。
- attention_mask (
torch.Tensorof 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.
- position_ids (
torch.LongTensorof 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]. - 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 thepast_key_valuesreturned by the model at a previous stage of decoding, whenuse_cache=Trueorconfig.use_cache=True.Only Cache instance is allowed as input, see our kv cache guide. If no
past_key_valuesare passed, DynamicCache will be initialized by default.The model will output the same cache format that is fed as input.
If
past_key_valuesare used, the user is expected to input only unprocessedinput_ids(those that don’t have their past key value states given to this model) of shape(batch_size, unprocessed_length)instead of allinput_idsof shape(batch_size, sequence_length). - inputs_embeds (
torch.FloatTensorof shape(batch_size, sequence_length, hidden_size), optional) — Optionally, instead of passinginput_idsyou can choose to directly pass an embedded representation. This is useful if you want more control over how to convertinput_idsindices into associated vectors than the model’s internal embedding lookup matrix. - labels (
torch.LongTensorof 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 (seeinput_idsdocstring). Tokens with indices set to-100are ignored (masked), the loss is only computed for the tokens with labels in[0, ..., config.vocab_size]. - use_cache (
bool, optional) — If set toTrue,past_key_valueskey value states are returned and can be used to speed up decoding (seepast_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.LongTensorof shape(sequence_length), optional) — Indices depicting the position of the input sequence tokens in the sequence. Contrarily toposition_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 to0) — If anint, compute logits for the lastlogits_to_keeptokens. If0, calculate logits for allinput_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 atorch.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.MoeCausalLMOutputWithPast 或 tuple(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=True且config.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'