@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / ecad851cbd

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 15, 03:30 UTC

[Model]Add Tencent HunYuanMoEV1 Model Support (#20114)

Signed-off-by: aiyiwang <aiyiwang@tencent.com> Signed-off-by: Jee Jee Li <pandaleefree@gmail.com> Co-authored-by: quinnrong <quinnrong@tencent.com> Co-authored-by: Jee Jee Li <pandaleefree@gmail.com>

aiyiwang2025 committed Jul 1, 2025 at 22:28 UTC ecad851cbd0a2a6f9922fc9f3e94bde3f8220176

6 files changed +949 -6

docs/models/supported_models.md

+2 -1

@@ -350,6 +350,7 @@ Specified using `--task generate`.

350 | `GraniteMoeSharedForCausalLM` | Granite MoE Shared | `ibm-research/moe-7b-1b-active-shared-experts` (test model) | ✅︎ | ✅︎ | ✅︎ |

351 | `GritLM` | GritLM | `parasail-ai/GritLM-7B-vllm`. | ✅︎ | ✅︎ | |

352 | `Grok1ModelForCausalLM` | Grok1 | `hpcai-tech/grok-1`. | ✅︎ | ✅︎ | ✅︎ |

354 | `InternLMForCausalLM` | InternLM | `internlm/internlm-7b`, `internlm/internlm-chat-7b`, etc. | ✅︎ | ✅︎ | ✅︎ |

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

356 | `InternLM3ForCausalLM` | InternLM3 | `internlm/internlm3-8b-instruct`, etc. | ✅︎ | ✅︎ | ✅︎ |

@@ -387,7 +388,7 @@ Specified using `--task generate`.

388 | `TeleChat2ForCausalLM` | TeleChat2 | `Tele-AI/TeleChat2-3B`, `Tele-AI/TeleChat2-7B`, `Tele-AI/TeleChat2-35B`, etc. | ✅︎ | ✅︎ | ✅︎ |

389 | `TeleFLMForCausalLM` | TeleFLM | `CofeAI/FLM-2-52B-Instruct-2407`, `CofeAI/Tele-FLM`, etc. | ✅︎ | ✅︎ | ✅︎ |

390 | `XverseForCausalLM` | XVERSE | `xverse/XVERSE-7B-Chat`, `xverse/XVERSE-13B-Chat`, `xverse/XVERSE-65B-Chat`, etc. | ✅︎ | ✅︎ | ✅︎ |

394

tests/models/registry.py

+3 -1

@@ -188,6 +188,8 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {

188 "GraniteMoeSharedForCausalLM": _HfExamplesInfo("ibm-research/moe-7b-1b-active-shared-experts"), # noqa: E501

189 "Grok1ModelForCausalLM": _HfExamplesInfo("hpcai-tech/grok-1",

190 trust_remote_code=True),

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

192 + trust_remote_code=True),

193 "InternLMForCausalLM": _HfExamplesInfo("internlm/internlm-chat-7b",

194 trust_remote_code=True),

195 "InternLM2ForCausalLM": _HfExamplesInfo("internlm/internlm2-chat-7b",

@@ -490,4 +492,4 @@ class HfExampleModels:

492 raise ValueError(f"No example model defined for {model_id}")

493

494

493 -HF_EXAMPLE_MODELS = HfExampleModels(_EXAMPLE_MODELS)

495 +HF_EXAMPLE_MODELS = HfExampleModels(_EXAMPLE_MODELS)

\ No newline at end of file

tests/models/test_initialization.py

+2 -1

@@ -33,7 +33,8 @@ def test_can_initialize(model_arch: str, monkeypatch: pytest.MonkeyPatch):

34 # Ensure at least 2 expert per group

35 # Since `grouped_topk` assums top-2

36 - num_experts = getattr(text_config, 'n_group', 1) * 2

36 + n_group = getattr(text_config, 'n_group', None)

37 + num_experts = n_group * 2 if n_group is not None else 2

39 text_config.update({

40 "num_layers": 1,

vllm/model_executor/layers/rotary_embedding.py

+44 -3

@@ -533,6 +533,41 @@ class DynamicNTKScalingRotaryEmbedding(RotaryEmbedding):

533 return cache

534

535

536 +class DynamicNTKAlphaRotaryEmbedding(RotaryEmbedding):

537 + """RotaryEmbedding extended with Dynamic NTK alpha.

538 +

539 + Based on the original RotaryEmbedding implementation.

540 + """

541 +

542 + def __init__(

543 + self,

544 + head_size: int,

545 + rotary_dim: int,

546 + max_position_embeddings: int,

547 + base: float,

548 + is_neox_style: bool,

549 + scaling_alpha: float,

550 + dtype: torch.dtype,

551 + ) -> None:

552 + self.scaling_alpha = scaling_alpha

553 + super().__init__(head_size, rotary_dim, max_position_embeddings, base,

554 + is_neox_style, dtype)

555 +

556 + def _compute_cos_sin_cache(self) -> torch.Tensor:

557 + # For Hunyuan DynamicNTKAlphaRotaryEmbedding

558 + max_len = self.max_position_embeddings

559 + base = self.base * self.scaling_alpha**(self.rotary_dim /

560 + (self.rotary_dim - 2))

561 + inv_freq = self._compute_inv_freq(base)

562 + t = torch.arange(max_len, dtype=torch.float)

563 +

564 + freqs = torch.einsum("i,j -> ij", t, inv_freq)

565 + cos = freqs.cos()

566 + sin = freqs.sin()

567 + cache = torch.cat((cos, sin), dim=-1)

568 + return cache

569 +

570 +

571 # Inverse dim formula to find dim based on number of rotations

572 def _yarn_find_correction_dim(num_rotations: int,

573 dim: int,

@@ -1929,9 +1964,15 @@ def get_rope(

1964 mixed_b)

1965 elif scaling_type == "dynamic":

1966 scaling_factor = rope_scaling["factor"]

1932 - rotary_emb = DynamicNTKScalingRotaryEmbedding(

1933 - head_size, rotary_dim, max_position, base, is_neox_style,

1934 - scaling_factor, dtype)

1967 + scaling_alpha = rope_scaling["alpha"]

1968 + if scaling_alpha:

1969 + rotary_emb = DynamicNTKAlphaRotaryEmbedding(

1970 + head_size, rotary_dim, max_position, base, is_neox_style,

1971 + scaling_alpha, dtype)

1972 + else:

1973 + rotary_emb = DynamicNTKScalingRotaryEmbedding(

1974 + head_size, rotary_dim, max_position, base, is_neox_style,

1975 + scaling_factor, dtype)

1976 elif scaling_type == "yarn":

1977 scaling_factor = rope_scaling["factor"]

1978 original_max_position = rope_scaling[

vllm/model_executor/models/hunyuan_v1_moe.py new

+897

@@ -0,0 +1,897 @@

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 +from collections.abc import Iterable

27 +from typing import Any, Optional, Union

28 +

29 +import regex as re

30 +import torch

31 +from torch import nn

32 +from transformers import PretrainedConfig

33 +

34 +from vllm.attention import Attention, AttentionType

35 +from vllm.compilation.decorators import support_torch_compile

36 +from vllm.config import CacheConfig, VllmConfig

37 +from vllm.distributed import (get_pp_group,

38 + get_tensor_model_parallel_world_size,

39 + tensor_model_parallel_all_reduce)

40 +from vllm.model_executor.layers.activation import SiluAndMul

41 +from vllm.model_executor.layers.fused_moe import FusedMoE

42 +from vllm.model_executor.layers.layernorm import RMSNorm

43 +from vllm.model_executor.layers.linear import (ColumnParallelLinear,

44 + MergedColumnParallelLinear,

45 + QKVParallelLinear,

46 + ReplicatedLinear,

47 + RowParallelLinear)

48 +from vllm.model_executor.layers.logits_processor import LogitsProcessor

49 +from vllm.model_executor.layers.quantization.base_config import (

50 + QuantizationConfig)

51 +from vllm.model_executor.layers.rotary_embedding import get_rope

52 +from vllm.model_executor.layers.sampler import SamplerOutput, get_sampler

53 +from vllm.model_executor.layers.vocab_parallel_embedding import (

54 + DEFAULT_VOCAB_PADDING_SIZE, ParallelLMHead, VocabParallelEmbedding)

55 +from vllm.model_executor.model_loader.weight_utils import (

56 + default_weight_loader, maybe_remap_kv_scale_name)

57 +from vllm.model_executor.sampling_metadata import SamplingMetadata

58 +from vllm.sequence import IntermediateTensors

59 +

60 +from .utils import PPMissingLayer, is_pp_missing_parameter, make_layers

61 +

62 +

63 +def _get_cla_factor(config: PretrainedConfig) -> int:

64 + if not getattr(config, "use_cla", False):

65 + return 1

66 + return getattr(config, "cla_share_factor", 1)

67 +

68 +

69 +class HunYuanMLP(nn.Module):

70 +

71 + def __init__(

72 + self,

73 + hidden_size: int,

74 + intermediate_size: int,

75 + hidden_act: str,

76 + quant_config: Optional[QuantizationConfig] = None,

77 + bias: bool = False,

78 + prefix: str = "",

79 + reduce_results: bool = True,

80 + ) -> None:

81 + super().__init__()

82 + self.gate_up_proj = MergedColumnParallelLinear(

83 + input_size=hidden_size,

84 + output_sizes=[intermediate_size] * 2,

85 + bias=bias,

86 + quant_config=quant_config,

87 + prefix=f"{prefix}.gate_up_proj",

88 + )

89 + self.down_proj = RowParallelLinear(

90 + input_size=intermediate_size,

91 + output_size=hidden_size,

92 + bias=bias,

93 + quant_config=quant_config,

94 + prefix=f"{prefix}.down_proj",

95 + reduce_results=reduce_results,

96 + )

97 + if hidden_act != "silu":

98 + raise ValueError(f"Unsupported activation: {hidden_act}. "

99 + "Only silu is supported for now.")

100 + self.act_fn = SiluAndMul()

101 +

102 + def forward(self, x):

103 + gate_up, _ = self.gate_up_proj(x)

104 + x = self.act_fn(gate_up)

105 + x, _ = self.down_proj(x)

106 + return x

107 +

108 +

109 +class HunYuanAttention(nn.Module):

110 +

111 + def __init__(

112 + self,

113 + config: PretrainedConfig,

114 + hidden_size: int,

115 + num_heads: int,

116 + num_kv_heads: int,

117 + rope_theta: float = 10000,

118 + rope_scaling: Optional[dict[str, Any]] = None,

119 + max_position_embeddings: int = 8192,

120 + quant_config: Optional[QuantizationConfig] = None,

121 + bias: bool = False,

122 + cache_config: Optional[CacheConfig] = None,

123 + prefix: str = "",

124 + layer_id: int = -1,

125 + ) -> None:

126 + super().__init__()

127 + self.hidden_size = hidden_size

128 + tp_size = get_tensor_model_parallel_world_size()

129 + self.total_num_heads = num_heads

130 + assert self.total_num_heads % tp_size == 0

131 + self.num_heads = self.total_num_heads // tp_size

132 + self.total_num_kv_heads = num_kv_heads

133 + if self.total_num_kv_heads >= tp_size:

134 + # Number of KV heads is greater than TP size, so we partition

135 + # the KV heads across multiple tensor parallel GPUs.

136 + assert self.total_num_kv_heads % tp_size == 0

137 + else:

138 + # Number of KV heads is less than TP size, so we replicate

139 + # the KV heads across multiple tensor parallel GPUs.

140 + assert tp_size % self.total_num_kv_heads == 0

141 + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)

142 + # MistralConfig has an optional head_dim introduced by Mistral-Nemo

143 + if hasattr(config, "head_dim"):

144 + self.head_dim = config.head_dim

145 + elif hasattr(config, "attention_head_dim"):

146 + self.head_dim = config.attention_head_dim

147 + else:

148 + self.head_dim = self.hidden_size // self.total_num_heads

149 + self.q_size = self.num_heads * self.head_dim

150 + self.kv_size = self.num_kv_heads * self.head_dim

151 + self.scaling = self.head_dim**-0.5

152 + self.rope_theta = rope_theta

153 + self.max_position_embeddings = max_position_embeddings

154 + self.use_qk_norm = getattr(config, "use_qk_norm", False)

155 + self.layer_id = layer_id

156 +

157 + self.qkv_proj = QKVParallelLinear(

158 + hidden_size=hidden_size,

159 + head_size=self.head_dim,

160 + total_num_heads=self.total_num_heads,

161 + total_num_kv_heads=self.total_num_kv_heads,

162 + bias=bias,

163 + quant_config=quant_config,

164 + prefix=f"{prefix}.qkv_proj",

165 + )

166 +

167 + self.o_proj = RowParallelLinear(

168 + input_size=self.total_num_heads * self.head_dim,

169 + output_size=hidden_size,

170 + bias=bias,

171 + quant_config=quant_config,

172 + prefix=f"{prefix}.o_proj",

173 + )

174 +

175 + self.rotary_emb = get_rope(

176 + self.head_dim,

177 + rotary_dim=self.head_dim,

178 + max_position=max_position_embeddings,

179 + base=rope_theta,

180 + rope_scaling=rope_scaling,

181 + is_neox_style=True,

182 + )

183 + self.attn = Attention(

184 + self.num_heads,

185 + self.head_dim,

186 + self.scaling,

187 + num_kv_heads=self.num_kv_heads,

188 + cache_config=cache_config,

189 + quant_config=quant_config,

190 + prefix=f"{prefix}.attn",

191 + )

192 +

193 + if self.use_qk_norm:

194 + self.query_layernorm = RMSNorm(self.head_dim,

195 + eps=config.rms_norm_eps)

196 + self.key_layernorm = RMSNorm(self.head_dim,

197 + eps=config.rms_norm_eps)

198 +

199 + def forward(

200 + self,

201 + positions: torch.Tensor,

202 + hidden_states: torch.Tensor,

203 + kv_states: Optional[tuple[torch.Tensor]] = None,

204 + ) -> torch.Tensor:

205 + qkv, _ = self.qkv_proj(hidden_states)

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

207 + q, k = self.rotary_emb(positions, q, k)

208 + ori_k = k

209 + if self.use_qk_norm:

210 + q = self.query_layernorm(

211 + q.view(-1, self.num_heads, self.head_dim).contiguous())

212 + k = self.key_layernorm(

213 + k.view(-1, self.num_kv_heads, self.head_dim).contiguous())

214 +

215 + attn_output = self.attn(q, k, v)

216 + # For o_proj

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

218 + output, _ = self.o_proj(attn_output)

219 + return output, (ori_k, v)

220 +

221 +

222 +class HunYuanCrossAttention(nn.Module):

223 +

224 + def __init__(

225 + self,

226 + config: PretrainedConfig,

227 + hidden_size: int,

228 + num_heads: int,

229 + num_kv_heads: int,

230 + rope_theta: float = 10000,

231 + rope_scaling: Optional[dict[str, Any]] = None,

232 + max_position_embeddings: int = 8192,

233 + quant_config: Optional[QuantizationConfig] = None,

234 + bias: bool = False,

235 + cache_config: Optional[CacheConfig] = None,

236 + prefix: str = "",

237 + layer_id: int = -1,

238 + ) -> None:

239 + super().__init__()

240 + self.hidden_size = hidden_size

241 + tp_size = get_tensor_model_parallel_world_size()

242 + self.total_num_heads = num_heads

243 + assert self.total_num_heads % tp_size == 0

244 + self.num_heads = self.total_num_heads // tp_size

245 + self.total_num_kv_heads = num_kv_heads

246 + if self.total_num_kv_heads >= tp_size:

247 + # Number of KV heads is greater than TP size, so we partition

248 + # the KV heads across multiple tensor parallel GPUs.

249 + assert self.total_num_kv_heads % tp_size == 0

250 + else:

251 + # Number of KV heads is less than TP size, so we replicate

252 + # the KV heads across multiple tensor parallel GPUs.

253 + assert tp_size % self.total_num_kv_heads == 0

254 + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)

255 + # MistralConfig has an optional head_dim introduced by Mistral-Nemo

256 + if hasattr(config, "head_dim"):

257 + self.head_dim = config.head_dim

258 + elif hasattr(config, "attention_head_dim"):

259 + self.head_dim = config.attention_head_dim

260 + else:

261 + self.head_dim = self.hidden_size // self.total_num_heads

262 + self.q_size = self.num_heads * self.head_dim

263 + self.kv_size = self.num_kv_heads * self.head_dim

264 + self.scaling = self.head_dim**-0.5

265 + self.rope_theta = rope_theta

266 + self.max_position_embeddings = max_position_embeddings

267 + self.use_qk_norm = getattr(config, "use_qk_norm", False)

268 + self.layer_id = layer_id

269 +

270 + self.q_proj = ColumnParallelLinear(

271 + hidden_size,

272 + hidden_size,

273 + bias=bias,

274 + quant_config=quant_config,

275 + prefix=f"{prefix}.q_proj",

276 + )

277 +

278 + self.o_proj = RowParallelLinear(

279 + input_size=self.total_num_heads * self.head_dim,

280 + output_size=hidden_size,

281 + bias=bias,

282 + quant_config=quant_config,

283 + prefix=f"{prefix}.o_proj",

284 + )

285 +

286 + self.rotary_emb = get_rope(

287 + self.head_dim,

288 + rotary_dim=self.head_dim,

289 + max_position=max_position_embeddings,

290 + base=rope_theta,

291 + rope_scaling=rope_scaling,

292 + is_neox_style=True,

293 + )

294 + self.attn = Attention(

295 + self.num_heads,

296 + self.head_dim,

297 + self.scaling,

298 + num_kv_heads=self.num_kv_heads,

299 + cache_config=cache_config,

300 + quant_config=quant_config,

301 + prefix=f"{prefix}.attn",

302 + attn_type=AttentionType.ENCODER_DECODER,

303 + )

304 +

305 + if self.use_qk_norm:

306 + self.query_layernorm = RMSNorm(self.head_dim,

307 + eps=config.rms_norm_eps)

308 + self.key_layernorm = RMSNorm(self.head_dim,

309 + eps=config.rms_norm_eps)

310 +

311 + def forward(

312 + self,

313 + positions: torch.Tensor,

314 + hidden_states: torch.Tensor,

315 + kv_states: Optional[tuple[torch.Tensor]] = None,

316 + ) -> torch.Tensor:

317 + assert kv_states is not None

318 + ori_k, v = kv_states # use last layer kv,

319 + k = ori_k

320 + q, _ = self.q_proj(hidden_states)

321 + k_tmp = torch.empty_like(k) # Todo: reduant rotary embedding

322 + q, _ = self.rotary_emb(positions, q, k_tmp)

323 + if self.use_qk_norm:

324 + q = self.query_layernorm(

325 + q.view(-1, self.num_heads, self.head_dim).contiguous())

326 + k = self.key_layernorm(

327 + k.view(-1, self.num_kv_heads, self.head_dim).contiguous())

328 +

329 + attn_output = self.attn(q, k, v)

330 + # For o_proj

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

332 + output, _ = self.o_proj(attn_output)

333 + return output, (ori_k, v)

334 +

335 +

336 +class HunYuanSparseMoeBlock(nn.Module):

337 +

338 + def __init__(

339 + self,

340 + config: PretrainedConfig,

341 + quant_config: Optional[QuantizationConfig] = None,

342 + layer_id: int = -1,

343 + prefix: str = "",

344 + ):

345 + super().__init__()

346 + self.tp_size = get_tensor_model_parallel_world_size()

347 +

348 + if self.tp_size > config.num_experts:

349 + raise ValueError(

350 + f"Tensor parallel size {self.tp_size} is greater than "

351 + f"the number of experts {config.num_experts}.")

352 +

353 + # Get layer_id topk if config.moe_topk is a list

354 + if isinstance(config.moe_topk, list):

355 + assert layer_id >= 0

356 + assert len(config.moe_topk) > layer_id

357 + top_k = config.moe_topk[layer_id]

358 + else:

359 + top_k = config.moe_topk

360 +

361 + # If it is moe, moe_intermediate_size is preferred

362 + intermediate_size = config.intermediate_size

363 + if config.moe_intermediate_size is not None:

364 + intermediate_size = (config.moe_intermediate_size if isinstance(

365 + config.moe_intermediate_size, int) else

366 + config.moe_intermediate_size[layer_id])

367 +

368 + self.experts = FusedMoE(

369 + num_experts=config.num_experts,

370 + top_k=top_k,

371 + hidden_size=config.hidden_size,

372 + intermediate_size=intermediate_size,

373 + reduce_results=False,

374 + renormalize=top_k > 1,

375 + quant_config=quant_config,

376 + prefix=f"{prefix}.experts",

377 + )

378 +

379 + self.gate = ReplicatedLinear(config.hidden_size,

380 + config.num_experts,

381 + bias=False,

382 + quant_config=None,

383 + prefix=f"{prefix}.gate")

384 + if config.use_mixed_mlp_moe > 0:

385 + # Get layer_id num_shared_expert if config.num_shared_expert is

386 + # a list.

387 + if isinstance(config.num_shared_expert, list):

388 + assert layer_id >= 0

389 + assert len(config.num_shared_expert) > layer_id

390 + num_shared_expert = config.num_shared_expert[layer_id]

391 + else:

392 + num_shared_expert = config.num_shared_expert

393 +

394 + self.shared_mlp = HunYuanMLP(

395 + hidden_size=config.hidden_size,

396 + intermediate_size=config.intermediate_size * num_shared_expert,

397 + hidden_act=config.hidden_act,

398 + quant_config=quant_config,

399 + reduce_results=False,

400 + )

401 + else:

402 + self.shared_mlp = None

403 +

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

405 + # NOTE: hidden_states can have either 1D or 2D shape.

406 + orig_shape = hidden_states.shape

407 + hidden_dim = hidden_states.shape[-1]

408 + hidden_states = hidden_states.view(-1, hidden_dim)

409 + shared_output = None

410 + if self.shared_mlp is not None:

411 + shared_output = self.shared_mlp(hidden_states)

412 +

413 + # router_logits: (num_tokens, n_experts)

414 + router_logits, _ = self.gate(hidden_states)

415 + final_hidden_states = self.experts(hidden_states=hidden_states,

416 + router_logits=router_logits)

417 + if shared_output is not None:

418 + final_hidden_states = final_hidden_states + shared_output

419 + if self.tp_size > 1:

420 + final_hidden_states = tensor_model_parallel_all_reduce(

421 + final_hidden_states)

422 +

423 + return final_hidden_states.view(orig_shape)

424 +

425 +

426 +class HunYuanDecoderLayer(nn.Module):

427 +

428 + def __init__(

429 + self,

430 + config: PretrainedConfig,

431 + cache_config: Optional[CacheConfig] = None,

432 + quant_config: Optional[QuantizationConfig] = None,

433 + prefix: str = "",

434 + layer_id: int = -1,

435 + ) -> None:

436 + super().__init__()

437 + assert layer_id >= 0

438 + self.layer_id = layer_id

439 + self.hidden_size = config.hidden_size

440 + self.intermediate_size = (config.intermediate_size if isinstance(

441 + config.intermediate_size, int) else

442 + config.intermediate_size[layer_id])

443 + rope_theta = getattr(config, "rope_theta", 10000)

444 + rope_scaling = getattr(config, "rope_scaling", None)

445 + if rope_scaling is not None and getattr(

446 + config, "original_max_position_embeddings", None):

447 + rope_scaling["original_max_position_embeddings"] = (

448 + config.original_max_position_embeddings)

449 + max_position_embeddings = getattr(config, "max_position_embeddings",

450 + 8192)

451 + attention_bias = getattr(config, "attention_bias", False) or getattr(

452 + config, "bias", False)

453 + cla_factor = _get_cla_factor(config)

454 + attention_type = (AttentionType.ENCODER_DECODER

455 + if layer_id >= 0 and layer_id % cla_factor != 0 else

456 + AttentionType.DECODER)

457 + if attention_type == AttentionType.DECODER:

458 + self.self_attn = HunYuanAttention(

459 + config=config,

460 + hidden_size=self.hidden_size,

461 + num_heads=config.num_attention_heads,

462 + num_kv_heads=getattr(config, "num_key_value_heads",

463 + config.num_attention_heads),

464 + rope_theta=rope_theta,

465 + rope_scaling=rope_scaling,

466 + max_position_embeddings=max_position_embeddings,

467 + quant_config=quant_config,

468 + bias=attention_bias,

469 + cache_config=cache_config,

470 + prefix=f"{prefix}.self_attn",

471 + layer_id=layer_id,

472 + )

473 + elif attention_type == AttentionType.ENCODER_DECODER:

474 + self.self_attn = HunYuanCrossAttention(

475 + config=config,

476 + hidden_size=self.hidden_size,

477 + num_heads=config.num_attention_heads,

478 + num_kv_heads=getattr(config, "num_key_value_heads",

479 + config.num_attention_heads),

480 + rope_theta=rope_theta,

481 + rope_scaling=rope_scaling,

482 + max_position_embeddings=max_position_embeddings,

483 + quant_config=quant_config,

484 + bias=attention_bias,

485 + cache_config=cache_config,

486 + prefix=f"{prefix}.self_attn",

487 + layer_id=layer_id,

488 + )

489 + else:

490 + raise RuntimeError(f"Unsupported attention type: {attention_type}")

491 +

492 + self.mlp = HunYuanSparseMoeBlock(

493 + config=config,

494 + quant_config=quant_config,

495 + layer_id=layer_id,

496 + prefix=f"{prefix}.mlp",

497 + )

498 + self.input_layernorm = RMSNorm(config.hidden_size,

499 + eps=config.rms_norm_eps)

500 + self.post_attention_layernorm = RMSNorm(config.hidden_size,

501 + eps=config.rms_norm_eps)

502 +

503 + def forward(

504 + self,

505 + positions: torch.Tensor,

506 + hidden_states: torch.Tensor,

507 + residual: Optional[torch.Tensor],

508 + kv_states: Optional[tuple[torch.Tensor]] = None,

509 + ) -> tuple[torch.Tensor, torch.Tensor]:

510 + # Self Attention

511 + if residual is None:

512 + residual = hidden_states

513 + hidden_states = self.input_layernorm(hidden_states)

514 + else:

515 + hidden_states, residual = self.input_layernorm(

516 + hidden_states, residual)

517 + hidden_states, ori_kv_states = self.self_attn(

518 + positions=positions,

519 + hidden_states=hidden_states,

520 + kv_states=kv_states,

521 + )

522 +

523 + # Fully Connected

524 + hidden_states, residual = self.post_attention_layernorm(

525 + hidden_states, residual)

526 + hidden_states = self.mlp(hidden_states)

527 + return hidden_states, residual, ori_kv_states

528 +

529 +

530 +@support_torch_compile

531 +class HunYuanModel(nn.Module):

532 +

533 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

534 + super().__init__()

535 +

536 + config = vllm_config.model_config.hf_config

537 + cache_config = vllm_config.cache_config

538 + quant_config = vllm_config.quant_config

539 + lora_config = vllm_config.lora_config

540 +

541 + self.config = config

542 + self.quant_config = quant_config

543 + self.padding_idx = config.pad_token_id

544 + lora_vocab = ((lora_config.lora_extra_vocab_size *

545 + (lora_config.max_loras or 1)) if lora_config else 0)

546 + self.vocab_size = config.vocab_size + lora_vocab

547 + self.org_vocab_size = config.vocab_size

548 + if get_pp_group().is_first_rank or (config.tie_word_embeddings

549 + and get_pp_group().is_last_rank):

550 + self.embed_tokens = VocabParallelEmbedding(

551 + self.vocab_size,

552 + config.hidden_size,

553 + org_num_embeddings=config.vocab_size,

554 + quant_config=quant_config,

555 + )

556 + else:

557 + self.embed_tokens = PPMissingLayer()

558 + self.start_layer, self.end_layer, self.layers = make_layers(

559 + config.num_hidden_layers,

560 + lambda prefix: HunYuanDecoderLayer(

561 + config=config,

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

563 + cache_config=cache_config,

564 + quant_config=quant_config,

565 + prefix=prefix,

566 + ),

567 + prefix=f"{prefix}.layers",

568 + )

569 + if get_pp_group().is_last_rank:

570 + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

571 + else:

572 + self.norm = PPMissingLayer()

573 +

574 + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:

575 + return self.embed_tokens(input_ids)

576 +

577 + def forward(

578 + self,

579 + input_ids: Optional[torch.Tensor],

580 + positions: torch.Tensor,

581 + intermediate_tensors: Optional[IntermediateTensors],

582 + inputs_embeds: Optional[torch.Tensor] = None,

583 + ) -> Union[torch.Tensor, IntermediateTensors]:

584 + if get_pp_group().is_first_rank:

585 + if inputs_embeds is not None:

586 + hidden_states = inputs_embeds

587 + else:

588 + hidden_states = self.get_input_embeddings(input_ids)

589 + residual = None

590 + else:

591 + assert intermediate_tensors is not None

592 + hidden_states = intermediate_tensors["hidden_states"]

593 + residual = intermediate_tensors["residual"]

594 +

595 + cla_factor = _get_cla_factor(self.config)

596 + prev_kv_states = None

597 + for i in range(self.start_layer, self.end_layer):

598 + layer = self.layers[i]

599 + hidden_states, residual, kv_states = layer(

600 + positions,

601 + hidden_states,

602 + residual,

603 + prev_kv_states,

604 + )

605 +

606 + if (getattr(self.config, "use_cla", False)

607 + and (i - self.start_layer) % cla_factor == 0):

608 + prev_kv_states = kv_states

609 + else:

610 + prev_kv_states = None

611 +

612 + if not get_pp_group().is_last_rank:

613 + return IntermediateTensors({

614 + "hidden_states": hidden_states,

615 + "residual": residual

616 + })

617 +

618 + hidden_states, _ = self.norm(hidden_states, residual)

619 + return hidden_states

620 +

621 +

622 +class HunYuanMoEV1ForCausalLM(nn.Module):

623 + packed_modules_mapping = {

624 + "qkv_proj": [

625 + "q_proj",

626 + "k_proj",

627 + "v_proj",

628 + ],

629 + "gate_up_proj": [

630 + "gate_proj",

631 + "up_proj",

632 + ],

633 + }

634 +

635 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

636 + super().__init__()

637 +

638 + config = vllm_config.model_config.hf_config

639 + quant_config = vllm_config.quant_config

640 + lora_config = vllm_config.lora_config

641 + self.config = config

642 + self.quant_config = quant_config

643 + self.lora_config = lora_config

644 +

645 + self.model = HunYuanModel(vllm_config=vllm_config, prefix="model")

646 + if get_pp_group().is_last_rank:

647 + self.unpadded_vocab_size = config.vocab_size

648 + if lora_config:

649 + self.unpadded_vocab_size += lora_config.lora_extra_vocab_size

650 + self.lm_head = ParallelLMHead(

651 + self.unpadded_vocab_size,

652 + config.hidden_size,

653 + org_num_embeddings=config.vocab_size,

654 + padding_size=DEFAULT_VOCAB_PADDING_SIZE,

655 + quant_config=quant_config,

656 + )

657 + if config.tie_word_embeddings:

658 + self.lm_head.weight = self.model.embed_tokens.weight

659 +

660 + logit_scale = getattr(config, "logit_scale", 1.0)

661 + self.logits_processor = LogitsProcessor(self.unpadded_vocab_size,

662 + config.vocab_size,

663 + logit_scale)

664 + self.sampler = get_sampler()

665 + else:

666 + self.lm_head = PPMissingLayer()

667 +

668 + def forward(

669 + self,

670 + input_ids: torch.Tensor,

671 + positions: torch.Tensor,

672 + intermediate_tensors: Optional[IntermediateTensors] = None,

673 + inputs_embeds: Optional[torch.Tensor] = None,

674 + ) -> Union[torch.Tensor, IntermediateTensors]:

675 + model_output = self.model(input_ids, positions, intermediate_tensors,

676 + inputs_embeds)

677 + return model_output

678 +

679 + def compute_logits(

680 + self,

681 + hidden_states: torch.Tensor,

682 + sampling_metadata: SamplingMetadata,

683 + ) -> Optional[torch.Tensor]:

684 + logits = self.logits_processor(self.lm_head, hidden_states,

685 + sampling_metadata)

686 + return logits

687 +

688 + def sample(

689 + self,

690 + logits: torch.Tensor,

691 + sampling_metadata: SamplingMetadata,

692 + ) -> Optional[SamplerOutput]:

693 + next_tokens = self.sampler(logits, sampling_metadata)

694 + return next_tokens

695 +

696 + def make_empty_intermediate_tensors(

697 + self, batch_size: int, dtype: torch.dtype,

698 + device: torch.device) -> IntermediateTensors:

699 + return IntermediateTensors({

700 + "hidden_states":

701 + torch.zeros((batch_size, self.config.hidden_size),

702 + dtype=dtype,

703 + device=device),

704 + "residual":

705 + torch.zeros((batch_size, self.config.hidden_size),

706 + dtype=dtype,

707 + device=device),

708 + })

709 +

710 + def _split_qkv_weight(self, qkv: torch.Tensor):

711 + num_attention_heads = self.config.num_attention_heads

712 + num_kv_heads = getattr(self.config, "num_key_value_heads",

713 + self.config.num_attention_heads)

714 + num_key_value_groups = num_attention_heads // num_kv_heads

715 + hidden_size = self.config.hidden_size

716 +

717 + if hasattr(self.config, "head_dim"):

718 + attention_head_dim = self.config.head_dim

719 + elif hasattr(self.config, "attention_head_dim"):

720 + attention_head_dim = self.config.attention_head_dim

721 + else:

722 + attention_head_dim = self.config.hidden_size // num_attention_heads

723 +

724 + qkv = qkv.reshape(num_kv_heads, num_key_value_groups + 2,

725 + attention_head_dim, hidden_size)

726 + q, k, v = torch.split(qkv, (num_key_value_groups, 1, 1), dim=1)

727 + q = q.reshape(-1, hidden_size)

728 + k = k.reshape(-1, hidden_size)

729 + v = v.reshape(-1, hidden_size)

730 + return torch.concat((q, k, v))

731 +

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

733 + cla_factor = _get_cla_factor(self.config)

734 + stacked_params_mapping = [

735 + # (param_name, shard_name, shard_id)

736 + (".qkv_proj", ".q_proj", "q"),

737 + (".qkv_proj", ".k_proj", "k"),

738 + (".qkv_proj", ".v_proj", "v"),

739 + (".gate_up_proj", ".gate_proj", 0),

740 + (".gate_up_proj", ".up_proj", 1),

741 + ]

742 +

743 + num_attention_heads = self.config.num_attention_heads

744 + num_kv_heads = getattr(self.config, "num_key_value_heads",

745 + self.config.num_attention_heads)

746 + split_params_mapping = [

747 + (".gate_up_proj", ".gate_and_up_proj", 2, [(1, 1), (0, 1)], None),

748 + (

749 + ".qkv_proj",

750 + ".qkv_proj",

751 + num_attention_heads + num_kv_heads * 2,

752 + [("q", num_attention_heads), ("k", num_kv_heads),

753 + ("v", num_kv_heads)],

754 + self._split_qkv_weight,

755 + ),

756 + ]

757 +

758 + # Params for weights, fp8 weight scales, fp8 activation scales

759 + # (param_name, weight_name, expert_id, shard_id)

760 + expert_params_mapping = FusedMoE.make_expert_params_mapping(

761 + ckpt_gate_proj_name="gate_proj",

762 + ckpt_down_proj_name="down_proj",

763 + ckpt_up_proj_name="up_proj",

764 + num_experts=self.config.num_experts,

765 + )

766 +

767 + params_dict = dict(self.named_parameters())

768 + for name, loaded_weight in weights:

769 + if "rotary_emb.inv_freq" in name:

770 + continue

771 + if "gate_proj_bias" in name:

772 + name = name.replace("gate_proj_bias", "gate_proj.bias")

773 + if "up_proj_bias" in name:

774 + name = name.replace("up_proj_bias", "up_proj.bias")

775 + if ("rotary_emb.cos_cached" in name

776 + or "rotary_emb.sin_cached" in name):

777 + # Models trained using ColossalAI may include these tensors in

778 + # the checkpoint. Skip them.

779 + continue

780 + # With tie_word_embeddings, we can skip lm_head.weight

781 + # The weight might appear unnecessarily in the files if the model is

782 + # processed with quantization, LoRA, fine-tuning, etc.

783 + if self.config.tie_word_embeddings and "lm_head.weight" in name:

784 + continue

785 + if self.quant_config is not None and (

786 + scale_name := self.quant_config.get_cache_scale(name)):

787 + # Loading kv cache scales for compressed-tensors quantization

788 + param = params_dict[scale_name]

789 + weight_loader = getattr(param, "weight_loader",

790 + default_weight_loader)

791 + loaded_weight = loaded_weight[0]

792 + weight_loader(param, loaded_weight)

793 + continue

794 +

795 + is_found = False

796 + for param_name, weight_name, shard_id in stacked_params_mapping:

797 + if weight_name not in name:

798 + continue

799 + if "mlp.experts" in name:

800 + continue

801 + # cross layer only have q_proj, skip qkv pack

802 + if weight_name == ".q_proj":

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

804 + if match:

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

806 + if cla_factor > 1 and layer_id % cla_factor != 0:

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 + param = params_dict[name]

817 + weight_loader = param.weight_loader

818 + weight_loader(param, loaded_weight, shard_id)

819 +

820 + is_found = True

821 + break

822 + if is_found:

823 + continue

824 +

825 + for (

826 + param_name,

827 + weight_name,

828 + den,

829 + split_param,

830 + func,

831 + ) in split_params_mapping:

832 + if weight_name not in name:

833 + continue

834 + name = name.replace(weight_name, param_name)

835 + # Skip loading extra bias for GPTQ models.

836 + if name.endswith(".bias") and name not in params_dict:

837 + continue

838 +

839 + if is_pp_missing_parameter(name, self):

840 + continue

841 +

842 + assert loaded_weight.shape[0] % den == 0

843 + units = loaded_weight.shape[0] // den

844 +

845 + param = params_dict[name]

846 + weight_loader = param.weight_loader

847 + offset = 0

848 + for shard_id, num in split_param:

849 + new_offset = offset + num * units

850 + if func:

851 + weight_loader(param,

852 + func(loaded_weight)[offset:new_offset],

853 + shard_id)

854 + else:

855 + weight_loader(param, loaded_weight[offset:new_offset],

856 + shard_id)

857 + offset = new_offset

858 +

859 + break

860 + else:

861 + # Skip loading extra bias for GPTQ models.

862 + if name.endswith(".bias") and name not in params_dict:

863 + continue

864 + for mapping in expert_params_mapping:

865 + param_name, weight_name, expert_id, shard_id = mapping

866 + if weight_name not in name:

867 + continue

868 + name = name.replace(weight_name, param_name)

869 + # Skip layers on other devices.

870 + if is_pp_missing_parameter(name, self):

871 + continue

872 + param = params_dict[name]

873 + weight_loader = param.weight_loader

874 + weight_loader(

875 + param,

876 + loaded_weight,

877 + name,

878 + shard_id=shard_id,

879 + expert_id=expert_id,

880 + )

881 + break

882 + else:

883 + # Remapping the name of FP8 kv-scale.

884 + name = maybe_remap_kv_scale_name(name, params_dict)

885 + if name is None:

886 + continue

887 +

888 + if is_pp_missing_parameter(name, self):

889 + continue

890 +

891 + if "mlp.gate.wg." in name:

892 + name = name.replace("wg.", "")

893 +

894 + param = params_dict[name]

895 + weight_loader = getattr(param, "weight_loader",

896 + default_weight_loader)

897 + weight_loader(param, loaded_weight)

vllm/model_executor/models/registry.py

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

73 "GraniteMoeSharedForCausalLM": ("granitemoeshared", "GraniteMoeSharedForCausalLM"), # noqa: E501

74 "GritLM": ("gritlm", "GritLM"),

75 "Grok1ModelForCausalLM": ("grok1", "Grok1ForCausalLM"),

76 + "HunYuanMoEV1ForCausalLM": ("hunyuan_v1_moe", "HunYuanMoEV1ForCausalLM"),

77 "InternLMForCausalLM": ("llama", "LlamaForCausalLM"),

78 "InternLM2ForCausalLM": ("internlm2", "InternLM2ForCausalLM"),

79 "InternLM2VEForCausalLM": ("internlm2_ve", "InternLM2VEForCausalLM"),