From b83be00cddf03387bf740acfe6b9ca07ec1a3c08 Mon Sep 17 00:00:00 2001 From: Harry Mellor <19981378+hmellor@users.noreply.github.com> Date: Thu, 9 Jul 2026 13:00:23 +0100 Subject: [PATCH] Migrate Olmo and Olmo2 to the Transformers modeling backend (#48100) Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> --- docs/models/supported_models.md | 4 +- vllm/model_executor/models/olmo.py | 378 ------------------ .../models/{olmo2.py => olmo3.py} | 34 +- vllm/model_executor/models/registry.py | 6 +- 4 files changed, 22 insertions(+), 400 deletions(-) delete mode 100644 vllm/model_executor/models/olmo.py rename vllm/model_executor/models/{olmo2.py => olmo3.py} (94%) diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index c3c2d3df058..f73ffd17c87 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -421,8 +421,6 @@ th { | `MPTForCausalLM` | MPT, MPT-Instruct, MPT-Chat, MPT-StoryWriter | `mosaicml/mpt-7b`, `mosaicml/mpt-7b-storywriter`, `mosaicml/mpt-30b`, etc. | | ✅︎ | | `NemotronForCausalLM` | Nemotron-3, Nemotron-4, Minitron | `nvidia/Minitron-8B-Base`, `mgoin/Nemotron-4-340B-Base-hf-FP8`, etc. | ✅︎ | ✅︎ | | `NemotronHForCausalLM` | Nemotron-H | `nvidia/Nemotron-H-8B-Base-8K`, `nvidia/Nemotron-H-47B-Base-8K`, `nvidia/Nemotron-H-56B-Base-8K`, etc. | ✅︎ | ✅︎ | -| `OlmoForCausalLM` | OLMo | `allenai/OLMo-1B-hf`, `allenai/OLMo-7B-hf`, etc. | ✅︎ | ✅︎ | -| `Olmo2ForCausalLM` | OLMo2 | `allenai/OLMo-2-0425-1B`, etc. | ✅︎ | ✅︎ | | `Olmo3ForCausalLM` | OLMo3 | `allenai/Olmo-3-7B-Instruct`, `allenai/Olmo-3-32B-Think`, etc. | ✅︎ | ✅︎ | | `OlmoHybridForCausalLM` | OLMo Hybrid | `allenai/Olmo-Hybrid-7B` | ✅︎ | ✅︎ | | `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 | Architecture | Models | Example HF Models | [LoRA](../features/lora.md) | [PP](../serving/parallelism_scaling.md) | | ------------ | ------ | ----------------- | -------------------- | ------------------------- | | `GPTBigCodeForCausalLM` | StarCoder, SantaCoder, WizardCoder | `bigcode/starcoder`, `bigcode/gpt_bigcode-santacoder`, `WizardLM/WizardCoder-15B-V1.0`, etc. | ✅︎ | | +| `OlmoForCausalLM` | OLMo | `allenai/OLMo-1B-hf`, `allenai/OLMo-7B-hf`, etc. | ✅︎ | ✅︎ | +| `Olmo2ForCausalLM` | OLMo2 | `allenai/OLMo-2-0425-1B`, etc. | ✅︎ | ✅︎ | | `SmolLM3ForCausalLM` | SmolLM3 | `HuggingFaceTB/SmolLM3-3B` | ✅︎ | ✅︎ | | `Starcoder2ForCausalLM` | Starcoder2 | `bigcode/starcoder2-3b`, `bigcode/starcoder2-7b`, `bigcode/starcoder2-15b`, etc. | ✅︎ | ✅︎ | diff --git a/vllm/model_executor/models/olmo.py b/vllm/model_executor/models/olmo.py deleted file mode 100644 index e62bd39238b..00000000000 --- a/vllm/model_executor/models/olmo.py +++ /dev/null @@ -1,378 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -# Adapted from -# https://github.com/huggingface/transformers/blob/v4.40.1/src/transformers/models/olmo/modeling_olmo.py -# Copyright 2024 The vLLM team. -# Copyright 2024 EleutherAI and the HuggingFace Inc. team. All rights reserved. -# -# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX -# and OPT implementations in this library. It has been modified from its -# original forms to accommodate minor architectural differences compared -# to GPT-NeoX and OPT used by the Meta AI team that trained the model. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""Inference-only OLMo model compatible with HuggingFace weights.""" - -from collections.abc import Iterable -from itertools import islice - -import torch -from torch import nn -from transformers import OlmoConfig - -from vllm.compilation.decorators import support_torch_compile -from vllm.config import CacheConfig, VllmConfig -from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size -from vllm.model_executor.layers.activation import SiluAndMul -from vllm.model_executor.layers.attention import Attention -from vllm.model_executor.layers.linear import ( - MergedColumnParallelLinear, - QKVParallelLinear, - RowParallelLinear, -) -from vllm.model_executor.layers.logits_processor import LogitsProcessor -from vllm.model_executor.layers.quantization import QuantizationConfig -from vllm.model_executor.layers.rotary_embedding import get_rope -from vllm.model_executor.layers.vocab_parallel_embedding import ( - ParallelLMHead, - VocabParallelEmbedding, -) -from vllm.sequence import IntermediateTensors - -from .interfaces import SupportsLoRA, SupportsPP -from .utils import ( - AutoWeightsLoader, - WeightsMapper, - make_empty_intermediate_tensors_factory, - make_layers, - maybe_prefix, -) - - -class OlmoAttention(nn.Module): - """ - This is the attention block where the output is computed as - `Attention(LN(x))` in `MLP(LN(x + Attention(LN(x))))` - (plus another skip connection). - """ - - def __init__( - self, - config: OlmoConfig, - cache_config: CacheConfig | None = None, - quant_config: QuantizationConfig | None = None, - prefix: str = "", - ): - super().__init__() - self.config = config - self.hidden_size = config.hidden_size - tensor_model_parallel_world_size = get_tensor_model_parallel_world_size() - self.total_num_heads = config.num_attention_heads - - assert self.hidden_size % self.total_num_heads == 0 - assert self.total_num_heads % tensor_model_parallel_world_size == 0 - - self.num_heads = self.total_num_heads // tensor_model_parallel_world_size - self.head_dim = self.hidden_size // self.total_num_heads - self.max_position_embeddings = config.max_position_embeddings - self.clip_qkv = config.clip_qkv - - # Attention input projection. Projects x -> (q, k, v) - self.qkv_proj = QKVParallelLinear( - self.hidden_size, - self.head_dim, - self.total_num_heads, - bias=config.attention_bias, - quant_config=quant_config, - prefix=f"{prefix}.qkv_proj", - ) - - # Rotary embeddings. - self.rotary_emb = get_rope( - self.head_dim, - max_position=self.max_position_embeddings, - rope_parameters=config.rope_parameters, - ) - self.scaling = self.head_dim**-0.5 - self.attn = Attention( - self.num_heads, - self.head_dim, - scale=self.scaling, - cache_config=cache_config, - quant_config=quant_config, - prefix=f"{prefix}.attn", - ) - - # Attention output projection. - self.o_proj = RowParallelLinear( - self.hidden_size, - self.hidden_size, - bias=config.attention_bias, - quant_config=quant_config, - prefix=f"{prefix}.o_proj", - ) - - def forward( - self, - positions: torch.Tensor, - hidden_states: torch.Tensor, - ) -> torch.Tensor: - qkv, _ = self.qkv_proj(hidden_states) - if self.clip_qkv is not None: - qkv.clamp_(min=-self.clip_qkv, max=self.clip_qkv) - q, k, v = qkv.chunk(chunks=3, dim=-1) - q, k = self.rotary_emb(positions, q, k) - attn_output = self.attn(q, k, v) - output, _ = self.o_proj(attn_output) - return output - - -class OlmoMLP(nn.Module): - """ - This is the MLP block where the output is computed as - `MLP(LN(x))` in `MLP(LN(x + Attention(LN(x))))` - (plus another skip connection). - """ - - def __init__( - self, - config: OlmoConfig, - quant_config: QuantizationConfig | None = None, - prefix: str = "", - ): - super().__init__() - self.config = config - self.hidden_size = config.hidden_size - self.intermediate_size = config.intermediate_size - - # Feed-forward input projection. - self.gate_up_proj = MergedColumnParallelLinear( - self.hidden_size, - [self.intermediate_size] * 2, - bias=False, - quant_config=quant_config, - prefix=f"{prefix}.gate_up_proj", - ) - - # Activation function. - self.act_fn = SiluAndMul() - - # Feed-forward output projection. - self.down_proj = RowParallelLinear( - self.intermediate_size, - self.hidden_size, - bias=False, - quant_config=quant_config, - prefix=f"{prefix}.down_proj", - ) - - def forward( - self, - x: torch.Tensor, - ) -> torch.Tensor: - gate_up, _ = self.gate_up_proj(x) - x = self.act_fn(gate_up) - x, _ = self.down_proj(x) - return x - - -class OlmoDecoderLayer(nn.Module): - """ - This is a typical transformer block where the output is - computed as `MLP(LN(x + Attention(LN(x))))` - (plus another skip connection). - """ - - def __init__( - self, - config: OlmoConfig, - cache_config: CacheConfig | None = None, - quant_config: QuantizationConfig | None = None, - prefix: str = "", - ): - super().__init__() - # Attention block. - self.self_attn = OlmoAttention( - config, cache_config, quant_config, prefix=f"{prefix}.self_attn" - ) - - # MLP block. - self.mlp = OlmoMLP(config, quant_config, prefix=f"{prefix}.mlp") - - # LayerNorm - self.input_layernorm = nn.LayerNorm( - config.hidden_size, elementwise_affine=False, bias=False - ) - self.post_attention_layernorm = nn.LayerNorm( - config.hidden_size, elementwise_affine=False, bias=False - ) - - def forward( - self, - positions: torch.Tensor, - hidden_states: torch.Tensor, - ) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor] | None]: - # Attention block. - residual = hidden_states - hidden_states = self.input_layernorm(hidden_states) - hidden_states = self.self_attn(positions, hidden_states) - hidden_states = hidden_states + residual - - # MLP block. - residual = hidden_states - hidden_states = self.post_attention_layernorm(hidden_states) - hidden_states = self.mlp(hidden_states) - hidden_states = residual + hidden_states - return hidden_states - - -@support_torch_compile -class OlmoModel(nn.Module): - def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): - super().__init__() - - config = vllm_config.model_config.hf_config - cache_config = vllm_config.cache_config - quant_config = vllm_config.quant_config - - self.config = config - - self.embed_tokens = VocabParallelEmbedding( - config.vocab_size, config.hidden_size - ) - self.start_layer, self.end_layer, self.layers = make_layers( - config.num_hidden_layers, - lambda prefix: OlmoDecoderLayer( - config, cache_config, quant_config, prefix=prefix - ), - prefix=f"{prefix}.layers", - ) - self.norm = nn.LayerNorm( - config.hidden_size, elementwise_affine=False, bias=False - ) - self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( - ["hidden_states"], config.hidden_size - ) - - def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: - return self.embed_tokens(input_ids) - - def forward( - self, - input_ids: torch.Tensor | None, - positions: torch.Tensor, - intermediate_tensors: IntermediateTensors | None, - inputs_embeds: torch.Tensor | None = None, - ) -> torch.Tensor | IntermediateTensors: - """ - Args: - input_ids: A tensor of shape `(batch_size, seq_len)`. - """ - if get_pp_group().is_first_rank: - if inputs_embeds is not None: - hidden_states = inputs_embeds - else: - hidden_states = self.embed_input_ids(input_ids) - else: - assert intermediate_tensors is not None - hidden_states = intermediate_tensors["hidden_states"] - - # Apply blocks one-by-one. - for layer in islice(self.layers, self.start_layer, self.end_layer): - # shape: (batch_size, seq_len, d_model) - hidden_states = layer(positions, hidden_states) - - if not get_pp_group().is_last_rank: - return IntermediateTensors({"hidden_states": hidden_states}) - # Apply final layer norm. - # shape: (batch_size, seq_len or 1, d_model) - hidden_states = self.norm(hidden_states) - return hidden_states - - -class OlmoForCausalLM(nn.Module, SupportsPP, SupportsLoRA): - """ - Extremely barebones HF model wrapper. - """ - - hf_to_vllm_mapper = WeightsMapper( - orig_to_new_stacked={ - # weight_name: (param_name, shard_id) - ".q_proj": (".qkv_proj", "q"), - ".k_proj": (".qkv_proj", "k"), - ".v_proj": (".qkv_proj", "v"), - ".gate_proj": (".gate_up_proj", 0), - ".up_proj": (".gate_up_proj", 1), - } - ) - packed_modules_mapping = { - "qkv_proj": ["q_proj", "k_proj", "v_proj"], - "gate_up_proj": ["gate_proj", "up_proj"], - } - - def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): - super().__init__() - config = vllm_config.model_config.hf_config - quant_config = vllm_config.quant_config - self.config = config - self.model = OlmoModel( - vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") - ) - if config.tie_word_embeddings: - self.lm_head = self.model.embed_tokens - else: - self.lm_head = ParallelLMHead( - config.vocab_size, - config.hidden_size, - quant_config=quant_config, - prefix=maybe_prefix(prefix, "lm_head"), - ) - self.logits_processor = LogitsProcessor(config.vocab_size) - self.make_empty_intermediate_tensors = ( - self.model.make_empty_intermediate_tensors - ) - - def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: - return self.model.embed_input_ids(input_ids) - - def forward( - self, - input_ids: torch.Tensor | None, - positions: torch.Tensor, - intermediate_tensors: IntermediateTensors | None = None, - inputs_embeds: torch.Tensor | None = None, - ) -> torch.Tensor | IntermediateTensors: - hidden_states = self.model( - input_ids=input_ids, - positions=positions, - intermediate_tensors=intermediate_tensors, - inputs_embeds=inputs_embeds, - ) - return hidden_states - - def compute_logits( - self, - hidden_states: torch.Tensor, - ) -> torch.Tensor | None: - logits = self.logits_processor(self.lm_head, hidden_states) - return logits - - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader( - self, - skip_prefixes=( - ["lm_head.weight"] if self.config.tie_word_embeddings else None - ), - ) - return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm/model_executor/models/olmo2.py b/vllm/model_executor/models/olmo3.py similarity index 94% rename from vllm/model_executor/models/olmo2.py rename to vllm/model_executor/models/olmo3.py index 489ec2616cb..3c36c059788 100644 --- a/vllm/model_executor/models/olmo2.py +++ b/vllm/model_executor/models/olmo3.py @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # Adapted from -# https://github.com/huggingface/transformers/blob/main/src/transformers/models/olmo2/modeling_olmo2.py +# https://github.com/huggingface/transformers/blob/main/src/transformers/models/olmo3/modeling_olmo3.py # Copyright 2024 The vLLM team. # Copyright 2024 EleutherAI and the HuggingFace Inc. team. All rights reserved. # @@ -22,7 +22,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -"""Inference-only OLMo2 model compatible with HuggingFace weights.""" +"""Inference-only OLMo3 model compatible with HuggingFace weights.""" from collections.abc import Iterable from functools import partial @@ -30,7 +30,7 @@ from itertools import islice import torch from torch import nn -from transformers import Olmo2Config, Olmo3Config +from transformers import Olmo3Config from vllm.compilation.decorators import support_torch_compile from vllm.config import VllmConfig @@ -64,7 +64,7 @@ from vllm.model_executor.models.utils import ( from vllm.sequence import IntermediateTensors -class Olmo2Attention(nn.Module): +class Olmo3Attention(nn.Module): """ This is the attention block where the output is computed as `Attention(LN(x))` in `MLP(LN(x + Attention(LN(x))))` @@ -74,7 +74,7 @@ class Olmo2Attention(nn.Module): def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() self.config = vllm_config.model_config.hf_config - assert isinstance(self.config, (Olmo2Config, Olmo3Config)) + assert isinstance(self.config, Olmo3Config) hidden_size = self.config.hidden_size self.tp_size = get_tensor_model_parallel_world_size() @@ -185,7 +185,7 @@ class Olmo2Attention(nn.Module): return output -class Olmo2MLP(nn.Module): +class Olmo3MLP(nn.Module): """ This is the MLP block where the output is computed as `MLP(x)` in `LN(MLP(x + LN(Attention(x))))` @@ -195,7 +195,7 @@ class Olmo2MLP(nn.Module): def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() config = vllm_config.model_config.hf_config - assert isinstance(config, (Olmo2Config, Olmo3Config)) + assert isinstance(config, Olmo3Config) hidden_size = config.hidden_size intermediate_size = config.intermediate_size @@ -230,7 +230,7 @@ class Olmo2MLP(nn.Module): return x -class Olmo2DecoderLayer(nn.Module): +class Olmo3DecoderLayer(nn.Module): """ This is a typical transformer block where the output is computed as `MLP(LN(x + Attention(LN(x))))` @@ -240,14 +240,14 @@ class Olmo2DecoderLayer(nn.Module): def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() config = vllm_config.model_config.hf_config - assert isinstance(config, (Olmo2Config, Olmo3Config)) + assert isinstance(config, Olmo3Config) # Attention block. - self.self_attn = Olmo2Attention( + self.self_attn = Olmo3Attention( vllm_config=vllm_config, prefix=f"{prefix}.self_attn" ) # MLP block. - self.mlp = Olmo2MLP(vllm_config=vllm_config, prefix=f"{prefix}.mlp") + self.mlp = Olmo3MLP(vllm_config=vllm_config, prefix=f"{prefix}.mlp") # LayerNorm self.post_attention_layernorm = RMSNorm( @@ -278,11 +278,11 @@ class Olmo2DecoderLayer(nn.Module): @support_torch_compile -class Olmo2Model(nn.Module): +class Olmo3Model(nn.Module): def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() self.config = vllm_config.model_config.hf_config - assert isinstance(self.config, (Olmo2Config, Olmo3Config)) + assert isinstance(self.config, Olmo3Config) self.embed_tokens = VocabParallelEmbedding( self.config.vocab_size, @@ -291,7 +291,7 @@ class Olmo2Model(nn.Module): ) self.start_layer, self.end_layer, self.layers = make_layers( self.config.num_hidden_layers, - lambda prefix: Olmo2DecoderLayer(vllm_config=vllm_config, prefix=prefix), + lambda prefix: Olmo3DecoderLayer(vllm_config=vllm_config, prefix=prefix), prefix=f"{prefix}.layers", ) self.norm = RMSNorm( @@ -343,7 +343,7 @@ class Olmo2Model(nn.Module): return hidden_states -class Olmo2ForCausalLM(nn.Module, SupportsPP, SupportsLoRA): +class Olmo3ForCausalLM(nn.Module, SupportsPP, SupportsLoRA): """ Extremely barebones HF model wrapper. """ @@ -366,9 +366,9 @@ class Olmo2ForCausalLM(nn.Module, SupportsPP, SupportsLoRA): def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() config = vllm_config.model_config.hf_config - assert isinstance(config, (Olmo2Config, Olmo3Config)) + assert isinstance(config, Olmo3Config) self.config = config - self.model = Olmo2Model( + self.model = Olmo3Model( vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") ) if config.tie_word_embeddings: diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index 8d5879cc3d5..6d70afc4f68 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -168,9 +168,7 @@ _TEXT_GENERATION_MODELS = { "NemotronForCausalLM": ("nemotron", "NemotronForCausalLM"), "NemotronHForCausalLM": ("nemotron_h", "NemotronHForCausalLM"), "NemotronHPuzzleForCausalLM": ("nemotron_h", "NemotronHForCausalLM"), - "OlmoForCausalLM": ("olmo", "OlmoForCausalLM"), - "Olmo2ForCausalLM": ("olmo2", "Olmo2ForCausalLM"), - "Olmo3ForCausalLM": ("olmo2", "Olmo2ForCausalLM"), + "Olmo3ForCausalLM": ("olmo3", "Olmo3ForCausalLM"), "OlmoHybridForCausalLM": ("olmo_hybrid", "OlmoHybridForCausalLM"), "OlmoeForCausalLM": ("olmoe", "OlmoeForCausalLM"), "OPTForCausalLM": ("opt", "OPTForCausalLM"), @@ -637,6 +635,8 @@ _SPECULATIVE_DECODING_MODELS = { _TRANSFORMERS_SUPPORTED_MODELS = { # Text generation models "GPTBigCodeForCausalLM": ("transformers", "TransformersForCausalLM"), + "OlmoForCausalLM": ("transformers", "TransformersForCausalLM"), + "Olmo2ForCausalLM": ("transformers", "TransformersForCausalLM"), "SmolLM3ForCausalLM": ("transformers", "TransformersForCausalLM"), "Starcoder2ForCausalLM": ("transformers", "TransformersForCausalLM"), # Multimodal models