@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 216 lines 10.7 KB
Raw
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 import PreTrainedConfig
6
7
8 class AXK1Config(PreTrainedConfig):
9 r"""This is the configuration class to store the configuration of a [`AXK1Model`].
10 It is used to instantiate an A.X model according to the specified arguments,
11 defining the model architecture. Instantiating a configuration with the defaults
12 will yield a similar configuration to that of the A.X K1.
13 Configuration objects inherit from [`PreTrainedConfig`] and can be used to control
14 the model outputs. Read the documentation from [`PreTrainedConfig`] for more
15 information.
16
17 Args:
18 vocab_size (`int`, *optional*, defaults to 163840):
19 Vocabulary size of the A.X K1 model. Defines the number of different
20 tokens that can be represented by the `inputs_ids` passed when calling
21 [`AXK1Model`]
22 hidden_size (`int`, *optional*, defaults to 7168):
23 Dimension of the hidden representations.
24 intermediate_size (`int`, *optional*, defaults to 18432):
25 Dimension of the MLP representations.
26 moe_intermediate_size (`int`, *optional*, defaults to 2048):
27 Dimension of the MoE representations.
28 num_hidden_layers (`int`, *optional*, defaults to 61):
29 Number of hidden layers in the Transformer decoder.
30 num_nextn_predict_layers (`int`, *optional*, defaults to 1):
31 Number of nextn predict layers in the AXK1 Model.
32 num_attention_heads (`int`, *optional*, defaults to 64):
33 Number of attention heads for each attention layer in the Transformer
34 decoder.
35 n_shared_experts (`int`, *optional*, defaults to 1):
36 Number of shared experts, None means dense model.
37 n_routed_experts (`int`, *optional*, defaults to 192):
38 Number of routed experts, None means dense model.
39 routed_scaling_factor (`float`, *optional*, defaults to 2.5):
40 Scaling factor or routed experts.
41 topk_method (`str`, *optional*, defaults to `noaux_tc`):
42 Topk method used in routed gate.
43 n_group (`int`, *optional*, defaults to 8):
44 Number of groups for routed experts.
45 topk_group (`int`, *optional*, defaults to 4):
46 Number of selected groups for each token(for each token, ensuring the
47 selected experts is only within `topk_group` groups).
48 num_experts_per_tok (`int`, *optional*, defaults to 8):
49 Number of selected experts, None means dense model.
50 moe_layer_freq (`int`, *optional*, defaults to 1):
51 The frequency of the MoE layer: one expert layer for every
52 `moe_layer_freq - 1` dense layers.
53 first_k_dense_replace (`int`, *optional*, defaults to 1):
54 Number of dense layers in shallow layers
55 (embed->dense->dense->...->dense->moe->moe...->lm_head).
56 \--k dense layers--/
57 norm_topk_prob (`bool`, *optional*, defaults to True):
58 Whether to normalize the weights of the routed experts.
59 scoring_func (`str`, *optional*, defaults to 'sigmoid'):
60 Method of computing expert weights.
61 aux_loss_alpha (`float`, *optional*, defaults to 0.0001):
62 Auxiliary loss weight coefficient.
63 seq_aux = (`bool`, *optional*, defaults to True):
64 Whether to compute the auxiliary loss for each individual sample.
65 num_key_value_heads (`int`, *optional*):
66 This is the number of key_value heads that should be used to implement
67 Grouped Query Attention. If `num_key_value_heads=num_attention_heads`,
68 the model will use Multi Head Attention (MHA), if `num_key_value_heads=1
69 the model will use Multi Query Attention (MQA) otherwise GQA is used. When
70 converting a multi-head checkpoint to a GQA checkpoint, each group key and
71 value head should be constructed by meanpooling all the original heads
72 within that group. For more details checkout
73 [this paper](https://arxiv.org/pdf/2305.13245.pdf).
74 If it is not specified, will default to `num_attention_heads`.
75 hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
76 The non-linear activation function (function or string) in the decoder.
77 max_position_embeddings (`int`, *optional*, defaults to 131072):
78 The maximum sequence length that this model might ever be used with.
79 initializer_range (`float`, *optional*, defaults to 0.02):
80 The standard deviation of the truncated_normal_initializer for
81 initializing all weight matrices.
82 rms_norm_eps (`float`, *optional*, defaults to 1e-06):
83 The epsilon used by the rms normalization layers.
84 use_cache (`bool`, *optional*, defaults to `True`):
85 Whether or not the model should return the last key/values attentions
86 (not used by all models). Only relevant if `config.is_decoder=True`.
87 pad_token_id (`int`, *optional*):
88 Padding token id.
89 bos_token_id (`int`, *optional*, defaults to 163691):
90 Beginning of stream token id.
91 eos_token_id (`int`, *optional*, defaults to 163691):
92 End of stream token id.
93 pretraining_tp (`int`, *optional*, defaults to 1):
94 Experimental feature. Tensor parallelism rank used during pretraining.
95 Please refer to
96 [this document](https://huggingface.co/docs/transformers/parallelism)
97 to understand more about it. This value is necessary to ensure exact
98 reproducibility of the pretraining results. Please refer to
99 [this issue](https://github.com/pytorch/pytorch/issues/76232).
100 tie_word_embeddings (`bool`, *optional*, defaults to `False`):
101 Whether to tie weight embeddings
102 rope_theta (`float`, *optional*, defaults to 10000.0):
103 The base period of the RoPE embeddings.
104 rope_scaling (`Dict`, *optional*):
105 Dictionary containing the scaling configuration for the RoPE embeddings.
106 Currently supports two scaling strategies: linear and dynamic.
107 Their scaling factor must be a float greater than 1. The expected format
108 is `{"type": strategy name, "factor": scaling factor}`. When using this
109 flag, don't update `max_position_embeddings` to the expected new maximum.
110 attention_bias (`bool`, defaults to `False`, *optional*, defaults to `False`):
111 Whether to use a bias in the query, key, value and output projection
112 layers during self-attention.
113 attention_dropout (`float`, *optional*, defaults to 0.0):
114 The dropout ratio for the attention probabilities.
115
116 """
117
118 model_type = "axk1"
119 keys_to_ignore_at_inference = ["past_key_values"]
120
121 def __init__(
122 self,
123 vocab_size: int = 163840,
124 hidden_size: int = 7168,
125 intermediate_size: int = 18432,
126 moe_intermediate_size: int = 2048,
127 num_hidden_layers: int = 61,
128 num_nextn_predict_layers: int | None = 1,
129 num_attention_heads: int = 64,
130 num_key_value_heads: int = 64,
131 n_shared_experts: int | None = 1,
132 n_routed_experts: int | None = 192,
133 ep_size: int | None = 8, ## Ignored - Expert parallel size
134 routed_scaling_factor: float | None = 2.5,
135 kv_lora_rank: int | None = 512,
136 q_lora_rank: int | None = 1536,
137 qk_rope_head_dim: int | None = 64,
138 v_head_dim: int | None = 128,
139 qk_nope_head_dim: int | None = 128,
140 topk_method: str | None = "noaux_tc",
141 n_group: int | None = 8,
142 topk_group: int | None = 4,
143 num_experts_per_tok: int | None = 8,
144 moe_layer_freq: int | None = 1,
145 first_k_dense_replace: int = 1,
146 norm_topk_prob: bool = True,
147 scoring_func: str | None = "sigmoid",
148 aux_loss_alpha: float | None = 0.0001,
149 seq_aux: float | None = True,
150 hidden_act: str | None = "silu",
151 max_position_embeddings: int | None = 131072,
152 initializer_range: float | None = 0.02,
153 rms_norm_eps: float = 1e-6,
154 use_cache: bool | None = True,
155 pad_token_id: int | None = None,
156 bos_token_id: int | None = 163691,
157 eos_token_id: int | None = 163691,
158 pretraining_tp: int | None = 1,
159 tie_word_embeddings: bool | None = False,
160 rope_theta: float | None = 10000.0,
161 rope_scaling: dict[str, Any] | None = None,
162 rope_parameters: dict[str, Any] | None = None,
163 attention_bias: bool | None = False,
164 attention_dropout: float | None = 0.0,
165 **kwargs,
166 ):
167 self.vocab_size = vocab_size
168 self.max_position_embeddings = max_position_embeddings
169 self.hidden_size = hidden_size
170 self.intermediate_size = intermediate_size
171 self.moe_intermediate_size = moe_intermediate_size
172 self.num_hidden_layers = num_hidden_layers
173 self.num_nextn_predict_layers = num_nextn_predict_layers
174 self.num_attention_heads = num_attention_heads
175 self.n_shared_experts = n_shared_experts
176 self.n_routed_experts = n_routed_experts
177 self.ep_size = ep_size
178 self.routed_scaling_factor = routed_scaling_factor
179 self.kv_lora_rank = kv_lora_rank
180 self.q_lora_rank = q_lora_rank
181 self.qk_rope_head_dim = qk_rope_head_dim
182 self.v_head_dim = v_head_dim
183 self.qk_nope_head_dim = qk_nope_head_dim
184 self.topk_method = topk_method
185 self.n_group = n_group
186 self.topk_group = topk_group
187 self.num_experts_per_tok = num_experts_per_tok
188 self.moe_layer_freq = moe_layer_freq
189 self.first_k_dense_replace = first_k_dense_replace
190 self.norm_topk_prob = norm_topk_prob
191 self.scoring_func = scoring_func
192 self.aux_loss_alpha = aux_loss_alpha
193 self.seq_aux = seq_aux
194 # for backward compatibility
195 if num_key_value_heads is None:
196 num_key_value_heads = num_attention_heads
197
198 self.num_key_value_heads = num_key_value_heads
199 self.hidden_act = hidden_act
200 self.initializer_range = initializer_range
201 self.rms_norm_eps = rms_norm_eps
202 self.pretraining_tp = pretraining_tp
203 self.use_cache = use_cache
204 self.rope_theta = rope_theta
205 self.rope_scaling = rope_scaling
206 self.rope_parameters = rope_parameters
207 self.attention_bias = attention_bias
208 self.attention_dropout = attention_dropout
209
210 super().__init__(
211 pad_token_id=pad_token_id,
212 bos_token_id=bos_token_id,
213 eos_token_id=eos_token_id,
214 tie_word_embeddings=tie_word_embeddings,
215 **kwargs,
216 )