Migrate Olmo and Olmo2 to the Transformers modeling backend (#48100)
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Harry Mellor committed
Jul 9, 2026 at 13:00 UTC
b83be00cddf03387bf740acfe6b9ca07ec1a3c08
4 files changed
+22
-400
docs/models/supported_models.md
+2
-2
@@ -421,8 +421,6 @@ th {
421
| `MPTForCausalLM` | MPT, MPT-Instruct, MPT-Chat, MPT-StoryWriter | `mosaicml/mpt-7b`, `mosaicml/mpt-7b-storywriter`, `mosaicml/mpt-30b`, etc. | | ✅︎ |
422
| `NemotronForCausalLM` | Nemotron-3, Nemotron-4, Minitron | `nvidia/Minitron-8B-Base`, `mgoin/Nemotron-4-340B-Base-hf-FP8`, etc. | ✅︎ | ✅︎ |
423
| `NemotronHForCausalLM` | Nemotron-H | `nvidia/Nemotron-H-8B-Base-8K`, `nvidia/Nemotron-H-47B-Base-8K`, `nvidia/Nemotron-H-56B-Base-8K`, etc. | ✅︎ | ✅︎ |
424
-| `OlmoForCausalLM` | OLMo | `allenai/OLMo-1B-hf`, `allenai/OLMo-7B-hf`, etc. | ✅︎ | ✅︎ |
425
-| `Olmo2ForCausalLM` | OLMo2 | `allenai/OLMo-2-0425-1B`, etc. | ✅︎ | ✅︎ |
424
| `Olmo3ForCausalLM` | OLMo3 | `allenai/Olmo-3-7B-Instruct`, `allenai/Olmo-3-32B-Think`, etc. | ✅︎ | ✅︎ |
425
| `OlmoHybridForCausalLM` | OLMo Hybrid | `allenai/Olmo-Hybrid-7B` | ✅︎ | ✅︎ |
426
| `OlmoeForCausalLM` | OLMoE | `allenai/OLMoE-1B-7B-0924`, `allenai/OLMoE-1B-7B-0924-Instruct`, etc. | | ✅︎ |
@@ -463,6 +461,8 @@ Some models are supported only via the [Transformers modeling backend](#transfor
461
| Architecture | Models | Example HF Models | [LoRA](../features/lora.md) | [PP](../serving/parallelism_scaling.md) |
462
| ------------ | ------ | ----------------- | -------------------- | ------------------------- |
463
| `GPTBigCodeForCausalLM` | StarCoder, SantaCoder, WizardCoder | `bigcode/starcoder`, `bigcode/gpt_bigcode-santacoder`, `WizardLM/WizardCoder-15B-V1.0`, etc. | ✅︎ | |
464
+| `OlmoForCausalLM` | OLMo | `allenai/OLMo-1B-hf`, `allenai/OLMo-7B-hf`, etc. | ✅︎ | ✅︎ |
465
+| `Olmo2ForCausalLM` | OLMo2 | `allenai/OLMo-2-0425-1B`, etc. | ✅︎ | ✅︎ |
466
| `SmolLM3ForCausalLM` | SmolLM3 | `HuggingFaceTB/SmolLM3-3B` | ✅︎ | ✅︎ |
467
| `Starcoder2ForCausalLM` | Starcoder2 | `bigcode/starcoder2-3b`, `bigcode/starcoder2-7b`, `bigcode/starcoder2-15b`, etc. | ✅︎ | ✅︎ |
468
vllm/model_executor/models/olmo.py
deleted
-378
@@ -1,378 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-# Adapted from
5
-# https://github.com/huggingface/transformers/blob/v4.40.1/src/transformers/models/olmo/modeling_olmo.py
6
-# Copyright 2024 The vLLM team.
7
-# Copyright 2024 EleutherAI and the HuggingFace Inc. team. All rights reserved.
8
-#
9
-# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
10
-# and OPT implementations in this library. It has been modified from its
11
-# original forms to accommodate minor architectural differences compared
12
-# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
13
-#
14
-# Licensed under the Apache License, Version 2.0 (the "License");
15
-# you may not use this file except in compliance with the License.
16
-# You may obtain a copy of the License at
17
-#
18
-# http://www.apache.org/licenses/LICENSE-2.0
19
-#
20
-# Unless required by applicable law or agreed to in writing, software
21
-# distributed under the License is distributed on an "AS IS" BASIS,
22
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
23
-# See the License for the specific language governing permissions and
24
-# limitations under the License.
25
-"""Inference-only OLMo model compatible with HuggingFace weights."""
26
-
27
-from collections.abc import Iterable
28
-from itertools import islice
29
-
30
-import torch
31
-from torch import nn
32
-from transformers import OlmoConfig
33
-
34
-from vllm.compilation.decorators import support_torch_compile
35
-from vllm.config import CacheConfig, VllmConfig
36
-from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
37
-from vllm.model_executor.layers.activation import SiluAndMul
38
-from vllm.model_executor.layers.attention import Attention
39
-from vllm.model_executor.layers.linear import (
40
- MergedColumnParallelLinear,
41
- QKVParallelLinear,
42
- RowParallelLinear,
43
-)
44
-from vllm.model_executor.layers.logits_processor import LogitsProcessor
45
-from vllm.model_executor.layers.quantization import QuantizationConfig
46
-from vllm.model_executor.layers.rotary_embedding import get_rope
47
-from vllm.model_executor.layers.vocab_parallel_embedding import (
48
- ParallelLMHead,
49
- VocabParallelEmbedding,
50
-)
51
-from vllm.sequence import IntermediateTensors
52
-
53
-from .interfaces import SupportsLoRA, SupportsPP
54
-from .utils import (
55
- AutoWeightsLoader,
56
- WeightsMapper,
57
- make_empty_intermediate_tensors_factory,
58
- make_layers,
59
- maybe_prefix,
60
-)
61
-
62
-
63
-class OlmoAttention(nn.Module):
64
- """
65
- This is the attention block where the output is computed as
66
- `Attention(LN(x))` in `MLP(LN(x + Attention(LN(x))))`
67
- (plus another skip connection).
68
- """
69
-
70
- def __init__(
71
- self,
72
- config: OlmoConfig,
73
- cache_config: CacheConfig | None = None,
74
- quant_config: QuantizationConfig | None = None,
75
- prefix: str = "",
76
- ):
77
- super().__init__()
78
- self.config = config
79
- self.hidden_size = config.hidden_size
80
- tensor_model_parallel_world_size = get_tensor_model_parallel_world_size()
81
- self.total_num_heads = config.num_attention_heads
82
-
83
- assert self.hidden_size % self.total_num_heads == 0
84
- assert self.total_num_heads % tensor_model_parallel_world_size == 0
85
-
86
- self.num_heads = self.total_num_heads // tensor_model_parallel_world_size
87
- self.head_dim = self.hidden_size // self.total_num_heads
88
- self.max_position_embeddings = config.max_position_embeddings
89
- self.clip_qkv = config.clip_qkv
90
-
91
- # Attention input projection. Projects x -> (q, k, v)
92
- self.qkv_proj = QKVParallelLinear(
93
- self.hidden_size,
94
- self.head_dim,
95
- self.total_num_heads,
96
- bias=config.attention_bias,
97
- quant_config=quant_config,
98
- prefix=f"{prefix}.qkv_proj",
99
- )
100
-
101
- # Rotary embeddings.
102
- self.rotary_emb = get_rope(
103
- self.head_dim,
104
- max_position=self.max_position_embeddings,
105
- rope_parameters=config.rope_parameters,
106
- )
107
- self.scaling = self.head_dim**-0.5
108
- self.attn = Attention(
109
- self.num_heads,
110
- self.head_dim,
111
- scale=self.scaling,
112
- cache_config=cache_config,
113
- quant_config=quant_config,
114
- prefix=f"{prefix}.attn",
115
- )
116
-
117
- # Attention output projection.
118
- self.o_proj = RowParallelLinear(
119
- self.hidden_size,
120
- self.hidden_size,
121
- bias=config.attention_bias,
122
- quant_config=quant_config,
123
- prefix=f"{prefix}.o_proj",
124
- )
125
-
126
- def forward(
127
- self,
128
- positions: torch.Tensor,
129
- hidden_states: torch.Tensor,
130
- ) -> torch.Tensor:
131
- qkv, _ = self.qkv_proj(hidden_states)
132
- if self.clip_qkv is not None:
133
- qkv.clamp_(min=-self.clip_qkv, max=self.clip_qkv)
134
- q, k, v = qkv.chunk(chunks=3, dim=-1)
135
- q, k = self.rotary_emb(positions, q, k)
136
- attn_output = self.attn(q, k, v)
137
- output, _ = self.o_proj(attn_output)
138
- return output
139
-
140
-
141
-class OlmoMLP(nn.Module):
142
- """
143
- This is the MLP block where the output is computed as
144
- `MLP(LN(x))` in `MLP(LN(x + Attention(LN(x))))`
145
- (plus another skip connection).
146
- """
147
-
148
- def __init__(
149
- self,
150
- config: OlmoConfig,
151
- quant_config: QuantizationConfig | None = None,
152
- prefix: str = "",
153
- ):
154
- super().__init__()
155
- self.config = config
156
- self.hidden_size = config.hidden_size
157
- self.intermediate_size = config.intermediate_size
158
-
159
- # Feed-forward input projection.
160
- self.gate_up_proj = MergedColumnParallelLinear(
161
- self.hidden_size,
162
- [self.intermediate_size] * 2,
163
- bias=False,
164
- quant_config=quant_config,
165
- prefix=f"{prefix}.gate_up_proj",
166
- )
167
-
168
- # Activation function.
169
- self.act_fn = SiluAndMul()
170
-
171
- # Feed-forward output projection.
172
- self.down_proj = RowParallelLinear(
173
- self.intermediate_size,
174
- self.hidden_size,
175
- bias=False,
176
- quant_config=quant_config,
177
- prefix=f"{prefix}.down_proj",
178
- )
179
-
180
- def forward(
181
- self,
182
- x: torch.Tensor,
183
- ) -> torch.Tensor:
184
- gate_up, _ = self.gate_up_proj(x)
185
- x = self.act_fn(gate_up)
186
- x, _ = self.down_proj(x)
187
- return x
188
-
189
-
190
-class OlmoDecoderLayer(nn.Module):
191
- """
192
- This is a typical transformer block where the output is
193
- computed as `MLP(LN(x + Attention(LN(x))))`
194
- (plus another skip connection).
195
- """
196
-
197
- def __init__(
198
- self,
199
- config: OlmoConfig,
200
- cache_config: CacheConfig | None = None,
201
- quant_config: QuantizationConfig | None = None,
202
- prefix: str = "",
203
- ):
204
- super().__init__()
205
- # Attention block.
206
- self.self_attn = OlmoAttention(
207
- config, cache_config, quant_config, prefix=f"{prefix}.self_attn"
208
- )
209
-
210
- # MLP block.
211
- self.mlp = OlmoMLP(config, quant_config, prefix=f"{prefix}.mlp")
212
-
213
- # LayerNorm
214
- self.input_layernorm = nn.LayerNorm(
215
- config.hidden_size, elementwise_affine=False, bias=False
216
- )
217
- self.post_attention_layernorm = nn.LayerNorm(
218
- config.hidden_size, elementwise_affine=False, bias=False
219
- )
220
-
221
- def forward(
222
- self,
223
- positions: torch.Tensor,
224
- hidden_states: torch.Tensor,
225
- ) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor] | None]:
226
- # Attention block.
227
- residual = hidden_states
228
- hidden_states = self.input_layernorm(hidden_states)
229
- hidden_states = self.self_attn(positions, hidden_states)
230
- hidden_states = hidden_states + residual
231
-
232
- # MLP block.
233
- residual = hidden_states
234
- hidden_states = self.post_attention_layernorm(hidden_states)
235
- hidden_states = self.mlp(hidden_states)
236
- hidden_states = residual + hidden_states
237
- return hidden_states
238
-
239
-
240
-@support_torch_compile
241
-class OlmoModel(nn.Module):
242
- def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
243
- super().__init__()
244
-
245
- config = vllm_config.model_config.hf_config
246
- cache_config = vllm_config.cache_config
247
- quant_config = vllm_config.quant_config
248
-
249
- self.config = config
250
-
251
- self.embed_tokens = VocabParallelEmbedding(
252
- config.vocab_size, config.hidden_size
253
- )
254
- self.start_layer, self.end_layer, self.layers = make_layers(
255
- config.num_hidden_layers,
256
- lambda prefix: OlmoDecoderLayer(
257
- config, cache_config, quant_config, prefix=prefix
258
- ),
259
- prefix=f"{prefix}.layers",
260
- )
261
- self.norm = nn.LayerNorm(
262
- config.hidden_size, elementwise_affine=False, bias=False
263
- )
264
- self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
265
- ["hidden_states"], config.hidden_size
266
- )
267
-
268
- def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
269
- return self.embed_tokens(input_ids)
270
-
271
- def forward(
272
- self,
273
- input_ids: torch.Tensor | None,
274
- positions: torch.Tensor,
275
- intermediate_tensors: IntermediateTensors | None,
276
- inputs_embeds: torch.Tensor | None = None,
277
- ) -> torch.Tensor | IntermediateTensors:
278
- """
279
- Args:
280
- input_ids: A tensor of shape `(batch_size, seq_len)`.
281
- """
282
- if get_pp_group().is_first_rank:
283
- if inputs_embeds is not None:
284
- hidden_states = inputs_embeds
285
- else:
286
- hidden_states = self.embed_input_ids(input_ids)
287
- else:
288
- assert intermediate_tensors is not None
289
- hidden_states = intermediate_tensors["hidden_states"]
290
-
291
- # Apply blocks one-by-one.
292
- for layer in islice(self.layers, self.start_layer, self.end_layer):
293
- # shape: (batch_size, seq_len, d_model)
294
- hidden_states = layer(positions, hidden_states)
295
-
296
- if not get_pp_group().is_last_rank:
297
- return IntermediateTensors({"hidden_states": hidden_states})
298
- # Apply final layer norm.
299
- # shape: (batch_size, seq_len or 1, d_model)
300
- hidden_states = self.norm(hidden_states)
301
- return hidden_states
302
-
303
-
304
-class OlmoForCausalLM(nn.Module, SupportsPP, SupportsLoRA):
305
- """
306
- Extremely barebones HF model wrapper.
307
- """
308
-
309
- hf_to_vllm_mapper = WeightsMapper(
310
- orig_to_new_stacked={
311
- # weight_name: (param_name, shard_id)
312
- ".q_proj": (".qkv_proj", "q"),
313
- ".k_proj": (".qkv_proj", "k"),
314
- ".v_proj": (".qkv_proj", "v"),
315
- ".gate_proj": (".gate_up_proj", 0),
316
- ".up_proj": (".gate_up_proj", 1),
317
- }
318
- )
319
- packed_modules_mapping = {
320
- "qkv_proj": ["q_proj", "k_proj", "v_proj"],
321
- "gate_up_proj": ["gate_proj", "up_proj"],
322
- }
323
-
324
- def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
325
- super().__init__()
326
- config = vllm_config.model_config.hf_config
327
- quant_config = vllm_config.quant_config
328
- self.config = config
329
- self.model = OlmoModel(
330
- vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
331
- )
332
- if config.tie_word_embeddings:
333
- self.lm_head = self.model.embed_tokens
334
- else:
335
- self.lm_head = ParallelLMHead(
336
- config.vocab_size,
337
- config.hidden_size,
338
- quant_config=quant_config,
339
- prefix=maybe_prefix(prefix, "lm_head"),
340
- )
341
- self.logits_processor = LogitsProcessor(config.vocab_size)
342
- self.make_empty_intermediate_tensors = (
343
- self.model.make_empty_intermediate_tensors
344
- )
345
-
346
- def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
347
- return self.model.embed_input_ids(input_ids)
348
-
349
- def forward(
350
- self,
351
- input_ids: torch.Tensor | None,
352
- positions: torch.Tensor,
353
- intermediate_tensors: IntermediateTensors | None = None,
354
- inputs_embeds: torch.Tensor | None = None,
355
- ) -> torch.Tensor | IntermediateTensors:
356
- hidden_states = self.model(
357
- input_ids=input_ids,
358
- positions=positions,
359
- intermediate_tensors=intermediate_tensors,
360
- inputs_embeds=inputs_embeds,
361
- )
362
- return hidden_states
363
-
364
- def compute_logits(
365
- self,
366
- hidden_states: torch.Tensor,
367
- ) -> torch.Tensor | None:
368
- logits = self.logits_processor(self.lm_head, hidden_states)
369
- return logits
370
-
371
- def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
372
- loader = AutoWeightsLoader(
373
- self,
374
- skip_prefixes=(
375
- ["lm_head.weight"] if self.config.tie_word_embeddings else None
376
- ),
377
- )
378
- return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
vllm/model_executor/models/olmo3.py
renamed
+17
-17
@@ -2,7 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
# Adapted from
5
-# https://github.com/huggingface/transformers/blob/main/src/transformers/models/olmo2/modeling_olmo2.py
5
+# https://github.com/huggingface/transformers/blob/main/src/transformers/models/olmo3/modeling_olmo3.py
6
# Copyright 2024 The vLLM team.
7
# Copyright 2024 EleutherAI and the HuggingFace Inc. team. All rights reserved.
8
#
@@ -22,7 +22,7 @@
22
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
23
# See the License for the specific language governing permissions and
24
# limitations under the License.
25
-"""Inference-only OLMo2 model compatible with HuggingFace weights."""
25
+"""Inference-only OLMo3 model compatible with HuggingFace weights."""
26
27
from collections.abc import Iterable
28
from functools import partial
@@ -30,7 +30,7 @@ from itertools import islice
30
31
import torch
32
from torch import nn
33
-from transformers import Olmo2Config, Olmo3Config
33
+from transformers import Olmo3Config
34
35
from vllm.compilation.decorators import support_torch_compile
36
from vllm.config import VllmConfig
@@ -64,7 +64,7 @@ from vllm.model_executor.models.utils import (
64
from vllm.sequence import IntermediateTensors
65
66
67
-class Olmo2Attention(nn.Module):
67
+class Olmo3Attention(nn.Module):
68
"""
69
This is the attention block where the output is computed as
70
`Attention(LN(x))` in `MLP(LN(x + Attention(LN(x))))`
@@ -74,7 +74,7 @@ class Olmo2Attention(nn.Module):
74
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
75
super().__init__()
76
self.config = vllm_config.model_config.hf_config
77
- assert isinstance(self.config, (Olmo2Config, Olmo3Config))
77
+ assert isinstance(self.config, Olmo3Config)
78
79
hidden_size = self.config.hidden_size
80
self.tp_size = get_tensor_model_parallel_world_size()
@@ -185,7 +185,7 @@ class Olmo2Attention(nn.Module):
185
return output
186
187
188
-class Olmo2MLP(nn.Module):
188
+class Olmo3MLP(nn.Module):
189
"""
190
This is the MLP block where the output is computed as
191
`MLP(x)` in `LN(MLP(x + LN(Attention(x))))`
@@ -195,7 +195,7 @@ class Olmo2MLP(nn.Module):
195
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
196
super().__init__()
197
config = vllm_config.model_config.hf_config
198
- assert isinstance(config, (Olmo2Config, Olmo3Config))
198
+ assert isinstance(config, Olmo3Config)
199
hidden_size = config.hidden_size
200
intermediate_size = config.intermediate_size
201
@@ -230,7 +230,7 @@ class Olmo2MLP(nn.Module):
230
return x
231
232
233
-class Olmo2DecoderLayer(nn.Module):
233
+class Olmo3DecoderLayer(nn.Module):
234
"""
235
This is a typical transformer block where the output is
236
computed as `MLP(LN(x + Attention(LN(x))))`
@@ -240,14 +240,14 @@ class Olmo2DecoderLayer(nn.Module):
240
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
241
super().__init__()
242
config = vllm_config.model_config.hf_config
243
- assert isinstance(config, (Olmo2Config, Olmo3Config))
243
+ assert isinstance(config, Olmo3Config)
244
# Attention block.
245
- self.self_attn = Olmo2Attention(
245
+ self.self_attn = Olmo3Attention(
246
vllm_config=vllm_config, prefix=f"{prefix}.self_attn"
247
)
248
249
# MLP block.
250
- self.mlp = Olmo2MLP(vllm_config=vllm_config, prefix=f"{prefix}.mlp")
250
+ self.mlp = Olmo3MLP(vllm_config=vllm_config, prefix=f"{prefix}.mlp")
251
252
# LayerNorm
253
self.post_attention_layernorm = RMSNorm(
@@ -278,11 +278,11 @@ class Olmo2DecoderLayer(nn.Module):
278
279
280
@support_torch_compile
281
-class Olmo2Model(nn.Module):
281
+class Olmo3Model(nn.Module):
282
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
283
super().__init__()
284
self.config = vllm_config.model_config.hf_config
285
- assert isinstance(self.config, (Olmo2Config, Olmo3Config))
285
+ assert isinstance(self.config, Olmo3Config)
286
287
self.embed_tokens = VocabParallelEmbedding(
288
self.config.vocab_size,
@@ -291,7 +291,7 @@ class Olmo2Model(nn.Module):
291
)
292
self.start_layer, self.end_layer, self.layers = make_layers(
293
self.config.num_hidden_layers,
294
- lambda prefix: Olmo2DecoderLayer(vllm_config=vllm_config, prefix=prefix),
294
+ lambda prefix: Olmo3DecoderLayer(vllm_config=vllm_config, prefix=prefix),
295
prefix=f"{prefix}.layers",
296
)
297
self.norm = RMSNorm(
@@ -343,7 +343,7 @@ class Olmo2Model(nn.Module):
343
return hidden_states
344
345
346
-class Olmo2ForCausalLM(nn.Module, SupportsPP, SupportsLoRA):
346
+class Olmo3ForCausalLM(nn.Module, SupportsPP, SupportsLoRA):
347
"""
348
Extremely barebones HF model wrapper.
349
"""
@@ -366,9 +366,9 @@ class Olmo2ForCausalLM(nn.Module, SupportsPP, SupportsLoRA):
366
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
367
super().__init__()
368
config = vllm_config.model_config.hf_config
369
- assert isinstance(config, (Olmo2Config, Olmo3Config))
369
+ assert isinstance(config, Olmo3Config)
370
self.config = config
371
- self.model = Olmo2Model(
371
+ self.model = Olmo3Model(
372
vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
373
)
374
if config.tie_word_embeddings:
vllm/model_executor/models/registry.py
+3
-3
@@ -168,9 +168,7 @@ _TEXT_GENERATION_MODELS = {
168
"NemotronForCausalLM": ("nemotron", "NemotronForCausalLM"),
169
"NemotronHForCausalLM": ("nemotron_h", "NemotronHForCausalLM"),
170
"NemotronHPuzzleForCausalLM": ("nemotron_h", "NemotronHForCausalLM"),
171
- "OlmoForCausalLM": ("olmo", "OlmoForCausalLM"),
172
- "Olmo2ForCausalLM": ("olmo2", "Olmo2ForCausalLM"),
173
- "Olmo3ForCausalLM": ("olmo2", "Olmo2ForCausalLM"),
171
+ "Olmo3ForCausalLM": ("olmo3", "Olmo3ForCausalLM"),
172
"OlmoHybridForCausalLM": ("olmo_hybrid", "OlmoHybridForCausalLM"),
173
"OlmoeForCausalLM": ("olmoe", "OlmoeForCausalLM"),
174
"OPTForCausalLM": ("opt", "OPTForCausalLM"),
@@ -637,6 +635,8 @@ _SPECULATIVE_DECODING_MODELS = {
635
_TRANSFORMERS_SUPPORTED_MODELS = {
636
# Text generation models
637
"GPTBigCodeForCausalLM": ("transformers", "TransformersForCausalLM"),
638
+ "OlmoForCausalLM": ("transformers", "TransformersForCausalLM"),
639
+ "Olmo2ForCausalLM": ("transformers", "TransformersForCausalLM"),
640
"SmolLM3ForCausalLM": ("transformers", "TransformersForCausalLM"),
641
"Starcoder2ForCausalLM": ("transformers", "TransformersForCausalLM"),
642
# Multimodal models