Transformers documentation

OLMo3

Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

该模型于 2025-09-16 贡献至 Hugging Face Transformers。

FlashAttention SDPA

OLMo3

Olmo3 是 OLMo2 的改进版本。更多详细信息将很快发布。

点击右侧边栏中的 OLMo3 模型,以查看将 OLMo3 应用于不同语言任务的更多示例。

下面的示例演示了如何使用 Pipeline、AutoModel 以及命令行来生成文本。

Pipeline
AutoModel
from transformers import pipeline


pipe = pipeline(
    task="text-generation",
    model="allenai/TBA",
    device=0,
)

result = pipe("Plants create energy through a process known as")
print(result)

量化通过以较低精度表示权重来降低大模型的内存负担。有关更多可用的量化后端,请参阅 量化 概述。

以下示例使用 torchao 仅将权重量化为 4 位。

#pip install torchao
from transformers import AutoModelForCausalLM, AutoTokenizer, TorchAoConfig


torchao_config = TorchAoConfig(
    "int4_weight_only",
    group_size=128
)

tokenizer = AutoTokenizer.from_pretrained(
    "allenai/TBA"
)

model = AutoModelForCausalLM.from_pretrained(
    "allenai/TBA",
    quantization_config=torchao_config,
    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))

备注

  • 通过向 from_pretrained() 添加 revision 参数,可以加载特定的中间检查点。

    from transformers import AutoModelForCausalLM
    
    model = AutoModelForCausalLM.from_pretrained("allenai/TBA", revision="stage1-step140000-tokens294B", device_map="auto")

Olmo3Config

class transformers.Olmo3Config

< >

( transformers_version: str | None = Nonearchitectures: list[str] | None = Noneoutput_hidden_states: bool | None = Falsereturn_dict: bool | None = Truedtype: str | torch.dtype | None = Nonechunk_size_feed_forward: int = 0is_encoder_decoder: bool = Falseid2label: dict[int, str] | dict[str, str] | None = Nonelabel2id: dict[str, int] | dict[str, str] | None = Noneproblem_type: Literal['regression', 'single_label_classification', 'multi_label_classification'] | None = Nonevocab_size: int = 50304hidden_size: int = 4096intermediate_size: int = 11008num_hidden_layers: int = 32num_attention_heads: int = 32num_key_value_heads: int | None = Nonehidden_act: str = 'silu'max_position_embeddings: int = 2048initializer_range: float = 0.02use_cache: bool = Truepad_token_id: int | None = 1bos_token_id: int | None = Noneeos_token_id: int | list[int] | None = 50279tie_word_embeddings: bool = Falserope_parameters: transformers.modeling_rope_utils.RopeParameters | dict | None = Noneattention_bias: bool = Falseattention_dropout: float | int = 0.0rms_norm_eps: float = 1e-05sliding_window: int | None = 4096layer_types: list[str] | None = None )

参数

  • vocab_size (int,可选,默认为 50304) — 模型的词汇表大小。定义了 input_ids 可以表示的不同 token 的数量。
  • hidden_size (int,可选,默认为 4096) — 隐藏表示的维度。
  • intermediate_size (int,可选,默认为 11008) — MLP 表示的维度。
  • num_hidden_layers (int,可选,默认为 32) — Transformer 解码器中的隐藏层数量。
  • num_attention_heads (int,可选,默认为 32) — 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,可选,默认为 silu) — 解码器中的非线性激活函数(函数或字符串)。例如,"gelu"、 "relu"、"silu" 等。
  • max_position_embeddings (int,可选,默认为 2048) — 该模型可能使用的最大序列长度。
  • initializer_range (float,可选,默认为 0.02) — 用于初始化所有权重矩阵的 truncated_normal_initializer 的标准差。
  • use_cache (bool,可选,默认为 True) — 模型是否应返回最后一次键/值注意力(并非所有模型都使用)。仅当 config.is_decoder=True 或模型为仅解码器生成式模型时相关。
  • pad_token_id (int,可选,默认为 1) — 词表中用于填充的 token id。
  • bos_token_id (int, 可选) — 词汇表中用于表示流起始的 Token id。
  • eos_token_id (Union[int, list[int]],可选,默认为 50279) — 词汇表中用于表示流结束(end-of-stream)的 token id。
  • tie_word_embeddings (bool,可选,默认为 False) — 是否根据模型的 tied_weights_keys 映射来绑定权重嵌入。
  • rope_parameters (Union[~modeling_rope_utils.RopeParameters, dict], 可选) — 包含 RoPE 嵌入配置参数的字典。该字典应包含 rope_theta 的值,并且如果希望使用 RoPE 处理更长的 max_position_embeddings,则可选地包含用于缩放的参数。
  • attention_bias (bool,可选,默认为 False) — 在自注意力机制中,是否在 query、key、value 和 output 投影层中使用偏置。
  • attention_dropout (Union[float, int],可选,默认为 0.0) — 注意力概率的 dropout 比率。
  • rms_norm_eps (float,可选,默认为 1e-05) — RMS 归一化层使用的 epsilon 值。
  • sliding_window (int,可选,默认为 4096) — 滑动窗口注意力的窗口大小。如果为 None,则不应用滑动窗口。
  • layer_types (list[str], 可选) — 一个列表,用于显式地将每个层索引映射到其层类型。如果未提供,将根据 config 值自动生成。

这是用于存储 Olmo3Model 配置的类。它根据指定的参数实例化 Olmo3 模型,并定义模型架构。使用默认值实例化配置,将得到与 allenai/Olmo-3-7B-Instruct 类似的配置。

配置对象继承自 PreTrainedConfig,可用于控制模型输出。有关更多信息,请参阅 PreTrainedConfig 的文档。

示例:

>>> from transformers import Olmo3Model, Olmo3Config

>>> # Initializing a Olmo3 7B style configuration
>>> configuration = Olmo3Config()

>>> # Initializing a model from the Olmo3 7B style configuration
>>> model = Olmo3Model(configuration)

>>> # Accessing the model configuration
>>> configuration = model.config

Olmo3ForCausalLM

类 transformers.Olmo3ForCausalLM

< >

( config: Olmo3Config )

参数

  • config (Olmo3Config) — 模型配置类,包含模型的所有参数。使用配置文件进行初始化时,不会加载与模型关联的权重,仅加载配置。请查看 from_pretrained() 方法以加载模型权重。

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

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

该模型也是 PyTorch torch.nn.Module 的子类。 请将其作为常规的 PyTorch Module 使用,有关通用用法和行为的所有事项,请参阅 PyTorch 文档。

forward

< >

( input_ids: typing.Optional[torch.LongTensor] = Noneattention_mask: typing.Optional[torch.Tensor] = Noneposition_ids: typing.Optional[torch.LongTensor] = Nonepast_key_values: transformers.cache_utils.Cache | None = Noneinputs_embeds: typing.Optional[torch.Tensor] = Nonelabels: typing.Optional[torch.LongTensor] = Noneuse_cache: bool | None = Nonelogits_to_keep: typing.Union[int, torch.Tensor] = 0**kwargs: *<class 'transformers.utils.generic.TransformersKwargs'> ) → CausalLMOutputWithPast or tuple(torch.Tensor)

参数

  • input_ids (torch.LongTensor 形状为 (batch_size, sequence_length), ,可选) — 输入序列中各 token 在词汇表中的索引。默认情况下,填充(padding)将被忽略。

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

    什么是输入 ID?

  • attention_mask (torch.Tensor 形状为 (batch_size, sequence_length), ,可选) — 用于避免对填充 token 索引执行注意力机制的掩码。掩码值取自 [0, 1]:
    • 1 表示未被屏蔽的 token,
    • 对于被掩码的 token,值为 0。

    什么是注意力掩码?

  • position_ids (torch.LongTensor 形状为 (batch_size, sequence_length), 可选) — 输入序列中每个 token 在位置嵌入中的位置索引。取值范围为 [0, config.n_positions - 1].

    什么是位置 ID?

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

    仅允许将 Cache 实例作为输入,请参阅我们的 kv cache 指南。 如果未传入 past_key_values,则默认会初始化 DynamicCache。

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

    如果使用了 past_key_values,用户应仅输入未处理的 input_ids(即未向该模型提供其过去键值状态的那些),其形状为 (batch_size, unprocessed_length),而不是形状为 (batch_size, sequence_length) 的全部 input_ids。

  • inputs_embeds(形状为 (batch_size, sequence_length, hidden_size) 的 torch.Tensor,可选) — 可选地,您可以选择直接传入嵌入表示,而不是传入 input_ids。如果您希望对 input_ids 索引如何转换为相应向量拥有比模型内部嵌入查找矩阵更多的控制权,这将非常有用。
  • labels (torch.LongTensor 形状为 (batch_size, sequence_length), 可选) — 用于计算因果语言建模损失的标签。索引应在 [-100, 0, ..., config.vocab_size] (参见 input_ids 文档字符串)。索引设置为 -100 的 token 将被忽略(掩码),损失仅针对标签在 [0, ..., config.vocab_size].

    标签在内部会进行移位,因此使用labels[..., 1:]来预测logits[..., :-1, :]。 请传入与input_ids匹配的未移位标签(即labels = input_ids)。在传入之前,请勿自行对标签进行移位。

    如果标签已经移位(例如在使用序列并行时),请将其作为 shift_labels 关键字参数传入,系统将直接使用该值。为了计算损失,仍然必须传入 labels。

  • use_cache (bool,可选) — 如果设置为 True,则会返回 past_key_values 键值状态,可用于加速解码(参见 past_key_values)。
  • logits_to_keep (Union[int, torch.Tensor], 可选,默认为 0) — 如果为 int,则计算最后 logits_to_keep 个 token 的 logits。如果为 0,则计算所有 input_ids 的 logits(特殊情况)。生成过程只需要最后一个 token 的 logits,仅针对该 token 进行计算可以节省内存,对于长序列或大词汇量而言,这种节省相当显著。 如果为 torch.Tensor,则必须是一维张量,对应于序列长度维度中需要保留的索引。 这在使用打包张量格式(将批次和序列长度合并为单一维度)时非常有用。

返回值

CausalLMOutputWithPast 或 tuple(torch.Tensor)

一个 CausalLMOutputWithPast 或 torch.Tensor 元组(如果传入了 return_dict=False 或 config.return_dict=False),其包含的元素取决于配置(Olmo3Config)和输入。

Olmo3ForCausalLM 的 forward 方法重写了 __call__ 特殊方法。

尽管前向传播的实现需要在此函数中定义,但之后应调用 Module 实例,而不是直接调用此函数,因为前者会负责执行预处理和后处理步骤,而后者则会静默地忽略这些步骤。

  • loss (torch.Tensor,形状为 (1,),可选,当提供 labels 时返回) — 语言建模损失(用于下一个 token 预测)。

  • logits(形状为 (batch_size, sequence_length, config.vocab_size) 的 torch.Tensor)— 语言建模头的预测分数(SoftMax 之前每个词汇表 token 的分数)。

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

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

  • hidden_states (tuple(torch.Tensor),可选,当传入 output_hidden_states=True 或 config.output_hidden_states=True 时返回) — 一个 torch.Tensor 元组(如果模型具有嵌入层,则包含嵌入层的输出,以及每一层的输出),形状为 (batch_size, sequence_length, hidden_size)。

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

  • attentions (tuple(torch.Tensor),可选,当传入 output_attentions=True 或 config.output_attentions=True 时返回) — 一个 torch.Tensor 元组(每层一个),形状为 (batch_size, num_heads, sequence_length, sequence_length)。

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

示例:

>>> from transformers import AutoTokenizer, Olmo3ForCausalLM

>>> model = Olmo3ForCausalLM.from_pretrained("meta-olmo3/Olmo3-2-7b-hf")
>>> tokenizer = AutoTokenizer.from_pretrained("meta-olmo3/Olmo3-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."

Olmo3ForSequenceClassification

class transformers.Olmo3ForSequenceClassification

< >

( config )

forward

< >

( input_ids: torch.LongTensor | None = Noneattention_mask: torch.Tensor | None = Noneposition_ids: torch.LongTensor | None = Nonepast_key_values: Cache | None = Noneinputs_embeds: torch.Tensor | None = Nonelabels: torch.LongTensor | None = Noneuse_cache: bool | None = None**kwargs: Unpack[TransformersKwargs] ) → SequenceClassifierOutputWithPast

参数

  • input_ids (torch.LongTensor 形状为 (batch_size, sequence_length), ,可选) — 输入序列中各 token 在词汇表中的索引。默认情况下,填充(padding)将被忽略。

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

    什么是输入 ID?

  • attention_mask (torch.Tensor 形状为 (batch_size, sequence_length), ,可选) — 用于避免对填充 token 索引执行注意力机制的掩码。掩码值取自 [0, 1]:
    • 1 表示未被屏蔽的 token,
    • 对于被掩码的 token,值为 0。

    什么是注意力掩码?

  • position_ids (torch.LongTensor 形状为 (batch_size, sequence_length), 可选) — 输入序列中每个 token 在位置嵌入中的位置索引。取值范围为 [0, config.n_positions - 1].

    什么是位置 ID?

  • past_key_values (Cache, 可选) — 预计算的隐藏状态(自注意力块和交叉注意力块中的键和值),可用于加速顺序解码。这通常由模型在解码的先前阶段返回的 past_key_values 组成,此时 use_cache=True 或 config.use_cache=True.

    仅允许将 Cache 实例作为输入,请参阅我们的 kv cache 指南。 如果未传入 past_key_values,则默认会初始化 DynamicCache。

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

    如果使用了 past_key_values,用户应仅输入未处理的 input_ids(即未向该模型提供其过去键值状态的那些),其形状为 (batch_size, unprocessed_length),而不是形状为 (batch_size, sequence_length) 的全部 input_ids。

  • inputs_embeds(形状为 (batch_size, sequence_length, hidden_size) 的 torch.Tensor,可选) — 可选地,您可以选择直接传入嵌入表示,而不是传入 input_ids。如果您希望对 input_ids 索引如何转换为相应向量拥有比模型内部嵌入查找矩阵更多的控制权,这将非常有用。
  • labels (torch.LongTensor,形状为 (batch_size,),可选) — 用于计算序列分类/回归损失的标签。索引应在 [0, ..., config.num_labels - 1] 范围内。如果 config.num_labels == 1,则计算回归损失(均方损失);如果 config.num_labels > 1,则计算分类损失(交叉熵)。
  • use_cache (bool,可选) — 如果设置为 True,则会返回 past_key_values 键值状态,可用于加速解码(参见 past_key_values)。

返回值

SequenceClassifierOutputWithPast

GenericForSequenceClassification 的 forward 方法重写了 __call__ 特殊方法。

尽管前向传播的实现需要在此函数中定义,但之后应调用 Module 实例,而不是直接调用此函数,因为前者会负责执行预处理和后处理步骤,而后者则会静默地忽略这些步骤。

Olmo3Model

class transformers.Olmo3Model

< >

( config: Olmo3Config )

参数

  • config (Olmo3Config) — 模型配置类,包含模型的所有参数。使用配置文件进行初始化时,不会加载与模型关联的权重,仅加载配置。请查看 from_pretrained() 方法以加载模型权重。

未附加任何特定头部的裸 Olmo3 模型,直接输出原始隐藏状态。

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

该模型也是 PyTorch torch.nn.Module 的子类。 请将其作为常规的 PyTorch Module 使用,有关通用用法和行为的所有事项,请参阅 PyTorch 文档。

forward

< >

( input_ids: typing.Optional[torch.LongTensor] = Noneattention_mask: typing.Optional[torch.Tensor] = Noneposition_ids: typing.Optional[torch.LongTensor] = Nonepast_key_values: transformers.cache_utils.Cache | None = Noneinputs_embeds: typing.Optional[torch.Tensor] = Noneuse_cache: bool | None = None**kwargs: *<class 'transformers.utils.generic.TransformersKwargs'> ) → BaseModelOutputWithPast or tuple(torch.Tensor)

参数

  • input_ids (torch.LongTensor 形状为 (batch_size, sequence_length), ,可选) — 输入序列中各 token 在词汇表中的索引。默认情况下,填充(padding)将被忽略。

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

    什么是输入 ID?

  • attention_mask (torch.Tensor 形状为 (batch_size, sequence_length), ,可选) — 用于避免对填充 token 索引执行注意力机制的掩码。掩码值取自 [0, 1]:
    • 1 表示未被屏蔽的 token,
    • 对于被掩码的 token,值为 0。

    什么是注意力掩码?

  • position_ids (torch.LongTensor 形状为 (batch_size, sequence_length), 可选) — 输入序列中每个 token 在位置嵌入中的位置索引。取值范围为 [0, config.n_positions - 1].

    什么是位置 ID?

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

    仅允许将 Cache 实例作为输入,请参阅我们的 kv cache 指南。 如果未传入 past_key_values,则默认会初始化 DynamicCache。

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

    如果使用了 past_key_values,用户应仅输入未处理的 input_ids(即未向该模型提供其过去键值状态的那些),其形状为 (batch_size, unprocessed_length),而不是形状为 (batch_size, sequence_length) 的全部 input_ids。

  • inputs_embeds(形状为 (batch_size, sequence_length, hidden_size) 的 torch.Tensor,可选) — 可选地,您可以选择直接传入嵌入表示,而不是传入 input_ids。如果您希望对 input_ids 索引如何转换为相应向量拥有比模型内部嵌入查找矩阵更多的控制权,这将非常有用。
  • use_cache (bool,可选) — 如果设置为 True,则会返回 past_key_values 键值状态,可用于加速解码(参见 past_key_values)。

返回值

BaseModelOutputWithPast 或 tuple(torch.Tensor)

一个 BaseModelOutputWithPast 或 torch.Tensor 元组(如果传入了 return_dict=False 或 config.return_dict=False),其包含的元素取决于配置(Olmo3Config)和输入。

Olmo3Model 的 forward 方法会重写 __call__ 特殊方法。

尽管前向传播的实现需要在此函数中定义,但之后应调用 Module 实例,而不是直接调用此函数,因为前者会负责执行预处理和后处理步骤,而后者则会静默地忽略这些步骤。

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

    如果使用了 past_key_values,则仅输出形状为 (batch_size, 1, hidden_size) 的序列的最后一个隐藏状态。

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

    包含预计算的隐藏状态(自注意力块中的键和值,以及如果 config.is_encoder_decoder=True 则还包括交叉注意力块中的键和值),可用于(参见 past_key_values 输入)加速顺序解码。

  • hidden_states (tuple(torch.Tensor),可选,当传入 output_hidden_states=True 或 config.output_hidden_states=True 时返回) — 一个 torch.Tensor 元组(如果模型具有嵌入层,则包含嵌入层的输出,以及每一层的输出),形状为 (batch_size, sequence_length, hidden_size)。

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

  • attentions (tuple(torch.Tensor),可选,当传入 output_attentions=True 或 config.output_attentions=True 时返回) — 一个 torch.Tensor 元组(每层一个),形状为 (batch_size, num_heads, sequence_length, sequence_length)。

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

Olmo3PreTrainedModel

class transformers.Olmo3PreTrainedModel

< >

( config: PreTrainedConfig*inputs**kwargs )

参数

  • config (PreTrainedConfig) — 模型配置类,包含模型的所有参数。使用配置文件进行初始化时,不会加载与模型关联的权重,仅加载配置。请查看 from_pretrained() 方法以加载模型权重。

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

该模型也是 PyTorch torch.nn.Module 的子类。 请将其作为常规的 PyTorch Module 使用,有关通用用法和行为的所有事项,请参阅 PyTorch 文档。

forward

( *args**kwargs )

点分路径(例如 torch.float32)的模拟值:属性访问链和调用行为如同透传装饰器,repr 返回点分路径,将其用作基类时会替换为普通的 type 基类(PEP 560 __mro_entries__),因此真实的子类保留正常的元类,且 inspect.signature 读取的是它们真实的 __init__,而非模拟对象的。

Update on GitHub