@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 185 lines 8.28 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.configuration_utils import PreTrainedConfig
6
7
8 class HYV3Config(PreTrainedConfig):
9 r"""This is the configuration class to store the configuration of a [`HYV3Model`].
10 It is used to instantiate a HYV3 model (HY V3 MoE language model) according to
11 the specified arguments.
12
13 Configuration objects inherit from [`PreTrainedConfig`] and can be used to
14 control the model outputs. Read the documentation from [`PreTrainedConfig`]
15 for more information.
16
17 Args:
18 vocab_size (`int`, *optional*, defaults to 120832):
19 Vocabulary size of the model.
20 hidden_size (`int`, *optional*, defaults to 4096):
21 Dimension of the hidden representations.
22 intermediate_size (`int`, *optional*, defaults to 13312):
23 Dimension of the dense FFN intermediate representations.
24 num_hidden_layers (`int`, *optional*, defaults to 80):
25 Number of hidden layers in the Transformer decoder.
26 num_attention_heads (`int`, *optional*, defaults to 64):
27 Number of attention heads for each attention layer.
28 num_key_value_heads (`int`, *optional*, defaults to 8):
29 Number of key-value heads for grouped-query attention.
30 head_dim (`int`, *optional*, defaults to 128):
31 Dimension per attention head.
32 hidden_act (`str`, *optional*, defaults to `"silu"`):
33 Activation function used in FFN layers.
34 max_position_embeddings (`int`, *optional*, defaults to 131072):
35 Maximum sequence length supported by the model.
36 initializer_range (`float`, *optional*, defaults to 0.006):
37 Standard deviation of the truncated normal initializer for weight
38 initialization.
39 rms_norm_eps (`float`, *optional*, defaults to 1e-5):
40 Epsilon for RMS normalization layers.
41 use_cache (`bool`, *optional*, defaults to `True`):
42 Whether to use KV cache for decoding.
43 pad_token_id (`int`, *optional*):
44 Padding token id.
45 bos_token_id (`int`, *optional*):
46 Beginning-of-sequence token id.
47 eos_token_id (`int` or `List[int]`, *optional*):
48 End-of-sequence token id(s).
49 rope_parameters (`dict`, *optional*):
50 The parameters of the RoPE embeddings.
51 qk_norm (`bool`, *optional*, defaults to `True`):
52 Whether to apply RMSNorm to query and key states before attention.
53 tie_word_embeddings (`bool`, *optional*, defaults to `False`):
54 Whether to tie input and output embedding weights.
55 enable_attention_fp32_softmax (`bool`, *optional*, defaults to `False`):
56 Whether to upcast attention softmax to float32. Note: the eager attention
57 path always computes softmax in float32 regardless of this setting; this
58 flag is reserved for future use with custom attention backends.
59 enable_lm_head_fp32 (`bool`, *optional*, defaults to `True`):
60 Whether to upcast the LM head computation to float32.
61 num_experts (`int`, *optional*, defaults to 192):
62 Total number of MoE experts.
63 num_experts_per_tok (`int`, *optional*, defaults to 8):
64 Number of experts selected per token (top-k routing).
65 num_shared_experts (`int`, *optional*, defaults to 1):
66 Number of always-active shared experts combined into a single MLP.
67 expert_hidden_dim (`int`, *optional*, defaults to 1536):
68 Intermediate dimension of each individual MoE expert.
69 moe_router_enable_expert_bias (`bool`, *optional*, defaults to `True`):
70 Whether to use per-expert load-balancing bias in the router.
71 moe_router_use_sigmoid (`bool`, *optional*, defaults to `True`):
72 Whether to use sigmoid (instead of softmax) for router scoring.
73 route_norm (`bool`, *optional*, defaults to `True`):
74 Whether to normalize routing scores when using sigmoid routing.
75 router_scaling_factor (`float`, *optional*):
76 Optional multiplicative scaling factor applied to routing scores.
77 use_grouped_mm (`bool`, *optional*, defaults to `False`):
78 Whether to use grouped GEMM for expert computation (not yet implemented).
79 enable_moe_fp32_combine (`bool`, *optional*, defaults to `False`):
80 Whether to accumulate expert outputs in float32.
81 first_k_dense_replace (`int`, *optional*, defaults to 1):
82 Number of initial decoder layers that use a dense FFN instead of MoE.
83 output_router_logits (`bool`, *optional*, defaults to `False`):
84 Whether to output router logits from each MoE layer. Useful for computing
85 auxiliary load-balancing loss during training. Disabled by default to avoid
86 the memory overhead of storing per-layer router tensors during inference.
87
88 Example:
89 ```python
90 >>> from transformers import HYV3Config, HYV3Model
91
92 >>> config = HYV3Config()
93 >>> model = HYV3Model(config)
94 ```
95
96 """
97
98 model_type = "hy_v3"
99 keys_to_ignore_at_inference = ["past_key_values"]
100
101 def __init__(
102 self,
103 vocab_size=120832,
104 hidden_size=4096,
105 intermediate_size=13312,
106 num_hidden_layers=80,
107 num_attention_heads=64,
108 num_key_value_heads=8,
109 head_dim=128,
110 hidden_act="silu",
111 max_position_embeddings=131072,
112 initializer_range=0.006,
113 rms_norm_eps=1e-5,
114 use_cache=True,
115 pad_token_id=None,
116 bos_token_id=None,
117 eos_token_id=None,
118 rope_parameters: dict[str, Any] | None = None,
119 qk_norm=True,
120 tie_word_embeddings=False,
121 enable_attention_fp32_softmax=False,
122 enable_lm_head_fp32=True,
123 # MoE specific
124 num_experts=192,
125 num_experts_per_tok=8,
126 num_shared_experts=1,
127 expert_hidden_dim=1536,
128 moe_router_enable_expert_bias=True,
129 moe_router_use_sigmoid=True,
130 route_norm=True,
131 router_scaling_factor=None,
132 use_grouped_mm=False,
133 enable_moe_fp32_combine=False,
134 # Dense/MoE layer control
135 first_k_dense_replace=1,
136 output_router_logits=False,
137 **kwargs,
138 ):
139 self.vocab_size = vocab_size
140 self.hidden_size = hidden_size
141 self.intermediate_size = intermediate_size
142 self.num_hidden_layers = num_hidden_layers
143 self.num_attention_heads = num_attention_heads
144 self.num_key_value_heads = num_key_value_heads
145 self.head_dim = head_dim
146 self.hidden_act = hidden_act
147 self.max_position_embeddings = max_position_embeddings
148 self.initializer_range = initializer_range
149 self.rms_norm_eps = rms_norm_eps
150 self.use_cache = use_cache
151 rope_theta = kwargs.pop("rope_theta", 11158840.0)
152 if rope_parameters is None:
153 rope_parameters = {"rope_type": "default", "rope_theta": rope_theta}
154 self.rope_parameters = rope_parameters
155 self.qk_norm = qk_norm
156 self.tie_word_embeddings = tie_word_embeddings
157 self.enable_lm_head_fp32 = enable_lm_head_fp32
158 self.enable_attention_fp32_softmax = enable_attention_fp32_softmax
159
160 # MoE specific
161 self.num_experts = num_experts
162 self.num_experts_per_tok = num_experts_per_tok
163 self.num_shared_experts = num_shared_experts
164 self.expert_hidden_dim = expert_hidden_dim
165 self.moe_router_enable_expert_bias = moe_router_enable_expert_bias
166 self.moe_router_use_sigmoid = moe_router_use_sigmoid
167 self.route_norm = route_norm
168 self.use_grouped_mm = use_grouped_mm
169 self.router_scaling_factor = router_scaling_factor
170 self.enable_moe_fp32_combine = enable_moe_fp32_combine
171
172 # Dense/MoE layer control
173 self.first_k_dense_replace = first_k_dense_replace
174 self.output_router_logits = output_router_logits
175
176 if eos_token_id is not None and isinstance(eos_token_id, int):
177 eos_token_id = [eos_token_id]
178
179 super().__init__(
180 pad_token_id=pad_token_id,
181 bos_token_id=bos_token_id,
182 eos_token_id=eos_token_id,
183 tie_word_embeddings=tie_word_embeddings,
184 **kwargs,
185 )