@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 277 lines 14.6 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 # Copyright 2025 The Qwen team, Alibaba Group and the HuggingFace Inc. team.
4 # All rights reserved.
5 #
6 # Licensed under the Apache License, Version 2.0 (the "License");
7 # you may not use this file except in compliance with the License.
8 # You may obtain a copy of the License at
9 #
10 # http://www.apache.org/licenses/LICENSE-2.0
11 #
12 # Unless required by applicable law or agreed to in writing, software
13 # distributed under the License is distributed on an "AS IS" BASIS,
14 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15 # See the License for the specific language governing permissions and
16 # limitations under the License.
17 """Qwen3-Next model configuration."""
18
19 from transformers.configuration_utils import PreTrainedConfig
20 from transformers.utils import logging
21
22 logger = logging.get_logger(__name__)
23
24
25 class Qwen3NextConfig(PreTrainedConfig):
26 r"""This is the configuration class to store the configuration of a [`Qwen3NextModel`]. It is used to instantiate a
27 Qwen3-Next model according to the specified arguments, defining the model architecture.
28 Instantiating a configuration with the defaults will yield a similar configuration to that of
29 Qwen3-Next-80B-A3B-Instruct [Qwen/Qwen3-Next-80B-A3B-Instruct](https://huggingface.co/Qwen/Qwen3-Next-80B-A3B-Instruct).
30
31 Configuration objects inherit from [`PreTrainedConfig`] and can be used to control the model outputs. Read the
32 documentation from [`PreTrainedConfig`] for more information.
33
34
35 Args:
36 vocab_size (`int`, *optional*, defaults to 151936):
37 Vocabulary size of the model. Defines the number of different tokens that can be represented by the
38 `inputs_ids`.
39 hidden_size (`int`, *optional*, defaults to 2048):
40 Dimension of the hidden representations.
41 intermediate_size (`int`, *optional*, defaults to 5632):
42 Dimension of the MLP representations.
43 num_hidden_layers (`int`, *optional*, defaults to 48):
44 Number of hidden layers in the Transformer encoder.
45 num_attention_heads (`int`, *optional*, defaults to 16):
46 Number of attention heads for each attention layer in the Transformer encoder.
47 num_key_value_heads (`int`, *optional*, defaults to 2):
48 This is the number of key_value heads that should be used to implement Grouped Query Attention. If
49 `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if
50 `num_key_value_heads=1` the model will use Multi Query Attention (MQA) otherwise GQA is used. When
51 converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed
52 by meanpooling all the original heads within that group. For more details checkout [this
53 paper](https://arxiv.org/pdf/2305.13245.pdf). If it is not specified, will default to `32`.
54 hidden_act (`str`, *optional*, defaults to `"silu"`):
55 The non-linear activation function in the decoder.
56 max_position_embeddings (`int`, *optional*, defaults to 32768):
57 The maximum sequence length that this model might ever be used with.
58 initializer_range (`float`, *optional*, defaults to 0.02):
59 The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
60 rms_norm_eps (`float`, *optional*, defaults to 1e-06):
61 The epsilon used by the rms normalization layers.
62 use_cache (`bool`, *optional*, defaults to `True`):
63 Whether or not the model should return the last key/values attentions (not used by all models). Only
64 relevant if `config.is_decoder=True`.
65 tie_word_embeddings (`bool`, *optional*, defaults to `False`):
66 Whether the model's input and output word embeddings should be tied.
67 rope_parameters (`dict`, *optional*):
68 Dictionary containing the scaling configuration for the RoPE embeddings. NOTE: if you apply new rope type
69 and you expect the model to work on longer `max_position_embeddings`, we recommend you to update this value
70 accordingly.
71 Expected contents:
72 `rope_theta` (`float`): The base period of the RoPE embeddings.
73 `rope_type` (`str`):
74 The sub-variant of RoPE to use. Can be one of ['default', 'linear', 'dynamic', 'yarn', 'longrope',
75 'llama3'], with 'default' being the original RoPE implementation.
76 `factor` (`float`, *optional*):
77 Used with all rope types except 'default'. The scaling factor to apply to the RoPE embeddings. In
78 most scaling types, a `factor` of x will enable the model to handle sequences of length x *
79 original maximum pre-trained length.
80 `original_max_position_embeddings` (`int`, *optional*):
81 Used with 'dynamic', 'longrope' and 'llama3'. The original max position embeddings used during
82 pretraining.
83 `attention_factor` (`float`, *optional*):
84 Used with 'yarn' and 'longrope'. The scaling factor to be applied on the attention
85 computation. If unspecified, it defaults to value recommended by the implementation, using the
86 `factor` field to infer the suggested value.
87 `beta_fast` (`float`, *optional*):
88 Only used with 'yarn'. Parameter to set the boundary for extrapolation (only) in the linear
89 ramp function. If unspecified, it defaults to 32.
90 `beta_slow` (`float`, *optional*):
91 Only used with 'yarn'. Parameter to set the boundary for interpolation (only) in the linear
92 ramp function. If unspecified, it defaults to 1.
93 `short_factor` (`List[float]`, *optional*):
94 Only used with 'longrope'. The scaling factor to be applied to short contexts (<
95 `original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden
96 size divided by the number of attention heads divided by 2
97 `long_factor` (`List[float]`, *optional*):
98 Only used with 'longrope'. The scaling factor to be applied to long contexts (<
99 `original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden
100 size divided by the number of attention heads divided by 2
101 `low_freq_factor` (`float`, *optional*):
102 Only used with 'llama3'. Scaling factor applied to low frequency components of the RoPE
103 `high_freq_factor` (`float`, *optional*):
104 Only used with 'llama3'. Scaling factor applied to high frequency components of the RoPE
105 `partial_rotary_factor` (`float`, *optional*, defaults to 0.25):
106 Percentage of the query and keys which will have rotary embedding.
107 attention_bias (`bool`, *optional*, defaults to `False`):
108 Whether to use a bias in the query, key, value and output projection layers during self-attention.
109 attention_dropout (`float`, *optional*, defaults to 0.0):
110 The dropout ratio for the attention probabilities.
111 head_dim (`int`, *optional*, defaults to 256):
112 Projection weights dimension in multi-head attention.
113 linear_conv_kernel_dim (`int`, *optional*, defaults to 4):
114 Kernel size of the convolution used in linear attention layers.
115 linear_key_head_dim (`int`, *optional*, defaults to 128):
116 Dimension of each key head in linear attention.
117 linear_value_head_dim (`int`, *optional*, defaults to 128):
118 Dimension of each value head in linear attention.
119 linear_num_key_heads (`int`, *optional*, defaults to 16):
120 Number of key heads used in linear attention layers.
121 linear_num_value_heads (`int`, *optional*, defaults to 32):
122 Number of value heads used in linear attention layers.
123 decoder_sparse_step (`int`, *optional*, defaults to 1):
124 The frequency of the MoE layer.
125 moe_intermediate_size (`int`, *optional*, defaults to 512):
126 Intermediate size of the routed expert.
127 shared_expert_intermediate_size (`int`, *optional*, defaults to 512):
128 Intermediate size of the shared expert.
129 num_experts_per_tok (`int`, *optional*, defaults to 10):
130 Number of selected experts.
131 num_experts (`int`, *optional*, defaults to 512):
132 Number of routed experts.
133 norm_topk_prob (`bool`, *optional*, defaults to `True`):
134 Whether to normalize the topk probabilities.
135 output_router_logits (`bool`, *optional*, defaults to `False`):
136 Whether or not the router logits should be returned by the model. Enabling this will also
137 allow the model to output the auxiliary loss, including load balancing loss and router z-loss.
138 router_aux_loss_coef (`float`, *optional*, defaults to 0.001):
139 The aux loss factor for the total loss.
140 mlp_only_layers (`list[int]`, *optional*, defaults to `[]`):
141 Indicate which layers use Qwen3NextMLP rather than Qwen3NextSparseMoeBlock
142 The list contains layer index, from 0 to num_layers-1 if we have num_layers layers
143 If `mlp_only_layers` is empty, `decoder_sparse_step` is used to determine the sparsity.
144 layer_types (`list[str]`, *optional*):
145 Types of each layer (attention or linear).
146
147 ```python
148 >>> from transformers import Qwen3NextModel, Qwen3NextConfig
149
150 >>> # Initializing a Qwen3Next style configuration
151 >>> configuration = Qwen3NextConfig()
152
153 >>> # Initializing a model from the Qwen3-Next-80B-A3B style configuration
154 >>> model = Qwen3NextModel(configuration)
155
156 >>> # Accessing the model configuration
157 >>> configuration = model.config
158 ```
159
160 """ # noqa: E501
161
162 model_type = "qwen3_next"
163 keys_to_ignore_at_inference = ["past_key_values"]
164
165 base_model_tp_plan = {
166 "layers.*.self_attn.q_proj": "colwise",
167 "layers.*.self_attn.k_proj": "colwise",
168 "layers.*.self_attn.v_proj": "colwise",
169 "layers.*.self_attn.o_proj": "rowwise",
170 "layers.*.mlp.experts.*.gate_proj": "colwise",
171 "layers.*.mlp.experts.*.up_proj": "colwise",
172 "layers.*.mlp.experts.*.down_proj": "rowwise",
173 "layers.*.mlp.shared_experts.gate_proj": "colwise",
174 "layers.*.mlp.shared_experts.up_proj": "colwise",
175 "layers.*.mlp.shared_experts.down_proj": "rowwise",
176 "layers.*.mlp.gate_proj": "colwise",
177 "layers.*.mlp.up_proj": "colwise",
178 "layers.*.mlp.down_proj": "rowwise",
179 }
180 base_model_pp_plan = {
181 "embed_tokens": (["input_ids"], ["inputs_embeds"]),
182 "layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
183 "norm": (["hidden_states"], ["hidden_states"]),
184 }
185
186 def __init__(
187 self,
188 vocab_size=151936,
189 hidden_size=2048,
190 intermediate_size=5632,
191 num_hidden_layers=48,
192 num_attention_heads=16,
193 num_key_value_heads=2,
194 hidden_act="silu",
195 max_position_embeddings=32768,
196 initializer_range=0.02,
197 rms_norm_eps=1e-6,
198 use_cache=True,
199 tie_word_embeddings=False,
200 rope_parameters=None,
201 attention_bias=False,
202 attention_dropout=0.0,
203 head_dim=256,
204 linear_conv_kernel_dim=4,
205 linear_key_head_dim=128,
206 linear_value_head_dim=128,
207 linear_num_key_heads=16,
208 linear_num_value_heads=32,
209 decoder_sparse_step=1,
210 moe_intermediate_size=512,
211 shared_expert_intermediate_size=512,
212 num_experts_per_tok=10,
213 num_experts=512,
214 norm_topk_prob=True,
215 output_router_logits=False,
216 router_aux_loss_coef=0.001,
217 mlp_only_layers=None,
218 layer_types=None,
219 **kwargs,
220 ):
221 if mlp_only_layers is None:
222 mlp_only_layers = []
223 self.vocab_size = vocab_size
224 self.max_position_embeddings = max_position_embeddings
225 self.hidden_size = hidden_size
226 self.intermediate_size = intermediate_size
227 self.num_hidden_layers = num_hidden_layers
228 self.num_attention_heads = num_attention_heads
229 self.num_key_value_heads = num_key_value_heads
230 self.hidden_act = hidden_act
231 self.initializer_range = initializer_range
232 self.rms_norm_eps = rms_norm_eps
233 self.use_cache = use_cache
234 # Try to set `rope_scaling` if available, otherwise use `rope_parameters`
235 rope_scaling = kwargs.pop("rope_scaling", None)
236 rope_parameters = rope_scaling or rope_parameters or {"rope_type": "default"}
237 rope_theta = kwargs.pop("rope_theta", 10000.0)
238 if "rope_theta" not in rope_parameters:
239 rope_parameters["rope_theta"] = rope_theta
240 partial_rotary_factor = kwargs.pop("partial_rotary_factor", 0.25)
241 if "partial_rotary_factor" not in rope_parameters:
242 rope_parameters["partial_rotary_factor"] = partial_rotary_factor
243 self.rope_parameters = rope_parameters
244 self.partial_rotary_factor = partial_rotary_factor
245 self.attention_bias = attention_bias
246 self.attention_dropout = attention_dropout
247 self.head_dim = head_dim
248
249 self.layer_types = layer_types
250 if self.layer_types is None:
251 self.layer_types = [
252 "linear_attention" if bool((i + 1) % 4) else "full_attention"
253 for i in range(self.num_hidden_layers)
254 ]
255 self.validate_layer_type()
256
257 # linear attention part
258 self.linear_conv_kernel_dim = linear_conv_kernel_dim
259 self.linear_key_head_dim = linear_key_head_dim
260 self.linear_value_head_dim = linear_value_head_dim
261 self.linear_num_key_heads = linear_num_key_heads
262 self.linear_num_value_heads = linear_num_value_heads
263
264 # MoE arguments
265 self.decoder_sparse_step = decoder_sparse_step
266 self.moe_intermediate_size = moe_intermediate_size
267 self.shared_expert_intermediate_size = shared_expert_intermediate_size
268 self.num_experts_per_tok = num_experts_per_tok
269 self.num_experts = num_experts
270 self.norm_topk_prob = norm_topk_prob
271 self.output_router_logits = output_router_logits
272 self.router_aux_loss_coef = router_aux_loss_coef
273 self.mlp_only_layers = mlp_only_layers
274 super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)
275
276
277 __all__ = ["Qwen3NextConfig"]