A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
[Model] Add Olmo3 model implementation (#24534)
Signed-off-by: Shane A <shanea@allenai.org> Co-authored-by: Isotr0py <mozf@mail2.sysu.edu.cn>
Shane A committed
Sep 12, 2025 at 20:26 UTC
89e08d6d180c76019daa5aa1bbf7759dfaedde2e
7 files changed
+114
-14
docs/models/supported_models.md
+1
@@ -389,6 +389,7 @@ th {
389
| `NemotronHForCausalLM` | Nemotron-H | `nvidia/Nemotron-H-8B-Base-8K`, `nvidia/Nemotron-H-47B-Base-8K`, `nvidia/Nemotron-H-56B-Base-8K`, etc. | ✅︎ | ✅︎ | ✅︎ |
390
| `OLMoForCausalLM` | OLMo | `allenai/OLMo-1B-hf`, `allenai/OLMo-7B-hf`, etc. | ✅︎ | ✅︎ | ✅︎ |
391
| `OLMo2ForCausalLM` | OLMo2 | `allenai/OLMo-2-0425-1B`, etc. | ✅︎ | ✅︎ | ✅︎ |
392
+| `OLMo3ForCausalLM` | OLMo3 | TBA | ✅︎ | ✅︎ | ✅︎ |
393
| `OLMoEForCausalLM` | OLMoE | `allenai/OLMoE-1B-7B-0924`, `allenai/OLMoE-1B-7B-0924-Instruct`, etc. | | ✅︎ | ✅︎ |
394
| `OPTForCausalLM` | OPT, OPT-IML | `facebook/opt-66b`, `facebook/opt-iml-max-30b`, etc. | | ✅︎ | ✅︎ |
395
| `OrionForCausalLM` | Orion | `OrionStarAI/Orion-14B-Base`, `OrionStarAI/Orion-14B-Chat`, etc. | | ✅︎ | ✅︎ |
tests/models/registry.py
+1
@@ -301,6 +301,7 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
301
trust_remote_code=True),
302
"OlmoForCausalLM": _HfExamplesInfo("allenai/OLMo-1B-hf"),
303
"Olmo2ForCausalLM": _HfExamplesInfo("allenai/OLMo-2-0425-1B"),
304
+ "Olmo3ForCausalLM": _HfExamplesInfo("shanearora/2025-sep-a-base-model"),
305
"OlmoeForCausalLM": _HfExamplesInfo("allenai/OLMoE-1B-7B-0924-Instruct"),
306
"OPTForCausalLM": _HfExamplesInfo("facebook/opt-125m",
307
{"1b": "facebook/opt-iml-max-1.3b"}),
vllm/model_executor/models/olmo2.py
+28
-14
@@ -52,10 +52,11 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
52
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
53
from vllm.model_executor.models.interfaces import SupportsLoRA, SupportsPP
54
from vllm.model_executor.models.utils import (
55
- AutoWeightsLoader, is_pp_missing_parameter,
55
+ AutoWeightsLoader, extract_layer_index, is_pp_missing_parameter,
56
make_empty_intermediate_tensors_factory, make_layers, maybe_prefix)
57
from vllm.model_executor.sampling_metadata import SamplingMetadata
58
from vllm.sequence import IntermediateTensors
59
+from vllm.transformers_utils.configs import Olmo3Config
60
61
62
class Olmo2Attention(nn.Module):
@@ -68,7 +69,7 @@ class Olmo2Attention(nn.Module):
69
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
70
super().__init__()
71
self.config = vllm_config.model_config.hf_config
71
- assert isinstance(self.config, Olmo2Config)
72
+ assert isinstance(self.config, (Olmo2Config, Olmo3Config))
73
74
hidden_size = self.config.hidden_size
75
self.tp_size = get_tensor_model_parallel_world_size()
@@ -111,14 +112,14 @@ class Olmo2Attention(nn.Module):
112
self.q_norm = RMSNorm(self.config.hidden_size,
113
eps=self.config.rms_norm_eps)
114
114
- # Rotary embeddings.
115
- self.rotary_emb = get_rope(
116
- self.head_dim,
117
- rotary_dim=self.head_dim,
118
- max_position=self.max_position_embeddings,
119
- base=self.rope_theta, # type: ignore
120
- )
115
self.scaling = self.head_dim**-0.5
116
+
117
+ layer_idx = extract_layer_index(prefix)
118
+ sliding_window = None
119
+ if ((layer_types := getattr(self.config, "layer_types", None))
120
+ is not None and layer_types[layer_idx] == "sliding_attention"):
121
+ sliding_window = self.config.sliding_window
122
+
123
self.attn = Attention(
124
self.num_heads,
125
self.head_dim,
@@ -126,7 +127,20 @@ class Olmo2Attention(nn.Module):
127
num_kv_heads=self.num_kv_heads,
128
cache_config=vllm_config.cache_config,
129
quant_config=vllm_config.quant_config,
129
- prefix=prefix,
130
+ per_layer_sliding_window=sliding_window,
131
+ prefix=f"{prefix}.attn",
132
+ )
133
+
134
+ # Rotary embeddings. Rope scaling is only applied on full attention
135
+ # layers.
136
+ self.rope_scaling = (self.config.rope_scaling
137
+ if sliding_window is None else None)
138
+ self.rotary_emb = get_rope(
139
+ self.head_dim,
140
+ rotary_dim=self.head_dim,
141
+ max_position=self.max_position_embeddings,
142
+ base=self.rope_theta, # type: ignore
143
+ rope_scaling=self.rope_scaling,
144
)
145
146
# Attention output projection.
@@ -176,7 +190,7 @@ class Olmo2MLP(nn.Module):
190
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
191
super().__init__()
192
config = vllm_config.model_config.hf_config
179
- assert isinstance(config, Olmo2Config)
193
+ assert isinstance(config, (Olmo2Config, Olmo3Config))
194
hidden_size = config.hidden_size
195
intermediate_size = config.intermediate_size
196
@@ -221,7 +235,7 @@ class Olmo2DecoderLayer(nn.Module):
235
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
236
super().__init__()
237
config = vllm_config.model_config.hf_config
224
- assert isinstance(config, Olmo2Config)
238
+ assert isinstance(config, (Olmo2Config, Olmo3Config))
239
# Attention block.
240
self.self_attn = Olmo2Attention(vllm_config=vllm_config,
241
prefix=f"{prefix}.self_attn")
@@ -261,7 +275,7 @@ class Olmo2Model(nn.Module):
275
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
276
super().__init__()
277
self.config = vllm_config.model_config.hf_config
264
- assert isinstance(self.config, Olmo2Config)
278
+ assert isinstance(self.config, (Olmo2Config, Olmo3Config))
279
280
self.embed_tokens = VocabParallelEmbedding(
281
self.config.vocab_size,
@@ -376,7 +390,7 @@ class Olmo2ForCausalLM(nn.Module, SupportsPP, SupportsLoRA):
390
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
391
super().__init__()
392
config = vllm_config.model_config.hf_config
379
- assert isinstance(config, Olmo2Config)
393
+ assert isinstance(config, (Olmo2Config, Olmo3Config))
394
self.config = config
395
self.model = Olmo2Model(vllm_config=vllm_config,
396
prefix=maybe_prefix(prefix, "model"))
vllm/model_executor/models/registry.py
+1
@@ -120,6 +120,7 @@ _TEXT_GENERATION_MODELS = {
120
"NemotronHForCausalLM": ("nemotron_h", "NemotronHForCausalLM"),
121
"OlmoForCausalLM": ("olmo", "OlmoForCausalLM"),
122
"Olmo2ForCausalLM": ("olmo2", "Olmo2ForCausalLM"),
123
+ "Olmo3ForCausalLM": ("olmo2", "Olmo2ForCausalLM"),
124
"OlmoeForCausalLM": ("olmoe", "OlmoeForCausalLM"),
125
"OPTForCausalLM": ("opt", "OPTForCausalLM"),
126
"OrionForCausalLM": ("orion", "OrionForCausalLM"),
vllm/transformers_utils/config.py
+1
@@ -75,6 +75,7 @@ _CONFIG_REGISTRY: dict[str, type[PretrainedConfig]] = LazyConfigDict(
75
eagle="EAGLEConfig",
76
speculators="SpeculatorsConfig",
77
nemotron="NemotronConfig",
78
+ olmo3="Olmo3Config",
79
ovis="OvisConfig",
80
ultravox="UltravoxConfig",
81
step3_vl="Step3VLConfig",
vllm/transformers_utils/configs/__init__.py
+2
@@ -23,6 +23,7 @@ from vllm.transformers_utils.configs.moonvit import MoonViTConfig
23
from vllm.transformers_utils.configs.nemotron import NemotronConfig
24
from vllm.transformers_utils.configs.nemotron_h import NemotronHConfig
25
from vllm.transformers_utils.configs.nemotron_vl import Nemotron_Nano_VL_Config
26
+from vllm.transformers_utils.configs.olmo3 import Olmo3Config
27
from vllm.transformers_utils.configs.ovis import OvisConfig
28
from vllm.transformers_utils.configs.qwen3_next import Qwen3NextConfig
29
from vllm.transformers_utils.configs.speculators.base import SpeculatorsConfig
@@ -45,6 +46,7 @@ __all__ = [
46
"NemotronConfig",
47
"NemotronHConfig",
48
"Nemotron_Nano_VL_Config",
49
+ "Olmo3Config",
50
"OvisConfig",
51
"SpeculatorsConfig",
52
"UltravoxConfig",
vllm/transformers_utils/configs/olmo3.py
new
+80
@@ -0,0 +1,80 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from transformers.configuration_utils import PretrainedConfig
5
+
6
+
7
+class Olmo3Config(PretrainedConfig):
8
+
9
+ model_type = "olmo3"
10
+ keys_to_ignore_at_inference = ["past_key_values"]
11
+
12
+ def __init__(
13
+ self,
14
+ vocab_size=50304,
15
+ hidden_size=4096,
16
+ intermediate_size=11008,
17
+ num_hidden_layers=32,
18
+ num_attention_heads=32,
19
+ num_key_value_heads=None,
20
+ hidden_act="silu",
21
+ max_position_embeddings=2048,
22
+ initializer_range=0.02,
23
+ use_cache=True,
24
+ pad_token_id=1,
25
+ bos_token_id=None,
26
+ eos_token_id=50279,
27
+ tie_word_embeddings=False,
28
+ rope_theta=10000.0,
29
+ rope_scaling=None,
30
+ attention_bias=False,
31
+ attention_dropout=0.0,
32
+ rms_norm_eps=1e-5,
33
+ sliding_window=4096,
34
+ layer_types=None,
35
+ **kwargs,
36
+ ):
37
+ # This model uses Olmo3ForCausalLM in transformers but Olmo2ForCausalLM
38
+ # in vLLM.
39
+ if "architectures" not in kwargs:
40
+ kwargs["architectures"] = ["Olmo2ForCausalLM"]
41
+ elif "Olmo3ForCausalLM" in kwargs["architectures"]:
42
+ kwargs["architectures"].remove("Olmo3ForCausalLM")
43
+ kwargs["architectures"].append("Olmo2ForCausalLM")
44
+
45
+ super().__init__(
46
+ pad_token_id=pad_token_id,
47
+ bos_token_id=bos_token_id,
48
+ eos_token_id=eos_token_id,
49
+ tie_word_embeddings=tie_word_embeddings,
50
+ **kwargs,
51
+ )
52
+ self.vocab_size = vocab_size
53
+ self.max_position_embeddings = max_position_embeddings
54
+ self.hidden_size = hidden_size
55
+ self.intermediate_size = intermediate_size
56
+ self.num_hidden_layers = num_hidden_layers
57
+ self.num_attention_heads = num_attention_heads
58
+
59
+ # for backward compatibility
60
+ if num_key_value_heads is None:
61
+ num_key_value_heads = num_attention_heads
62
+
63
+ self.num_key_value_heads = num_key_value_heads
64
+ self.hidden_act = hidden_act
65
+ self.initializer_range = initializer_range
66
+ self.use_cache = use_cache
67
+ self.rope_theta = rope_theta
68
+ self.rope_scaling = rope_scaling
69
+ self.attention_bias = attention_bias
70
+ self.attention_dropout = attention_dropout
71
+
72
+ self.rms_norm_eps = rms_norm_eps
73
+
74
+ self.sliding_window = sliding_window
75
+ self.layer_types = layer_types
76
+ if self.layer_types is None:
77
+ self.layer_types = [
78
+ "sliding_attention" if (i + 1) % 4 != 0 else "full_attention"
79
+ for i in range(self.num_hidden_layers)
80
+ ]