@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 48d7132962

[Model] Migrate FlexOlmo, Olmo3 and Hunyuan V1/VL to the Transformers modeling backend (#53615)

Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Signed-off-by: Cyrus Leung <tlleungac@connect.ust.hk> Co-authored-by: Xianbao QIAN <xianbao.qian@gmail.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>

Harry Mellor committed Aug 24, 2026 at 22:39 UTC 48d7132962ef3e2e6982a316013294c0029e3a64

14 files changed +47 -3028

docs/models/supported_models.md

+5 -5

@@ -384,7 +384,6 @@ th {

386 | `FalconH1ForCausalLM` | Falcon-H1 | `tiiuae/Falcon-H1-34B-Base`, `tiiuae/Falcon-H1-34B-Instruct`, etc. | ✅︎ | ✅︎ |

387 | `GemmaForCausalLM` | Gemma | `google/gemma-2b`, `google/gemma-1.1-2b-it`, etc. | ✅︎ | ✅︎ |

388 | `Gemma2ForCausalLM` | Gemma 2 | `google/gemma-2-9b`, `google/gemma-2-27b`, etc. | ✅︎ | ✅︎ |

389 | `Gemma3ForCausalLM` | Gemma 3 | `google/gemma-3-1b-it`, etc. | ✅︎ | ✅︎ |

@@ -406,8 +405,6 @@ th {

405 | `GraniteMoeSWAForCausalLM` | Granite MoE SWA | `ibm-granite/granite-swash-3b-a600m` | ✅︎ | ✅︎ |

406 | `GraniteSWAForCausalLM` | Granite SWA | `ibm-granite/granite-swash-2b` | ✅︎ | ✅︎ |

409 -| `HunYuanDenseV1ForCausalLM` | Hunyuan Dense | `tencent/Hunyuan-7B-Instruct` | ✅︎ | ✅︎ |

410 -| `HunYuanMoEV1ForCausalLM` | Hunyuan-A13B | `tencent/Hunyuan-A13B-Instruct`, `tencent/Hunyuan-A13B-Pretrain`, `tencent/Hunyuan-A13B-Instruct-FP8`, etc. | ✅︎ | ✅︎ |

408 | `HYV3ForCausalLM` | HY3 | `tencent/Hy3-preview-Base`, `tencent/Hy3-preview` | ✅︎ | ✅︎ |

409 | `HyperCLOVAXForCausalLM` | HyperCLOVAX-SEED-Think-14B | `naver-hyperclovax/HyperCLOVAX-SEED-Think-14B` | ✅︎ | ✅︎ |

410 | `InternLM2ForCausalLM` | InternLM2 | `internlm/internlm2-7b`, `internlm/internlm2-chat-7b`, etc. | ✅︎ | ✅︎ |

@@ -436,7 +433,6 @@ th {

433 | `MixtralForCausalLM` | Mixtral-8x7B, Mixtral-8x7B-Instruct | `mistralai/Mixtral-8x7B-v0.1`, `mistralai/Mixtral-8x7B-Instruct-v0.1`, `mistral-community/Mixtral-8x22B-v0.1`, etc. | ✅︎ | ✅︎ |

434 | `NemotronForCausalLM` | Nemotron-3, Nemotron-4, Minitron | `nvidia/Minitron-8B-Base`, `mgoin/Nemotron-4-340B-Base-hf-FP8`, etc. | ✅︎ | ✅︎ |

435 | `NemotronHForCausalLM` | Nemotron-H | `nvidia/Nemotron-H-8B-Base-8K`, `nvidia/Nemotron-H-47B-Base-8K`, `nvidia/Nemotron-H-56B-Base-8K`, etc. | ✅︎ | ✅︎ |

439 -| `Olmo3ForCausalLM` | OLMo3 | `allenai/Olmo-3-7B-Instruct`, `allenai/Olmo-3-32B-Think`, etc. | ✅︎ | ✅︎ |

436 | `OlmoHybridForCausalLM` | OLMo Hybrid | `allenai/Olmo-Hybrid-7B` | ✅︎ | ✅︎ |

438 | `OPTForCausalLM` | OPT, OPT-IML | `facebook/opt-66b`, `facebook/opt-iml-max-30b`, etc. | ✅︎ | ✅︎ |

@@ -471,9 +467,13 @@ Some models are supported only via the [Transformers modeling backend](#transfor

467

469 | ------------ | ------ | ----------------- | -------------------- | ------------------------- |

472 +| `HunYuanDenseV1ForCausalLM` | Hunyuan Dense | `tencent/Hunyuan-7B-Instruct` | ✅︎ | ✅︎ |

473 +| `HunYuanMoEV1ForCausalLM` | Hunyuan-A13B | `tencent/Hunyuan-A13B-Instruct`, `tencent/Hunyuan-A13B-Pretrain`, `tencent/Hunyuan-A13B-Instruct-FP8`, etc. | ✅︎ | ✅︎ |

474 | `OlmoForCausalLM` | OLMo | `allenai/OLMo-1B-hf`, `allenai/OLMo-7B-hf`, etc. | ✅︎ | ✅︎ |

475 | `Olmo2ForCausalLM` | OLMo2 | `allenai/OLMo-2-0425-1B`, etc. | ✅︎ | ✅︎ |

476 +| `Olmo3ForCausalLM` | OLMo3 | `allenai/Olmo-3-7B-Instruct`, `allenai/Olmo-3-32B-Think`, etc. | ✅︎ | ✅︎ |

477 | `SmolLM3ForCausalLM` | SmolLM3 | `HuggingFaceTB/SmolLM3-3B` | ✅︎ | ✅︎ |

478 | `Starcoder2ForCausalLM` | Starcoder2 | `bigcode/starcoder2-3b`, `bigcode/starcoder2-7b`, `bigcode/starcoder2-15b`, etc. | ✅︎ | ✅︎ |

479 | `VaultGemmaForCausalLM` | VaultGemma | `google/vaultgemma-1b` | ✅︎ | ✅︎ |

@@ -543,7 +543,6 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen

543 | `GraniteSpeechPlusForConditionalGeneration` | Granite Speech Plus | T + A | `ibm-granite/granite-speech-4.1-2b-plus` | ✅︎ | ✅︎ |

545 | `H2OVLChatModel` | H2OVL | T + IE+ | `h2oai/h2ovl-mississippi-800m`, `h2oai/h2ovl-mississippi-2b`, etc. | ✅︎ | ✅︎ |

546 -| `HunYuanVLForConditionalGeneration` | HunyuanOCR | T + IE+ | `tencent/HunyuanOCR`, etc. | ✅︎ | ✅︎ |

546 | `Idefics3ForConditionalGeneration` | Idefics3 | T + I | `HuggingFaceM4/Idefics3-8B-Llama3`, etc. | ✅︎ | |

547 | `IsaacForConditionalGeneration` | Isaac | T + I+ | `PerceptronAI/Isaac-0.1` | ✅︎ | ✅︎ |

548 | `InternS1ForConditionalGeneration` | Intern-S1 | T + IE+ + VE+ | `internlm/Intern-S1`, `internlm/Intern-S1-mini`, etc. | ✅︎ | ✅︎ |

@@ -620,6 +619,7 @@ Some models are supported only via the [Transformers modeling backend](#transfor

620 | ------------ | ------ | ------ | ----------------- | --------------------------- | --------------------------------------- |

621 | `Emu3ForConditionalGeneration` | Emu3 | T + I+ | `BAAI/Emu3-Chat-hf` | ✅︎ | ✅︎ |

622 +| `HunYuanVLForConditionalGeneration` | HunyuanOCR | T + IE+ | `tencent/HunyuanOCR`, etc. | ✅︎ | ✅︎ |

623 | `VibeVoiceAsrForConditionalGeneration` | VibeVoice-ASR | T + A+ | `microsoft/VibeVoice-ASR-HF` | ✅︎ | ✅︎ |

624

625 ^ You need to set the architecture name via `--hf-overrides` to match the one in vLLM.

tests/models/registry.py

+1 -4

@@ -329,9 +329,7 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {

329 min_transformers_version="5.9.0",

330 ),

331 "HunYuanDenseV1ForCausalLM": _HfExamplesInfo("tencent/Hunyuan-7B-Instruct"),

332 - "HunYuanMoEV1ForCausalLM": _HfExamplesInfo(

333 - "tencent/Hunyuan-A13B-Instruct", trust_remote_code=True

334 - ),

332 + "HunYuanMoEV1ForCausalLM": _HfExamplesInfo("tencent/Hunyuan-A13B-Instruct"),

333 "HYV3ForCausalLM": _HfExamplesInfo("tencent/Hy3-preview", trust_remote_code=True),

334 "HyperCLOVAXForCausalLM": _HfExamplesInfo(

335 "naver-hyperclovax/HyperCLOVAX-SEED-Think-14B",

@@ -938,7 +936,6 @@ _MULTIMODAL_EXAMPLE_MODELS = {

936 ),

937 "HunYuanVLForConditionalGeneration": _HfExamplesInfo(

938 "tencent/HunyuanOCR",

941 - hf_overrides={"num_experts": 0},

939 min_transformers_version="5.13",

940 ),

941 "Idefics3ForConditionalGeneration": _HfExamplesInfo(

vllm/model_executor/models/flex_olmo.py deleted

-156

@@ -1,156 +0,0 @@

1 -# SPDX-License-Identifier: Apache-2.0

2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 -

4 -# Licensed under the Apache License, Version 2.0 (the "License");

5 -# you may not use this file except in compliance with the License.

6 -# You may obtain a copy of the License at

7 -#

8 -# http://www.apache.org/licenses/LICENSE-2.0

9 -#

10 -# Unless required by applicable law or agreed to in writing, software

11 -# distributed under the License is distributed on an "AS IS" BASIS,

12 -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

13 -# See the License for the specific language governing permissions and

14 -# limitations under the License.

15 -"""Inference-only FlexOlmo model compatible with HuggingFace weights."""

16 -

17 -import torch

18 -from torch import nn

19 -

20 -from vllm.config import VllmConfig

21 -from vllm.distributed import get_tensor_model_parallel_world_size

22 -from vllm.logger import init_logger

23 -from vllm.model_executor.layers.fused_moe import (

24 - FusedMoEFactory,

25 -)

26 -from vllm.model_executor.layers.layernorm import RMSNorm

27 -from vllm.model_executor.layers.linear import ReplicatedLinear

28 -from vllm.model_executor.models.olmoe import OlmoeAttention, OlmoeForCausalLM

29 -from vllm.transformers_utils.configs.flex_olmo import FlexOlmoConfig

30 -

31 -logger = init_logger(__name__)

32 -

33 -

34 -class FlexOlmoAttention(OlmoeAttention):

35 - def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

36 - super().__init__(vllm_config=vllm_config, prefix=prefix)

37 -

38 - hf_config = vllm_config.model_config.hf_config

39 - assert isinstance(hf_config, FlexOlmoConfig)

40 -

41 - self.k_norm = RMSNorm(

42 - self.total_num_kv_heads * self.head_dim, eps=hf_config.rms_norm_eps

43 - )

44 - self.q_norm = RMSNorm(

45 - self.total_num_heads * self.head_dim, eps=hf_config.rms_norm_eps

46 - )

47 -

48 -

49 -class FlexOlmoMoE(nn.Module):

50 - """A tensor-parallel MoE implementation for FlexOlmo that shards each expert

51 - across all ranks.

52 -

53 - Each expert's weights are sharded across all ranks and a fused MoE

54 - kernel is used for the forward pass, and finally we reduce the outputs

55 - across ranks.

56 - """

57 -

58 - def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

59 - super().__init__()

60 -

61 - hf_config = vllm_config.model_config.hf_config

62 - assert isinstance(hf_config, FlexOlmoConfig)

63 -

64 - tp_size = get_tensor_model_parallel_world_size()

65 -

66 - # Gate always runs at half / full precision for now.

67 - self.gate = ReplicatedLinear(

68 - hf_config.hidden_size,

69 - hf_config.num_experts,

70 - bias=False,

71 - return_bias=False,

72 - quant_config=None,

73 - prefix=f"{prefix}.gate",

74 - )

75 -

76 - self.experts = FusedMoEFactory(

77 - num_experts=hf_config.num_experts,

78 - top_k=hf_config.num_experts_per_tok,

79 - hidden_size=hf_config.hidden_size,

80 - intermediate_size=hf_config.intermediate_size,

81 - renormalize=False,

82 - quant_config=None,

83 - tp_size=tp_size,

84 - prefix=f"{prefix}.experts",

85 - router_logits_dtype=torch.float32,

86 - )

87 -

88 - self.top_k = hf_config.num_experts_per_tok

89 -

90 - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:

91 - # NOTE: hidden_states can have either 1D or 2D shape.

92 - orig_shape = hidden_states.shape

93 - hidden_dim = hidden_states.shape[-1]

94 - hidden_states = hidden_states.view(-1, hidden_dim)

95 -

96 - # router_logits: (num_tokens, n_experts)

97 - router_logits = self.gate(hidden_states)

98 - # Warning: The experts mutate the hidden state input! This messes up

99 - # basic things like the residual stream.

100 - final_hidden_states = self.experts(

101 - hidden_states=hidden_states.detach().clone(),

102 - router_logits=router_logits.float(),

103 - )

104 -

105 - return final_hidden_states.view(orig_shape)

106 -

107 -

108 -class FlexOlmoDecoderLayer(nn.Module):

109 - def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:

110 - super().__init__()

111 - hf_config = vllm_config.model_config.hf_config

112 - assert isinstance(hf_config, FlexOlmoConfig)

113 -

114 - self.self_attn = FlexOlmoAttention(

115 - vllm_config=vllm_config, prefix=f"{prefix}.self_attn"

116 - )

117 - self.post_attention_layernorm = RMSNorm(

118 - hf_config.hidden_size, eps=hf_config.rms_norm_eps

119 - )

120 - self.post_feedforward_layernorm = RMSNorm(

121 - hf_config.hidden_size, eps=hf_config.rms_norm_eps

122 - )

123 -

124 - self.mlp = FlexOlmoMoE(vllm_config=vllm_config, prefix=f"{prefix}.mlp")

125 -

126 - def forward(

127 - self,

128 - positions: torch.Tensor,

129 - hidden_states: torch.Tensor,

130 - residual: torch.Tensor | None,

131 - ) -> tuple[torch.Tensor, torch.Tensor | None]:

132 - # Attention block.

133 - residual = hidden_states

134 - hidden_states = self.self_attn(positions, hidden_states)

135 - hidden_states = self.post_attention_layernorm(hidden_states)

136 - hidden_states = hidden_states + residual

137 -

138 - # MLP block.

139 - residual = hidden_states

140 - hidden_states = self.mlp(hidden_states)

141 - hidden_states = self.post_feedforward_layernorm(hidden_states)

142 - hidden_states = residual + hidden_states

143 - return hidden_states, None

144 -

145 -

146 -class FlexOlmoForCausalLM(OlmoeForCausalLM):

147 - fall_back_to_pt_during_load = False

148 -

149 - def __init__(

150 - self,

151 - *,

152 - vllm_config: VllmConfig,

153 - prefix: str = "",

154 - layer_type: type[nn.Module] = FlexOlmoDecoderLayer,

155 - ):

156 - super().__init__(vllm_config=vllm_config, prefix=prefix, layer_type=layer_type)

vllm/model_executor/models/hunyuan_v1.py deleted

-1038

@@ -1,1038 +0,0 @@

1 -# SPDX-License-Identifier: Apache-2.0

2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 -

4 -# coding=utf-8

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 HunYuan model compatible with HuggingFace weights."""

26 -

27 -import typing

28 -from collections.abc import Callable, Iterable

29 -from itertools import islice

30 -

31 -import regex as re

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, get_current_vllm_config

38 -from vllm.distributed import (

39 - get_ep_group,

40 - get_pp_group,

41 - get_tensor_model_parallel_world_size,

42 -)

43 -from vllm.model_executor.layers.activation import SiluAndMul

44 -from vllm.model_executor.layers.attention import Attention

45 -from vllm.model_executor.layers.fused_moe import (

46 - FusedMoEFactory,

47 - fused_moe_make_expert_params_mapping,

48 -)

49 -from vllm.model_executor.layers.layernorm import RMSNorm

50 -from vllm.model_executor.layers.linear import (

51 - ColumnParallelLinear,

52 - MergedColumnParallelLinear,

53 - QKVParallelLinear,

54 - ReplicatedLinear,

55 - RowParallelLinear,

56 -)

57 -from vllm.model_executor.layers.logits_processor import LogitsProcessor

58 -from vllm.model_executor.layers.quantization import QuantizationConfig

59 -from vllm.model_executor.layers.rotary_embedding import get_rope

60 -from vllm.model_executor.layers.vocab_parallel_embedding import (

61 - ParallelLMHead,

62 - VocabParallelEmbedding,

63 -)

64 -from vllm.model_executor.model_loader.weight_utils import (

65 - default_weight_loader,

66 - maybe_remap_kv_scale_name,

67 -)

68 -from vllm.sequence import IntermediateTensors

69 -from vllm.v1.attention.backend import AttentionType

70 -

71 -from .interfaces import (

72 - EagleModelMixin,

73 - MixtureOfExperts,

74 - SupportsEagle,

75 - SupportsEagle3,

76 - SupportsLoRA,

77 - SupportsPP,

78 -)

79 -from .utils import (

80 - AutoWeightsLoader,

81 - PPMissingLayer,

82 - is_pp_missing_parameter,

83 - make_layers,

84 - maybe_prefix,

85 -)

86 -

87 -

88 -def _is_moe(config: PretrainedConfig) -> bool:

89 - num_experts = getattr(config, "num_experts", None)

90 - if isinstance(num_experts, int):

91 - return num_experts > 1

92 - if isinstance(num_experts, list) and num_experts:

93 - # Ensure all elements are integers before calling max.

94 - if all(isinstance(e, int) for e in num_experts):

95 - return max(num_experts) > 1

96 - else:

97 - return False

98 - return False

99 -

100 -

101 -def _get_cla_factor(config: PretrainedConfig) -> int:

102 - if not getattr(config, "use_cla", False):

103 - return 1

104 - return getattr(config, "cla_share_factor", 1)

105 -

106 -

107 -class HunYuanMLP(nn.Module):

108 - def __init__(

109 - self,

110 - hidden_size: int,

111 - intermediate_size: int,

112 - hidden_act: str,

113 - quant_config: QuantizationConfig | None = None,

114 - bias: bool = False,

115 - prefix: str = "",

116 - reduce_results: bool = True,

117 - ) -> None:

118 - super().__init__()

119 - self.gate_up_proj = MergedColumnParallelLinear(

120 - input_size=hidden_size,

121 - output_sizes=[intermediate_size] * 2,

122 - bias=bias,

123 - quant_config=quant_config,

124 - prefix=f"{prefix}.gate_up_proj",

125 - )

126 - self.down_proj = RowParallelLinear(

127 - input_size=intermediate_size,

128 - output_size=hidden_size,

129 - bias=bias,

130 - quant_config=quant_config,

131 - prefix=f"{prefix}.down_proj",

132 - reduce_results=reduce_results,

133 - )

134 - if hidden_act != "silu":

135 - raise ValueError(

136 - f"Unsupported activation: {hidden_act}. Only silu is supported for now."

137 - )

138 - self.act_fn = SiluAndMul()

139 -

140 - def forward(self, x):

141 - gate_up, _ = self.gate_up_proj(x)

142 - x = self.act_fn(gate_up)

143 - x, _ = self.down_proj(x)

144 - return x

145 -

146 -

147 -class HunYuanAttention(nn.Module):

148 - def __init__(

149 - self,

150 - config: PretrainedConfig,

151 - hidden_size: int,

152 - num_heads: int,

153 - num_kv_heads: int,

154 - max_position_embeddings: int = 8192,

155 - quant_config: QuantizationConfig | None = None,

156 - bias: bool = False,

157 - cache_config: CacheConfig | None = None,

158 - prefix: str = "",

159 - layer_id: int = -1,

160 - ) -> None:

161 - super().__init__()

162 - self.hidden_size = hidden_size

163 - tp_size = get_tensor_model_parallel_world_size()

164 - self.total_num_heads = num_heads

165 - assert self.total_num_heads % tp_size == 0

166 - self.num_heads = self.total_num_heads // tp_size

167 - self.total_num_kv_heads = num_kv_heads

168 - if self.total_num_kv_heads >= tp_size:

169 - # Number of KV heads is greater than TP size, so we partition

170 - # the KV heads across multiple tensor parallel GPUs.

171 - assert self.total_num_kv_heads % tp_size == 0

172 - else:

173 - # Number of KV heads is less than TP size, so we replicate

174 - # the KV heads across multiple tensor parallel GPUs.

175 - assert tp_size % self.total_num_kv_heads == 0

176 - self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)

177 -

178 - if hasattr(config, "head_dim") and config.head_dim:

179 - self.head_dim = config.head_dim

180 - elif hasattr(config, "attention_head_dim"):

181 - self.head_dim = config.attention_head_dim

182 - else:

183 - self.head_dim = self.hidden_size // self.total_num_heads

184 - self.q_size = self.num_heads * self.head_dim

185 - self.kv_size = self.num_kv_heads * self.head_dim

186 - self.scaling = self.head_dim**-0.5

187 - self.max_position_embeddings = max_position_embeddings

188 - self.use_qk_norm = getattr(config, "use_qk_norm", False)

189 - self.layer_id = layer_id

190 -

191 - self.qkv_proj = QKVParallelLinear(

192 - hidden_size=hidden_size,

193 - head_size=self.head_dim,

194 - total_num_heads=self.total_num_heads,

195 - total_num_kv_heads=self.total_num_kv_heads,

196 - bias=bias,

197 - quant_config=quant_config,

198 - prefix=f"{prefix}.qkv_proj",

199 - )

200 -

201 - self.o_proj = RowParallelLinear(

202 - input_size=self.total_num_heads * self.head_dim,

203 - output_size=hidden_size,

204 - bias=bias,

205 - quant_config=quant_config,

206 - prefix=f"{prefix}.o_proj",

207 - )

208 -

209 - self.rotary_emb = get_rope(

210 - self.head_dim,

211 - max_position=max_position_embeddings,

212 - rope_parameters=config.rope_parameters,

213 - is_neox_style=True,

214 - )

215 - self.attn = Attention(

216 - self.num_heads,

217 - self.head_dim,

218 - self.scaling,

219 - num_kv_heads=self.num_kv_heads,

220 - cache_config=cache_config,

221 - quant_config=quant_config,

222 - prefix=f"{prefix}.attn",

223 - )

224 -

225 - if self.use_qk_norm:

226 - self.query_layernorm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)

227 - self.key_layernorm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)

228 -

229 - def forward(

230 - self,

231 - positions: torch.Tensor,

232 - hidden_states: torch.Tensor,

233 - kv_states: tuple[torch.Tensor] | None = None,

234 - ) -> torch.Tensor:

235 - qkv, _ = self.qkv_proj(hidden_states)

236 - q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)

237 - q, k = self.rotary_emb(positions, q, k)

238 - ori_k = k

239 - if self.use_qk_norm:

240 - q = self.query_layernorm(

241 - q.view(-1, self.num_heads, self.head_dim),

242 - )

243 - k = self.key_layernorm(

244 - k.view(-1, self.num_kv_heads, self.head_dim),

245 - )

246 -

247 - attn_output = self.attn(q, k, v)

248 - # For o_proj

249 - attn_output = attn_output.view(q.shape[0], -1)

250 - output, _ = self.o_proj(attn_output)

251 - return output, (ori_k, v)

252 -

253 -

254 -class HunYuanCrossAttention(nn.Module):

255 - def __init__(

256 - self,

257 - config: PretrainedConfig,

258 - hidden_size: int,

259 - num_heads: int,

260 - num_kv_heads: int,

261 - max_position_embeddings: int = 8192,

262 - quant_config: QuantizationConfig | None = None,

263 - bias: bool = False,

264 - cache_config: CacheConfig | None = None,

265 - prefix: str = "",

266 - layer_id: int = -1,

267 - ) -> None:

268 - super().__init__()

269 - self.hidden_size = hidden_size

270 - tp_size = get_tensor_model_parallel_world_size()

271 - self.total_num_heads = num_heads

272 - assert self.total_num_heads % tp_size == 0

273 - self.num_heads = self.total_num_heads // tp_size

274 - self.total_num_kv_heads = num_kv_heads

275 - if self.total_num_kv_heads >= tp_size:

276 - # Number of KV heads is greater than TP size, so we partition

277 - # the KV heads across multiple tensor parallel GPUs.

278 - assert self.total_num_kv_heads % tp_size == 0

279 - else:

280 - # Number of KV heads is less than TP size, so we replicate

281 - # the KV heads across multiple tensor parallel GPUs.

282 - assert tp_size % self.total_num_kv_heads == 0

283 - self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)

284 - # MistralConfig has an optional head_dim introduced by Mistral-Nemo

285 - if hasattr(config, "head_dim"):

286 - self.head_dim = config.head_dim

287 - elif hasattr(config, "attention_head_dim"):

288 - self.head_dim = config.attention_head_dim

289 - else:

290 - self.head_dim = self.hidden_size // self.total_num_heads

291 - self.q_size = self.num_heads * self.head_dim

292 - self.kv_size = self.num_kv_heads * self.head_dim

293 - self.scaling = self.head_dim**-0.5

294 - self.max_position_embeddings = max_position_embeddings

295 - self.use_qk_norm = getattr(config, "use_qk_norm", False)

296 - self.layer_id = layer_id

297 -

298 - self.q_proj = ColumnParallelLinear(

299 - hidden_size,

300 - hidden_size,

301 - bias=bias,

302 - quant_config=quant_config,

303 - prefix=f"{prefix}.q_proj",

304 - )

305 -

306 - self.o_proj = RowParallelLinear(

307 - input_size=self.total_num_heads * self.head_dim,

308 - output_size=hidden_size,

309 - bias=bias,

310 - quant_config=quant_config,

311 - prefix=f"{prefix}.o_proj",

312 - )

313 -

314 - self.rotary_emb = get_rope(

315 - self.head_dim,

316 - max_position=max_position_embeddings,

317 - rope_parameters=config.rope_parameters,

318 - is_neox_style=True,

319 - )

320 - self.attn = Attention(

321 - self.num_heads,

322 - self.head_dim,

323 - self.scaling,

324 - num_kv_heads=self.num_kv_heads,

325 - cache_config=cache_config,

326 - quant_config=quant_config,

327 - prefix=f"{prefix}.attn",

328 - attn_type=AttentionType.ENCODER_DECODER,

329 - )

330 -

331 - if self.use_qk_norm:

332 - self.query_layernorm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)

333 - self.key_layernorm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)

334 -

335 - def forward(

336 - self,

337 - positions: torch.Tensor,

338 - hidden_states: torch.Tensor,

339 - kv_states: tuple[torch.Tensor] | None = None,

340 - ) -> torch.Tensor:

341 - assert kv_states is not None

342 - ori_k, v = kv_states # use last layer kv,

343 - k = ori_k

344 - q, _ = self.q_proj(hidden_states)

345 - k_tmp = torch.empty_like(k) # Todo: reduant rotary embedding

346 - q, _ = self.rotary_emb(positions, q, k_tmp)

347 - if self.use_qk_norm:

348 - q = self.query_layernorm(

349 - q.view(-1, self.num_heads, self.head_dim),

350 - )

351 - k = self.key_layernorm(

352 - k.view(-1, self.num_kv_heads, self.head_dim),

353 - )

354 -

355 - attn_output = self.attn(q, k, v)

356 - # For o_proj

357 - attn_output = attn_output.view(q.shape[0], -1)

358 - output, _ = self.o_proj(attn_output)

359 - return output, (ori_k, v)

360 -

361 -

362 -class HunYuanSparseMoeBlock(nn.Module):

363 - def __init__(

364 - self,

365 - config: PretrainedConfig,

366 - quant_config: QuantizationConfig | None = None,

367 - layer_id: int = -1,

368 - prefix: str = "",

369 - enable_eplb: bool = False,

370 - ):

371 - super().__init__()

372 - self.tp_size = get_tensor_model_parallel_world_size()

373 -

374 - self.ep_group = get_ep_group().device_group

375 - self.ep_size = self.ep_group.size()

376 - self.n_routed_experts = config.num_experts

377 -

378 - if self.tp_size > config.num_experts:

379 - raise ValueError(

380 - f"Tensor parallel size {self.tp_size} is greater than "

381 - f"the number of experts {config.num_experts}."

382 - )

383 -

384 - # Get layer_id topk if config.moe_topk is a list

385 - if isinstance(config.moe_topk, list):

386 - assert layer_id >= 0

387 - assert len(config.moe_topk) > layer_id

388 - top_k = config.moe_topk[layer_id]

389 - else:

390 - top_k = config.moe_topk

391 -

392 - # If it is moe, moe_intermediate_size is preferred

393 - intermediate_size = config.intermediate_size

394 - if config.moe_intermediate_size is not None:

395 - intermediate_size = (

396 - config.moe_intermediate_size

397 - if isinstance(config.moe_intermediate_size, int)

398 - else config.moe_intermediate_size[layer_id]

399 - )

400 -

401 - # Load balancing settings.

402 - vllm_config = get_current_vllm_config()

403 - eplb_config = vllm_config.parallel_config.eplb_config

404 - self.enable_eplb = enable_eplb

405 -

406 - self.n_logical_experts = self.n_routed_experts

407 - self.n_redundant_experts = eplb_config.num_redundant_experts

408 - self.n_physical_experts = self.n_logical_experts + self.n_redundant_experts

409 - self.n_local_physical_experts = self.n_physical_experts // self.ep_size

410 - self.gate = ReplicatedLinear(

411 - config.hidden_size,

412 - config.num_experts,

413 - bias=False,

414 - quant_config=None,

415 - prefix=f"{prefix}.gate",

416 - )

417 - if config.use_mixed_mlp_moe > 0:

418 - # Get layer_id num_shared_expert if config.num_shared_expert is

419 - # a list.

420 - if isinstance(config.num_shared_expert, list):

421 - assert layer_id >= 0

422 - assert len(config.num_shared_expert) > layer_id

423 - num_shared_expert = config.num_shared_expert[layer_id]

424 - else:

425 - num_shared_expert = config.num_shared_expert

426 -

427 - self.shared_mlp = HunYuanMLP(

428 - hidden_size=config.hidden_size,

429 - intermediate_size=config.intermediate_size * num_shared_expert,

430 - hidden_act=config.hidden_act,

431 - quant_config=quant_config,

432 - reduce_results=False,

433 - prefix=f"{prefix}.shared_mlp",

434 - )

435 - else:

436 - self.shared_mlp = None

437 -

438 - self.experts = FusedMoEFactory(

439 - shared_experts=self.shared_mlp,

440 - num_experts=self.n_routed_experts,

441 - top_k=top_k,

442 - hidden_size=config.hidden_size,

443 - intermediate_size=intermediate_size,

444 - renormalize=top_k > 1,

445 - quant_config=quant_config,

446 - prefix=f"{prefix}.experts",

447 - enable_eplb=self.enable_eplb,

448 - num_redundant_experts=self.n_redundant_experts,

449 - )

450 -

451 - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:

452 - # NOTE: hidden_states can have either 1D or 2D shape.

453 - orig_shape = hidden_states.shape

454 - hidden_dim = hidden_states.shape[-1]

455 - hidden_states = hidden_states.view(-1, hidden_dim)

456 -

457 - # router_logits: (num_tokens, n_experts)

458 - router_logits, _ = self.gate(hidden_states)

459 - final_hidden_states = self.experts(

460 - hidden_states=hidden_states, router_logits=router_logits

461 - )

462 -

463 - return final_hidden_states.view(orig_shape)

464 -

465 -

466 -class HunYuanDecoderLayer(nn.Module):

467 - def __init__(

468 - self,

469 - config: PretrainedConfig,

470 - cache_config: CacheConfig | None = None,

471 - quant_config: QuantizationConfig | None = None,

472 - prefix: str = "",

473 - layer_id: int = -1,

474 - enable_eplb: bool = False,

475 - ) -> None:

476 - super().__init__()

477 - assert layer_id >= 0

478 - self.layer_id = layer_id

479 - self.hidden_size = config.hidden_size

480 - self.intermediate_size = (

481 - config.intermediate_size

482 - if isinstance(config.intermediate_size, int)

483 - else config.intermediate_size[layer_id]

484 - )

485 - max_position_embeddings = getattr(config, "max_position_embeddings", 8192)

486 - attention_bias = getattr(config, "attention_bias", False) or getattr(

487 - config, "bias", False

488 - )

489 - cla_factor = _get_cla_factor(config)

490 - attention_type = (

491 - AttentionType.ENCODER_DECODER

492 - if layer_id >= 0 and layer_id % cla_factor != 0

493 - else AttentionType.DECODER

494 - )

495 - if attention_type == AttentionType.DECODER:

496 - self.self_attn = HunYuanAttention(

497 - config=config,

498 - hidden_size=self.hidden_size,

499 - num_heads=config.num_attention_heads,

500 - num_kv_heads=getattr(

501 - config, "num_key_value_heads", config.num_attention_heads

502 - ),

503 - max_position_embeddings=max_position_embeddings,

504 - quant_config=quant_config,

505 - bias=attention_bias,

506 - cache_config=cache_config,

507 - prefix=f"{prefix}.self_attn",

508 - layer_id=layer_id,

509 - )

510 - elif attention_type == AttentionType.ENCODER_DECODER:

511 - self.self_attn = HunYuanCrossAttention(

512 - config=config,

513 - hidden_size=self.hidden_size,

514 - num_heads=config.num_attention_heads,

515 - num_kv_heads=getattr(

516 - config, "num_key_value_heads", config.num_attention_heads

517 - ),

518 - max_position_embeddings=max_position_embeddings,

519 - quant_config=quant_config,

520 - bias=attention_bias,

521 - cache_config=cache_config,

522 - prefix=f"{prefix}.self_attn",

523 - layer_id=layer_id,

524 - )

525 - else:

526 - raise RuntimeError(f"Unsupported attention type: {attention_type}")

527 -

528 - if _is_moe(config):

529 - self.mlp = HunYuanSparseMoeBlock(

530 - config=config,

531 - quant_config=quant_config,

532 - layer_id=layer_id,

533 - prefix=f"{prefix}.mlp",

534 - enable_eplb=enable_eplb,

535 - )

536 - else:

537 - self.mlp = HunYuanMLP(

538 - hidden_size=self.hidden_size,

539 - intermediate_size=self.intermediate_size,

540 - hidden_act=config.hidden_act,

541 - quant_config=quant_config,

542 - bias=getattr(config, "mlp_bias", False),

543 - prefix=f"{prefix}.mlp",

544 - )

545 -

546 - self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

547 - self.post_attention_layernorm = RMSNorm(

548 - config.hidden_size, eps=config.rms_norm_eps

549 - )

550 -

551 - def forward(

552 - self,

553 - positions: torch.Tensor,

554 - hidden_states: torch.Tensor,

555 - residual: torch.Tensor | None,

556 - kv_states: tuple[torch.Tensor] | None = None,

557 - ) -> tuple[torch.Tensor, torch.Tensor]:

558 - # Self Attention

559 - if residual is None:

560 - residual = hidden_states

561 - hidden_states = self.input_layernorm(hidden_states)

562 - else:

563 - hidden_states, residual = self.input_layernorm(hidden_states, residual)

564 - hidden_states, ori_kv_states = self.self_attn(

565 - positions=positions,

566 - hidden_states=hidden_states,

567 - kv_states=kv_states,

568 - )

569 -

570 - # Fully Connected

571 - hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)

572 - hidden_states = self.mlp(hidden_states)

573 - return hidden_states, residual, ori_kv_states

574 -

575 -

576 -@support_torch_compile(

577 - dynamic_arg_dims={

578 - "input_ids": 0,

579 - # positions is of shape (xd, seq_len) if xdrope is enabled for hunyuan-vl,

580 - # otherwise (seq_len, ).

581 - "positions": -1,

582 - "intermediate_tensors": 0,

583 - "inputs_embeds": 0,

584 - }

585 -)

586 -class HunYuanModel(nn.Module, EagleModelMixin):

587 - def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

588 - super().__init__()

589 -

590 - config = vllm_config.model_config.hf_config

591 - cache_config = vllm_config.cache_config

592 - quant_config = vllm_config.quant_config

593 -

594 - eplb_config = vllm_config.parallel_config.eplb_config

595 - enable_eplb = vllm_config.parallel_config.enable_eplb

596 - self.num_redundant_experts = eplb_config.num_redundant_experts

597 -

598 - self.config = config

599 - self.quant_config = quant_config

600 -

601 - self.vocab_size = config.vocab_size

602 -

603 - if get_pp_group().is_first_rank or (

604 - config.tie_word_embeddings and get_pp_group().is_last_rank

605 - ):

606 - self.embed_tokens = VocabParallelEmbedding(

607 - self.vocab_size,

608 - config.hidden_size,

609 - quant_config=quant_config,

610 - )

611 - else:

612 - self.embed_tokens = PPMissingLayer()

613 - self.start_layer, self.end_layer, self.layers = make_layers(

614 - config.num_hidden_layers,

615 - lambda prefix: HunYuanDecoderLayer(

616 - config=config,

617 - layer_id=int(prefix.split(".")[-1]),

618 - cache_config=cache_config,

619 - quant_config=quant_config,

620 - prefix=prefix,

621 - enable_eplb=enable_eplb,

622 - ),

623 - prefix=f"{prefix}.layers",

624 - )

625 - if get_pp_group().is_last_rank:

626 - self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

627 - else:

628 - self.norm = PPMissingLayer()

629 -

630 - def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:

631 - return self.embed_tokens(input_ids)

632 -

633 - def forward(

634 - self,

635 - input_ids: torch.Tensor | None,

636 - positions: torch.Tensor,

637 - intermediate_tensors: IntermediateTensors | None,

638 - inputs_embeds: torch.Tensor | None = None,

639 - ) -> torch.Tensor | IntermediateTensors:

640 - if get_pp_group().is_first_rank:

641 - if inputs_embeds is not None:

642 - hidden_states = inputs_embeds

643 - else:

644 - hidden_states = self.embed_input_ids(input_ids)

645 - residual = None

646 - else:

647 - assert intermediate_tensors is not None

648 - hidden_states = intermediate_tensors["hidden_states"]

649 - residual = intermediate_tensors["residual"]

650 -

651 - cla_factor = _get_cla_factor(self.config)

652 - prev_kv_states = None

653 - aux_hidden_states = self._maybe_add_hidden_state([], 0, hidden_states, residual)

654 - for i, layer in enumerate(

655 - islice(self.layers, self.start_layer, self.end_layer)

656 - ):

657 - hidden_states, residual, kv_states = layer(

658 - positions,

659 - hidden_states,

660 - residual,

661 - prev_kv_states,

662 - )

663 -

664 - if getattr(self.config, "use_cla", False) and i % cla_factor == 0:

665 - prev_kv_states = kv_states

666 - else:

667 - prev_kv_states = None

668 -

669 - self._maybe_add_hidden_state(

670 - aux_hidden_states, i + 1, hidden_states, residual

671 - )

672 -

673 - if not get_pp_group().is_last_rank:

674 - return IntermediateTensors(

675 - {"hidden_states": hidden_states, "residual": residual}

676 - )

677 -

678 - hidden_states, _ = self.norm(hidden_states, residual)

679 -

680 - if len(aux_hidden_states) > 0:

681 - return hidden_states, aux_hidden_states

682 - return hidden_states

683 -

684 - def _split_qkv_weight(self, qkv: torch.Tensor):

685 - num_attention_heads = self.config.num_attention_heads

686 - num_kv_heads = getattr(

687 - self.config, "num_key_value_heads", self.config.num_attention_heads

688 - )

689 - num_key_value_groups = num_attention_heads // num_kv_heads

690 - hidden_size = self.config.hidden_size

691 -

692 - if hasattr(self.config, "head_dim"):

693 - attention_head_dim = self.config.head_dim

694 - elif hasattr(self.config, "attention_head_dim"):

695 - attention_head_dim = self.config.attention_head_dim

696 - else:

697 - attention_head_dim = self.config.hidden_size // num_attention_heads

698 -

699 - qkv = qkv.reshape(

700 - num_kv_heads, num_key_value_groups + 2, attention_head_dim, hidden_size

701 - )

702 - q, k, v = torch.split(qkv, (num_key_value_groups, 1, 1), dim=1)

703 - q = q.reshape(-1, hidden_size)

704 - k = k.reshape(-1, hidden_size)

705 - v = v.reshape(-1, hidden_size)

706 - return torch.concat((q, k, v))

707 -

708 - def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:

709 - if _is_moe(self.config):

710 - # Params for weights, fp8 weight scales, fp8 activation scales

711 - # (param_name, weight_name, expert_id, shard_id)

712 - return fused_moe_make_expert_params_mapping(

713 - self,

714 - ckpt_gate_proj_name="gate_proj",

715 - ckpt_down_proj_name="down_proj",

716 - ckpt_up_proj_name="up_proj",

717 - num_experts=self.config.num_experts,

718 - num_redundant_experts=self.num_redundant_experts,

719 - )

720 - else:

721 - return []

722 -

723 - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):

724 - cla_factor = _get_cla_factor(self.config)

725 - stacked_params_mapping = [

726 - # (param_name, shard_name, shard_id)

727 - (".qkv_proj", ".q_proj", "q"),

728 - (".qkv_proj", ".k_proj", "k"),

729 - (".qkv_proj", ".v_proj", "v"),

730 - (".gate_up_proj", ".gate_proj", 0),

731 - (".gate_up_proj", ".up_proj", 1),

732 - ]

733 -

734 - num_attention_heads = self.config.num_attention_heads

735 - num_kv_heads = getattr(

736 - self.config, "num_key_value_heads", self.config.num_attention_heads

737 - )

738 - split_params_mapping = [

739 - (".gate_up_proj", ".gate_and_up_proj", 2, [(1, 1), (0, 1)], None),

740 - (

741 - ".qkv_proj",

742 - ".qkv_proj",

743 - num_attention_heads + num_kv_heads * 2,

744 - [("q", num_attention_heads), ("k", num_kv_heads), ("v", num_kv_heads)],

745 - self._split_qkv_weight,

746 - ),

747 - ]

748 -

749 - params_dict = dict(self.named_parameters())

750 - loaded_params: set[str] = set()

751 - expert_params_mapping = self.get_expert_mapping()

752 - for name, loaded_weight in weights:

753 - if "rotary_emb.inv_freq" in name:

754 - continue

755 - if "gate_proj_bias" in name:

756 - name = name.replace("gate_proj_bias", "gate_proj.bias")

757 - if "up_proj_bias" in name:

758 - name = name.replace("up_proj_bias", "up_proj.bias")

759 - if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name:

760 - # Models trained using ColossalAI may include these tensors in

761 - # the checkpoint. Skip them.

762 - continue

763 - # With tie_word_embeddings, we can skip lm_head.weight

764 - # The weight might appear unnecessarily in the files if the model is

765 - # processed with quantization, LoRA, fine-tuning, etc.

766 - if self.config.tie_word_embeddings and "lm_head.weight" in name:

767 - continue

768 -

769 - is_found = False

770 - for param_name, weight_name, shard_id in stacked_params_mapping:

771 - if weight_name not in name:

772 - continue

773 - if "mlp.experts" in name:

774 - continue

775 - # cross layer only have q_proj, skip qkv pack

776 - if weight_name == ".q_proj":

777 - match = re.search(r"layers\.\d+", name)

778 - if match:

779 - layer_id = int(match.group(0).split(".")[-1])

780 - if cla_factor > 1 and layer_id % cla_factor != 0:

781 - continue

782 - name = name.replace(weight_name, param_name)

783 - # Skip loading extra bias for GPTQ models.

784 - if name.endswith(".bias") and name not in params_dict:

785 - continue

786 -

787 - if is_pp_missing_parameter(name, self):

788 - continue

789 -

790 - param = params_dict[name]

791 - weight_loader = param.weight_loader

792 - weight_loader(param, loaded_weight, shard_id)

793 - loaded_params.add(name)

794 - is_found = True

795 - break

796 - if is_found:

797 - continue

798 -

799 - for (

800 - param_name,

801 - weight_name,

802 - den,

803 - split_param,

804 - func,

805 - ) in split_params_mapping:

806 - if weight_name not in name:

807 - continue

808 - name = name.replace(weight_name, param_name)

809 - # Skip loading extra bias for GPTQ models.

810 - if name.endswith(".bias") and name not in params_dict:

811 - continue

812 -

813 - if is_pp_missing_parameter(name, self):

814 - continue

815 -

816 - assert loaded_weight.shape[0] % den == 0

817 - units = loaded_weight.shape[0] // den

818 -

819 - param = params_dict[name]

820 - weight_loader = param.weight_loader

821 - offset = 0

822 - for shard_id, num in split_param:

823 - new_offset = offset + num * units

824 - if func:

825 - weight_loader(

826 - param, func(loaded_weight)[offset:new_offset], shard_id

827 - )

828 - else:

829 - weight_loader(param, loaded_weight[offset:new_offset], shard_id)

830 - offset = new_offset

831 -

832 - break

833 - else:

834 - # Skip loading extra bias for GPTQ models.

835 - if name.endswith(".bias") and name not in params_dict:

836 - continue

837 - is_expert_weight = False

838 - for mapping in expert_params_mapping:

839 - param_name, weight_name, expert_id, shard_id = mapping

840 - if weight_name not in name:

841 - continue

842 - # this is an expert weight and should not be

843 - # attempted to load as other weights later

844 - is_expert_weight = True

845 -

846 - # Do not modify `name` since the loop may continue here

847 - # Instead, create a new variable

848 - name_mapped = name.replace(weight_name, param_name)

849 - if is_pp_missing_parameter(name_mapped, self):

850 - continue

851 - param = params_dict[name_mapped]

852 - # We should ask the weight loader to return success or not

853 - # here since otherwise we may skip experts with other

854 - # available replicas.

855 - weight_loader = typing.cast(

856 - Callable[..., bool], param.weight_loader

857 - )

858 - success = weight_loader(

859 - param,

860 - loaded_weight,

861 - name_mapped,

862 - shard_id=shard_id,

863 - expert_id=expert_id,

864 - return_success=True,

865 - )

866 - if success:

867 - name = name_mapped

868 - break

869 - else:

870 - if is_expert_weight:

871 - # We've checked that this is an expert weight

872 - # However it's not mapped locally to this rank

873 - # So we simply skip it

874 - continue

875 - # Remapping the name of FP8 kv-scale.

876 - name = maybe_remap_kv_scale_name(name, params_dict)

877 - if name is None:

878 - continue

879 -

880 - if is_pp_missing_parameter(name, self):

881 - continue

882 -

883 - if "mlp.gate.wg." in name:

884 - name = name.replace("wg.", "")

885 -

886 - param = params_dict[name]

887 - weight_loader = getattr(

888 - param, "weight_loader", default_weight_loader

889 - )

890 - weight_loader(param, loaded_weight)

891 - loaded_params.add(name)

892 - return loaded_params

893 -

894 -

895 -class HunyuanV1ModelBase(

896 - nn.Module, SupportsLoRA, SupportsPP, SupportsEagle, SupportsEagle3

897 -):

898 - packed_modules_mapping = {

899 - "qkv_proj": [

900 - "q_proj",

901 - "k_proj",

902 - "v_proj",

903 - ],

904 - "gate_up_proj": [

905 - "gate_proj",

906 - "up_proj",

907 - ],

908 - }

909 -

910 - def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

911 - super().__init__()

912 -

913 - config = vllm_config.model_config.hf_config

914 - quant_config = vllm_config.quant_config

915 - self.config = config

916 - self.quant_config = quant_config

917 -

918 - self.model = HunYuanModel(

919 - vllm_config=vllm_config,

920 - prefix=maybe_prefix(prefix, "model"),

921 - )

922 - if get_pp_group().is_last_rank:

923 - self.lm_head = ParallelLMHead(

924 - config.vocab_size,

925 - config.hidden_size,

926 - quant_config=quant_config,

927 - prefix=maybe_prefix(prefix, "lm_head"),

928 - )

929 - if config.tie_word_embeddings:

930 - self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens)

931 -

932 - logit_scale = getattr(config, "logit_scale", 1.0)

933 - self.logits_processor = LogitsProcessor(

934 - config.vocab_size, scale=logit_scale

935 - )

936 - else:

937 - self.lm_head = PPMissingLayer()

938 -

939 - def forward(

940 - self,

941 - input_ids: torch.Tensor | None,

942 - positions: torch.Tensor,

943 - intermediate_tensors: IntermediateTensors | None = None,

944 - inputs_embeds: torch.Tensor | None = None,

945 - ) -> torch.Tensor | IntermediateTensors:

946 - model_output = self.model(

947 - input_ids, positions, intermediate_tensors, inputs_embeds

948 - )

949 - return model_output

950 -

951 - def compute_logits(

952 - self,

953 - hidden_states: torch.Tensor,

954 - ) -> torch.Tensor | None:

955 - logits = self.logits_processor(self.lm_head, hidden_states)

956 - return logits

957 -

958 - def make_empty_intermediate_tensors(

959 - self, batch_size: int, dtype: torch.dtype, device: torch.device

960 - ) -> IntermediateTensors:

961 - return IntermediateTensors(

962 - {

963 - "hidden_states": torch.zeros(

964 - (batch_size, self.config.hidden_size), dtype=dtype, device=device

965 - ),

966 - "residual": torch.zeros(

967 - (batch_size, self.config.hidden_size), dtype=dtype, device=device

968 - ),

969 - }

970 - )

971 -

972 - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

973 - loader = AutoWeightsLoader(self)

974 - return loader.load_weights(weights)

975 -

976 - def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:

977 - return self.model.embed_input_ids(input_ids)

978 -

979 -

980 -class HunYuanMoEV1Base(HunyuanV1ModelBase, MixtureOfExperts):

981 - def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

982 - super().__init__(vllm_config=vllm_config, prefix=prefix)

983 -

984 - # Set MoE hyperparameters

985 - self.num_expert_groups = 1

986 - self.moe_layers = []

987 - example_layer = None

988 - for layer in self.model.layers:

989 - if isinstance(layer, PPMissingLayer):

990 - continue

991 -

992 - assert isinstance(layer, HunYuanDecoderLayer)

993 - if isinstance(layer.mlp, HunYuanSparseMoeBlock):

994 - example_layer = layer.mlp

995 - self.moe_layers.append(layer.mlp.experts)

996 -

997 - if example_layer is None:

998 - raise RuntimeError("No HunYuanMoE layer found in model.layers.")

999 -

1000 - self.num_moe_layers = len(self.moe_layers)

1001 - self.num_logical_experts = example_layer.n_logical_experts

1002 - self.num_physical_experts = example_layer.n_physical_experts

1003 - self.num_local_physical_experts = example_layer.n_local_physical_experts

1004 - self.num_routed_experts = example_layer.n_routed_experts

1005 - self.num_redundant_experts = example_layer.n_redundant_experts

1006 -

1007 - def update_physical_experts_metadata(

1008 - self,

1009 - num_physical_experts: int,

1010 - num_local_physical_experts: int,

1011 - ) -> None:

1012 - assert self.num_local_physical_experts == num_local_physical_experts

1013 - self.num_physical_experts = num_physical_experts

1014 - self.num_local_physical_experts = num_local_physical_experts

1015 - self.num_redundant_experts = num_physical_experts - self.num_logical_experts

1016 - for layer in self.model.layers:

1017 - if isinstance(layer.mlp, HunYuanSparseMoeBlock):

1018 - moe = layer.mlp

1019 - moe.n_local_physical_experts = num_local_physical_experts

1020 - moe.n_physical_experts = num_physical_experts

1021 - moe.n_redundant_experts = self.num_redundant_experts

1022 - moe.experts.update_expert_map()

1023 -

1024 - def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:

1025 - return self.model.get_expert_mapping()

1026 -

1027 -

1028 -class HunYuanDenseV1Base(HunyuanV1ModelBase):

1029 - def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

1030 - super().__init__(vllm_config=vllm_config, prefix=prefix)

1031 -

1032 -

1033 -class HunYuanDenseV1ForCausalLM(HunYuanDenseV1Base):

1034 - pass

1035 -

1036 -

1037 -class HunYuanMoEV1ForCausalLM(HunYuanMoEV1Base):

1038 - pass

vllm/model_executor/models/hunyuan_vision.py deleted

-1025

@@ -1,1025 +0,0 @@

1 -# SPDX-License-Identifier: Apache-2.0

2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 -

4 -# coding=utf-8

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 HunYuan-VL model compatible with HuggingFace weights."""

26 -

27 -from collections.abc import Callable, Iterable, Mapping, Sequence

28 -from functools import partial

29 -from typing import Annotated, Any, Literal, TypeAlias

30 -

31 -import torch

32 -import torch.nn as nn

33 -import torch.nn.functional as F

34 -from transformers import BatchFeature, HunYuanVLProcessor

35 -from transformers.models.hunyuan_vl.image_processing_hunyuan_vl import (

36 - HunYuanVLImageProcessor,

37 - smart_resize,

38 -)

39 -

40 -from vllm.config import VllmConfig

41 -from vllm.config.multimodal import BaseDummyOptions

42 -from vllm.distributed import parallel_state

43 -from vllm.distributed import utils as dist_utils

44 -from vllm.inputs import ModalityData, MultiModalDataDict

45 -from vllm.logger import init_logger

46 -from vllm.model_executor.layers.activation import get_act_fn

47 -from vllm.model_executor.layers.attention import MMEncoderAttention

48 -from vllm.model_executor.layers.layernorm import RMSNorm

49 -from vllm.model_executor.layers.linear import (

50 - ColumnParallelLinear,

51 - QKVParallelLinear,

52 - RowParallelLinear,

53 -)

54 -from vllm.model_executor.layers.quantization import QuantizationConfig

55 -from vllm.model_executor.models.module_mapping import MultiModelKeys

56 -from vllm.multimodal import MULTIMODAL_REGISTRY

57 -from vllm.multimodal.inputs import (

58 - ImageItem,

59 - MultiModalFeatureSpec,

60 - MultiModalFieldConfig,

61 - MultiModalKwargsItems,

62 -)

63 -from vllm.multimodal.parse import (

64 - DictEmbeddingItems,

65 - ImageSize,

66 - ModalityDataItems,

67 - MultiModalDataItems,

68 - MultiModalDataParser,

69 -)

70 -from vllm.multimodal.processing import (

71 - BaseDummyInputsBuilder,

72 - BaseMultiModalProcessor,

73 - BaseProcessingInfo,

74 - PromptReplacement,

75 - PromptUpdate,

76 - PromptUpdateDetails,

77 -)

78 -from vllm.sequence import IntermediateTensors

79 -from vllm.transformers_utils.configs.hunyuan_vl import (

80 - HunYuanVLConfig,

81 - HunYuanVLVisionConfig,

82 -)

83 -from vllm.utils.tensor_schema import TensorSchema, TensorShape

84 -

85 -from .interfaces import (

86 - MultiModalEmbeddings,

87 - SupportsEagle,

88 - SupportsEagle3,

89 - SupportsLoRA,

90 - SupportsMultiModal,

91 - SupportsPP,

92 - SupportsQuant,

93 - SupportsXDRoPE,

94 -)

95 -from .utils import (

96 - AutoWeightsLoader,

97 - WeightsMapper,

98 - init_vllm_registered_model,

99 - maybe_prefix,

100 -)

101 -from .vision import is_vit_use_data_parallel

102 -

103 -logger = init_logger(__name__)

104 -

105 -# === Vision Inputs === #

106 -

107 -

108 -class HunYuanVLImagePixelInputs(TensorSchema):

109 - """

110 - Dimensions:

111 - - np: Number of patches

112 - - ni: Number of images

113 - - cps: Number of channels * patch_size * patch_size

114 - """

115 -

116 - type: Literal["pixel_values"]

117 -

118 - pixel_values: Annotated[

119 - torch.Tensor,

120 - TensorShape("np", "cps"),

121 - ]

122 -

123 - image_grid_thw: Annotated[

124 - torch.Tensor,

125 - TensorShape("ni", 3),

126 - ]

127 -

128 -

129 -class HunYuanVLImageEmbeddingInputs(TensorSchema):

130 - """

131 - Dimensions:

132 - - nf: Number of image features

133 - - hs: Hidden size

134 - - ni: Number of images

135 - """

136 -

137 - type: Literal["image_embeds"]

138 -

139 - image_embeds: Annotated[

140 - torch.Tensor,

141 - TensorShape("nf", "hs"),

142 - ]

143 -

144 - image_grid_thw: Annotated[

145 - torch.Tensor,

146 - TensorShape("ni", 3),

147 - ]

148 -

149 -

150 -HunYuanVLImageInputs: TypeAlias = (

151 - HunYuanVLImagePixelInputs | HunYuanVLImageEmbeddingInputs

152 -)

153 -

154 -# === Vision Encoder === #

155 -

156 -

157 -class HunYuanVisionMLP(nn.Module):

158 - def __init__(

159 - self,

160 - in_features: int,

161 - hidden_features: int,

162 - bias: bool = True,

163 - act_fn: Callable[[torch.Tensor], torch.Tensor] = F.gelu,

164 - quant_config: QuantizationConfig | None = None,

165 - prefix: str = "",

166 - ):

167 - super().__init__()

168 - use_data_parallel = is_vit_use_data_parallel()

169 - self.dense_h_to_4h = ColumnParallelLinear(

170 - in_features,

171 - hidden_features,

172 - bias=bias,

173 - quant_config=quant_config,

174 - prefix=f"{prefix}.dense_h_to_4h",

175 - disable_tp=use_data_parallel,

176 - )

177 - self.dense_4h_to_h = RowParallelLinear(

178 - hidden_features,

179 - in_features,

180 - bias=bias,

181 - quant_config=quant_config,

182 - prefix=f"{prefix}.dense_4h_to_h",

183 - disable_tp=use_data_parallel,

184 - )

185 - self.act_fn = act_fn

186 -

187 - def forward(self, x: torch.Tensor):

188 - x_up, _ = self.dense_h_to_4h(x)

189 - x_down, _ = self.dense_4h_to_h(self.act_fn(x_up))

190 - return x_down

191 -

192 -

193 -class HunYuanVisionAttention(nn.Module):

194 - def __init__(

195 - self,

196 - embed_dim: int,

197 - num_heads: int,

198 - projection_size: int,

199 - quant_config: QuantizationConfig | None = None,

200 - prefix: str = "",

201 - ) -> None:

202 - super().__init__()

203 - # Per attention head and per partition values.

204 - use_data_parallel = is_vit_use_data_parallel()

205 - self.tp_size = (

206 - 1

207 - if use_data_parallel

208 - else parallel_state.get_tensor_model_parallel_world_size()

209 - )

210 - self.hidden_size_per_attention_head = dist_utils.divide(

211 - projection_size, num_heads

212 - )

213 - self.num_attention_heads_per_partition = dist_utils.divide(

214 - num_heads, self.tp_size

215 - )

216 -

217 - self.qkv = QKVParallelLinear(

218 - hidden_size=embed_dim,

219 - head_size=self.hidden_size_per_attention_head,

220 - total_num_heads=num_heads,

221 - total_num_kv_heads=num_heads,

222 - bias=True,

223 - quant_config=quant_config,

224 - prefix=f"{prefix}.qkv",

225 - disable_tp=use_data_parallel,

226 - )

227 -

228 - self.o_proj = RowParallelLinear(

229 - input_size=projection_size,

230 - output_size=embed_dim,

231 - quant_config=quant_config,

232 - prefix=f"{prefix}.o_proj",

233 - disable_tp=use_data_parallel,

234 - )

235 -

236 - self.scale = self.hidden_size_per_attention_head**-0.5

237 - self.attn = MMEncoderAttention(

238 - self.num_attention_heads_per_partition,

239 - self.hidden_size_per_attention_head,

240 - self.scale,

241 - prefix=f"{prefix}.attn",

242 - )

243 -

244 - def forward(

245 - self,

246 - x: torch.Tensor,

247 - ) -> torch.Tensor:

248 - qkv, _ = self.qkv(x)

249 - q, k, v = qkv.chunk(3, dim=-1)

250 - out = self.attn(q, k, v)

251 - output, _ = self.o_proj(out)

252 - return output

253 -

254 -

255 -class HunYuanVisionBlock(nn.Module):

256 - def __init__(

257 - self,

258 - dim: int,

259 - num_heads: int,

260 - mlp_hidden_dim: int,

261 - act_fn: Callable[[torch.Tensor], torch.Tensor] = F.gelu,

262 - norm_layer: Callable[[int], nn.Module] | None = None,

263 - quant_config: QuantizationConfig | None = None,

264 - prefix: str = "",

265 - ) -> None:

266 - super().__init__()

267 - if norm_layer is None:

268 - norm_layer = partial(nn.LayerNorm, eps=1e-6)

269 - self.input_layernorm = norm_layer(dim)

270 - self.post_attention_layernorm = norm_layer(dim)

271 - self.self_attn = HunYuanVisionAttention(

272 - embed_dim=dim,

273 - num_heads=num_heads,

274 - projection_size=dim,

275 - quant_config=quant_config,

276 - prefix=f"{prefix}.self_attn",

277 - )

278 - self.mlp = HunYuanVisionMLP(

279 - dim,

280 - mlp_hidden_dim,

281 - act_fn=act_fn,

282 - bias=True,

283 - quant_config=quant_config,

284 - prefix=f"{prefix}.mlp",

285 - )

286 -

287 - def forward(

288 - self,

289 - x: torch.Tensor,

290 - ) -> torch.Tensor:

291 - x = x + self.self_attn(self.input_layernorm(x))

292 - x = x + self.mlp(self.post_attention_layernorm(x))

293 - return x

294 -

295 -

296 -class HunYuanVisionPatchEmbed(nn.Module):

297 - def __init__(self, config: HunYuanVLVisionConfig):

298 - super().__init__()

299 -

300 - self.config = config

301 - self.embed_dim = config.hidden_size

302 - self.patch_size = config.patch_size

303 - self.num_channels = config.num_channels

304 - self.spatial_merge_size = config.spatial_merge_size

305 - self.interpolate_mode = config.interpolate_mode

306 -

307 - self.patch_embedding = nn.Conv2d(

308 - in_channels=config.num_channels,

309 - out_channels=self.embed_dim,

310 - kernel_size=self.patch_size,

311 - stride=self.patch_size,

312 - bias=True,

313 - )

314 -

315 - self.max_num_patches = (config.max_image_size // self.patch_size) ** 2

316 -

317 - self.num_positions = self.max_num_patches + 1

318 - self.position_edge = int(self.num_positions**0.5)

319 - # first token is cls token, skip it

320 - self.position_embedding = nn.Embedding(self.num_positions, self.embed_dim)

321 -

322 - self.patch_pos_embed = None

323 -

324 - def forward(

325 - self, pixel_values: torch.Tensor, grid_thw: list[list[int]]

326 - ) -> torch.Tensor:

327 - num_patches = pixel_values.size(0)

328 - pixel_values = pixel_values.reshape(

329 - num_patches, self.num_channels, self.patch_size, self.patch_size

330 - )

331 -

332 - patch_embeds = self.patch_embedding(pixel_values)

333 - patch_embeds = patch_embeds.squeeze(-1).squeeze(-1).unsqueeze(0)

334 -

335 - if self.patch_pos_embed is None:

336 - patch_pos_shape = (

337 - 1,

338 - self.position_edge,

339 - self.position_edge,

340 - self.embed_dim,

341 - )

342 - self.patch_pos_embed = (

343 - self.position_embedding.weight[1:, :]

344 - .reshape(patch_pos_shape)

345 - .permute(0, 3, 1, 2)

346 - .float()

347 - )

348 -

349 - patch_pos_embed_list = []

350 - for grid in grid_thw:

351 - _, h0, w0 = grid

352 - # we add a small number to avoid floating point error in the interpolation

353 - # see discussion at https://github.com/facebookresearch/dino/issues/8

354 - h0, w0 = h0 + 0.1, w0 + 0.1

355 - patch_pos_embed = nn.functional.interpolate(

356 - self.patch_pos_embed,

357 - scale_factor=(h0 / self.position_edge, w0 / self.position_edge),

358 - mode=self.interpolate_mode,

359 - align_corners=False,

360 - )

361 -

362 - patch_pos_embed = (

363 - patch_pos_embed.reshape(self.embed_dim, -1)

364 - .transpose(0, 1)

365 - .unsqueeze(0)

366 - .to(patch_embeds.dtype)

367 - )

368 - patch_pos_embed_list.append(patch_pos_embed)

369 -

370 - patch_pos_embed = torch.cat(patch_pos_embed_list, dim=1)

371 - embeddings = patch_embeds + patch_pos_embed

372 -

373 - return embeddings

374 -

375 -

376 -class HunYuanVisionPatchMerger(nn.Module):

377 - def __init__(

378 - self,

379 - in_channels,

380 - out_channels,

381 - spatial_merge_size=2,

382 - rms_norm_eps=1e-5,

383 - prefix="",

384 - ):

385 - super().__init__()

386 - self.spatial_merge_size = spatial_merge_size

387 - embed_std = out_channels**-0.5

388 -

389 - self.proj = nn.Sequential(

390 - nn.Conv2d(

391 - in_channels,

392 - in_channels * 2,

393 - kernel_size=spatial_merge_size,

394 - stride=spatial_merge_size,

395 - ),

396 - nn.GELU(),

397 - nn.Conv2d(in_channels * 2, in_channels * 4, kernel_size=1),

398 - )

399 - self.mlp = nn.Linear(in_channels * 4, out_channels)

400 -

401 - self.image_newline = nn.Parameter(torch.randn(in_channels * 4) * embed_std)

402 - self.image_begin = nn.Parameter(torch.randn(out_channels) * embed_std)

403 - self.image_end = nn.Parameter(torch.randn(out_channels) * embed_std)

404 - self.image_sep = nn.Parameter(torch.randn(out_channels) * embed_std)

405 -

406 - self.before_rms = RMSNorm(in_channels, eps=rms_norm_eps)

407 - self.after_rms = RMSNorm(out_channels, eps=rms_norm_eps)

408 -

409 - def forward(self, x, size=(16, 16)):

410 - x = self.before_rms(x)

411 -

412 - h, w = size

413 - dtype = x.dtype

414 - x = x.permute(0, 2, 1).reshape(x.shape[0], -1, h, w)

415 -

416 - x = self.proj(x) # b,c,h,w

417 - b, c, h, w = x.shape

418 - x = torch.cat(

419 - [x, self.image_newline.reshape(1, c, 1, 1).expand(b, c, h, 1).to(dtype)],

420 - dim=-1,

421 - )

422 - x = x.reshape(b, c, -1).permute(0, 2, 1)

423 - x = self.mlp(x)

424 -

425 - begin = self.image_begin.reshape(1, 1, -1).expand(b, 1, x.shape[-1]).to(dtype)

426 - end = self.image_end.reshape(1, 1, -1).expand(b, 1, x.shape[-1]).to(dtype)

427 - x = torch.cat([begin, x, end], dim=1)

428 -

429 - return self.after_rms(x)

430 -

431 -

432 -class HunYuanVisionTransformer(nn.Module):

433 - hf_to_vllm_mapper = WeightsMapper(

434 - orig_to_new_stacked={

435 - ".q_proj": (".qkv", "q"),

436 - ".k_proj": (".qkv", "k"),

437 - ".v_proj": (".qkv", "v"),

438 - }

439 - )

440 -

441 - def __init__(

442 - self,

443 - vision_config: HunYuanVLVisionConfig,

444 - quant_config: QuantizationConfig | None = None,

445 - prefix: str = "",

446 - ) -> None:

447 - super().__init__()

448 -

449 - num_hidden_layers = vision_config.num_hidden_layers

450 - self.hidden_size = vision_config.hidden_size

451 - self.num_heads = vision_config.num_attention_heads

452 - self.spatial_merge_size = vision_config.spatial_merge_size

453 -

454 - from vllm.compilation.backends import set_model_tag

455 -

456 - with set_model_tag("HunYuanVisionPatchEmbed"):

457 - self.embeddings = HunYuanVisionPatchEmbed(vision_config)

458 -

459 - norm_layer = partial(nn.LayerNorm, eps=vision_config.rms_norm_eps)

460 -

461 - with set_model_tag("HunYuanVisionBlock"):

462 - self.layers = nn.ModuleList(

463 - [

464 - HunYuanVisionBlock(

465 - dim=vision_config.hidden_size,

466 - num_heads=vision_config.num_attention_heads,

467 - mlp_hidden_dim=vision_config.intermediate_size,

468 - act_fn=get_act_fn(vision_config.hidden_act),

469 - norm_layer=norm_layer,

470 - quant_config=quant_config,

471 - prefix=f"{prefix}.layers.{layer_idx}",

472 - )

473 - for layer_idx in range(num_hidden_layers)

474 - ]

475 - )

476 -

477 - with set_model_tag("HunYuanVisionPatchMerger"):

478 - self.perceive = HunYuanVisionPatchMerger(

479 - vision_config.hidden_size,

480 - vision_config.out_hidden_size,

481 - spatial_merge_size=vision_config.spatial_merge_size,

482 - rms_norm_eps=vision_config.rms_norm_eps,

483 - prefix=f"{prefix}.perceive",

484 - )

485 -

486 - @property

487 - def dtype(self) -> torch.dtype:

488 - return self.embeddings.patch_embedding.weight.dtype

489 -

490 - @property

491 - def device(self) -> torch.device:

492 - return self.embeddings.patch_embedding.weight.device

493 -

494 - def forward(

495 - self,

496 - x: torch.Tensor,

497 - grid_thw: list[list[int]],

498 - ) -> torch.Tensor:

499 - # patchify

500 - seq_len = x.size(0)

501 - cu_seqlens: list = [0]

502 -

503 - hidden_states = x.to(device=self.device, dtype=self.dtype)

504 - # embeddings = patch_embeds + patch_pos_embed

505 - hidden_states = self.embeddings(hidden_states, grid_thw)

506 -

507 - for t, h, w in grid_thw:

508 - t, h, w = int(t), int(h), int(w)

509 - cu_seqlens.append(h * w)

510 -

511 - cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int32)

512 - cu_seqlens = torch.cumsum(cu_seqlens, dim=0, dtype=torch.int32)

513 -

514 - cu_seqlens = cu_seqlens.to(device=self.device, non_blocking=True)

515 -

516 - hidden_states = hidden_states.reshape(seq_len, -1)

517 - hidden_states = hidden_states.unsqueeze(0)

518 -

519 - # build per-image lengths once

520 - split_lengths = [int(h) * int(w) for (_, h, w) in grid_thw]

521 - for layer in self.layers:

522 - # hidden_states: (1, T_total, D)

523 - parts = hidden_states.split(split_lengths, dim=1) # list of (1, L_i, D)

524 - parts = [layer(p) for p in parts]

525 - hidden_states = torch.cat(parts, dim=1)

526 -

527 - # adapter

528 - split_lengths = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist()

529 - split_items = hidden_states.split(split_lengths, dim=1)

530 - image_embeds_list = []

531 - for grid, split_item in zip(grid_thw, split_items):

532 - image_embeds_list.append(

533 - self.perceive(split_item.contiguous(), size=grid[1:]).squeeze(0)

534 - )

535 -

536 - return image_embeds_list

537 -

538 - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

539 - loader = AutoWeightsLoader(self)

540 - return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

541 -

542 -

543 -def _hunyuan_vl_field_config(hf_inputs: Mapping[str, torch.Tensor]):

544 - image_grid_thw = hf_inputs.get("image_grid_thw", torch.empty((0, 3)))

545 - image_grid_sizes = image_grid_thw.prod(-1)

546 - return dict(

547 - pixel_values=MultiModalFieldConfig.flat_from_sizes("image", image_grid_sizes),

548 - image_embeds=MultiModalFieldConfig.flat_from_sizes("image", image_grid_sizes),

549 - image_grid_thw=MultiModalFieldConfig.batched("image", keep_on_cpu=True),

550 - )

551 -

552 -

553 -class HunYuanVLMultiModalDataParser(MultiModalDataParser):

554 - # The patch grid is what sizes the placeholder range.

555 - embedding_fields = {

556 - "image": {"image_embeds": "values", "image_grid_thw": "metadata"},

557 - }

558 -

559 - def _parse_image_data(

560 - self,

561 - data: dict[str, torch.Tensor] | ModalityData[ImageItem],

562 - ) -> ModalityDataItems[Any, Any] | None:

563 - if isinstance(data, dict):

564 - required, optional = self.embedding_field_sets("image")

565 - return DictEmbeddingItems(

566 - data,

567 - modality="image",

568 - required_fields=required,

569 - optional_fields=optional,

570 - fields_factory=_hunyuan_vl_field_config,

571 - )

572 -

573 - return super()._parse_image_data(data)

574 -

575 -

576 -class HunYuanVLProcessingInfo(BaseProcessingInfo):

577 - def get_hf_config(self):

578 - return self.ctx.get_hf_config(HunYuanVLConfig)

579 -

580 - def get_hf_processor(

581 - self,

582 - **kwargs: object,

583 - ) -> HunYuanVLProcessor:

584 - # transformers>=5.13 replaced `use_fast` with `backend`; pin the

585 - # PIL backend to match the released HunyuanOCR checkpoint packing.

586 - kwargs.pop("use_fast", None)

587 - kwargs.setdefault("backend", "pil")

588 - return self.ctx.get_hf_processor(

589 - HunYuanVLProcessor,

590 - **kwargs,

591 - )

592 -

593 - def get_image_processor(

594 - self,

595 - **kwargs: object,

596 - ) -> HunYuanVLImageProcessor:

597 - return self.get_hf_processor(**kwargs).image_processor

598 -

599 - def get_data_parser(self):

600 - return HunYuanVLMultiModalDataParser(

601 - expected_hidden_size=self._get_expected_hidden_size(),

602 - allow_missing_mm_embeddings=self.allow_missing_mm_embeddings,

603 - )

604 -

605 - def get_supported_mm_limits(self) -> Mapping[str, int | None]:

606 - return {"image": None}

607 -

608 - def get_mm_max_tokens_per_item(

609 - self,

610 - seq_len: int,

611 - mm_counts: Mapping[str, int],

612 - ) -> Mapping[str, int]:

613 - max_image_tokens = self.get_max_image_tokens()

614 - # TODO: support video

615 - max_video_tokens = 0

616 - return {"image": max_image_tokens, "video": max_video_tokens}

617 -

618 - def _get_vision_info(

619 - self,

620 - *,

621 - image_width: int,

622 - image_height: int,

623 - num_frames: int = 1,

624 - do_resize: bool = True,

625 - image_processor: HunYuanVLImageProcessor,

626 - mm_kwargs: Mapping[str, object],

627 - ) -> tuple[ImageSize, int]:

628 - hf_config = self.get_hf_config()

629 - vision_config = hf_config.vision_config

630 - patch_size = vision_config.patch_size

631 - spatial_merge_size = vision_config.spatial_merge_size

632 -

633 - mm_kwargs = self.ctx.get_merged_mm_kwargs(mm_kwargs)

634 - size = image_processor.size

635 - if override_size := mm_kwargs.get("size"):

636 - size = size | override_size

637 - if (override_min_pixels := mm_kwargs.get("min_pixels")) is not None:

638 - size = size | {"shortest_edge": override_min_pixels}

639 - if (override_max_pixels := mm_kwargs.get("max_pixels")) is not None:

640 - size = size | {"longest_edge": override_max_pixels}

641 -

642 - if do_resize:

643 - resized_height, resized_width = smart_resize(

644 - height=image_height,

645 - width=image_width,

646 - factor=patch_size * spatial_merge_size,

647 - min_pixels=size["shortest_edge"],

648 - max_pixels=size["longest_edge"],

649 - )

650 - preprocessed_size = ImageSize(width=resized_width, height=resized_height)

651 - else:

652 - preprocessed_size = ImageSize(width=image_width, height=image_height)

653 -

654 - grid_t = 1

655 - grid_h = preprocessed_size.height // patch_size

656 - grid_w = preprocessed_size.width // patch_size

657 -

658 - num_vision_tokens = (

659 - grid_t * grid_h // spatial_merge_size * (grid_w // spatial_merge_size + 1)

660 - + 2

661 - )

662 -

663 - return preprocessed_size, num_vision_tokens

664 -

665 - def get_num_image_tokens(

666 - self,

667 - *,

668 - image_width: int,

669 - image_height: int,

670 - image_processor: HunYuanVLImageProcessor,

671 - mm_kwargs: Mapping[str, object],

672 - ) -> int:

673 - _, num_image_tokens = self._get_vision_info(

674 - image_width=image_width,

675 - image_height=image_height,

676 - image_processor=image_processor,

677 - mm_kwargs=mm_kwargs,

678 - )

679 - return num_image_tokens

680 -

681 - def get_image_size_with_most_features(self) -> ImageSize:

682 - image_processor = self.get_image_processor()

683 -

684 - max_image_size, _ = self._get_vision_info(

685 - image_width=512,

686 - image_height=8192,

687 - image_processor=image_processor,

688 - mm_kwargs={},

689 - )

690 - return max_image_size

691 -

692 - def get_max_image_tokens(self) -> int:

693 - image_processor = self.get_image_processor()

694 - target_width, target_height = self.get_image_size_with_most_features()

695 -

696 - return self.get_num_image_tokens(

697 - image_width=target_width,

698 - image_height=target_height,

699 - image_processor=image_processor,

700 - mm_kwargs={},

701 - )

702 -

703 -

704 -class HunYuanVLDummyInputsBuilder(BaseDummyInputsBuilder[HunYuanVLProcessingInfo]):

705 - def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:

706 - num_images = mm_counts.get("image", 0)

707 -

708 - hf_processor = self.info.get_hf_processor(typ=HunYuanVLProcessor)

709 - image_placeholder = (

710 - f"{hf_processor.image_start_token}{hf_processor.image_token}"

711 - f"{hf_processor.image_end_token}"

712 - )

713 -

714 - return image_placeholder * num_images

715 -

716 - def get_dummy_mm_data(

717 - self,

718 - seq_len: int,

719 - mm_counts: Mapping[str, int],

720 - mm_options: Mapping[str, BaseDummyOptions],

721 - ) -> MultiModalDataDict:

722 - num_images = mm_counts.get("image", 1)

723 -

724 - target_width, target_height = self.info.get_image_size_with_most_features()

725 -

726 - return {

727 - "image": self._get_dummy_images(

728 - width=target_width, height=target_height, num_images=num_images

729 - ),

730 - }

731 -

732 -

733 -class HunYuanVLMultiModalProcessor(BaseMultiModalProcessor[HunYuanVLProcessingInfo]):

734 - def _get_prompt_updates(

735 - self,

736 - mm_items: MultiModalDataItems,

737 - hf_processor_mm_kwargs: Mapping[str, Any],

738 - out_mm_kwargs: MultiModalKwargsItems,

739 - ) -> Sequence[PromptUpdate]:

740 - hf_processor = self.info.get_hf_processor(**hf_processor_mm_kwargs)

741 - image_processor = self.info.get_image_processor(**hf_processor_mm_kwargs)

742 -

743 - token_ids = {

744 - "image": hf_processor.image_token_id,

745 - "image_start": hf_processor.image_start_token_id,

746 - "image_end": hf_processor.image_end_token_id,

747 - }

748 -

749 - merge_size = image_processor.merge_size

750 -

751 - def get_replacement_hunyuan_vl(item_idx: int, modality: str):

752 - out_item = out_mm_kwargs[modality][item_idx]

753 - grid_thw = out_item[f"{modality}_grid_thw"].data

754 - assert isinstance(grid_thw, torch.Tensor)

755 -

756 - _, grid_h, grid_w = grid_thw

757 - num_tokens = (int(grid_h) // merge_size) * (

758 - int(grid_w) // merge_size + 1

759 - ) + 2

760 - tokens = (

761 - [token_ids[f"{modality}_start"]]

762 - + [token_ids[modality]] * num_tokens

763 - + [token_ids[f"{modality}_end"]]

764 - )

765 - return PromptUpdateDetails.select_token_id(tokens, token_ids[modality])

766 -

767 - return [

768 - PromptReplacement(

769 - modality=modality,

770 - target=[

771 - token_ids[f"{modality}_start"],

772 - token_ids[modality],

773 - token_ids[f"{modality}_end"],

774 - ],

775 - replacement=partial(get_replacement_hunyuan_vl, modality=modality),

776 - )

777 - for modality in ("image",)

778 - ]

779 -

780 - def _get_mm_fields_config(

781 - self,

782 - hf_inputs: BatchFeature,

783 - hf_processor_mm_kwargs: Mapping[str, object],

784 - ) -> Mapping[str, MultiModalFieldConfig]:

785 - return _hunyuan_vl_field_config(hf_inputs)

786 -

787 -

788 -@MULTIMODAL_REGISTRY.register_processor(

789 - HunYuanVLMultiModalProcessor,

790 - info=HunYuanVLProcessingInfo,

791 - dummy_inputs=HunYuanVLDummyInputsBuilder,

792 -)

793 -class HunYuanVLForConditionalGeneration(

794 - nn.Module,

795 - SupportsMultiModal,

796 - SupportsLoRA,

797 - SupportsPP,

798 - SupportsQuant,

799 - SupportsXDRoPE,

800 - SupportsEagle,

801 - SupportsEagle3,

802 -):

803 - # To ensure correct weight loading and mapping.

804 - hf_to_vllm_mapper = WeightsMapper(

805 - orig_to_new_prefix={

806 - # mapping for new names in checkpoint saved after transformers v4.52

807 - "vit.vit.": "visual.",

808 - "vit.": "visual.",

809 - "model.": "language_model.model.",

810 - }

811 - )

812 -

813 - supports_encoder_tp_data = True

814 -

815 - def get_xdrope_input_positions(

816 - self,

817 - input_tokens: list[int],

818 - mm_features: list[MultiModalFeatureSpec],

819 - ) -> torch.Tensor:

820 - kwargs = MultiModalFeatureSpec.gather_kwargs(

821 - mm_features,

822 - {"image_grid_thw"},

823 - )

824 - image_grid_thw = [item.tolist() for item in kwargs.get("image_grid_thw", [])]

825 -

826 - hf_config = self.config

827 - image_start_token_id = hf_config.image_start_token_id

828 - spatial_merge_size = hf_config.vision_config.spatial_merge_size

829 - xd_num = len(hf_config.rope_scaling["xdrope_section"])

830 -

831 - input_tokens_tensor = torch.tensor(input_tokens)

832 - image_start_indices = torch.argwhere(

833 - input_tokens_tensor == image_start_token_id

834 - ).squeeze(1)

835 -

836 - p_index = torch.arange(len(input_tokens_tensor))

837 - w_index = torch.arange(len(input_tokens_tensor))

838 - h_index = torch.arange(len(input_tokens_tensor))

839 - t_index = torch.arange(len(input_tokens_tensor))

840 - for image_index in range(len(image_start_indices)):

841 - # +1 : first image_token, +2: for xdrope positions

842 - pos = image_start_indices[image_index] + 2

843 - t, h, w = image_grid_thw[image_index]

844 - _, llm_grid_h, llm_grid_w = (

845 - t,

846 - h // spatial_merge_size,

847 - w // spatial_merge_size,

848 - )

849 -

850 - token_num = (llm_grid_w + 1) * llm_grid_h

851 - w_index[pos : pos + token_num].copy_(

852 - torch.arange(0, llm_grid_w + 1)

853 - .reshape(1, -1)

854 - .expand(llm_grid_h, -1)

855 - .reshape(-1)

856 - )

857 - h_index[pos : pos + token_num].copy_(

858 - torch.arange(0, llm_grid_h)

859 - .reshape(-1, 1)

860 - .expand(-1, llm_grid_w + 1)

861 - .reshape(-1)

862 - )

863 - t_index[pos : pos + token_num] = image_index

864 -

865 - if xd_num == 4:

866 - llm_positions = torch.stack([p_index, w_index, h_index, t_index])

867 - elif xd_num == 3:

868 - llm_positions = torch.stack([w_index, h_index, t_index])

869 -

870 - return llm_positions

871 -

872 - @classmethod

873 - def get_placeholder_str(cls, modality: str, i: int) -> str | None:

874 - if modality.startswith("image"):

875 - return "<｜hy_place▁holder▁no▁100｜><｜hy_place▁holder▁no▁102｜><｜hy_place▁holder▁no▁101｜>" # noqa: E501

876 -

877 - raise ValueError("Only image modality is supported")

878 -

879 - def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

880 - super().__init__()

881 - config: HunYuanVLConfig = vllm_config.model_config.hf_config

882 -

883 - self.config = config

884 -

885 - with self._mark_tower_model(vllm_config, {"image"}):

886 - self.visual = HunYuanVisionTransformer(

887 - config.vision_config,

888 - quant_config=vllm_config.quant_config,

889 - prefix=maybe_prefix(prefix, "visual"),

890 - )

891 -

892 - with self._mark_language_model(vllm_config):

893 - self.language_model = init_vllm_registered_model(

894 - vllm_config=vllm_config,

895 - prefix=maybe_prefix(prefix, "language_model.model"),

896 - architectures=[

897 - "HunYuanDenseV1ForCausalLM",

898 - "HunYuanMoEV1ForCausalLM",

899 - ],

900 - )

901 -

902 - self.make_empty_intermediate_tensors = (

903 - self.language_model.make_empty_intermediate_tensors

904 - )

905 -

906 - def _parse_and_validate_image_input(

907 - self, **kwargs: object

908 - ) -> HunYuanVLImageInputs | None:

909 - pixel_values = kwargs.pop("pixel_values", None)

910 - image_embeds = kwargs.pop("image_embeds", None)

911 - image_grid_thw = kwargs.pop("image_grid_thw", None)

912 -

913 - if pixel_values is None and image_embeds is None:

914 - return None

915 -

916 - # TODO: refine

917 - if isinstance(pixel_values, list):

918 - pixel_values = torch.cat(pixel_values, dim=0)

919 - if len(pixel_values.shape) == 3:

920 - last_dim = pixel_values.shape[-1]

921 - pixel_values = pixel_values.reshape(-1, last_dim)

922 - image_grid_thw = image_grid_thw.reshape(-1, 3)

923 -

924 - if pixel_values is not None:

925 - return HunYuanVLImagePixelInputs(

926 - type="pixel_values",

927 - pixel_values=pixel_values,

928 - image_grid_thw=image_grid_thw,

929 - )

930 -

931 - if image_embeds is not None:

932 - return HunYuanVLImageEmbeddingInputs(

933 - type="image_embeds",

934 - image_embeds=image_embeds,

935 - image_grid_thw=image_grid_thw,

936 - )

937 -

938 - def _process_image_input(

939 - self, image_input: HunYuanVLImageInputs

940 - ) -> tuple[torch.Tensor, ...]:

941 - grid_thw = image_input["image_grid_thw"]

942 - assert grid_thw.ndim == 2

943 - grid_thw_list = grid_thw.tolist()

944 -

945 - if image_input["type"] == "image_embeds":

946 - image_embeds = image_input["image_embeds"].type(self.visual.dtype)

947 - else:

948 - pixel_values = image_input["pixel_values"]

949 -

950 - # TODO: use_data_parallel (split image_embeds in visual)

951 - image_embeds = self.visual(pixel_values, grid_thw=grid_thw_list)

952 -

953 - return image_embeds

954 -

955 - def _parse_and_validate_multimodal_inputs(self, **kwargs: object) -> dict:

956 - mm_input_by_modality = {}

957 -

958 - # Preserve the order of modalities if there are multiple of them

959 - # from the order of kwargs.

960 - for input_key in kwargs:

961 - if (

962 - input_key in ("pixel_values", "image_embeds")

963 - and "image" not in mm_input_by_modality

964 - ):

965 - mm_input_by_modality["image"] = self._parse_and_validate_image_input(

966 - **kwargs

967 - )

968 - return mm_input_by_modality

969 -

970 - def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:

971 - mm_input_by_modality = self._parse_and_validate_multimodal_inputs(**kwargs)

972 - if not mm_input_by_modality:

973 - return []

974 -

975 - # The result multimodal_embeddings is tuple of tensors, with each

976 - # tensor correspoending to a multimodal data item (image or video).

977 - multimodal_embeddings: tuple[torch.Tensor, ...] = ()

978 -

979 - # NOTE: It is important to iterate over the keys in this dictionary

980 - # to preserve the order of the modalities.

981 - for modality in mm_input_by_modality:

982 - multimodal_input = mm_input_by_modality[modality]

983 - if modality == "image":

984 - image_embeddings = self._process_image_input(multimodal_input)

985 - multimodal_embeddings += tuple(image_embeddings)

986 - return multimodal_embeddings

987 -

988 - def forward(

989 - self,

990 - input_ids: torch.Tensor | None,

991 - positions: torch.Tensor,

992 - intermediate_tensors: IntermediateTensors | None,

993 - inputs_embeds: torch.Tensor | None,

994 - **kwargs: object,

995 - ) -> torch.Tensor | IntermediateTensors:

996 - if intermediate_tensors is not None:

997 - inputs_embeds = None

998 -

999 - hidden_states = self.language_model(

1000 - input_ids=input_ids,

1001 - positions=positions,

1002 - intermediate_tensors=intermediate_tensors,

1003 - inputs_embeds=inputs_embeds,

1004 - )

1005 - return hidden_states

1006 -

1007 - def compute_logits(

1008 - self,

1009 - hidden_states: torch.Tensor,

1010 - ) -> torch.Tensor | None:

1011 - return self.language_model.compute_logits(hidden_states)

1012 -

1013 - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

1014 - loader = AutoWeightsLoader(self)

1015 - return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

1016 -

1017 - def get_mm_mapping(self) -> MultiModelKeys:

1018 - """

1019 - Get the module prefix in multimodal models

1020 - """

1021 - return MultiModelKeys.from_string_field(

1022 - language_model="language_model.model",

1023 - connector="visual.perceive",

1024 - tower_model="visual",

1025 - )

vllm/model_executor/models/olmo3.py deleted

-412

@@ -1,412 +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/main/src/transformers/models/olmo3/modeling_olmo3.py

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 OLMo3 model compatible with HuggingFace weights."""

26 -

27 -from collections.abc import Iterable

28 -from functools import partial

29 -from itertools import islice

30 -

31 -import torch

32 -from torch import nn

33 -from transformers import Olmo3Config

34 -

35 -from vllm.compilation.decorators import support_torch_compile

36 -from vllm.config import VllmConfig

37 -from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size

38 -from vllm.distributed.communication_op import tensor_model_parallel_all_gather

39 -from vllm.distributed.parallel_state import get_tensor_model_parallel_rank

40 -from vllm.distributed.utils import split_tensor_along_last_dim

41 -from vllm.model_executor.layers.activation import SiluAndMul

42 -from vllm.model_executor.layers.attention import Attention

43 -from vllm.model_executor.layers.layernorm import RMSNorm

44 -from vllm.model_executor.layers.linear import (

45 - MergedColumnParallelLinear,

46 - QKVParallelLinear,

47 - RowParallelLinear,

48 -)

49 -from vllm.model_executor.layers.logits_processor import LogitsProcessor

50 -from vllm.model_executor.layers.rotary_embedding import get_rope

51 -from vllm.model_executor.layers.vocab_parallel_embedding import (

52 - ParallelLMHead,

53 - VocabParallelEmbedding,

54 -)

55 -from vllm.model_executor.models.interfaces import SupportsLoRA, SupportsPP

56 -from vllm.model_executor.models.utils import (

57 - AutoWeightsLoader,

58 - WeightsMapper,

59 - extract_layer_index,

60 - make_empty_intermediate_tensors_factory,

61 - make_layers,

62 - maybe_prefix,

63 -)

64 -from vllm.sequence import IntermediateTensors

65 -

66 -

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))))`

71 - (plus another skip connection).

72 - """

73 -

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, Olmo3Config)

78 -

79 - hidden_size = self.config.hidden_size

80 - self.tp_size = get_tensor_model_parallel_world_size()

81 - self.total_num_heads = self.config.num_attention_heads

82 -

83 - assert hidden_size % self.total_num_heads == 0

84 - assert self.total_num_heads % self.tp_size == 0

85 -

86 - self.num_heads = self.total_num_heads // self.tp_size

87 - self.total_num_kv_heads = (

88 - self.config.num_key_value_heads or self.total_num_heads

89 - )

90 - if self.total_num_kv_heads >= self.tp_size:

91 - assert self.total_num_kv_heads % self.tp_size == 0

92 - else:

93 - assert self.tp_size % self.total_num_kv_heads == 0

94 -

95 - self.num_kv_heads = max(1, self.total_num_kv_heads // self.tp_size)

96 - self.head_dim = hidden_size // self.total_num_heads

97 - self.q_size = self.num_heads * self.head_dim

98 - self.kv_size = self.num_kv_heads * self.head_dim

99 - self.max_position_embeddings = self.config.max_position_embeddings

100 -

101 - # Attention input projection. Projects x -> (q, k, v)

102 - self.qkv_proj = QKVParallelLinear(

103 - hidden_size,

104 - self.head_dim,

105 - self.total_num_heads,

106 - self.total_num_kv_heads,

107 - bias=False,

108 - quant_config=vllm_config.quant_config,

109 - prefix=f"{prefix}.qkv_proj",

110 - )

111 -

112 - self.tp_rank = get_tensor_model_parallel_rank()

113 - self.k_norm = RMSNorm(

114 - self.total_num_kv_heads * self.head_dim,

115 - eps=self.config.rms_norm_eps,

116 - )

117 - self.q_norm = RMSNorm(self.config.hidden_size, eps=self.config.rms_norm_eps)

118 -

119 - self.scaling = self.head_dim**-0.5

120 -

121 - layer_idx = extract_layer_index(prefix)

122 - sliding_window = None

123 - if (

124 - layer_types := getattr(self.config, "layer_types", None)

125 - ) is not None and layer_types[layer_idx] == "sliding_attention":

126 - sliding_window = self.config.sliding_window

127 -

128 - self.attn = Attention(

129 - self.num_heads,

130 - self.head_dim,

131 - self.scaling,

132 - num_kv_heads=self.num_kv_heads,

133 - cache_config=vllm_config.cache_config,

134 - quant_config=vllm_config.quant_config,

135 - per_layer_sliding_window=sliding_window,

136 - prefix=f"{prefix}.attn",

137 - )

138 -

139 - # Rotary embeddings. Rope scaling is only applied on full attention layers.

140 - rope_parameters = self.config.rope_parameters

141 - attn_type = "full_attention" if sliding_window is None else "sliding_attention"

142 - rope_parameters = rope_parameters.get(attn_type, rope_parameters)

143 - self.rotary_emb = get_rope(

144 - self.head_dim,

145 - max_position=self.max_position_embeddings,

146 - rope_parameters=rope_parameters,

147 - )

148 -

149 - # Attention output projection.

150 - self.o_proj = RowParallelLinear(

151 - self.total_num_heads * self.head_dim,

152 - hidden_size,

153 - bias=False,

154 - quant_config=vllm_config.quant_config,

155 - prefix=f"{prefix}.o_proj",

156 - )

157 -

158 - def _apply_qk_norm(

159 - self, q: torch.Tensor, k: torch.Tensor

160 - ) -> tuple[torch.Tensor, torch.Tensor]:

161 - if self.tp_size > 1:

162 - q = tensor_model_parallel_all_gather(q.contiguous())

163 - k = tensor_model_parallel_all_gather(k.contiguous())

164 - q = self.q_norm(q)

165 - k = self.k_norm(k)

166 - if self.tp_size > 1:

167 - splitter = partial(split_tensor_along_last_dim, num_partitions=self.tp_size)

168 - q = splitter(q)[self.tp_rank]

169 - k = splitter(k)[self.tp_rank]

170 - return q, k

171 -

172 - def forward(

173 - self,

174 - positions: torch.Tensor,

175 - hidden_states: torch.Tensor,

176 - ) -> torch.Tensor:

177 - qkv, _ = self.qkv_proj(hidden_states)

178 - q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)

179 - q, k = self._apply_qk_norm(q, k)

180 - q, k = self.rotary_emb(positions, q, k)

181 - attn_output = self.attn(q, k, v)

182 - output, _ = self.o_proj(attn_output)

183 - return output

184 -

185 -

186 -class Olmo3MLP(nn.Module):

187 - """

188 - This is the MLP block where the output is computed as

189 - `MLP(x)` in `LN(MLP(x + LN(Attention(x))))`

190 - (plus another skip connection).

191 - """

192 -

193 - def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

194 - super().__init__()

195 - config = vllm_config.model_config.hf_config

196 - assert isinstance(config, Olmo3Config)

197 - hidden_size = config.hidden_size

198 - intermediate_size = config.intermediate_size

199 -

200 - # Feed-forward input projection.

201 - self.gate_up_proj = MergedColumnParallelLinear(

202 - hidden_size,

203 - [intermediate_size] * 2,

204 - bias=False,

205 - quant_config=vllm_config.quant_config,

206 - prefix=f"{prefix}.gate_up_proj",

207 - )

208 -

209 - # Activation function.

210 - self.act_fn = SiluAndMul()

211 -

212 - # Feed-forward output projection.

213 - self.down_proj = RowParallelLinear(

214 - intermediate_size,

215 - hidden_size,

216 - bias=False,

217 - quant_config=vllm_config.quant_config,

218 - prefix=f"{prefix}.down_proj",

219 - )

220 -

221 - def forward(

222 - self,

223 - x: torch.Tensor,

224 - ) -> torch.Tensor:

225 - gate_up, _ = self.gate_up_proj(x)

226 - x = self.act_fn(gate_up)

227 - x, _ = self.down_proj(x)

228 - return x

229 -

230 -

231 -class Olmo3DecoderLayer(nn.Module):

232 - """

233 - This is a typical transformer block where the output is

234 - computed as `MLP(LN(x + Attention(LN(x))))`

235 - (plus another skip connection).

236 - """

237 -

238 - def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

239 - super().__init__()

240 - config = vllm_config.model_config.hf_config

241 - assert isinstance(config, Olmo3Config)

242 - # Attention block.

243 - self.self_attn = Olmo3Attention(

244 - vllm_config=vllm_config, prefix=f"{prefix}.self_attn"

245 - )

246 -

247 - # MLP block.

248 - self.mlp = Olmo3MLP(vllm_config=vllm_config, prefix=f"{prefix}.mlp")

249 -

250 - # LayerNorm

251 - self.post_attention_layernorm = RMSNorm(

252 - config.hidden_size, eps=config.rms_norm_eps

253 - )

254 -

255 - self.post_feedforward_layernorm = RMSNorm(

256 - config.hidden_size, eps=config.rms_norm_eps

257 - )

258 -

259 - def forward(

260 - self,

261 - positions: torch.Tensor,

262 - hidden_states: torch.Tensor,

263 - ) -> torch.Tensor:

264 - # Attention block.

265 - residual = hidden_states

266 - hidden_states = self.self_attn(positions, hidden_states)

267 - hidden_states = self.post_attention_layernorm(hidden_states)

268 - hidden_states = hidden_states + residual

269 -

270 - # MLP block.

271 - residual = hidden_states

272 - hidden_states = self.mlp(hidden_states)

273 - hidden_states = self.post_feedforward_layernorm(hidden_states)

274 - hidden_states = residual + hidden_states

275 - return hidden_states

276 -

277 -

278 -@support_torch_compile

279 -class Olmo3Model(nn.Module):

280 - def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

281 - super().__init__()

282 - self.config = vllm_config.model_config.hf_config

283 - assert isinstance(self.config, Olmo3Config)

284 -

285 - self.embed_tokens = VocabParallelEmbedding(

286 - self.config.vocab_size,

287 - self.config.hidden_size,

288 - prefix=f"{prefix}.embed_tokens",

289 - )

290 - self.start_layer, self.end_layer, self.layers = make_layers(

291 - self.config.num_hidden_layers,

292 - lambda prefix: Olmo3DecoderLayer(vllm_config=vllm_config, prefix=prefix),

293 - prefix=f"{prefix}.layers",

294 - )

295 - self.norm = RMSNorm(

296 - self.config.hidden_size,

297 - eps=self.config.rms_norm_eps,

298 - )

299 - self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(

300 - ["hidden_states"], self.config.hidden_size

301 - )

302 -

303 - def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:

304 - return self.embed_tokens(input_ids)

305 -

306 - def forward(

307 - self,

308 - input_ids: torch.Tensor | None,

309 - positions: torch.Tensor,

310 - intermediate_tensors: IntermediateTensors | None,

311 - inputs_embeds: torch.Tensor | None = None,

312 - ) -> torch.Tensor | IntermediateTensors:

313 - """

314 - Args:

315 - input_ids: A tensor of shape `(batch_size, seq_len)`.

316 - """

317 - if get_pp_group().is_first_rank:

318 - if inputs_embeds is not None:

319 - hidden_states = inputs_embeds

320 - # Get embeddings of input.

321 - # shape: (batch_size, seq_len, d_model)

322 - else:

323 - hidden_states = self.embed_tokens(input_ids)

324 -

325 - else:

326 - assert intermediate_tensors is not None

327 - hidden_states = intermediate_tensors["hidden_states"]

328 - assert isinstance(hidden_states, torch.Tensor)

329 -

330 - # Apply blocks one-by-one.

331 - for layer in islice(self.layers, self.start_layer, self.end_layer):

332 - # shape: (batch_size, seq_len, d_model)

333 - hidden_states = layer(positions, hidden_states)

334 -

335 - if not get_pp_group().is_last_rank:

336 - return IntermediateTensors({"hidden_states": hidden_states})

337 -

338 - # Apply final layer norm.

339 - # shape: (batch_size, seq_len or 1, d_model)

340 - hidden_states = self.norm(hidden_states)

341 - return hidden_states

342 -

343 -

344 -class Olmo3ForCausalLM(nn.Module, SupportsPP, SupportsLoRA):

345 - """

346 - Extremely barebones HF model wrapper.

347 - """

348 -

349 - hf_to_vllm_mapper = WeightsMapper(

350 - orig_to_new_stacked={

351 - # weight_name: (param_name, shard_id)

352 - ".q_proj": (".qkv_proj", "q"),

353 - ".k_proj": (".qkv_proj", "k"),

354 - ".v_proj": (".qkv_proj", "v"),

355 - ".gate_proj": (".gate_up_proj", 0),

356 - ".up_proj": (".gate_up_proj", 1),

357 - }

358 - )

359 - packed_modules_mapping = {

360 - "qkv_proj": ["q_proj", "k_proj", "v_proj"],

361 - "gate_up_proj": ["gate_proj", "up_proj"],

362 - }

363 -

364 - def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

365 - super().__init__()

366 - config = vllm_config.model_config.hf_config

367 - assert isinstance(config, Olmo3Config)

368 - self.config = config

369 - self.model = Olmo3Model(

370 - vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")

371 - )

372 - self.lm_head = ParallelLMHead(

373 - config.vocab_size,

374 - config.hidden_size,

375 - quant_config=vllm_config.quant_config,

376 - prefix=maybe_prefix(prefix, "lm_head"),

377 - )

378 - if config.tie_word_embeddings:

379 - self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens)

380 - self.logits_processor = LogitsProcessor(config.vocab_size)

381 - self.make_empty_intermediate_tensors = (

382 - self.model.make_empty_intermediate_tensors

383 - )

384 -

385 - def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:

386 - return self.model.embed_input_ids(input_ids)

387 -

388 - def forward(

389 - self,

390 - input_ids: torch.Tensor | None,

391 - positions: torch.Tensor,

392 - intermediate_tensors: IntermediateTensors | None = None,

393 - inputs_embeds: torch.Tensor | None = None,

394 - ) -> torch.Tensor | IntermediateTensors:

395 - hidden_states = self.model(

396 - input_ids=input_ids,

397 - positions=positions,

398 - intermediate_tensors=intermediate_tensors,

399 - inputs_embeds=inputs_embeds,

400 - )

401 - return hidden_states

402 -

403 - def compute_logits(

404 - self,

405 - hidden_states: torch.Tensor,

406 - ) -> torch.Tensor | None:

407 - logits = self.logits_processor(self.lm_head, hidden_states)

408 - return logits

409 -

410 - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):

411 - loader = AutoWeightsLoader(self)

412 - return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

vllm/model_executor/models/registry.py

+8 -8

@@ -104,7 +104,6 @@ _TEXT_GENERATION_MODELS = {

104 "FalconForCausalLM": ("falcon", "FalconForCausalLM"),

105 "FalconMambaForCausalLM": ("mamba", "MambaForCausalLM"),

106 "FalconH1ForCausalLM": ("falcon_h1", "FalconH1ForCausalLM"),

107 - "FlexOlmoForCausalLM": ("flex_olmo", "FlexOlmoForCausalLM"),

107 "GemmaForCausalLM": ("gemma", "GemmaForCausalLM"),

108 "Gemma2ForCausalLM": ("gemma2", "Gemma2ForCausalLM"),

109 "Gemma3ForCausalLM": ("gemma3", "Gemma3ForCausalLM"),

@@ -128,8 +127,6 @@ _TEXT_GENERATION_MODELS = {

127 "GraniteMoeSWAForCausalLM": ("granitemoeshared", "GraniteMoeSharedForCausalLM"),

128 "GraniteSWAForCausalLM": ("granite", "GraniteForCausalLM"),

129 "HrmTextForCausalLM": ("hrm_text", "HrmTextForCausalLM"),

131 - "HunYuanMoEV1ForCausalLM": ("hunyuan_v1", "HunYuanMoEV1ForCausalLM"),

132 - "HunYuanDenseV1ForCausalLM": ("hunyuan_v1", "HunYuanDenseV1ForCausalLM"),

130 "HYV3ForCausalLM": ("hy_v3", "HYV3ForCausalLM"),

131 "HCXVisionV2ForCausalLM": ("hyperclovax_vision_v2", "HCXVisionV2ForCausalLM"),

132 "HyperCLOVAXForCausalLM": ("hyperclovax", "HyperCLOVAXForCausalLM"),

@@ -180,7 +177,6 @@ _TEXT_GENERATION_MODELS = {

177 "NemotronForCausalLM": ("nemotron", "NemotronForCausalLM"),

178 "NemotronHForCausalLM": ("nemotron_h", "NemotronHForCausalLM"),

179 "NemotronHPuzzleForCausalLM": ("nemotron_h", "NemotronHForCausalLM"),

183 - "Olmo3ForCausalLM": ("olmo3", "Olmo3ForCausalLM"),

180 "MuseGlimmerForCausalLM": ("muse_glimmer", "MuseGlimmerForCausalLM"),

181 "OlmoHybridForCausalLM": ("olmo_hybrid", "OlmoHybridForCausalLM"),

182 "OlmoeForCausalLM": ("olmoe", "OlmoeForCausalLM"),

@@ -420,10 +416,6 @@ _MULTIMODAL_MODELS = {

416 "Granite4VisionForConditionalGeneration",

417 ),

418 "H2OVLChatModel": ("h2ovl", "H2OVLChatModel"),

423 - "HunYuanVLForConditionalGeneration": (

424 - "hunyuan_vision",

425 - "HunYuanVLForConditionalGeneration",

426 - ),

419 "InternVLChatModel": ("internvl", "InternVLChatModel"),

420 "InternS1ForConditionalGeneration": (

421 "interns1",

@@ -679,9 +671,13 @@ _SPECULATIVE_DECODING_MODELS = {

671

672 _TRANSFORMERS_SUPPORTED_MODELS = {

673 # Text generation models

674 + "FlexOlmoForCausalLM": ("transformers", "TransformersMoEForCausalLM"),

675 "GPTBigCodeForCausalLM": ("transformers", "TransformersForCausalLM"),

676 + "HunYuanDenseV1ForCausalLM": ("transformers", "TransformersForCausalLM"),

677 + "HunYuanMoEV1ForCausalLM": ("transformers", "TransformersMoEForCausalLM"),

678 "OlmoForCausalLM": ("transformers", "TransformersForCausalLM"),

679 "Olmo2ForCausalLM": ("transformers", "TransformersForCausalLM"),

680 + "Olmo3ForCausalLM": ("transformers", "TransformersForCausalLM"),

681 "SmolLM3ForCausalLM": ("transformers", "TransformersForCausalLM"),

682 "Starcoder2ForCausalLM": ("transformers", "TransformersForCausalLM"),

683 "VaultGemmaForCausalLM": ("transformers", "TransformersForCausalLM"),

@@ -690,6 +686,10 @@ _TRANSFORMERS_SUPPORTED_MODELS = {

686 "transformers",

687 "TransformersMultiModalForCausalLM",

688 ),

689 + "HunYuanVLForConditionalGeneration": (

690 + "transformers",

691 + "TransformersMultiModalForCausalLM",

692 + ),

693 "VibeVoiceAsrForConditionalGeneration": (

694 "transformers",

695 "TransformersMultiModalForCausalLM",

vllm/model_executor/models/transformers/moe.py

+19 -6

@@ -37,10 +37,10 @@ from vllm.model_executor.layers.fused_moe import (

37 from vllm.model_executor.models.interfaces import MixtureOfExperts

38 from vllm.model_executor.models.transformers.fuser import get_fuser

39 from vllm.model_executor.models.transformers.fusers.moe import MoEBlockFuser

40 -from vllm.model_executor.models.utils import maybe_prefix

40 +from vllm.model_executor.models.utils import extract_layer_index, maybe_prefix

41 from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE, direct_register_custom_op

43 -from .utils import log_replacement

43 +from .utils import log_replacement, maybe_per_layer

45 if TYPE_CHECKING:

46 from vllm.config import VllmConfig

@@ -183,7 +183,7 @@ class MoEMixin(MixtureOfExperts):

183 )

184

185 # Common kwargs

186 - renormalize = getattr(text_config, "norm_topk_prob", top_k > 1)

186 + norm_topk_prob = getattr(text_config, "norm_topk_prob", None)

187

188 # Routed scaling factor kwargs

189 routed_scaling_factor = getattr(text_config, "routed_scaling_factor", 1.0)

@@ -266,12 +266,25 @@ class MoEMixin(MixtureOfExperts):

266 self.num_shared_experts = 1

267 break

268

269 + # Only pay for the layer index if something is per-layer;

270 + # `extract_layer_index` raises for unnumbered module paths.

271 + per_layer = (top_k, intermediate_size, norm_topk_prob)

272 + layer_idx = (

273 + extract_layer_index(qual_name)

274 + if any(isinstance(v, list) for v in per_layer)

275 + else 0

276 + )

277 + layer_top_k = maybe_per_layer(top_k, layer_idx)

278 + layer_renormalize = maybe_per_layer(norm_topk_prob, layer_idx)

279 + if layer_renormalize is None:

280 + layer_renormalize = layer_top_k > 1

281 +

282 kwargs: dict[str, Any] = dict(

283 num_experts=num_experts,

271 - top_k=top_k,

284 + top_k=layer_top_k,

285 hidden_size=hidden_size,

273 - intermediate_size=intermediate_size,

274 - renormalize=renormalize,

286 + intermediate_size=maybe_per_layer(intermediate_size, layer_idx),

287 + renormalize=layer_renormalize,

288 quant_config=self.quant_config,

289 prefix=qual_name,

290 activation=activation,

vllm/model_executor/models/transformers/multimodal.py

@@ -173,6 +173,11 @@ class MultiModalDummyInputsBuilder(BaseDummyInputsBuilder[MultiModalProcessingIn

173 image_token = processor.boi_token

174 else:

175 image_token = getattr(processor, "image_token", "")

176 + # Some processors (e.g. HunYuanVL) reject a bare image token and

177 + # require each one to be wrapped in its start/end markers.

178 + start_token = getattr(processor, "image_start_token", "")

179 + end_token = getattr(processor, "image_end_token", "")

180 + image_token = f"{start_token}{image_token}{end_token}"

181 text += image_token * num_images

182 return text

183

vllm/model_executor/models/transformers/utils.py

+9 -1

@@ -22,7 +22,7 @@ from functools import lru_cache

22 from itertools import chain

23 from operator import attrgetter

24 from pathlib import Path

25 -from typing import TYPE_CHECKING, Literal

25 +from typing import TYPE_CHECKING, Literal, TypeVar

27 import torch

28 from torch import nn

@@ -306,6 +306,14 @@ def recursive_replace_linear(

306 _recursive_replace(model, prefix=prefix)

307

308

309 +T = TypeVar("T")

310 +

311 +

312 +def maybe_per_layer(value: T | list[T], layer_idx: int) -> T:

313 + """Pick `layer_idx`'s entry from a config field that may be sized per layer."""

314 + return value[layer_idx] if isinstance(value, list) else value

315 +

316 +

317 def named_state(module: nn.Module) -> Iterator[tuple[str, torch.Tensor]]:

318 """`module`'s own state (i.e. named parameters and buffers)."""

319 return chain(module.named_parameters(), module.named_buffers())

vllm/transformers_utils/config.py

-2

@@ -87,11 +87,9 @@ _CONFIG_REGISTRY: dict[str, type[PretrainedConfig]] = LazyConfigDict(

87 deepseek_v4="DeepseekV4Config",

88 dots3_note="Dots3NoteConfig",

89 k3_dspark="K3DSparkConfig",

90 - flex_olmo="FlexOlmoConfig",

90 funaudiochat="FunAudioChatConfig",

91 granite4_vision="Granite4VisionConfig",

92 hyperclovax="HyperCLOVAXConfig",

94 - hunyuan_vl="HunYuanVLConfig",

93 hy_v3="HYV3Config",

94 isaac="IsaacConfig",

95 kimi_k2="DeepseekV3Config", # Kimi K2 uses same architecture as DeepSeek V3

vllm/transformers_utils/configs/__init__.py

-8

@@ -37,13 +37,9 @@ _CLASS_TO_MODULE: dict[str, str] = {

37 "K3DSparkConfig": "vllm.transformers_utils.configs.k3_dspark",

38 "DotsOCRConfig": "vllm.transformers_utils.configs.dotsocr",

39 "EAGLEConfig": "vllm.transformers_utils.configs.eagle",

40 - "FlexOlmoConfig": "vllm.transformers_utils.configs.flex_olmo",

40 "FunAudioChatConfig": "vllm.transformers_utils.configs.funaudiochat",

41 "FunAudioChatAudioEncoderConfig": "vllm.transformers_utils.configs.funaudiochat",

42 "Granite4VisionConfig": "vllm.transformers_utils.configs.granite4_vision",

44 - "HunYuanVLConfig": "vllm.transformers_utils.configs.hunyuan_vl",

45 - "HunYuanVLTextConfig": "vllm.transformers_utils.configs.hunyuan_vl",

46 - "HunYuanVLVisionConfig": "vllm.transformers_utils.configs.hunyuan_vl",

43 "HYV3Config": "vllm.transformers_utils.configs.hy_v3",

44 "HyperCLOVAXConfig": "vllm.transformers_utils.configs.hyperclovax",

45 "IsaacConfig": "vllm.transformers_utils.configs.isaac",

@@ -130,13 +126,9 @@ __all__ = [

126 "K3DSparkConfig",

127 "DotsOCRConfig",

128 "EAGLEConfig",

133 - "FlexOlmoConfig",

129 "FunAudioChatConfig",

130 "FunAudioChatAudioEncoderConfig",

131 "Granite4VisionConfig",

137 - "HunYuanVLConfig",

138 - "HunYuanVLTextConfig",

139 - "HunYuanVLVisionConfig",

132 "HYV3Config",

133 "HyperCLOVAXConfig",

134 "IsaacConfig",

vllm/transformers_utils/configs/flex_olmo.py deleted

-82

@@ -1,82 +0,0 @@

1 -# SPDX-License-Identifier: Apache-2.0

2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 -from typing import Any

4 -

5 -from transformers.configuration_utils import PretrainedConfig

6 -

7 -

8 -class FlexOlmoConfig(PretrainedConfig):

9 - model_type = "flex_olmo"

10 - keys_to_ignore_at_inference = ["past_key_values"]

11 -

12 - def __init__(

13 - self,

14 - vocab_size=100352,

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=4096,

22 - initializer_range=0.02,

23 - rms_norm_eps=1e-06,

24 - use_cache=True,

25 - pad_token_id=100277,

26 - bos_token_id=None,

27 - eos_token_id=100257,

28 - tie_word_embeddings=False,

29 - rope_parameters: dict[str, Any] | None = None,

30 - attention_bias=False,

31 - attention_dropout=0.0,

32 - num_experts_per_tok=5,

33 - num_experts=7,

34 - output_router_logits=False,

35 - router_aux_loss_coef=0.01,

36 - norm_topk_prob=False,

37 - **kwargs,

38 - ):

39 - if "architectures" not in kwargs:

40 - kwargs["architectures"] = ["FlexOlmoForCausalLM"]

41 -

42 - self.vocab_size = vocab_size

43 - self.max_position_embeddings = max_position_embeddings

44 - self.hidden_size = hidden_size

45 - self.intermediate_size = intermediate_size

46 - self.num_hidden_layers = num_hidden_layers

47 - self.num_attention_heads = num_attention_heads

48 -

49 - # for backward compatibility

50 - if num_key_value_heads is None:

51 - num_key_value_heads = num_attention_heads

52 -

53 - self.num_key_value_heads = num_key_value_heads

54 - self.hidden_act = hidden_act

55 - self.initializer_range = initializer_range

56 - self.rms_norm_eps = rms_norm_eps

57 - self.use_cache = use_cache

58 - # Try to set `rope_scaling` if available, otherwise use `rope_parameters`

59 - rope_scaling = kwargs.pop("rope_scaling", None)

60 - rope_parameters = rope_scaling or rope_parameters or {"rope_type": "default"}

61 - rope_theta = kwargs.pop("rope_theta", 500000.0)

62 - if "rope_theta" not in rope_parameters:

63 - rope_parameters["rope_theta"] = rope_theta

64 - self.rope_parameters = rope_parameters

65 - self.attention_bias = attention_bias

66 - self.attention_dropout = attention_dropout

67 - self.num_experts_per_tok = num_experts_per_tok

68 - self.num_experts = num_experts

69 - self.output_router_logits = output_router_logits

70 - self.router_aux_loss_coef = router_aux_loss_coef

71 - self.norm_topk_prob = norm_topk_prob

72 - # Validate the correctness of rotary position embeddings parameters

73 - # BC: if there is a 'type' field, move it to 'rope_type'.

74 - if self.rope_parameters is not None and "type" in self.rope_parameters:

75 - self.rope_parameters["rope_type"] = self.rope_parameters["type"]

76 - super().__init__(

77 - pad_token_id=pad_token_id,

78 - bos_token_id=bos_token_id,

79 - eos_token_id=eos_token_id,

80 - tie_word_embeddings=tie_word_embeddings,

81 - **kwargs,

82 - )

vllm/transformers_utils/configs/hunyuan_vl.py deleted

-281

@@ -1,281 +0,0 @@

1 -# SPDX-License-Identifier: Apache-2.0

2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 -

4 -# adapted from https://github.com/ManaEstras/transformers/blob/v4.57.1.hyvl/src/transformers/models/hunyuan_vl/configuration_hunyuan_vl.py

5 -

6 -from transformers import PretrainedConfig

7 -

8 -

9 -class HunYuanVLVisionConfig(PretrainedConfig):

10 - model_type = "hunyuan_vl"

11 - base_config_key = "vision_config"

12 -

13 - def __init__(

14 - self,

15 - hidden_act="gelu",

16 - hidden_size=1152,

17 - intermediate_size=4304,

18 - interpolate_mode="bilinear",

19 - rms_norm_eps=1e-05,

20 - learnable_mlp_pooling_size=0,

21 - num_attention_heads=16,

22 - num_key_value_heads=None,

23 - num_channels=3,

24 - num_hidden_layers=27,

25 - out_hidden_size=4096,

26 - patch_size=16,

27 - remove_prenorm=True,

28 - spatial_merge_size=2,

29 - temporal_patch_size=1,

30 - resize_resolution=2048,

31 - img_max_token_num=4096,

32 - max_image_size=2048,

33 - video_max_image_size=768,

34 - video_min_image_size=256,

35 - min_image_size=512,

36 - anyres_vit_max_image_size=2048,

37 - max_vit_seq_len=16384,

38 - text_hidden_size=3072,

39 - **kwargs,

40 - ):

41 - super().__init__(**kwargs)

42 -

43 - self.hidden_act = hidden_act

44 - self.hidden_size = hidden_size

45 - self.intermediate_size = intermediate_size

46 - self.interpolate_mode = interpolate_mode

47 - self.learnable_mlp_pooling_size = learnable_mlp_pooling_size

48 - self.num_attention_heads = num_attention_heads

49 - if not num_key_value_heads:

50 - self.num_key_value_heads = num_attention_heads

51 - else:

52 - self.num_key_value_heads = num_key_value_heads

53 - self.num_channels = num_channels

54 - self.num_hidden_layers = num_hidden_layers

55 - self.out_hidden_size = out_hidden_size

56 - self.patch_size = patch_size

57 - self.remove_prenorm = remove_prenorm

58 - self.spatial_merge_size = spatial_merge_size

59 - self.temporal_patch_size = temporal_patch_size

60 - self.rms_norm_eps = rms_norm_eps

61 -

62 - self.resize_resolution = resize_resolution

63 - self.img_max_token_num = img_max_token_num

64 - self.max_image_size = max_image_size

65 - self.min_image_size = min_image_size

66 - self.video_max_image_size = video_max_image_size

67 - self.video_min_image_size = video_min_image_size

68 - self.anyres_vit_max_image_size = anyres_vit_max_image_size

69 - self.max_vit_seq_len = max_vit_seq_len

70 - self.text_hidden_size = text_hidden_size

71 -

72 -

73 -class HunYuanVLTextConfig(PretrainedConfig):

74 - r"""

75 - This is the configuration class to store the configuration of a [`HunYuanVLTextConfig`]. It is used to instantiate an

76 - HunYuan model according to the specified arguments, defining the model architecture. Instantiating a configuration

77 - with the defaults will yield a similar configuration to that of the HunYuan-7B.

78 - Hunyuan-7B-Instruct [tencent/Hunyuan-7B-Instruct](https://huggingface.co/tencent/Hunyuan-7B-Instruct).

79 -

80 - Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the

81 - documentation from [`PretrainedConfig`] for more information.

82 -

83 -

84 - Args:

85 - vocab_size (`int`, *optional*, defaults to 290943):

86 - Vocabulary size of the HunYuan model. Defines the number of different tokens that can be represented by the

87 - `inputs_ids` passed when calling [`HunYuanVLTextConfig`]

88 - hidden_size (`int`, *optional*, defaults to 4096):

89 - Dimension of the hidden representations.

90 - intermediate_size (`int`, *optional*, defaults to 11008):

91 - Dimension of the MLP representations or shared MLP representations.

92 - num_hidden_layers (`int`, *optional*, defaults to 32):

93 - Number of hidden layers in the Transformer decoder.

94 - num_attention_heads (`int`, *optional*, defaults to 32):

95 - Number of attention heads for each attention layer in the Transformer decoder.

96 - num_key_value_heads (`int`, *optional*):

97 - This is the number of key_value heads that should be used to implement Grouped Query Attention. If

98 - `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if

99 - `num_key_value_heads=1 the model will use Multi Query Attention (MQA) otherwise GQA is used. When

100 - converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed

101 - by meanpooling all the original heads within that group. For more details checkout [this

102 - paper](https://huggingface.co/papers/2305.13245). If it is not specified, will default to

103 - `num_attention_heads`.

104 - hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):

105 - The non-linear activation function (function or string) in the decoder.

106 - max_position_embeddings (`int`, *optional*, defaults to 2048):

107 - The maximum sequence length that this model might ever be used with.

108 - initializer_range (`float`, *optional*, defaults to 0.02):

109 - The standard deviation of the truncated_normal_initializer for initializing all weight matrices.

110 - rms_norm_eps (`float`, *optional*, defaults to 1e-05):

111 - The epsilon used by the rms normalization layers.

112 - use_cache (`bool`, *optional*, defaults to `True`):

113 - Whether or not the model should return the last key/values attentions (not used by all models). Only

114 - relevant if `config.is_decoder=True`.

115 - pad_token_id (`int`, *optional*, defaults to 0):

116 - Padding token id.

117 - bos_token_id (`int`, *optional*, defaults to 1):

118 - Beginning of stream token id.

119 - eos_token_id (`int`, *optional*, defaults to 2):

120 - End of stream token id.

121 - eod_token_id (int, *optional*, defaults to 3):

122 - Token ID representing the end-of-document marker. Used to indicate the termination of a text sequence.

123 - Example: In multi-document processing, this token helps the model distinguish between separate documents.

124 - pretraining_tp (`int`, *optional*, defaults to 1):

125 - Experimental feature. Tensor parallelism rank used during pretraining. Please refer to [this

126 - document](https://huggingface.co/docs/transformers/parallelism) to understand more about it. This value is

127 - necessary to ensure exact reproducibility of the pretraining results. Please refer to [this

128 - issue](https://github.com/pytorch/pytorch/issues/76232).

129 - tie_word_embeddings (`bool`, *optional*, defaults to `False`):

130 - Whether to tie weight embeddings

131 - rope_theta (`float`, *optional*, defaults to 10000.0):

132 - The base period of the RoPE embeddings.

133 - rope_scaling (`Dict`, *optional*):

134 - Dictionary containing the scaling configuration for the RoPE embeddings. Currently supports two scaling

135 - strategies: linear and dynamic. Their scaling factor must be a float greater than 1. The expected format is

136 - `{"type": strategy name, "factor": scaling factor}`. When using this flag, don't update

137 - `max_position_embeddings` to the expected new maximum. See the following thread for more information on how

138 - these scaling strategies behave:

139 - https://www.reddit.com/r/LocalLLaMA/comments/14mrgpr/dynamically_scaled_rope_further_increases/. This is an

140 - experimental feature, subject to breaking API changes in future versions.

141 - attention_bias (`bool`, defaults to `False`, *optional*, defaults to `False`):

142 - Whether to use a bias in the query, key, value and output projection layers during self-attention.

143 - attention_dropout (`float`, *optional*, defaults to 0.0):

144 - The dropout ratio for the attention probabilities.

145 - head_dim (`int`, *optional*, defaults to 128):

146 - The attention head dimension.

147 - """ # noqa: E501

148 -

149 - model_type = "hunyuan_vl_text"

150 - keys_to_ignore_at_inference = ["past_key_values"]

151 -

152 - def __init__(

153 - self,

154 - vocab_size=290943,

155 - hidden_size=4096,

156 - intermediate_size: int = 11008,

157 - num_hidden_layers=32,

158 - num_attention_heads=32,

159 - num_key_value_heads=None,

160 - hidden_act="silu",

161 - max_position_embeddings=2048,

162 - initializer_range=0.02,

163 - rms_norm_eps=1e-5,

164 - use_cache=True,

165 - pad_token_id=0,

166 - bos_token_id=1,

167 - eos_token_id=2,

168 - eod_token_id=3,

169 - pretraining_tp=1,

170 - tie_word_embeddings=False,

171 - rope_theta=10000.0,

172 - rope_scaling=None,

173 - attention_bias=False,

174 - attention_dropout=0.0,

175 - head_dim=None,

176 - **kwargs,

177 - ):

178 - self.vocab_size = vocab_size

179 - self.max_position_embeddings = max_position_embeddings

180 - self.hidden_size = hidden_size

181 - self.intermediate_size = intermediate_size

182 - self.num_hidden_layers = num_hidden_layers

183 - self.num_attention_heads = num_attention_heads

184 - self.head_dim = head_dim

185 - # for backward compatibility

186 - if num_key_value_heads is None:

187 - num_key_value_heads = num_attention_heads

188 -

189 - self.num_key_value_heads = num_key_value_heads

190 - self.hidden_act = hidden_act

191 - self.initializer_range = initializer_range

192 - self.rms_norm_eps = rms_norm_eps

193 - self.pretraining_tp = pretraining_tp

194 - self.use_cache = use_cache

195 - self.rope_theta = rope_theta

196 - self.rope_scaling = rope_scaling

197 - self.attention_bias = attention_bias

198 - self.attention_dropout = attention_dropout

199 -

200 - super().__init__(

201 - pad_token_id=pad_token_id,

202 - bos_token_id=bos_token_id,

203 - eos_token_id=eos_token_id,

204 - tie_word_embeddings=tie_word_embeddings,

205 - **kwargs,

206 - )

207 -

208 -

209 -class HunYuanVLConfig(PretrainedConfig):

210 - model_type = "hunyuan_vl"

211 - sub_configs = {

212 - "vision_config": HunYuanVLVisionConfig,

213 - "text_config": HunYuanVLTextConfig,

214 - }

215 - keys_to_ignore_at_inference = ["past_key_values"]

216 -

217 - def __init__(

218 - self,

219 - text_config=None,

220 - vision_config=None,

221 - im_start_id=120118,

222 - im_end_id=120119,

223 - image_token_id=120120,

224 - im_newline_id=120121,

225 - video_start_id=120122,

226 - video_end_id=120123,

227 - **kwargs,

228 - ):

229 - # We need to init super() here so that it does not reset values

230 - # that are in text config to the BaseClass defaults. The Base

231 - # config has many text related defaults and not all defaults are

232 - # same as for `HunYuanVLTextConfig`.

233 - super().__init__(**kwargs)

234 -

235 - if isinstance(vision_config, dict):

236 - self.vision_config = self.sub_configs["vision_config"](**vision_config)

237 - elif vision_config is None:

238 - self.vision_config = self.sub_configs["vision_config"]()

239 -

240 - if isinstance(text_config, dict):

241 - self.text_config = self.sub_configs["text_config"](**text_config)

242 - elif text_config is None:

243 - # For BC use all kwargs to init `TextConfig`

244 - self.text_config = self.sub_configs["text_config"](**kwargs)

245 -

246 - self.image_token_id = image_token_id

247 - self.im_start_id = im_start_id

248 - self.im_end_id = im_end_id

249 - self.im_newline_id = im_newline_id

250 - self.video_start_id = video_start_id

251 - self.video_end_id = video_end_id

252 -

253 - self.vision_config.text_hidden_size = self.text_config.hidden_size

254 -

255 - # Attention implementation to use. It sets it recursively on sub-configs

256 - # so we call it again in the end.

257 - self._attn_implementation = kwargs.pop("attn_implementation", None)

258 -

259 - def __setattr__(self, key, value):

260 - if (

261 - (text_config := super().__getattribute__("__dict__").get("text_config"))

262 - is not None

263 - and key not in ["dtype", "_attn_implementation_internal"]

264 - and key in text_config.__dict__

265 - ):

266 - setattr(text_config, key, value)

267 - else:

268 - super().__setattr__(key, value)

269 -

270 - def __getattribute__(self, key):

271 - if "text_config" in super().__getattribute__("__dict__") and key not in [

272 - "_name_or_path",

273 - "model_type",

274 - "dtype",

275 - "_attn_implementation_internal",

276 - ]:

277 - text_config = super().__getattribute__("text_config")

278 - if key in text_config.__dict__:

279 - return getattr(text_config, key)

280 -

281 - return super().__getattribute__(key)