Transformers 文档
OLMo2
并获得增强的文档体验
开始使用
该模型于 2024-12-31 发布，并于 2024-11-25 添加到 Hugging Face Transformers。
OLMo2
OLMo2 通过改变原始模型的架构和训练方法，在 OLMo 的基础上进行了改进。这包括排除所有偏置项以提高训练稳定性、非参数层归一化、SwiGLU 激活函数、旋转位置嵌入以及一种改进的基于 BPE 的分词器，该分词器会屏蔽个人身份信息。它在 Dolma（一个包含 3T 词元的 dataset）上进行了预训练。
您可以在 OLMo2 集合下找到所有原始 OLMo2 检查点。
点击侧边栏右侧的 OLMo2 模型，以获取更多关于如何将 OLMo2 应用于不同语言任务的示例。
下面的示例演示了如何使用 Pipeline、AutoModel 以及从命令行生成文本。
import torch
from transformers import pipeline
pipe = pipeline(
task="text-generation",
model="allenai/OLMo-2-0425-1B",
dtype=torch.float16,
device=0,
)
result = pipe("Plants create energy through a process known as")
print(result)量化通过以较低精度表示权重来减少大型模型的内存负担。有关更多可用量化后端，请参阅量化概述。
下面的示例使用 torchao 将权重量化为 4 位。
#pip install torchao
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, TorchAoConfig
torchao_config = TorchAoConfig(
"int4_weight_only",
group_size=128
)
tokenizer = AutoTokenizer.from_pretrained(
"allenai/OLMo-2-0425-1B"
)
model = AutoModelForCausalLM.from_pretrained(
"allenai/OLMo-2-0425-1B",
quantization_config=torchao_config,
dtype=torch.bfloat16,
device_map="auto",
attn_implementation="sdpa"
)
input_ids = tokenizer("Plants create energy through a process known as", return_tensors="pt").to(model.device)
output = model.generate(**input_ids, max_length=50, cache_implementation="static")
print(tokenizer.decode(output[0], skip_special_tokens=True))
注意事项
OLMo2 使用 RMSNorm 而非标准的层归一化 (LayerNorm)。RMSNorm 应用于注意力查询和键，并且它是在注意力层和前馈层之后应用，而不是之前。
OLMo2 需要 Transformers v4.48 或更高版本。
通过向 from_pretrained() 添加
revision参数来加载特定的中间检查点。from transformers import AutoModelForCausalLM model = AutoModelForCausalLM.from_pretrained("allenai/OLMo-2-0425-1B", revision="stage1-step140000-tokens294B")
Olmo2Config
class transformers.Olmo2Config
< source >( vocab_size: int | None = 50304 hidden_size: int | None = 4096 intermediate_size: int | None = 11008 num_hidden_layers: int | None = 32 num_attention_heads: int | None = 32 num_key_value_heads: int | None = None hidden_act: str | None = 'silu' max_position_embeddings: int | None = 2048 initializer_range: float | None = 0.02 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: bool | 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 rms_norm_eps: int | None = 1e-05 **kwargs )
参数
- vocab_size (
int, optional, defaults to 50304) — Olmo2 模型的分词表大小。定义了在调用 Olmo2Model 时传入的inputs_ids可以表示的不同词元数量。 - hidden_size (
int, optional, defaults to 4096) — 隐藏表示的维度。 - intermediate_size (
int, optional, defaults to 11008) — MLP 表示的维度。 - num_hidden_layers (
int, optional, defaults to 32) — Transformer 解码器中的隐藏层数量。 - num_attention_heads (
int, optional, defaults to 32) — 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 检查点时，每个组的键和值头都应该通过对该组内所有原始头进行平均池化来构建。更多详细信息，请参阅 这篇论文。如果未指定，将默认为num_attention_heads。 - hidden_act (
str或function, optional, defaults to"silu") — 解码器中的非线性激活函数（函数或字符串）。 - max_position_embeddings (
int, optional, defaults to 2048) — 该模型可能使用的最大序列长度。 - initializer_range (
float, optional, defaults to 0.02) — 用于初始化所有权重矩阵的截断正态分布初始化器的标准差。 - use_cache (
bool, optional, defaults toTrue) — 模型是否应返回最后的键/值注意力（并非所有模型都使用）。仅在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 比率。 - rms_norm_eps (
float, optional, defaults to 1e-05) — RMS 归一化层使用的 epsilon。
这是用于存储 Olmo2Model 配置的配置类。它用于根据指定的参数实例化一个 OLMo2 模型，从而定义模型的架构。使用默认值实例化配置将得到一个与 allenai/Olmo2-7B-1124-hf 类似的配置。
配置对象继承自 PreTrainedConfig，可用于控制模型输出。有关更多信息，请阅读 PreTrainedConfig 的文档。
>>> from transformers import Olmo2Model, Olmo2Config
>>> # Initializing a Olmo2 7B style configuration
>>> configuration = Olmo2Config()
>>> # Initializing a model from the Olmo2 7B style configuration
>>> model = Olmo2Model(configuration)
>>> # Accessing the model configuration
>>> configuration = model.configOlmo2Model
class transformers.Olmo2Model
< source >( config: Olmo2Config )
参数
- config (Olmo2Config) — 模型配置类，包含模型的所有参数。使用配置文件初始化不会加载模型权重，仅加载配置。请查看 from_pretrained() 方法来加载模型权重。
The bare Olmo2 Model outputting raw hidden-states without any specific head on top。
此模型继承自 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 cache_position: torch.LongTensor | None = None use_cache: bool | None = None **kwargs: typing_extensions.Unpack[transformers.utils.generic.TransformersKwargs] ) → transformers.modeling_outputs.BaseModelOutputWithPast or 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) — 预计算的隐藏状态（自注意力块和交叉注意力块中的键和值），可用于加速顺序解码。这通常由模型在解码的早期阶段返回的past_key_values组成，当use_cache=True或config.use_cache=True时。只有 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索引转换为关联向量，这将非常有用。 - cache_position (
torch.LongTensorof shape(sequence_length), optional) — 指示输入序列 token 在序列中位置的索引。与position_ids相反，此张量不受填充的影响。它用于在正确位置更新缓存并推断完整序列长度。 - use_cache (
bool, optional) — 如果设置为True，则返回past_key_values键值状态，可用于加速解码（请参阅past_key_values）。
返回
transformers.modeling_outputs.BaseModelOutputWithPast or tuple(torch.FloatTensor)
A transformers.modeling_outputs.BaseModelOutputWithPast 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 (Olmo2Config) and inputs。
-
last_hidden_state (
torch.FloatTensor, 形状为(batch_size, sequence_length, hidden_size)) — 模型最后一层输出的隐藏状态序列。如果使用了
past_key_values，则只输出形状为(batch_size, 1, 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 后的注意力权重，用于计算自注意力头中的加权平均值。
The Olmo2Model forward method, overrides the __call__ special method。
虽然 forward pass 的实现需要在此函数中定义，但你应该在之后调用
Module实例而不是这个，因为前者负责运行预处理和后处理步骤，而后者会静默地忽略它们。
Olmo2ForCausalLM
class transformers.Olmo2ForCausalLM
< 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 (Olmo2ForCausalLM) — 模型配置类，包含模型的所有参数。使用配置文件初始化不会加载模型权重，仅加载配置。请查看 from_pretrained() 方法来加载模型权重。
The Olmo2 Model for causal language modeling。
此模型继承自 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 cache_position: torch.LongTensor | None = None logits_to_keep: int | torch.Tensor = 0 **kwargs: typing_extensions.Unpack[transformers.utils.generic.TransformersKwargs] ) → transformers.modeling_outputs.CausalLMOutputWithPast or 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) — 预计算的隐藏状态（自注意力块和交叉注意力块中的键和值），可用于加速顺序解码。这通常由模型在解码的早期阶段返回的past_key_values组成，当use_cache=True或config.use_cache=True时。只有 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索引转换为关联向量，这将非常有用。 - labels (
torch.LongTensorof shape(batch_size, sequence_length), optional) — 用于计算掩码语言模型损失的标签。索引应为[0, ..., config.vocab_size]或 -100（请参阅input_ids文档字符串）。索引设置为-100的 token 将被忽略（掩码），损失仅为标签在[0, ..., config.vocab_size]中的 token 计算。 - use_cache (
bool, optional) — 如果设置为True，则返回past_key_values键值状态，可用于加速解码（请参阅past_key_values）。 - cache_position (
torch.LongTensorof shape(sequence_length), optional) — 指示输入序列 token 在序列中位置的索引。与position_ids相反，此张量不受填充的影响。它用于在正确位置更新缓存并推断完整序列长度。 - logits_to_keep (
Union[int, torch.Tensor], optional, defaults to0) — 如果是int，则计算最后logits_to_keep个 token 的 logits。如果是0，则计算所有input_ids的 logits（特殊情况）。仅生成需要最后一个 token 的 logits，并且仅为该 token 计算 logits 可以节省内存，这对于长序列或大词汇量非常重要。如果是torch.Tensor，则必须是 1D 张量，对应于序列长度维度的索引。当使用打包张量格式（批次和序列长度的单个维度）时，这很有用。
返回
transformers.modeling_outputs.CausalLMOutputWithPast or tuple(torch.FloatTensor)
A transformers.modeling_outputs.CausalLMOutputWithPast 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 (Olmo2Config) and inputs。
-
loss (
torch.FloatTensor形状为(1,)，可选，当提供labels时返回) — 语言建模损失（用于下一个 token 预测）。 -
logits (形状为
(batch_size, sequence_length, config.vocab_size)的torch.FloatTensor) — 语言建模头部的预测分数（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 Olmo2ForCausalLM forward method, overrides the __call__ special method。
虽然 forward pass 的实现需要在此函数中定义，但你应该在之后调用
Module实例而不是这个，因为前者负责运行预处理和后处理步骤，而后者会静默地忽略它们。
示例
>>> from transformers import AutoTokenizer, Olmo2ForCausalLM
>>> model = Olmo2ForCausalLM.from_pretrained("meta-olmo2/Olmo2-2-7b-hf")
>>> tokenizer = AutoTokenizer.from_pretrained("meta-olmo2/Olmo2-2-7b-hf")
>>> 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 conscious, but I can talk to you."