@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 38efa28278
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, 02:15 UTC

[Model] Add Ling implementation (#20680)

Signed-off-by: vito.yy <vito.yy@antgroup.com>

ant-yy committed Jul 14, 2025 at 22:10 UTC 38efa28278b4accf8eb2a7258f9f999fdbdd9f63
4 files changed +534
docs/models/supported_models.md
+1
@@ -316,6 +316,7 @@ Specified using `--task generate`.
316 | `AquilaForCausalLM` | Aquila, Aquila2 | `BAAI/Aquila-7B`, `BAAI/AquilaChat-7B`, etc. | ✅︎ | ✅︎ | ✅︎ |
317 | `ArcticForCausalLM` | Arctic | `Snowflake/snowflake-arctic-base`, `Snowflake/snowflake-arctic-instruct`, etc. | | ✅︎ | ✅︎ |
318 | `BaiChuanForCausalLM` | Baichuan2, Baichuan | `baichuan-inc/Baichuan2-13B-Chat`, `baichuan-inc/Baichuan-7B`, etc. | ✅︎ | ✅︎ | ✅︎ |
319 +| `BailingMoeForCausalLM` | Ling | `inclusionAI/Ling-lite-1.5`, `inclusionAI/Ling-plus`, etc. | | ✅︎ | ✅︎ |
320 | `BambaForCausalLM` | Bamba | `ibm-ai-platform/Bamba-9B-fp8`, `ibm-ai-platform/Bamba-9B` | ✅︎ | ✅︎ | ✅︎ |
321 | `BloomForCausalLM` | BLOOM, BLOOMZ, BLOOMChat | `bigscience/bloom`, `bigscience/bloomz`, etc. | | ✅︎ | |
322 | `BartForConditionalGeneration` | BART | `facebook/bart-base`, `facebook/bart-large-cnn`, etc. | | | |
tests/models/registry.py
+2
@@ -141,6 +141,8 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
141 trust_remote_code=True),
142 "BaichuanForCausalLM": _HfExamplesInfo("baichuan-inc/Baichuan2-7B-chat",
143 trust_remote_code=True),
144 + "BailingMoeForCausalLM": _HfExamplesInfo("inclusionAI/Ling-lite-1.5",
145 + trust_remote_code=True),
146 "BambaForCausalLM": _HfExamplesInfo("ibm-ai-platform/Bamba-9B",
147 extras={"tiny": "hmellor/tiny-random-BambaForCausalLM"}), # noqa: E501
148 "BloomForCausalLM": _HfExamplesInfo("bigscience/bloom-560m",
vllm/model_executor/models/bailing_moe.py new
+530
@@ -0,0 +1,530 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +# Adapted from
5 +# https://github.com/inclusionAI/Ling/blob/master/models/modeling_bailing_moe.py
6 +# Copyright 2023 The vLLM team.
7 +# Copyright 2023 Antgroup and The HuggingFace Inc. team. All rights reserved.
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 BailingMoE model compatible with HuggingFace weights."""
26 +from collections.abc import Iterable
27 +from typing import Optional, Union
28 +
29 +import torch
30 +import torch.nn.functional as F
31 +from torch import nn
32 +from transformers.configuration_utils import PretrainedConfig
33 +
34 +from vllm.attention import Attention
35 +from vllm.config import CacheConfig, VllmConfig
36 +from vllm.distributed import (get_pp_group, get_tensor_model_parallel_rank,
37 + get_tensor_model_parallel_world_size,
38 + tensor_model_parallel_all_reduce)
39 +from vllm.model_executor.layers.activation import SiluAndMul
40 +from vllm.model_executor.layers.fused_moe import FusedMoE
41 +from vllm.model_executor.layers.layernorm import RMSNorm
42 +from vllm.model_executor.layers.linear import (MergedColumnParallelLinear,
43 + QKVParallelLinear,
44 + ReplicatedLinear,
45 + RowParallelLinear)
46 +from vllm.model_executor.layers.logits_processor import LogitsProcessor
47 +from vllm.model_executor.layers.quantization.base_config import (
48 + QuantizationConfig)
49 +from vllm.model_executor.layers.rotary_embedding import get_rope
50 +from vllm.model_executor.layers.sampler import SamplerOutput, get_sampler
51 +from vllm.model_executor.layers.vocab_parallel_embedding import (
52 + ParallelLMHead, VocabParallelEmbedding)
53 +from vllm.model_executor.model_loader.weight_utils import default_weight_loader
54 +from vllm.model_executor.sampling_metadata import SamplingMetadata
55 +from vllm.sequence import IntermediateTensors
56 +
57 +from .interfaces import SupportsPP
58 +from .utils import (AutoWeightsLoader, PPMissingLayer, is_pp_missing_parameter,
59 + make_empty_intermediate_tensors_factory, make_layers,
60 + maybe_prefix)
61 +
62 +
63 +class BailingAttention(nn.Module):
64 +
65 + def __init__(
66 + self,
67 + config: PretrainedConfig,
68 + cache_config: Optional[CacheConfig] = None,
69 + quant_config: Optional[QuantizationConfig] = None,
70 + prefix: str = "",
71 + ):
72 + super().__init__()
73 + self.hidden_size = config.hidden_size
74 + self.total_num_heads = config.num_attention_heads
75 + self.total_kv_heads = config.num_key_value_heads
76 + tp_size = get_tensor_model_parallel_world_size()
77 +
78 + assert self.total_num_heads % tp_size == 0
79 + assert self.total_kv_heads % tp_size == 0
80 + assert self.total_num_heads >= self.total_kv_heads
81 +
82 + self.num_heads = self.total_num_heads // tp_size
83 + self.head_dim = config.head_dim or (self.hidden_size //
84 + self.total_num_heads)
85 + self.q_size_per_rank = self.head_dim * self.num_heads
86 +
87 + self.num_kv_heads = self.total_kv_heads // tp_size
88 + self.kv_size_per_rank = self.num_kv_heads * self.head_dim
89 + self.scale = self.head_dim**-0.5
90 +
91 + self.query_key_value = QKVParallelLinear(
92 + self.hidden_size,
93 + self.head_dim,
94 + self.total_num_heads,
95 + self.total_kv_heads,
96 + bias=(config.use_bias or config.use_qkv_bias),
97 + quant_config=quant_config,
98 + prefix=f"{prefix}.query_key_value",
99 + )
100 +
101 + self.dense = RowParallelLinear(
102 + self.total_num_heads * self.head_dim,
103 + self.hidden_size,
104 + bias=config.use_bias,
105 + quant_config=quant_config,
106 + prefix=f"{prefix}.dense",
107 + )
108 +
109 + self.attn = Attention(self.num_heads,
110 + self.head_dim,
111 + self.scale,
112 + num_kv_heads=self.num_kv_heads,
113 + cache_config=cache_config,
114 + prefix=f"{prefix}.attn")
115 +
116 + self.rotary_emb = get_rope(
117 + self.head_dim,
118 + rotary_dim=self.head_dim,
119 + max_position=config.max_position_embeddings,
120 + base=config.rope_theta,
121 + is_neox_style=True,
122 + rope_scaling=config.rope_scaling,
123 + )
124 +
125 + def forward(
126 + self,
127 + hidden_states: torch.Tensor,
128 + position_ids: torch.Tensor,
129 + ) -> torch.Tensor:
130 +
131 + qkv, _ = self.query_key_value(hidden_states)
132 + q, k, v = qkv.split([
133 + self.q_size_per_rank, self.kv_size_per_rank, self.kv_size_per_rank
134 + ],
135 + dim=-1)
136 +
137 + q, k = self.rotary_emb(position_ids, q, k)
138 +
139 + context_layer = self.attn(q, k, v)
140 +
141 + attn_output, _ = self.dense(context_layer)
142 + return attn_output
143 +
144 +
145 +class BailingMLP(nn.Module):
146 +
147 + def __init__(
148 + self,
149 + intermediate_size: int,
150 + config: PretrainedConfig,
151 + quant_config: Optional[QuantizationConfig] = None,
152 + reduce_results: Optional[bool] = True,
153 + prefix: str = "",
154 + ) -> None:
155 + super().__init__()
156 + self.gate_up_proj = MergedColumnParallelLinear(
157 + config.hidden_size,
158 + [intermediate_size] * 2,
159 + bias=config.use_bias,
160 + quant_config=quant_config,
161 + prefix=f"{prefix}.gate_up_proj",
162 + )
163 + self.down_proj = RowParallelLinear(
164 + intermediate_size,
165 + config.hidden_size,
166 + bias=config.use_bias,
167 + quant_config=quant_config,
168 + reduce_results=reduce_results,
169 + prefix=f"{prefix}.down_proj",
170 + )
171 + self.act_fn = SiluAndMul()
172 +
173 + def forward(self, x):
174 + x, _ = self.gate_up_proj(x)
175 + x = self.act_fn(x)
176 + x, _ = self.down_proj(x)
177 + return x
178 +
179 +
180 +class BailingMoE(nn.Module):
181 +
182 + def __init__(
183 + self,
184 + intermediate_size: int,
185 + config: PretrainedConfig,
186 + quant_config: Optional[QuantizationConfig] = None,
187 + reduce_results: Optional[bool] = True,
188 + prefix: str = "",
189 + ):
190 + super().__init__()
191 +
192 + self.tp_size = get_tensor_model_parallel_world_size()
193 + self.tp_rank = get_tensor_model_parallel_rank()
194 + self.num_experts = config.num_experts
195 + self.top_k = config.num_experts_per_tok
196 + self.norm_expert_prob = config.norm_topk_prob
197 + self.hidden_size = config.hidden_size
198 + self.quant_config = quant_config
199 + self.num_shared_experts = config.num_shared_experts
200 + # Gate always runs at half / full precision for now.
201 + self.gate = ReplicatedLinear(self.hidden_size,
202 + self.num_experts,
203 + bias=False,
204 + quant_config=None)
205 +
206 + self.experts = FusedMoE(num_experts=self.num_experts,
207 + top_k=self.top_k,
208 + hidden_size=self.hidden_size,
209 + intermediate_size=config.moe_intermediate_size,
210 + reduce_results=False,
211 + renormalize=self.norm_expert_prob,
212 + quant_config=quant_config,
213 + prefix=f"{prefix}.experts")
214 +
215 + if self.num_shared_experts > 0:
216 + intermediate_size = (config.moe_intermediate_size *
217 + self.num_shared_experts)
218 + self.shared_experts = BailingMLP(
219 + intermediate_size=intermediate_size,
220 + config=config,
221 + quant_config=quant_config,
222 + reduce_results=False,
223 + prefix=f"{prefix}.shared_experts")
224 + else:
225 + self.shared_experts = None
226 +
227 + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
228 + num_tokens, hidden_size = hidden_states.shape
229 + hidden_states = hidden_states.view(-1, hidden_size)
230 + if self.num_shared_experts > 0:
231 + shared_output = self.shared_experts(hidden_states)
232 + # router_logits: (num_tokens, n_experts)
233 + router_logits, _ = self.gate(hidden_states)
234 + final_hidden_states = self.experts(hidden_states=hidden_states,
235 + router_logits=router_logits)
236 +
237 + if self.num_shared_experts > 0:
238 + final_hidden_states = final_hidden_states + shared_output
239 +
240 + if self.tp_size > 1:
241 + final_hidden_states = tensor_model_parallel_all_reduce(
242 + final_hidden_states)
243 + return final_hidden_states.view(num_tokens, hidden_size)
244 +
245 +
246 +class BailingMoeBlock(nn.Module):
247 +
248 + def __init__(
249 + self,
250 + config: PretrainedConfig,
251 + cache_config: Optional[CacheConfig] = None,
252 + quant_config: Optional[QuantizationConfig] = None,
253 + prefix: str = "",
254 + ):
255 + super().__init__()
256 + hidden_size = config.hidden_size
257 + intermediate_size = config.intermediate_size
258 + self.input_layernorm = RMSNorm(hidden_size, eps=config.rms_norm_eps)
259 + self.attention = BailingAttention(config,
260 + cache_config,
261 + quant_config,
262 + prefix=f"{prefix}.attention")
263 + self.post_attention_layernorm = RMSNorm(hidden_size,
264 + eps=config.rms_norm_eps)
265 + self.mlp = BailingMoE(intermediate_size,
266 + config,
267 + quant_config,
268 + True,
269 + prefix=f"{prefix}.mlp")
270 +
271 + def forward(
272 + self,
273 + hidden_states: torch.Tensor,
274 + position_ids: torch.Tensor,
275 + residual: Optional[torch.Tensor],
276 + ) -> torch.Tensor:
277 + if residual is None:
278 + residual = hidden_states
279 + hidden_states = self.input_layernorm(hidden_states)
280 + else:
281 + hidden_states, residual = self.input_layernorm(
282 + hidden_states, residual)
283 +
284 + hidden_states = self.attention(
285 + hidden_states=hidden_states,
286 + position_ids=position_ids,
287 + )
288 +
289 + hidden_states, residual = self.post_attention_layernorm(
290 + hidden_states, residual)
291 + hidden_states = self.mlp(hidden_states)
292 + return hidden_states, residual
293 +
294 +
295 +class BailingMoeModel(nn.Module):
296 +
297 + def __init__(
298 + self,
299 + *,
300 + vllm_config: VllmConfig,
301 + prefix: str = "",
302 + ):
303 + super().__init__()
304 + config = vllm_config.model_config.hf_config
305 + cache_config = vllm_config.cache_config
306 + quant_config = vllm_config.quant_config
307 +
308 + self.config = config
309 + self.vocab_size = config.vocab_size
310 + self.embed_dim = config.hidden_size
311 +
312 + if get_pp_group().is_first_rank or (config.tie_word_embeddings
313 + and get_pp_group().is_last_rank):
314 + self.word_embeddings = VocabParallelEmbedding(
315 + self.vocab_size, self.embed_dim)
316 + else:
317 + self.word_embeddings = PPMissingLayer()
318 +
319 + self.embedding_dropout = torch.nn.Dropout(config.embedding_dropout)
320 +
321 + self.start_layer, self.end_layer, self.layers = make_layers(
322 + config.num_hidden_layers,
323 + lambda prefix: BailingMoeBlock(
324 + config=config,
325 + cache_config=cache_config,
326 + quant_config=quant_config,
327 + prefix=prefix,
328 + ),
329 + prefix=f"{prefix}.layers")
330 +
331 + self.make_empty_intermediate_tensors = (
332 + make_empty_intermediate_tensors_factory(
333 + ["hidden_states", "residual"], config.hidden_size))
334 +
335 + if get_pp_group().is_last_rank:
336 + self.norm = RMSNorm(self.embed_dim, eps=config.rms_norm_eps)
337 + else:
338 + self.norm = PPMissingLayer()
339 +
340 + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
341 + return self.word_embeddings(input_ids)
342 +
343 + def forward(
344 + self,
345 + input_ids: torch.Tensor,
346 + position_ids: torch.Tensor,
347 + intermediate_tensors: Optional[IntermediateTensors],
348 + inputs_embeds: Optional[torch.Tensor] = None,
349 + ) -> Union[torch.Tensor, IntermediateTensors]:
350 + if get_pp_group().is_first_rank:
351 + if inputs_embeds is not None:
352 + hidden_states = inputs_embeds
353 + else:
354 + hidden_states = self.get_input_embeddings(input_ids)
355 + residual = None
356 + else:
357 + assert intermediate_tensors is not None
358 + hidden_states = intermediate_tensors["hidden_states"]
359 + residual = intermediate_tensors["residual"]
360 +
361 + for i in range(self.start_layer, self.end_layer):
362 + layer = self.layers[i]
363 + hidden_states, residual = layer(
364 + hidden_states,
365 + position_ids,
366 + residual,
367 + )
368 +
369 + if not get_pp_group().is_last_rank:
370 + return IntermediateTensors({
371 + "hidden_states": hidden_states,
372 + "residual": residual
373 + })
374 +
375 + hidden_states, _ = self.norm(hidden_states, residual)
376 + return hidden_states
377 +
378 + def load_weights(self, weights: Iterable[tuple[str,
379 + torch.Tensor]]) -> set[str]:
380 + stacked_params_mapping = [
381 + # (param_name, shard_name, shard_id)
382 + ("gate_up_proj", "gate_proj", 0),
383 + ("gate_up_proj", "up_proj", 1),
384 + ]
385 + expert_params_mapping = FusedMoE.make_expert_params_mapping(
386 + ckpt_gate_proj_name="gate_proj",
387 + ckpt_down_proj_name="down_proj",
388 + ckpt_up_proj_name="up_proj",
389 + num_experts=self.config.num_experts)
390 +
391 + params_dict = dict(self.named_parameters(remove_duplicate=False))
392 + loaded_params: set[str] = set()
393 + for name, loaded_weight in weights:
394 + if self.config.norm_head and "lm_head.weight" in name:
395 + loaded_weight = F.normalize(loaded_weight,
396 + dim=0,
397 + p=2,
398 + eps=1e-7)
399 +
400 + for (param_name, weight_name, shard_id) in stacked_params_mapping:
401 + if weight_name not in name:
402 + continue
403 + if "mlp.experts" in name:
404 + continue
405 + name = name.replace(weight_name, param_name)
406 + # Skip loading extra bias for GPTQ models.
407 + if name.endswith(".bias") and name not in params_dict:
408 + continue
409 + if name not in params_dict:
410 + continue
411 +
412 + if is_pp_missing_parameter(name, self):
413 + continue
414 +
415 + param = params_dict[name]
416 + weight_loader = param.weight_loader
417 + weight_loader(param, loaded_weight, shard_id)
418 + break
419 + else:
420 + for mapping in expert_params_mapping:
421 + param_name, weight_name, expert_id, shard_id = mapping
422 + if weight_name not in name:
423 + continue
424 + name = name.replace(weight_name, param_name)
425 +
426 + if is_pp_missing_parameter(name, self):
427 + continue
428 + param = params_dict[name]
429 + weight_loader = param.weight_loader
430 + weight_loader(param,
431 + loaded_weight,
432 + name,
433 + shard_id=shard_id,
434 + expert_id=expert_id)
435 + break
436 + else:
437 + if name.endswith(".bias") and name not in params_dict:
438 + continue
439 + if name not in params_dict:
440 + continue
441 +
442 + if is_pp_missing_parameter(name, self):
443 + continue
444 +
445 + param = params_dict[name]
446 + weight_loader = getattr(param, "weight_loader",
447 + default_weight_loader)
448 + weight_loader(param, loaded_weight)
449 + loaded_params.add(name)
450 + return loaded_params
451 +
452 +
453 +class BailingMoeForCausalLM(nn.Module, SupportsPP):
454 +
455 + packed_modules_mapping = {
456 + "query_key_value": ["query_key_value"],
457 + "gate_up_proj": [
458 + "gate_proj",
459 + "up_proj",
460 + ],
461 + }
462 +
463 + def __init__(
464 + self,
465 + *,
466 + vllm_config: VllmConfig,
467 + prefix: str = "",
468 + ) -> None:
469 + super().__init__()
470 +
471 + config = vllm_config.model_config.hf_config
472 + quant_config = vllm_config.quant_config
473 +
474 + self.config = config
475 + self.quant_config = quant_config
476 + self.max_position_embeddings = config.max_position_embeddings
477 + self.model = BailingMoeModel(vllm_config=vllm_config,
478 + prefix=maybe_prefix(prefix, "model"))
479 + if get_pp_group().is_last_rank:
480 + self.lm_head = (self.word_embeddings if config.tie_word_embeddings
481 + else ParallelLMHead(config.vocab_size,
482 + config.hidden_size,
483 + quant_config=quant_config))
484 + self.logits_processor = LogitsProcessor(config.vocab_size)
485 + else:
486 + self.lm_head = PPMissingLayer()
487 +
488 + self.sampler = get_sampler()
489 + self.make_empty_intermediate_tensors = (
490 + self.model.make_empty_intermediate_tensors)
491 +
492 + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
493 + return self.model.get_input_embeddings(input_ids)
494 +
495 + def forward(
496 + self,
497 + input_ids: torch.Tensor,
498 + positions: torch.Tensor,
499 + intermediate_tensors: Optional[IntermediateTensors] = None,
500 + inputs_embeds: Optional[torch.Tensor] = None,
501 + ) -> Union[torch.Tensor, IntermediateTensors]:
502 + model_output = self.model(input_ids, positions, intermediate_tensors,
503 + inputs_embeds)
504 + return model_output
505 +
506 + def compute_logits(
507 + self,
508 + hidden_states: torch.Tensor,
509 + sampling_metadata: SamplingMetadata,
510 + ) -> Optional[torch.Tensor]:
511 + logits = self.logits_processor(self.lm_head, hidden_states,
512 + sampling_metadata)
513 + return logits
514 +
515 + def sample(
516 + self,
517 + logits: torch.Tensor,
518 + sampling_metadata: SamplingMetadata,
519 + ) -> Optional[SamplerOutput]:
520 + next_tokens = self.sampler(logits, sampling_metadata)
521 + return next_tokens
522 +
523 + def load_weights(self, weights: Iterable[tuple[str,
524 + torch.Tensor]]) -> set[str]:
525 + loader = AutoWeightsLoader(
526 + self,
527 + skip_prefixes=(["lm_head."]
528 + if self.config.tie_word_embeddings else None),
529 + )
530 + return loader.load_weights(weights)
vllm/model_executor/models/registry.py
+1
@@ -41,6 +41,7 @@ _TEXT_GENERATION_MODELS = {
41 "BaiChuanForCausalLM": ("baichuan", "BaiChuanForCausalLM"),
42 # baichuan-13b, lower case 'c' in the class name
43 "BaichuanForCausalLM": ("baichuan", "BaichuanForCausalLM"),
44 + "BailingMoeForCausalLM": ("bailing_moe", "BailingMoeForCausalLM"),
45 "BambaForCausalLM": ("bamba", "BambaForCausalLM"),
46 "BloomForCausalLM": ("bloom", "BloomForCausalLM"),
47 "ChatGLMModel": ("chatglm", "ChatGLMForCausalLM"),