@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 219 lines 9.52 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 # Copyright 2024 HuggingFace Inc. team. All rights reserved.
5 # Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
6 #
7 # Licensed under the Apache License, Version 2.0 (the "License");
8 # you may not use this file except in compliance with the License.
9 # You may obtain a copy of the License at
10 #
11 # http://www.apache.org/licenses/LICENSE-2.0
12 #
13 # Unless required by applicable law or agreed to in writing, software
14 # distributed under the License is distributed on an "AS IS" BASIS,
15 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
16 # See the License for the specific language governing permissions and
17 # limitations under the License.
18 """Nemotron model configuration."""
19
20 from transformers import PreTrainedConfig
21 from transformers.utils import logging
22
23 logger = logging.get_logger(__name__)
24
25
26 class NemotronConfig(PreTrainedConfig):
27 r"""This is the configuration class to store the configuration of a
28 [`NemotronModel`]. It is used to instantiate a Nemotron model
29 according to the specified arguments, defining the model architecture.
30 Instantiating a configuration with the defaults will yield a similar
31 configuration to that of the Nemotron-8B.
32
33 Configuration objects inherit from [`PreTrainedConfig`] and can be
34 used to control the model outputs. Read the documentation from
35 [`PreTrainedConfig`] for more information.
36
37
38 Args:
39 vocab_size (`int`, *optional*, defaults to 256000):
40 Vocabulary size of the Nemotron model. Defines the number of
41 different tokens that can be represented by the
42 `inputs_ids` passed when calling [`NemotronModel`]
43 hidden_size (`int`, *optional*, defaults to 6144):
44 Dimension of the hidden representations.
45 intermediate_size (`int`, *optional*, defaults to 24576):
46 Dimension of the MLP representations.
47 num_hidden_layers (`int`, *optional*, defaults to 32):
48 Number of hidden layers in the Transformer decoder.
49 num_attention_heads (`int`, *optional*, defaults to 48):
50 Number of attention heads for each attention layer in the
51 Transformer decoder.
52 head_dim (`int`, *optional*):
53 Projection weights dimension in multi-head attention. Set to
54 hidden_size // num_attention_heads if None
55 num_key_value_heads (`int`, *optional*):
56 This is the number of key_value heads that should be used to
57 implement Grouped Query Attention. If
58 `num_key_value_heads=num_attention_heads`, the model will use
59 Multi Head Attention (MHA), if
60 `num_key_value_heads=1 the model will use Multi Query Attention
61 (MQA) otherwise GQA is used. When converting a multi-head
62 checkpoint to a GQA checkpoint, each group key and value
63 head should be constructed by meanpooling all the original
64 heads within that group. For more details checkout
65 [this paper](https://arxiv.org/pdf/2305.13245.pdf). If it
66 is not specified, will default to `num_attention_heads`.
67 hidden_act (`str` or `function`, *optional*, defaults to `"relu2"`):
68 The non-linear activation function (function or string) in the
69 decoder.
70 max_position_embeddings (`int`, *optional*, defaults to 4096):
71 The maximum sequence length that this model might ever be used
72 with.
73 initializer_range (`float`, *optional*, defaults to 0.0134):
74 The standard deviation of the truncated_normal_initializer for
75 initializing all weight matrices.
76 norm_eps (`float`, *optional*, defaults to 1e-05):
77 The epsilon used by the normalization layers.
78 use_cache (`bool`, *optional*, defaults to `True`):
79 Whether or not the model should return the last key/values
80 attentions (not used by all models). Only relevant if
81 `config.is_decoder=True`.
82 pad_token_id (`int`, *optional*):
83 Padding token id.
84 bos_token_id (`int`, *optional*, defaults to 2):
85 Beginning of stream token id.
86 eos_token_id (`int`, *optional*, defaults to 3):
87 End of stream token id.
88 tie_word_embeddings (`bool`, *optional*, defaults to `False`):
89 Whether to tie weight embeddings
90 rope_parameters (`dict`, *optional*):
91 The parameters of the RoPE embeddings. Expected contents:
92 `rope_theta` (`float`): The base period of the RoPE embeddings.
93 `rope_type` (`str`):
94 The sub-variant of RoPE to use. Can be one of ['default', 'linear',
95 'dynamic', 'yarn', 'longrope', 'llama3'], with 'default' being the
96 original RoPE implementation.
97 `partial_rotary_factor` (`float`, *optional*, defaults to 0.5):
98 Percentage of the query and keys which will have rotary embedding.
99 attention_bias (`bool`, *optional*, defaults to `False`):
100 Whether to use a bias in the query, key, value and output
101 projection layers during self-attention.
102 attention_dropout (`float`, *optional*, defaults to 0.0):
103 The dropout ratio for the attention probabilities.
104 mlp_bias (`bool`, *optional*, defaults to `False`):
105 Whether to use a bias in up_proj and down_proj layers in the MLP
106 layers.
107
108 ```python
109 >>> from transformers import NemotronModel, NemotronConfig
110 >>> # Initializing a Nemotron nemotron-15b style configuration
111 >>> configuration = NemotronConfig()
112 >>> # Initializing a model from the nemotron-15b style configuration
113 >>> model = NemotronModel(configuration)
114 >>> # Accessing the model configuration
115 >>> configuration = model.config
116 ```
117
118 """
119
120 model_type = "nemotron"
121 keys_to_ignore_at_inference = ["past_key_values"]
122
123 def __init__(
124 self,
125 vocab_size=256000,
126 hidden_size=6144,
127 intermediate_size=24576,
128 num_hidden_layers=32,
129 num_attention_heads=48,
130 head_dim=None,
131 num_key_value_heads=None,
132 hidden_act="relu2",
133 max_position_embeddings=4096,
134 initializer_range=0.0134,
135 norm_eps=1e-5,
136 use_cache=True,
137 pad_token_id=None,
138 bos_token_id=2,
139 eos_token_id=3,
140 tie_word_embeddings=False,
141 rope_parameters=None,
142 attention_bias=False,
143 attention_dropout=0.0,
144 mlp_bias=False,
145 **kwargs,
146 ):
147 self.vocab_size = vocab_size
148 self.max_position_embeddings = max_position_embeddings
149 self.hidden_size = hidden_size
150 self.intermediate_size = intermediate_size
151 self.num_hidden_layers = num_hidden_layers
152 self.num_attention_heads = num_attention_heads
153 head_dim = head_dim or kwargs.get("kv_channels")
154 self.head_dim = (
155 head_dim if head_dim is not None else (hidden_size // num_attention_heads)
156 )
157
158 # for backward compatibility
159 if num_key_value_heads is None:
160 num_key_value_heads = num_attention_heads
161
162 self.num_key_value_heads = num_key_value_heads
163 self.hidden_act = hidden_act
164 self.initializer_range = initializer_range
165 self.norm_eps = norm_eps
166 self.use_cache = use_cache
167 # Try to set `rope_scaling` if available, otherwise use `rope_parameters`
168 rope_scaling = kwargs.pop("rope_scaling", None)
169 rope_parameters = rope_scaling or rope_parameters or {"rope_type": "default"}
170 rope_theta = kwargs.pop("rope_theta", 10000.0)
171 if "rope_theta" not in rope_parameters:
172 rope_parameters["rope_theta"] = rope_theta
173 # for backward compatibility
174 partial_rotary_factor = (
175 kwargs.get("rope_percent")
176 or kwargs.get("rope_percentage")
177 or kwargs.get("partial_rotary_factor")
178 or 0.5
179 )
180 if "partial_rotary_factor" not in rope_parameters:
181 rope_parameters["partial_rotary_factor"] = partial_rotary_factor
182 self.rope_parameters = rope_parameters
183 self._rope_parameters_validation()
184 self.attention_bias = attention_bias
185 self.attention_dropout = attention_dropout
186 self.mlp_bias = mlp_bias
187
188 super().__init__(
189 pad_token_id=pad_token_id,
190 bos_token_id=bos_token_id,
191 eos_token_id=eos_token_id,
192 tie_word_embeddings=tie_word_embeddings,
193 **kwargs,
194 )
195
196 def _rope_parameters_validation(self):
197 """Validate the `rope_parameters` configuration."""
198 if self.rope_parameters is None:
199 return
200
201 rope_type: str | None = self.rope_parameters.get("rope_type", None)
202 factor: float | None = self.rope_parameters.get("factor", None)
203
204 if rope_type not in {"default", "linear", "dynamic"}:
205 raise ValueError(
206 "`rope_type` must be one of ['default', 'linear', 'dynamic'], "
207 f"got {rope_type}"
208 )
209 if rope_type != "default":
210 if factor is None:
211 raise ValueError(
212 "If `rope_type` is not 'default', `rope_parameters` "
213 "must include a `factor` field. Got `None`."
214 )
215 if not isinstance(factor, float) or factor <= 1.0:
216 raise ValueError(
217 "`rope_parameters`'s factor field must be a float > 1, got "
218 f"{factor}"
219 )