[Model] Remove Ouro (#49786)
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
Harry Mellor committed
Jul 25, 2026 at 10:50 UTC
dbcc1cdd0a7b41a75f14e3e432cce52072862a6b
4 files changed
+1
-451
docs/models/supported_models.md
-1
@@ -427,7 +427,6 @@ th {
427
| `OlmoeForCausalLM` | OLMoE | `allenai/OLMoE-1B-7B-0924`, `allenai/OLMoE-1B-7B-0924-Instruct`, etc. | | ✅︎ |
428
| `OPTForCausalLM` | OPT, OPT-IML | `facebook/opt-66b`, `facebook/opt-iml-max-30b`, etc. | ✅︎ | ✅︎ |
429
| `OrionForCausalLM` | Orion | `OrionStarAI/Orion-14B-Base`, `OrionStarAI/Orion-14B-Chat`, etc. | | ✅︎ |
430
-| `OuroForCausalLM` | ouro | `ByteDance/Ouro-1.4B`, `ByteDance/Ouro-2.6B`, etc. | ✅︎ | |
430
| `PanguEmbeddedForCausalLM` | openPangu-Embedded-7B | `FreedomIntelligence/openPangu-Embedded-7B-V1.1` | ✅︎ | ✅︎ |
431
| `PanguProMoEV2ForCausalLM` | openpangu-pro-moe-v2 | | ✅︎ | ✅︎ |
432
| `PanguUltraMoEForCausalLM` | openpangu-ultra-moe-718b-model | `FreedomIntelligence/openPangu-Ultra-MoE-718B-V1.1` | ✅︎ | ✅︎ |
tests/models/registry.py
-1
@@ -461,7 +461,6 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
461
"OrionForCausalLM": _HfExamplesInfo(
462
"OrionStarAI/Orion-14B-Chat", trust_remote_code=True
463
),
464
- "OuroForCausalLM": _HfExamplesInfo("ByteDance/Ouro-1.4B", trust_remote_code=True),
464
"PanguEmbeddedForCausalLM": _HfExamplesInfo(
465
"FreedomIntelligence/openPangu-Embedded-7B-V1.1", trust_remote_code=True
466
),
vllm/model_executor/models/ouro.py
deleted
-448
@@ -1,448 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
5
-# Adapted from
6
-# https://github.com/huggingface/transformers/blob/v4.28.0/src/transformers/models/qwen2/modeling_qwen2.py
7
-# Copyright 2024 The Qwen team.
8
-# Copyright 2023 The vLLM team.
9
-# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.
10
-#
11
-# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
12
-# and OPT implementations in this library. It has been modified from its
13
-# original forms to accommodate minor architectural differences compared
14
-# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
15
-#
16
-# Licensed under the Apache License, Version 2.0 (the "License");
17
-# you may not use this file except in compliance with the License.
18
-# You may obtain a copy of the License at
19
-#
20
-# http://www.apache.org/licenses/LICENSE-2.0
21
-#
22
-# Unless required by applicable law or agreed to in writing, software
23
-# distributed under the License is distributed on an "AS IS" BASIS,
24
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
25
-# See the License for the specific language governing permissions and
26
-# limitations under the License.
27
-"""Inference-only Ouro model compatible with HuggingFace weights."""
28
-
29
-from collections.abc import Iterable
30
-from typing import Any
31
-
32
-import torch
33
-from torch import nn
34
-from transformers import PretrainedConfig
35
-
36
-from vllm.compilation.decorators import support_torch_compile
37
-from vllm.config import CacheConfig, VllmConfig
38
-from vllm.distributed import get_tensor_model_parallel_world_size
39
-from vllm.model_executor.layers.activation import SiluAndMul
40
-from vllm.model_executor.layers.attention import Attention
41
-from vllm.model_executor.layers.layernorm import RMSNorm
42
-from vllm.model_executor.layers.linear import (
43
- MergedColumnParallelLinear,
44
- QKVParallelLinear,
45
- RowParallelLinear,
46
-)
47
-from vllm.model_executor.layers.logits_processor import LogitsProcessor
48
-from vllm.model_executor.layers.quantization import QuantizationConfig
49
-from vllm.model_executor.layers.rotary_embedding import get_rope
50
-from vllm.model_executor.layers.vocab_parallel_embedding import (
51
- ParallelLMHead,
52
- VocabParallelEmbedding,
53
-)
54
-from vllm.sequence import IntermediateTensors
55
-from vllm.v1.attention.backend import AttentionType
56
-
57
-from .interfaces import SupportsLoRA
58
-from .utils import (
59
- AutoWeightsLoader,
60
- WeightsMapper,
61
- extract_layer_index,
62
- make_empty_intermediate_tensors_factory,
63
- make_layers,
64
- maybe_prefix,
65
-)
66
-
67
-
68
-class OuroMLP(nn.Module):
69
- def __init__(
70
- self,
71
- hidden_size: int,
72
- intermediate_size: int,
73
- hidden_act: str,
74
- quant_config: QuantizationConfig | None = None,
75
- prefix: str = "",
76
- ) -> None:
77
- super().__init__()
78
- self.gate_up_proj = MergedColumnParallelLinear(
79
- hidden_size,
80
- [intermediate_size] * 2,
81
- bias=False,
82
- quant_config=quant_config,
83
- prefix=f"{prefix}.gate_up_proj",
84
- )
85
- self.down_proj = RowParallelLinear(
86
- intermediate_size,
87
- hidden_size,
88
- bias=False,
89
- quant_config=quant_config,
90
- prefix=f"{prefix}.down_proj",
91
- )
92
- if hidden_act != "silu":
93
- raise ValueError(
94
- f"Unsupported activation: {hidden_act}. Only silu is supported for now."
95
- )
96
- self.act_fn = SiluAndMul()
97
-
98
- def forward(self, x):
99
- gate_up, _ = self.gate_up_proj(x)
100
- x = self.act_fn(gate_up)
101
- x, _ = self.down_proj(x)
102
- return x
103
-
104
-
105
-class OuroAttention(nn.Module):
106
- def __init__(
107
- self,
108
- config: PretrainedConfig,
109
- hidden_size: int,
110
- num_heads: int,
111
- num_kv_heads: int,
112
- max_position: int = 4096 * 32,
113
- cache_config: CacheConfig | None = None,
114
- quant_config: QuantizationConfig | None = None,
115
- prefix: str = "",
116
- attn_type: str = AttentionType.DECODER,
117
- dual_chunk_attention_config: dict[str, Any] | None = None,
118
- ) -> None:
119
- super().__init__()
120
- self.hidden_size = hidden_size
121
- tp_size = get_tensor_model_parallel_world_size()
122
- self.total_num_heads = num_heads
123
- assert self.total_num_heads % tp_size == 0
124
- self.num_heads = self.total_num_heads // tp_size
125
- self.total_num_kv_heads = num_kv_heads
126
- if self.total_num_kv_heads >= tp_size:
127
- # Number of KV heads is greater than TP size, so we partition
128
- # the KV heads across multiple tensor parallel GPUs.
129
- assert self.total_num_kv_heads % tp_size == 0
130
- else:
131
- # Number of KV heads is less than TP size, so we replicate
132
- # the KV heads across multiple tensor parallel GPUs.
133
- assert tp_size % self.total_num_kv_heads == 0
134
- self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
135
- self.head_dim = hidden_size // self.total_num_heads
136
- self.q_size = self.num_heads * self.head_dim
137
- self.kv_size = self.num_kv_heads * self.head_dim
138
- self.scaling = self.head_dim**-0.5
139
- self.dual_chunk_attention_config = dual_chunk_attention_config
140
-
141
- # Get total_ut_steps from config, default to 4 if not specified
142
- total_ut_steps = getattr(config, "total_ut_steps", 4)
143
-
144
- # Use total number of hidden layers instead of hardcoded 24
145
- total_layers = config.num_hidden_layers
146
-
147
- self.qkv_proj = QKVParallelLinear(
148
- hidden_size,
149
- self.head_dim,
150
- self.total_num_heads,
151
- self.total_num_kv_heads,
152
- bias=False,
153
- quant_config=quant_config,
154
- prefix=f"{prefix}.qkv_proj",
155
- )
156
- self.o_proj = RowParallelLinear(
157
- self.total_num_heads * self.head_dim,
158
- hidden_size,
159
- bias=False,
160
- quant_config=quant_config,
161
- prefix=f"{prefix}.o_proj",
162
- )
163
-
164
- self.rotary_emb = get_rope(
165
- self.head_dim,
166
- max_position=max_position,
167
- rope_parameters=config.rope_parameters,
168
- dual_chunk_attention_config=dual_chunk_attention_config,
169
- )
170
- self.attn = nn.ModuleList()
171
- for ut_step in range(total_ut_steps):
172
- base_layer_idx = extract_layer_index(prefix)
173
- unique_layer_idx = ut_step * total_layers + base_layer_idx
174
-
175
- unique_prefix = prefix.replace(
176
- f"layers.{base_layer_idx}", f"layers.{unique_layer_idx}"
177
- )
178
-
179
- self.attn.append(
180
- Attention(
181
- self.num_heads,
182
- self.head_dim,
183
- self.scaling,
184
- num_kv_heads=self.num_kv_heads,
185
- cache_config=cache_config,
186
- quant_config=quant_config,
187
- attn_type=attn_type,
188
- prefix=f"{unique_prefix}.attn",
189
- **{
190
- "layer_idx": unique_layer_idx,
191
- "dual_chunk_attention_config": dual_chunk_attention_config,
192
- }
193
- if dual_chunk_attention_config
194
- else {},
195
- )
196
- )
197
-
198
- def forward(
199
- self,
200
- positions: torch.Tensor,
201
- hidden_states: torch.Tensor,
202
- current_ut: int,
203
- ) -> torch.Tensor:
204
- qkv, _ = self.qkv_proj(hidden_states)
205
- q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
206
- q, k = self.rotary_emb(positions, q, k)
207
- attn_output = self.attn[current_ut](q, k, v)
208
- output, _ = self.o_proj(attn_output)
209
- return output
210
-
211
-
212
-class OuroDecoderLayer(nn.Module):
213
- def __init__(
214
- self,
215
- config: PretrainedConfig,
216
- cache_config: CacheConfig | None = None,
217
- quant_config: QuantizationConfig | None = None,
218
- prefix: str = "",
219
- ) -> None:
220
- super().__init__()
221
- self.hidden_size = config.hidden_size
222
- dual_chunk_attention_config = getattr(
223
- config, "dual_chunk_attention_config", None
224
- )
225
-
226
- if getattr(config, "is_causal", True):
227
- attn_type = AttentionType.DECODER
228
- else:
229
- attn_type = AttentionType.ENCODER_ONLY
230
-
231
- self.self_attn = OuroAttention(
232
- config=config,
233
- hidden_size=self.hidden_size,
234
- num_heads=config.num_attention_heads,
235
- max_position=config.max_position_embeddings,
236
- num_kv_heads=config.num_key_value_heads,
237
- cache_config=cache_config,
238
- quant_config=quant_config,
239
- prefix=f"{prefix}.self_attn",
240
- attn_type=attn_type,
241
- dual_chunk_attention_config=dual_chunk_attention_config,
242
- )
243
- self.mlp = OuroMLP(
244
- hidden_size=self.hidden_size,
245
- intermediate_size=config.intermediate_size,
246
- hidden_act=config.hidden_act,
247
- quant_config=quant_config,
248
- prefix=f"{prefix}.mlp",
249
- )
250
- self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
251
- self.input_layernorm_2 = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
252
-
253
- self.post_attention_layernorm = RMSNorm(
254
- config.hidden_size, eps=config.rms_norm_eps
255
- )
256
- self.post_attention_layernorm_2 = RMSNorm(
257
- config.hidden_size, eps=config.rms_norm_eps
258
- )
259
-
260
- def forward(
261
- self,
262
- positions: torch.Tensor,
263
- hidden_states: torch.Tensor,
264
- current_ut: int,
265
- residual: torch.Tensor | None = None,
266
- ) -> tuple[torch.Tensor, torch.Tensor]:
267
- if residual is None:
268
- residual = hidden_states
269
- hidden_states = self.input_layernorm(hidden_states)
270
- else:
271
- hidden_states, residual = self.input_layernorm(hidden_states, residual)
272
- hidden_states = self.self_attn(
273
- positions=positions, hidden_states=hidden_states, current_ut=current_ut
274
- )
275
- hidden_states = self.input_layernorm_2(hidden_states)
276
-
277
- hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
278
- hidden_states = self.mlp(hidden_states)
279
- hidden_states = self.post_attention_layernorm_2(hidden_states)
280
-
281
- return hidden_states, residual
282
-
283
-
284
-@support_torch_compile(
285
- dynamic_arg_dims={
286
- "input_ids": 0,
287
- "positions": -1,
288
- "intermediate_tensors": 0,
289
- "inputs_embeds": 0,
290
- }
291
-)
292
-class OuroModel(nn.Module):
293
- def __init__(
294
- self,
295
- *,
296
- vllm_config: VllmConfig,
297
- prefix: str = "",
298
- decoder_layer_type: type[nn.Module] = OuroDecoderLayer,
299
- ):
300
- super().__init__()
301
-
302
- config = vllm_config.model_config.hf_config
303
- cache_config = vllm_config.cache_config
304
- quant_config = vllm_config.quant_config
305
-
306
- # TODO (@robertgshaw2): see if this can be moved out
307
- if cache_config.sliding_window is not None and hasattr(
308
- config, "max_window_layers"
309
- ):
310
- assert config.max_window_layers == config.num_hidden_layers, (
311
- "Sliding window for some but all layers is not supported. "
312
- "This model uses sliding window but `max_window_layers` = {} "
313
- "is less than `num_hidden_layers` = {}. Please open an issue "
314
- "to discuss this feature.".format(
315
- config.max_window_layers,
316
- config.num_hidden_layers,
317
- )
318
- )
319
-
320
- self.config = config
321
- self.quant_config = quant_config
322
- self.vocab_size = config.vocab_size
323
-
324
- self.embed_tokens = VocabParallelEmbedding(
325
- config.vocab_size,
326
- config.hidden_size,
327
- quant_config=quant_config,
328
- prefix=f"{prefix}.embed_tokens",
329
- )
330
-
331
- # Use the provided decoder layer type or default to OuroDecoderLayer
332
- decoder_layer_type = decoder_layer_type or OuroDecoderLayer
333
- self.start_layer, self.end_layer, self.layers = make_layers(
334
- config.num_hidden_layers,
335
- lambda prefix: decoder_layer_type(
336
- config=config,
337
- cache_config=cache_config,
338
- quant_config=quant_config,
339
- prefix=prefix,
340
- ),
341
- prefix=f"{prefix}.layers",
342
- )
343
-
344
- self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
345
- ["hidden_states", "residual"], config.hidden_size
346
- )
347
- self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
348
- self.early_exit_gate = RowParallelLinear(config.hidden_size, 1, bias=True)
349
-
350
- self.total_ut_steps = getattr(self.config, "total_ut_steps", 4)
351
-
352
- def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
353
- return self.embed_tokens(input_ids)
354
-
355
- def forward(
356
- self,
357
- input_ids: torch.Tensor | None,
358
- positions: torch.Tensor,
359
- intermediate_tensors: IntermediateTensors | None = None,
360
- inputs_embeds: torch.Tensor | None = None,
361
- ) -> torch.Tensor | IntermediateTensors:
362
- if inputs_embeds is not None:
363
- hidden_states = inputs_embeds
364
- else:
365
- hidden_states = self.embed_input_ids(input_ids)
366
-
367
- for current_ut in range(self.total_ut_steps):
368
- residual = None
369
- for layer in self.layers[self.start_layer : self.end_layer]:
370
- hidden_states, residual = layer(
371
- positions, hidden_states, current_ut, residual
372
- )
373
- hidden_states, _ = self.norm(hidden_states, residual)
374
- return hidden_states
375
-
376
-
377
-class OuroForCausalLM(nn.Module, SupportsLoRA):
378
- hf_to_vllm_mapper = WeightsMapper(
379
- orig_to_new_stacked={
380
- # weight_name: (param_name, shard_id)
381
- ".q_proj": (".qkv_proj", "q"),
382
- ".k_proj": (".qkv_proj", "k"),
383
- ".v_proj": (".qkv_proj", "v"),
384
- ".gate_proj": (".gate_up_proj", 0),
385
- ".up_proj": (".gate_up_proj", 1),
386
- }
387
- )
388
- packed_modules_mapping = {
389
- "qkv_proj": ["q_proj", "k_proj", "v_proj"],
390
- "gate_up_proj": ["gate_proj", "up_proj"],
391
- }
392
-
393
- def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
394
- super().__init__()
395
- config = vllm_config.model_config.hf_config
396
- quant_config = vllm_config.quant_config
397
-
398
- self.config = config
399
-
400
- self.quant_config = quant_config
401
- self.model = OuroModel(
402
- vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
403
- )
404
-
405
- if config.tie_word_embeddings:
406
- self.lm_head = self.model.embed_tokens
407
- else:
408
- self.lm_head = ParallelLMHead(
409
- config.vocab_size,
410
- config.hidden_size,
411
- quant_config=quant_config,
412
- prefix=maybe_prefix(prefix, "lm_head"),
413
- )
414
-
415
- self.logits_processor = LogitsProcessor(config.vocab_size)
416
-
417
- self.make_empty_intermediate_tensors = (
418
- self.model.make_empty_intermediate_tensors
419
- )
420
-
421
- def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
422
- return self.model.embed_input_ids(input_ids)
423
-
424
- def forward(
425
- self,
426
- input_ids: torch.Tensor | None,
427
- positions: torch.Tensor,
428
- intermediate_tensors: IntermediateTensors | None = None,
429
- inputs_embeds: torch.Tensor | None = None,
430
- ) -> torch.Tensor | IntermediateTensors:
431
- hidden_states = self.model(
432
- input_ids, positions, intermediate_tensors, inputs_embeds
433
- )
434
- return hidden_states
435
-
436
- def compute_logits(
437
- self,
438
- hidden_states: torch.Tensor,
439
- ) -> torch.Tensor | None:
440
- logits = self.logits_processor(self.lm_head, hidden_states)
441
- return logits
442
-
443
- def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
444
- loader = AutoWeightsLoader(
445
- self,
446
- skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
447
- )
448
- return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
vllm/model_executor/models/registry.py
+1
-1
@@ -183,7 +183,6 @@ _TEXT_GENERATION_MODELS = {
183
"OlmoeForCausalLM": ("olmoe", "OlmoeForCausalLM"),
184
"OPTForCausalLM": ("opt", "OPTForCausalLM"),
185
"OrionForCausalLM": ("orion", "OrionForCausalLM"),
186
- "OuroForCausalLM": ("ouro", "OuroForCausalLM"),
186
"PanguEmbeddedForCausalLM": ("openpangu", "PanguEmbeddedForCausalLM"),
187
"PanguProMoEV2ForCausalLM": ("openpangu", "PanguProMoEV2ForCausalLM"),
188
"PanguUltraMoEForCausalLM": ("openpangu", "PanguUltraMoEForCausalLM"),
@@ -760,6 +759,7 @@ _PREVIOUSLY_SUPPORTED_MODELS = {
759
"PersimmonForCausalLM": "0.25.0",
760
"FuyuForCausalLM": "0.25.0",
761
"Plamo2ForCausalLM": "0.26.0",
762
+ "OuroForCausalLM": "0.26.0",
763
}
764
765
_OOT_SUPPORTED_MODELS = {