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'