@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 285 lines 12.5 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
5 from transformers.configuration_utils import PreTrainedConfig
6
7
8 class OlmoHybridConfig(PreTrainedConfig):
9 r"""Configuration class for [`OlmoHybridModel`]. It is used to
10 instantiate an OLMo Hybrid model according to the specified
11 arguments, defining the model architecture. Instantiating a
12 configuration with the defaults will yield a similar
13 configuration to that of the
14 [allenai/Olmo-Hybrid-7B](https://huggingface.co/allenai/Olmo-Hybrid-7B)
15 model.
16
17 Configuration objects inherit from [`PreTrainedConfig`] and
18 can be used to control the model outputs. Read the
19 documentation from [`PreTrainedConfig`] for more information.
20
21 Args:
22 vocab_size (`int`, *optional*, defaults to 100352):
23 Vocabulary size of the OlmoHybrid model. Defines
24 the number of different tokens that can be
25 represented by the `inputs_ids` passed when
26 calling [`OlmoHybridModel`].
27 hidden_size (`int`, *optional*, defaults to 3840):
28 Dimension of the hidden representations.
29 intermediate_size (`int`, *optional*,
30 defaults to 11008):
31 Dimension of the MLP representations.
32 num_hidden_layers (`int`, *optional*,
33 defaults to 32):
34 Number of hidden layers in the Transformer
35 decoder.
36 num_attention_heads (`int`, *optional*,
37 defaults to 30):
38 Number of attention heads for each attention
39 layer in the Transformer decoder.
40 num_key_value_heads (`int`, *optional*):
41 This is the number of key_value heads that
42 should be used to implement Grouped Query
43 Attention. If
44 `num_key_value_heads=num_attention_heads`,
45 the model will use Multi Head Attention (MHA),
46 if `num_key_value_heads=1` the model will use
47 Multi Query Attention (MQA) otherwise GQA is
48 used. When converting a multi-head checkpoint
49 to a GQA checkpoint, each group key and value
50 head should be constructed by meanpooling all
51 the original heads within that group. For more
52 details, check out
53 [this paper](https://huggingface.co/papers/2305.13245).
54 If it is not specified, will default to
55 `num_attention_heads`.
56 hidden_act (`str` or `function`, *optional*,
57 defaults to `"silu"`):
58 The non-linear activation function (function
59 or string) in the decoder.
60 max_position_embeddings (`int`, *optional*,
61 defaults to 65536):
62 The maximum sequence length that this model
63 might ever be used with.
64 initializer_range (`float`, *optional*,
65 defaults to 0.02):
66 The standard deviation of the
67 truncated_normal_initializer for initializing
68 all weight matrices.
69 use_cache (`bool`, *optional*, defaults to `True`):
70 Whether or not the model should return the last
71 key/values attentions (not used by all models).
72 Only relevant if `config.is_decoder=True`.
73 pad_token_id (`int`, *optional*,
74 defaults to 100277):
75 Padding token id.
76 bos_token_id (`int`, *optional*):
77 Beginning of stream token id.
78 eos_token_id (`int`, *optional*,
79 defaults to 100257):
80 End of stream token id.
81 tie_word_embeddings (`bool`, *optional*,
82 defaults to `False`):
83 Whether to tie weight embeddings.
84 rope_parameters (`RopeParameters`, *optional*):
85 Dictionary containing the configuration
86 parameters for the RoPE embeddings. Can be
87 `None` to disable RoPE.
88 attention_bias (`bool`, *optional*,
89 defaults to `False`):
90 Whether to use a bias in the query, key, value
91 and output projection layers during
92 self-attention.
93 attention_dropout (`float`, *optional*,
94 defaults to 0.0):
95 The dropout ratio for the attention
96 probabilities.
97 rms_norm_eps (`float`, *optional*,
98 defaults to 1e-06):
99 The epsilon used by the rms normalization
100 layers.
101 layer_types (`list`, *optional*):
102 Attention pattern for each layer. Can contain
103 `"full_attention"` or `"linear_attention"`.
104 Defaults to linear attention for most layers
105 with full attention for every 4th layer.
106 linear_num_key_heads (`int`, *optional*):
107 Number of key heads for the linear attention
108 layers. Defaults to `num_attention_heads`.
109 linear_num_value_heads (`int`, *optional*):
110 Number of value heads for the linear attention
111 layers. Defaults to `num_attention_heads`.
112 linear_key_head_dim (`int`, *optional*):
113 Dimension of each key head in linear attention
114 layers. Defaults to
115 `0.75 * hidden_size / linear_num_key_heads`.
116 linear_value_head_dim (`int`, *optional*):
117 Dimension of each value head in linear
118 attention layers. Defaults to
119 `2 * linear_key_head_dim`.
120 linear_a_log_min (`float`, *optional*,
121 defaults to 0.0):
122 Minimum value for uniform initialization of
123 A_log in GatedDeltaNet layers.
124 linear_a_log_max (`float`, *optional*,
125 defaults to 16.0):
126 Maximum value for uniform initialization of
127 A_log in GatedDeltaNet layers.
128 linear_dt_min (`float`, *optional*,
129 defaults to 0.001):
130 Minimum value for dt initialization in
131 GatedDeltaNet layers.
132 linear_dt_max (`float`, *optional*,
133 defaults to 0.1):
134 Maximum value for dt initialization in
135 GatedDeltaNet layers.
136 linear_dt_init_floor (`float`, *optional*,
137 defaults to 0.0001):
138 Floor value for clamping dt during
139 initialization in GatedDeltaNet layers.
140 linear_conv_kernel_dim (`int`, *optional*,
141 defaults to 4):
142 Kernel size for the short convolution applied
143 to queries, keys, and values in linear
144 attention layers.
145 linear_allow_neg_eigval (`bool`, *optional*,
146 defaults to `True`):
147 Whether to allow negative eigenvalues in the
148 GatedDeltaNet recurrence. When `True`, the
149 beta parameter is scaled by 2.0 to allow
150 values in range [0, 2] instead of [0, 1].
151 ```python
152 >>> from transformers import (
153 ... OlmoHybridModel,
154 ... OlmoHybridConfig,
155 ... )
156
157 >>> configuration = OlmoHybridConfig()
158 >>> model = OlmoHybridModel(configuration)
159 >>> configuration = model.config
160 ```
161
162 """
163
164 model_type = "olmo_hybrid"
165 keys_to_ignore_at_inference = ["past_key_values"]
166 base_model_tp_plan = {
167 "layers.*.self_attn.q_proj": "colwise_gather_output",
168 "layers.*.self_attn.k_proj": "colwise_gather_output",
169 "layers.*.self_attn.v_proj": "colwise_gather_output",
170 "layers.*.self_attn.o_proj": "rowwise_split_input",
171 "layers.*.mlp.gate_proj": "colwise",
172 "layers.*.mlp.up_proj": "colwise",
173 "layers.*.mlp.down_proj": "rowwise",
174 }
175 base_model_pp_plan = {
176 "embed_tokens": (["input_ids"], ["inputs_embeds"]),
177 "layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
178 "norm": (["hidden_states"], ["hidden_states"]),
179 }
180
181 def __init__(
182 self,
183 vocab_size: int | None = 100352,
184 hidden_size: int | None = 3840,
185 intermediate_size: int | None = 11008,
186 num_hidden_layers: int | None = 32,
187 num_attention_heads: int | None = 30,
188 num_key_value_heads: int | None = None,
189 hidden_act: str | None = "silu",
190 max_position_embeddings: int | None = 65536,
191 initializer_range: float | None = 0.02,
192 use_cache: bool | None = True,
193 pad_token_id: int | None = 100277,
194 bos_token_id: int | None = None,
195 eos_token_id: int | None = 100257,
196 tie_word_embeddings: bool | None = False,
197 rope_parameters=None,
198 attention_bias: bool | None = False,
199 attention_dropout: float | None = 0.0,
200 rms_norm_eps: float | None = 1e-06,
201 layer_types: list[str] | None = None,
202 linear_num_key_heads: int | None = None,
203 linear_num_value_heads: int | None = None,
204 linear_key_head_dim: int | None = None,
205 linear_value_head_dim: int | None = None,
206 linear_a_log_min: float = 0.0,
207 linear_a_log_max: float = 16.0,
208 linear_dt_min: float = 0.001,
209 linear_dt_max: float = 0.1,
210 linear_dt_init_floor: float = 1e-4,
211 linear_conv_kernel_dim: int = 4,
212 linear_allow_neg_eigval: bool = True,
213 **kwargs,
214 ):
215 super().__init__(**kwargs)
216
217 assert num_hidden_layers is not None
218 assert hidden_size is not None
219 assert num_attention_heads is not None
220
221 if layer_types is None:
222 # Default: linear attention for most layers, full attention every 4th layer
223 layer_types = ["linear_attention"] * int(num_hidden_layers)
224 for i in range(int(num_hidden_layers)):
225 if i % 4 == 3:
226 layer_types[i] = "full_attention"
227 # Ensure at least one full attention layer for small num_hidden_layers
228 if "full_attention" not in layer_types:
229 layer_types[-1] = "full_attention"
230
231 self.layer_types = layer_types
232 self.validate_layer_type()
233 if "linear_attention" not in layer_types:
234 raise ValueError(
235 "OLMoHybrid expects at least one 'linear_attention' layer."
236 )
237 if all(t == "linear_attention" for t in layer_types):
238 raise ValueError("OLMoHybrid expects at least one attention layer.")
239
240 self.layer_types = layer_types
241
242 if linear_num_key_heads is None:
243 linear_num_key_heads = num_attention_heads
244 if linear_num_value_heads is None:
245 linear_num_value_heads = num_attention_heads
246 if linear_key_head_dim is None:
247 linear_key_head_dim = int(0.75 * hidden_size / linear_num_key_heads)
248 if linear_value_head_dim is None:
249 linear_value_head_dim = 2 * linear_key_head_dim
250
251 self.linear_num_key_heads = linear_num_key_heads
252 self.linear_num_value_heads = linear_num_value_heads
253 self.linear_key_head_dim = linear_key_head_dim
254 self.linear_value_head_dim = linear_value_head_dim
255 self.linear_a_log_min = linear_a_log_min
256 self.linear_a_log_max = linear_a_log_max
257 self.linear_dt_min = linear_dt_min
258 self.linear_dt_max = linear_dt_max
259 self.linear_dt_init_floor = linear_dt_init_floor
260 self.linear_conv_kernel_dim = linear_conv_kernel_dim
261 self.linear_allow_neg_eigval = linear_allow_neg_eigval
262 self.vocab_size = vocab_size
263 self.max_position_embeddings = max_position_embeddings
264 self.hidden_size = hidden_size
265 self.intermediate_size = intermediate_size
266 self.num_hidden_layers = num_hidden_layers
267 self.num_attention_heads = num_attention_heads
268
269 # for backward compatibility
270 if num_key_value_heads is None:
271 num_key_value_heads = num_attention_heads
272
273 self.num_key_value_heads = num_key_value_heads
274 self.hidden_act = hidden_act
275 self.initializer_range = initializer_range
276 self.rms_norm_eps = rms_norm_eps
277 self.use_cache = use_cache
278 self.attention_bias = attention_bias
279 self.attention_dropout = attention_dropout
280 self.rope_parameters = rope_parameters
281
282 self.tie_word_embeddings = tie_word_embeddings
283 self.pad_token_id = pad_token_id
284 self.bos_token_id = bos_token_id
285 self.eos_token_id = eos_token_id