karylab_agents/vllm
1
0
Fork 0
forked from Karylab-cklius/vllm
Code Pull Requests 2 Actions 1 Packages Activity

Migrate Olmo and Olmo2 to the Transformers modeling backend (#48100)

Browse Source 
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
This commit is contained in:
Harry Mellor
2026-07-09 05:00:23 -07:00
committed by GitHub
parent 412414d8e0
commit b83be00cdd
4 changed files with 22 additions and 400 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/models/supported_models.md
+2 -2
View File
@@ -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. | ✅︎ | ✅︎ |
vllm/model_executor/models/olmo.py
-378
View File
@@ -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)
vllm/model_executor/models/olmo2.py → vllm/model_executor/models/olmo3.py
+17 -17
View File
@@ -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:
vllm/model_executor/models/registry.py
+3 -3
View File
@@ -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