Skip to content

Commit 68937f8

Browse files
rjg-lyhrjg-lyh
committed
[V0.9.1] Use AddRmsNormQuant ops in the custom model to optimize Qwen3's
performance Signed-off-by: rjg-lyh <1318825571@qq.com>
1 parent 9acc082 commit 68937f8

File tree

4 files changed

+198
-5
lines changed

4 files changed

+198
-5
lines changed

‎vllm_ascend/models/__init__.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ def register_model():
1111
from .qwen2_5_vl import \
1212
AscendQwen2_5_VLForConditionalGeneration # noqa: F401
1313
from .qwen2_vl import AscendQwen2VLForConditionalGeneration # noqa: F401
14+
from .qwen3 import CustomQwen3ForCausalLM # noqa: F401
1415

1516
ModelRegistry.register_model(
1617
"DeepSeekMTPModel",
@@ -52,3 +53,7 @@ def register_model():
5253
ModelRegistry.register_model(
5354
"Qwen3MoeForCausalLM",
5455
"vllm_ascend.models.qwen3_moe:CustomQwen3MoeForCausalLM")
56+
57+
ModelRegistry.register_model(
58+
"Qwen3ForCausalLM",
59+
"vllm_ascend.models.qwen3:CustomQwen3ForCausalLM")

‎vllm_ascend/models/qwen3.py

Lines changed: 144 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,144 @@
1+
from collections.abc import Iterable
2+
from typing import Optional, Union
3+
4+
import torch
5+
from torch import nn
6+
from transformers import Qwen3Config
7+
8+
from vllm.compilation.decorators import support_torch_compile
9+
from vllm.config import CacheConfig, VllmConfig
10+
from vllm.distributed import get_pp_group
11+
from vllm.model_executor.layers.logits_processor import LogitsProcessor
12+
from vllm.model_executor.layers.quantization import QuantizationConfig
13+
from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
14+
from vllm.model_executor.sampling_metadata import SamplingMetadata
15+
from vllm.sequence import IntermediateTensors
16+
17+
from vllm.model_executor.models.interfaces import SupportsLoRA, SupportsPP
18+
from vllm.model_executor.models.qwen2 import Qwen2Model
19+
from vllm.model_executor.models.qwen3 import Qwen3DecoderLayer
20+
from vllm.model_executor.models.utils import AutoWeightsLoader, PPMissingLayer, maybe_prefix
21+
22+
from vllm_ascend.ops.layernorm import AddRMSNormQuant
23+
24+
class CustomQwen3DecoderLayer(Qwen3DecoderLayer):
25+
26+
def __init__(
27+
self,
28+
config: Qwen3Config,
29+
cache_config: Optional[CacheConfig] = None,
30+
quant_config: Optional[QuantizationConfig] = None,
31+
prefix: str = "",
32+
) -> None:
33+
super().__init__(config=config,
34+
cache_config=cache_config,
35+
quant_config=quant_config,
36+
prefix=prefix)
37+
if quant_config is not None:
38+
from vllm_ascend.quantization.quant_config import AscendQuantConfig
39+
assert isinstance(quant_config, AscendQuantConfig)
40+
self.input_layernorm = AddRMSNormQuant(config.hidden_size,
41+
layer=self.self_attn.qkv_proj,
42+
eps=config.rms_norm_eps)
43+
self.post_attention_layernorm = AddRMSNormQuant(config.hidden_size,
44+
layer=self.mlp.gate_up_proj,
45+
eps=config.rms_norm_eps)
46+
47+
48+
ALL_DECODER_LAYER_TYPES = {
49+
"attention": CustomQwen3DecoderLayer,
50+
}
51+
52+
53+
@support_torch_compile(
54+
dynamic_arg_dims={
55+
"input_ids": 0,
56+
# positions is of shape (3, seq_len) if mrope is enabled for qwen2-vl,
57+
# otherwise (seq_len, ).
58+
"positions": -1,
59+
"intermediate_tensors": 0,
60+
"inputs_embeds": 0,
61+
})
62+
class CustomQwen3Model(Qwen2Model):
63+
64+
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
65+
super().__init__(vllm_config=vllm_config,
66+
prefix=prefix,
67+
decoder_layer_type=CustomQwen3DecoderLayer)
68+
69+
70+
class CustomQwen3ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
71+
# add `CustomQwen3Model` to init self.model
72+
packed_modules_mapping = {
73+
"qkv_proj": [
74+
"q_proj",
75+
"k_proj",
76+
"v_proj",
77+
],
78+
"gate_up_proj": [
79+
"gate_proj",
80+
"up_proj",
81+
],
82+
}
83+
84+
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
85+
super().__init__()
86+
config = vllm_config.model_config.hf_config
87+
quant_config = vllm_config.quant_config
88+
lora_config = vllm_config.lora_config
89+
90+
self.config = config
91+
self.lora_config = lora_config
92+
93+
self.quant_config = quant_config
94+
self.model = CustomQwen3Model(vllm_config=vllm_config,
95+
prefix=maybe_prefix(prefix, "model"))
96+
97+
if get_pp_group().is_last_rank:
98+
if config.tie_word_embeddings:
99+
self.lm_head = self.model.embed_tokens
100+
else:
101+
self.lm_head = ParallelLMHead(config.vocab_size,
102+
config.hidden_size,
103+
quant_config=quant_config,
104+
prefix=maybe_prefix(
105+
prefix, "lm_head"))
106+
else:
107+
self.lm_head = PPMissingLayer()
108+
109+
self.logits_processor = LogitsProcessor(config.vocab_size)
110+
111+
self.make_empty_intermediate_tensors = (
112+
self.model.make_empty_intermediate_tensors)
113+
114+
def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
115+
return self.model.get_input_embeddings(input_ids)
116+
117+
def forward(
118+
self,
119+
input_ids: torch.Tensor,
120+
positions: torch.Tensor,
121+
intermediate_tensors: Optional[IntermediateTensors] = None,
122+
inputs_embeds: Optional[torch.Tensor] = None,
123+
) -> Union[torch.Tensor, IntermediateTensors]:
124+
hidden_states = self.model(input_ids, positions, intermediate_tensors,
125+
inputs_embeds)
126+
return hidden_states
127+
128+
def compute_logits(
129+
self,
130+
hidden_states: torch.Tensor,
131+
sampling_metadata: SamplingMetadata,
132+
) -> Optional[torch.Tensor]:
133+
logits = self.logits_processor(self.lm_head, hidden_states,
134+
sampling_metadata)
135+
return logits
136+
137+
def load_weights(self, weights: Iterable[tuple[str,
138+
torch.Tensor]]) -> set[str]:
139+
loader = AutoWeightsLoader(
140+
self,
141+
skip_prefixes=(["lm_head."]
142+
if self.config.tie_word_embeddings else None),
143+
)
144+
return loader.load_weights(weights)

‎vllm_ascend/ops/layernorm.py

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,45 @@
2121
from vllm.model_executor.layers.layernorm import RMSNorm
2222

2323

24+
class AddRMSNormQuant(RMSNorm):
25+
"""Root mean square normalization.
26+
27+
Computes x -> w * x / sqrt(E[x^2] + eps) where w is the learned weight.
28+
Refer to https://arxiv.org/abs/1910.07467
29+
"""
30+
def __init__(
31+
self,
32+
hidden_size: int,
33+
layer: torch.nn.Module,
34+
eps: float = 1e-6,
35+
var_hidden_size: Optional[int] = None,
36+
has_weight: bool = True,
37+
dtype: Optional[torch.dtype] = None,
38+
) -> None:
39+
super().__init__(hidden_size,
40+
eps,
41+
var_hidden_size,
42+
has_weight,
43+
dtype)
44+
self.layer = layer
45+
46+
def forward(
47+
self,
48+
x: torch.Tensor,
49+
residual: Optional[torch.Tensor] = None,
50+
) -> Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
51+
import torch_npu
52+
53+
if residual is not None:
54+
x, _, residual = torch_npu.npu_add_rms_norm_quant(x, residual, self.weight,
55+
self.layer.aclnn_input_scale,
56+
self.layer.aclnn_input_offset,
57+
epsilon=self.variance_epsilon)
58+
return x, residual
59+
60+
x, residual = torch_npu.npu_rms_norm(x, self.weight, self.variance_epsilon)
61+
return x
62+
2463
def forward_oot(
2564
self,
2665
x: torch.Tensor,

‎vllm_ascend/quantization/w8a8.py

Lines changed: 10 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@ class AscendW8A8LinearMethod:
3636
Args:
3737
w_sym: whether the linear weight is symmetrically quantized.
3838
"""
39+
params_dtype: torch.dtype = torch.bfloat16
3940

4041
def __init__(self) -> None:
4142
# aclnn quant matmul requires to transpose matrix B, set to true by default.
@@ -59,6 +60,7 @@ def get_pertensor_param(params_dtype: torch.dtype) -> Dict[str, Any]:
5960
params_dict = {}
6061
params_dict["input_scale"] = torch.empty(1, dtype=params_dtype)
6162
params_dict["input_offset"] = torch.empty(1, dtype=torch.int8)
63+
AscendW8A8LinearMethod.params_dtype = params_dtype
6264
return params_dict
6365

6466
@staticmethod
@@ -80,6 +82,7 @@ def get_perchannel_param(
8082
params_dict["weight_offset"] = torch.empty(output_size,
8183
1,
8284
dtype=params_dtype)
85+
AscendW8A8LinearMethod.params_dtype = params_dtype
8386
return params_dict
8487

8588
def get_pergroup_param(self, input_size: int, output_size: int,
@@ -93,11 +96,10 @@ def apply(
9396
bias: Optional[torch.Tensor] = None,
9497
tp_rank: Optional[int] = 0,
9598
) -> torch.Tensor:
96-
original_dtype = x.dtype
97-
if original_dtype != torch.int8:
99+
if x.dtype != torch.int8:
98100
x = quant_per_tensor(
99101
x,
100-
layer.aclnn_input_scale,
102+
layer.aclnn_input_scale_reciprocal,
101103
layer.aclnn_input_offset,
102104
)
103105
quant_bias = layer.quant_bias if tp_rank == 0 else None
@@ -106,12 +108,15 @@ def apply(
106108
layer.weight,
107109
layer.deq_scale,
108110
bias=quant_bias,
109-
output_dtype=original_dtype,
111+
output_dtype=AscendW8A8LinearMethod.params_dtype,
110112
)
111113

112114
def process_weights_after_loading(self, layer):
113115
expanding_factor = layer.weight.data.shape[1]
114-
layer.aclnn_input_scale = 1 / torch.nn.Parameter(
116+
layer.aclnn_input_scale = torch.nn.Parameter(
117+
layer.input_scale.data.repeat(expanding_factor),
118+
requires_grad=False)
119+
layer.aclnn_input_scale_reciprocal = 1 / torch.nn.Parameter(
115120
layer.input_scale.data.repeat(expanding_factor),
116121
requires_grad=False)
117122
layer.aclnn_input_offset = torch.nn.Parameter(

0 commit comments

Comments
 (0)