vllm-npu 0.4.2__py3-none-any.whl

Sign up to get free protection for your applications and to get access to all the features.
Files changed (219) hide show
  1. vllm/__init__.py +23 -0
  2. vllm/_custom_ops.py +251 -0
  3. vllm/attention/__init__.py +13 -0
  4. vllm/attention/backends/__init__.py +0 -0
  5. vllm/attention/backends/abstract.py +127 -0
  6. vllm/attention/backends/flash_attn.py +271 -0
  7. vllm/attention/backends/flashinfer.py +220 -0
  8. vllm/attention/backends/rocm_flash_attn.py +374 -0
  9. vllm/attention/backends/torch_sdpa.py +250 -0
  10. vllm/attention/backends/xformers.py +393 -0
  11. vllm/attention/layer.py +56 -0
  12. vllm/attention/ops/__init__.py +0 -0
  13. vllm/attention/ops/paged_attn.py +216 -0
  14. vllm/attention/ops/prefix_prefill.py +792 -0
  15. vllm/attention/ops/triton_flash_attention.py +810 -0
  16. vllm/attention/selector.py +91 -0
  17. vllm/block.py +84 -0
  18. vllm/config.py +1225 -0
  19. vllm/core/__init__.py +0 -0
  20. vllm/core/block/__init__.py +0 -0
  21. vllm/core/block/block_table.py +295 -0
  22. vllm/core/block/common.py +199 -0
  23. vllm/core/block/cpu_gpu_block_allocator.py +228 -0
  24. vllm/core/block/interfaces.py +205 -0
  25. vllm/core/block/naive_block.py +318 -0
  26. vllm/core/block/prefix_caching_block.py +606 -0
  27. vllm/core/block_manager_v1.py +625 -0
  28. vllm/core/block_manager_v2.py +258 -0
  29. vllm/core/evictor_v1.py +105 -0
  30. vllm/core/evictor_v2.py +127 -0
  31. vllm/core/interfaces.py +113 -0
  32. vllm/core/policy.py +45 -0
  33. vllm/core/scheduler.py +1163 -0
  34. vllm/distributed/__init__.py +3 -0
  35. vllm/distributed/communication_op.py +237 -0
  36. vllm/distributed/device_communicators/__init__.py +0 -0
  37. vllm/distributed/device_communicators/custom_all_reduce.py +274 -0
  38. vllm/distributed/device_communicators/pynccl.py +287 -0
  39. vllm/distributed/device_communicators/pynccl_utils.py +66 -0
  40. vllm/distributed/parallel_state.py +339 -0
  41. vllm/distributed/utils.py +136 -0
  42. vllm/engine/__init__.py +0 -0
  43. vllm/engine/arg_utils.py +649 -0
  44. vllm/engine/async_llm_engine.py +737 -0
  45. vllm/engine/llm_engine.py +784 -0
  46. vllm/engine/metrics.py +368 -0
  47. vllm/engine/output_processor/__init__.py +0 -0
  48. vllm/engine/output_processor/interfaces.py +76 -0
  49. vllm/engine/output_processor/multi_step.py +142 -0
  50. vllm/engine/output_processor/single_step.py +284 -0
  51. vllm/engine/output_processor/stop_checker.py +101 -0
  52. vllm/engine/output_processor/util.py +19 -0
  53. vllm/entrypoints/__init__.py +0 -0
  54. vllm/entrypoints/api_server.py +119 -0
  55. vllm/entrypoints/llm.py +259 -0
  56. vllm/entrypoints/openai/__init__.py +0 -0
  57. vllm/entrypoints/openai/api_server.py +186 -0
  58. vllm/entrypoints/openai/cli_args.py +115 -0
  59. vllm/entrypoints/openai/protocol.py +460 -0
  60. vllm/entrypoints/openai/serving_chat.py +392 -0
  61. vllm/entrypoints/openai/serving_completion.py +347 -0
  62. vllm/entrypoints/openai/serving_engine.py +234 -0
  63. vllm/envs.py +217 -0
  64. vllm/executor/__init__.py +0 -0
  65. vllm/executor/cpu_executor.py +152 -0
  66. vllm/executor/distributed_gpu_executor.py +115 -0
  67. vllm/executor/executor_base.py +115 -0
  68. vllm/executor/gpu_executor.py +150 -0
  69. vllm/executor/multiproc_worker_utils.py +263 -0
  70. vllm/executor/neuron_executor.py +91 -0
  71. vllm/executor/ray_gpu_executor.py +327 -0
  72. vllm/executor/ray_utils.py +119 -0
  73. vllm/logger.py +153 -0
  74. vllm/logging/__init__.py +5 -0
  75. vllm/logging/formatter.py +15 -0
  76. vllm/lora/__init__.py +0 -0
  77. vllm/lora/fully_sharded_layers.py +262 -0
  78. vllm/lora/layers.py +1181 -0
  79. vllm/lora/lora.py +167 -0
  80. vllm/lora/models.py +645 -0
  81. vllm/lora/punica.py +213 -0
  82. vllm/lora/request.py +32 -0
  83. vllm/lora/utils.py +98 -0
  84. vllm/lora/worker_manager.py +251 -0
  85. vllm/model_executor/__init__.py +7 -0
  86. vllm/model_executor/guided_decoding/__init__.py +25 -0
  87. vllm/model_executor/guided_decoding/lm_format_enforcer_decoding.py +70 -0
  88. vllm/model_executor/guided_decoding/outlines_decoding.py +130 -0
  89. vllm/model_executor/guided_decoding/outlines_logits_processors.py +184 -0
  90. vllm/model_executor/layers/__init__.py +0 -0
  91. vllm/model_executor/layers/activation.py +173 -0
  92. vllm/model_executor/layers/fused_moe/__init__.py +7 -0
  93. vllm/model_executor/layers/fused_moe/configs/E=16,N=1344,device_name=NVIDIA_A100-SXM4-40GB.json +146 -0
  94. vllm/model_executor/layers/fused_moe/configs/E=16,N=1344,device_name=NVIDIA_A100-SXM4-80GB.json +146 -0
  95. vllm/model_executor/layers/fused_moe/configs/E=16,N=1344,device_name=NVIDIA_H100_80GB_HBM3.json +146 -0
  96. vllm/model_executor/layers/fused_moe/configs/E=16,N=2688,device_name=NVIDIA_A100-SXM4-80GB.json +146 -0
  97. vllm/model_executor/layers/fused_moe/configs/E=16,N=2688,device_name=NVIDIA_H100_80GB_HBM3.json +146 -0
  98. vllm/model_executor/layers/fused_moe/configs/E=8,N=1792,device_name=NVIDIA_A100-SXM4-40GB.json +146 -0
  99. vllm/model_executor/layers/fused_moe/configs/E=8,N=1792,device_name=NVIDIA_A100-SXM4-80GB.json +146 -0
  100. vllm/model_executor/layers/fused_moe/configs/E=8,N=1792,device_name=NVIDIA_H100_80GB_HBM3.json +146 -0
  101. vllm/model_executor/layers/fused_moe/configs/E=8,N=2048,device_name=NVIDIA_A100-SXM4-80GB.json +146 -0
  102. vllm/model_executor/layers/fused_moe/configs/E=8,N=2048,device_name=NVIDIA_H100_80GB_HBM3.json +146 -0
  103. vllm/model_executor/layers/fused_moe/configs/E=8,N=3584,device_name=NVIDIA_A100-SXM4-40GB.json +146 -0
  104. vllm/model_executor/layers/fused_moe/configs/E=8,N=3584,device_name=NVIDIA_A100-SXM4-80GB.json +146 -0
  105. vllm/model_executor/layers/fused_moe/configs/E=8,N=3584,device_name=NVIDIA_H100_80GB_HBM3,dtype=float8.json +140 -0
  106. vllm/model_executor/layers/fused_moe/configs/E=8,N=3584,device_name=NVIDIA_H100_80GB_HBM3.json +146 -0
  107. vllm/model_executor/layers/fused_moe/configs/E=8,N=4096,device_name=NVIDIA_A100-SXM4-80GB.json +146 -0
  108. vllm/model_executor/layers/fused_moe/configs/E=8,N=4096,device_name=NVIDIA_H100_80GB_HBM3.json +146 -0
  109. vllm/model_executor/layers/fused_moe/configs/E=8,N=7168,device_name=NVIDIA_A100-SXM4-80GB.json +146 -0
  110. vllm/model_executor/layers/fused_moe/configs/E=8,N=7168,device_name=NVIDIA_H100_80GB_HBM3,dtype=float8.json +146 -0
  111. vllm/model_executor/layers/fused_moe/configs/E=8,N=7168,device_name=NVIDIA_H100_80GB_HBM3.json +146 -0
  112. vllm/model_executor/layers/fused_moe/fused_moe.py +479 -0
  113. vllm/model_executor/layers/layernorm.py +71 -0
  114. vllm/model_executor/layers/linear.py +709 -0
  115. vllm/model_executor/layers/logits_processor.py +115 -0
  116. vllm/model_executor/layers/ops/__init__.py +0 -0
  117. vllm/model_executor/layers/ops/rand.py +157 -0
  118. vllm/model_executor/layers/ops/sample.py +406 -0
  119. vllm/model_executor/layers/quantization/__init__.py +35 -0
  120. vllm/model_executor/layers/quantization/aqlm.py +376 -0
  121. vllm/model_executor/layers/quantization/awq.py +175 -0
  122. vllm/model_executor/layers/quantization/base_config.py +97 -0
  123. vllm/model_executor/layers/quantization/fp8.py +265 -0
  124. vllm/model_executor/layers/quantization/gptq.py +224 -0
  125. vllm/model_executor/layers/quantization/gptq_marlin.py +438 -0
  126. vllm/model_executor/layers/quantization/marlin.py +227 -0
  127. vllm/model_executor/layers/quantization/schema.py +84 -0
  128. vllm/model_executor/layers/quantization/squeezellm.py +137 -0
  129. vllm/model_executor/layers/rejection_sampler.py +405 -0
  130. vllm/model_executor/layers/rotary_embedding.py +525 -0
  131. vllm/model_executor/layers/sampler.py +1051 -0
  132. vllm/model_executor/layers/vocab_parallel_embedding.py +155 -0
  133. vllm/model_executor/model_loader/__init__.py +30 -0
  134. vllm/model_executor/model_loader/loader.py +362 -0
  135. vllm/model_executor/model_loader/neuron.py +136 -0
  136. vllm/model_executor/model_loader/tensorizer.py +368 -0
  137. vllm/model_executor/model_loader/utils.py +41 -0
  138. vllm/model_executor/model_loader/weight_utils.py +372 -0
  139. vllm/model_executor/models/__init__.py +119 -0
  140. vllm/model_executor/models/baichuan.py +410 -0
  141. vllm/model_executor/models/bloom.py +327 -0
  142. vllm/model_executor/models/chatglm.py +386 -0
  143. vllm/model_executor/models/commandr.py +373 -0
  144. vllm/model_executor/models/dbrx.py +413 -0
  145. vllm/model_executor/models/decilm.py +122 -0
  146. vllm/model_executor/models/deepseek.py +438 -0
  147. vllm/model_executor/models/falcon.py +444 -0
  148. vllm/model_executor/models/gemma.py +393 -0
  149. vllm/model_executor/models/gpt2.py +266 -0
  150. vllm/model_executor/models/gpt_bigcode.py +274 -0
  151. vllm/model_executor/models/gpt_j.py +281 -0
  152. vllm/model_executor/models/gpt_neox.py +295 -0
  153. vllm/model_executor/models/internlm2.py +323 -0
  154. vllm/model_executor/models/jais.py +333 -0
  155. vllm/model_executor/models/llama.py +442 -0
  156. vllm/model_executor/models/llava.py +239 -0
  157. vllm/model_executor/models/minicpm.py +531 -0
  158. vllm/model_executor/models/mixtral.py +583 -0
  159. vllm/model_executor/models/mixtral_quant.py +404 -0
  160. vllm/model_executor/models/mpt.py +295 -0
  161. vllm/model_executor/models/olmo.py +356 -0
  162. vllm/model_executor/models/opt.py +349 -0
  163. vllm/model_executor/models/orion.py +319 -0
  164. vllm/model_executor/models/phi.py +300 -0
  165. vllm/model_executor/models/qwen.py +284 -0
  166. vllm/model_executor/models/qwen2.py +367 -0
  167. vllm/model_executor/models/qwen2_moe.py +447 -0
  168. vllm/model_executor/models/stablelm.py +301 -0
  169. vllm/model_executor/models/starcoder2.py +302 -0
  170. vllm/model_executor/models/xverse.py +366 -0
  171. vllm/model_executor/sampling_metadata.py +588 -0
  172. vllm/model_executor/utils.py +35 -0
  173. vllm/outputs.py +150 -0
  174. vllm/py.typed +2 -0
  175. vllm/sampling_params.py +340 -0
  176. vllm/sequence.py +766 -0
  177. vllm/spec_decode/__init__.py +0 -0
  178. vllm/spec_decode/batch_expansion.py +397 -0
  179. vllm/spec_decode/interfaces.py +73 -0
  180. vllm/spec_decode/metrics.py +191 -0
  181. vllm/spec_decode/multi_step_worker.py +203 -0
  182. vllm/spec_decode/ngram_worker.py +176 -0
  183. vllm/spec_decode/spec_decode_worker.py +472 -0
  184. vllm/spec_decode/top1_proposer.py +200 -0
  185. vllm/spec_decode/util.py +228 -0
  186. vllm/test_utils.py +41 -0
  187. vllm/transformers_utils/__init__.py +0 -0
  188. vllm/transformers_utils/config.py +58 -0
  189. vllm/transformers_utils/configs/__init__.py +16 -0
  190. vllm/transformers_utils/configs/chatglm.py +68 -0
  191. vllm/transformers_utils/configs/dbrx.py +278 -0
  192. vllm/transformers_utils/configs/falcon.py +87 -0
  193. vllm/transformers_utils/configs/jais.py +236 -0
  194. vllm/transformers_utils/configs/mpt.py +178 -0
  195. vllm/transformers_utils/detokenizer.py +313 -0
  196. vllm/transformers_utils/tokenizer.py +149 -0
  197. vllm/transformers_utils/tokenizer_group/__init__.py +33 -0
  198. vllm/transformers_utils/tokenizer_group/base_tokenizer_group.py +55 -0
  199. vllm/transformers_utils/tokenizer_group/ray_tokenizer_group.py +169 -0
  200. vllm/transformers_utils/tokenizer_group/tokenizer_group.py +78 -0
  201. vllm/transformers_utils/tokenizers/__init__.py +5 -0
  202. vllm/transformers_utils/tokenizers/baichuan.py +255 -0
  203. vllm/usage/__init__.py +0 -0
  204. vllm/usage/usage_lib.py +209 -0
  205. vllm/utils.py +677 -0
  206. vllm/worker/__init__.py +0 -0
  207. vllm/worker/cache_engine.py +105 -0
  208. vllm/worker/cpu_model_runner.py +346 -0
  209. vllm/worker/cpu_worker.py +321 -0
  210. vllm/worker/model_runner.py +1168 -0
  211. vllm/worker/neuron_model_runner.py +196 -0
  212. vllm/worker/neuron_worker.py +98 -0
  213. vllm/worker/worker.py +345 -0
  214. vllm/worker/worker_base.py +146 -0
  215. vllm_npu-0.4.2.dist-info/LICENSE +201 -0
  216. vllm_npu-0.4.2.dist-info/METADATA +173 -0
  217. vllm_npu-0.4.2.dist-info/RECORD +219 -0
  218. vllm_npu-0.4.2.dist-info/WHEEL +5 -0
  219. vllm_npu-0.4.2.dist-info/top_level.txt +1 -0
@@ -0,0 +1,87 @@
1
+ # Adapted from
2
+ # https://huggingface.co/tiiuae/falcon-7b/blob/main/configuration_RW.py
3
+ # Copyright 2023 The vLLM team.
4
+ # Copyright 2022 the Big Science Workshop and HuggingFace Inc. team.
5
+ # 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
+ """Falcon configuration"""
19
+ from transformers.configuration_utils import PretrainedConfig
20
+
21
+
22
+ class RWConfig(PretrainedConfig):
23
+ model_type = "falcon"
24
+ keys_to_ignore_at_inference = ["past_key_values"]
25
+ attribute_map = {
26
+ "num_hidden_layers": "n_layer",
27
+ "num_attention_heads": "n_head",
28
+ "num_kv_heads": "n_head_kv",
29
+ }
30
+
31
+ def __init__(
32
+ self,
33
+ vocab_size=250880,
34
+ hidden_size=64,
35
+ n_layer=2,
36
+ n_head=8,
37
+ layer_norm_epsilon=1e-5,
38
+ initializer_range=0.02,
39
+ use_cache=True,
40
+ bos_token_id=1,
41
+ eos_token_id=2,
42
+ hidden_dropout=0.0,
43
+ attention_dropout=0.0,
44
+ multi_query=True,
45
+ n_head_kv=None,
46
+ alibi=False,
47
+ bias=False,
48
+ parallel_attn=False,
49
+ new_decoder_architecture=False,
50
+ **kwargs,
51
+ ) -> None:
52
+ self.vocab_size = vocab_size
53
+ # Backward compatibility with n_embed kwarg
54
+ n_embed = kwargs.pop("n_embed", None)
55
+ self.hidden_size = hidden_size if n_embed is None else n_embed
56
+ self.n_layer = n_layer
57
+ self.n_head = n_head
58
+ self.layer_norm_epsilon = layer_norm_epsilon
59
+ self.initializer_range = initializer_range
60
+ self.use_cache = use_cache
61
+ self.hidden_dropout = hidden_dropout
62
+ self.attention_dropout = attention_dropout
63
+
64
+ self.bos_token_id = bos_token_id
65
+ self.eos_token_id = eos_token_id
66
+ self.multi_query = multi_query
67
+ self.n_head_kv = 1 if n_head_kv is None else n_head_kv
68
+ self.alibi = alibi
69
+ self.bias = bias
70
+ self.parallel_attn = parallel_attn
71
+ self.new_decoder_architecture = new_decoder_architecture
72
+
73
+ if self.hidden_size == 8192:
74
+ # Hack for falcon-40b
75
+ self.new_decoder_architecture = True
76
+
77
+ super().__init__(bos_token_id=bos_token_id,
78
+ eos_token_id=eos_token_id,
79
+ **kwargs)
80
+
81
+ @property
82
+ def head_dim(self):
83
+ return self.hidden_size // self.n_head
84
+
85
+ @property
86
+ def rotary(self):
87
+ return not self.alibi
@@ -0,0 +1,236 @@
1
+ # coding=utf-8
2
+ # Copyright 2023 The OpenAI Team Authors and HuggingFace Inc. team.
3
+ # Copyright (c) 2018, NVIDIA CORPORATION. All rights reserved.
4
+ # Copyright 2023 Cerebras Systems.
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
+ """JAIS 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 JAISConfig(PretrainedConfig):
26
+ """
27
+ This is the configuration class to store the configuration of a
28
+ [`JAISModel`]. It is used to instantiate a JAIS model according to the
29
+ specified arguments, defining the model architecture.
30
+
31
+ Configuration objects inherit from [`PretrainedConfig`] and can be used
32
+ to control the model outputs. Read the documentation from
33
+ [`PretrainedConfig`] for more information.
34
+
35
+
36
+ Args:
37
+ vocab_size (`int`, *optional*, defaults to 50257):
38
+ Vocabulary size of the JAIS model. Defines the number of different
39
+ tokens that can be represented by the
40
+ `inputs_ids` passed when calling [`JAISModel`].
41
+ n_positions (`int`, *optional*, defaults to 1024):
42
+ The maximum sequence length that this model might ever be used
43
+ with. Typically set this to something large just in case
44
+ (e.g., 512 or 1024 or 2048).
45
+ n_embd (`int`, *optional*, defaults to 768):
46
+ Dimensionality of the embeddings and hidden states.
47
+ n_layer (`int`, *optional*, defaults to 12):
48
+ Number of hidden layers in the Transformer encoder.
49
+ n_head (`int`, *optional*, defaults to 12):
50
+ Number of attention heads for each attention layer in the
51
+ Transformer encoder.
52
+ n_inner (`int`, *optional*, defaults to None):
53
+ Dimensionality of the inner feed-forward layers. `None` will set
54
+ it to 4 times n_embd
55
+ activation_function (`str`, *optional*, defaults to `"gelu"`):
56
+ Activation function, to be selected in the list
57
+ `["relu", "silu", "gelu", "tanh", "gelu_new", "swiglu"]`.
58
+ resid_pdrop (`float`, *optional*, defaults to 0.1):
59
+ The dropout probability for all fully connected layers in
60
+ the embeddings, encoder, and pooler.
61
+ embd_pdrop (`float`, *optional*, defaults to 0.1):
62
+ The dropout ratio for the embeddings.
63
+ attn_pdrop (`float`, *optional*, defaults to 0.1):
64
+ The dropout ratio for the attention.
65
+ layer_norm_epsilon (`float`, *optional*, defaults to 1e-5):
66
+ The epsilon to use in the layer normalization layers.
67
+ initializer_range (`float`, *optional*, defaults to 0.02):
68
+ The standard deviation of the truncated_normal_initializer for
69
+ initializing all weight matrices.
70
+ scale_attn_weights (`bool`, *optional*, defaults to `True`):
71
+ Scale attention weights by dividing by sqrt(hidden_size)..
72
+ use_cache (`bool`, *optional*, defaults to `True`):
73
+ Whether or not the model should return the last key/values
74
+ attentions (not used by all models).
75
+ scale_attn_by_inverse_layer_idx (`bool`, *optional*,
76
+ defaults to `False`):
77
+ Whether to additionally scale attention weights by
78
+ `1 / layer_idx + 1`.
79
+ reorder_and_upcast_attn (`bool`, *optional*, defaults to `False`):
80
+ Whether to scale keys (K) prior to computing attention
81
+ (dot-product)
82
+ and upcast attention dot-product/softmax to float() when training
83
+ with mixed precision.
84
+ position_embedding_type (`str`, *optional*, defaults to `"learned"`):
85
+ Positional embedding can be either `"alibi"` or `"learned"`.
86
+ mup_width_scale (`float`, *optional*, defaults to 1.0):
87
+ muP parameter to scale learning rate and initializers. Calculated
88
+ as (`d_model,0 / d_model`), where
89
+ `d_model` is the model's width and `d_model,0` is the proxy
90
+ model's width.
91
+ mup_embeddings_scale (`float`, *optional*, defaults to 1.0):
92
+ muP parameter to scale token and position embeddings.
93
+ mup_output_alpha (`float`, *optional*, defaults to 1.0):
94
+ muP parameter to scale output logits
95
+ (`output_logits_scale = mup_output_alpha * mup_width_scale`).
96
+ mup_scale_qk_dot_by_d (`bool`, *optional*, defaults to `False`):
97
+ Scale attention weights by dividing by hidden_size instead of
98
+ sqrt(hidden_size). Need to set scale_attn_weights to `True` as
99
+ well.
100
+ alibi_scaling (`Dict`, *optional*):
101
+ Dictionary containing the scaling configuration for ALiBi
102
+ embeddings. Currently only supports linear
103
+ scaling strategy. Can specify either the scaling `factor` (must be
104
+ a float greater than 1) for fixed scaling
105
+ or `train_seq_len` for dynamic scaling on input samples with
106
+ sequence length > `train_seq_len`. The expected
107
+ formats are `{"type": strategy name, "factor": scaling factor}` or
108
+ `{"type": strategy name,
109
+ "train_seq_len": training sequence length}`.
110
+ architectures (`List`, *optional*, defaults to ['JAISLMHeadModel']):
111
+ architecture names for Jais.
112
+
113
+ Example:
114
+
115
+ ```python
116
+ >>> from transformers import JAISConfig, JAISModel
117
+
118
+ >>> # Initializing a JAIS configuration
119
+ >>> configuration = JAISConfig()
120
+
121
+ >>> # Initializing a model (with random weights) from the configuration
122
+ >>> model = JAISModel(configuration)
123
+
124
+ >>> # Accessing the model configuration
125
+ >>> configuration = model.config
126
+ ```"""
127
+
128
+ model_type = "jais"
129
+ keys_to_ignore_at_inference = ["past_key_values"]
130
+ attribute_map = {
131
+ "hidden_size": "n_embd",
132
+ "max_position_embeddings": "n_positions",
133
+ "num_attention_heads": "n_head",
134
+ "num_hidden_layers": "n_layer",
135
+ }
136
+
137
+ def __init__(
138
+ self,
139
+ vocab_size=50257,
140
+ n_positions=1024,
141
+ n_embd=768,
142
+ n_layer=12,
143
+ n_head=12,
144
+ n_inner=None,
145
+ activation_function="gelu_new",
146
+ resid_pdrop=0.1,
147
+ embd_pdrop=0.1,
148
+ attn_pdrop=0.1,
149
+ layer_norm_epsilon=1e-5,
150
+ initializer_range=0.02,
151
+ scale_attn_weights=True,
152
+ use_cache=True,
153
+ bos_token_id=50256,
154
+ eos_token_id=50256,
155
+ scale_attn_by_inverse_layer_idx=False,
156
+ reorder_and_upcast_attn=False,
157
+ position_embedding_type="learned",
158
+ mup_width_scale=1.0,
159
+ mup_embeddings_scale=1.0,
160
+ mup_output_alpha=1.0,
161
+ mup_scale_qk_dot_by_d=False,
162
+ alibi_scaling=None,
163
+ architectures=None,
164
+ **kwargs,
165
+ ):
166
+ self.vocab_size = vocab_size
167
+ self.n_positions = n_positions
168
+ self.n_embd = n_embd
169
+ self.n_layer = n_layer
170
+ self.n_head = n_head
171
+ self.n_inner = n_inner
172
+ self.activation_function = activation_function
173
+ self.resid_pdrop = resid_pdrop
174
+ self.embd_pdrop = embd_pdrop
175
+ self.attn_pdrop = attn_pdrop
176
+ self.layer_norm_epsilon = layer_norm_epsilon
177
+ self.initializer_range = initializer_range
178
+ self.scale_attn_weights = scale_attn_weights
179
+ self.use_cache = use_cache
180
+ self.scale_attn_by_inverse_layer_idx = scale_attn_by_inverse_layer_idx
181
+ self.reorder_and_upcast_attn = reorder_and_upcast_attn
182
+
183
+ self.bos_token_id = bos_token_id
184
+ self.eos_token_id = eos_token_id
185
+
186
+ self.position_embedding_type = position_embedding_type
187
+ self.mup_width_scale = mup_width_scale
188
+ self.mup_embeddings_scale = mup_embeddings_scale
189
+ self.mup_output_alpha = mup_output_alpha
190
+ self.mup_scale_qk_dot_by_d = mup_scale_qk_dot_by_d
191
+
192
+ self.alibi_scaling = alibi_scaling
193
+ self._alibi_scaling_validation()
194
+ if architectures is None:
195
+ architectures = ["JAISLMHeadModel"]
196
+
197
+ super().__init__(
198
+ bos_token_id=bos_token_id,
199
+ eos_token_id=eos_token_id,
200
+ architectures=architectures,
201
+ **kwargs,
202
+ )
203
+
204
+ def _alibi_scaling_validation(self):
205
+ """
206
+ Validate the `alibi_scaling` configuration.
207
+ """
208
+ if self.alibi_scaling is None:
209
+ return
210
+
211
+ if (not isinstance(self.alibi_scaling, dict)
212
+ or len(self.alibi_scaling) != 2):
213
+ raise ValueError(
214
+ "`alibi_scaling` must be a dictionary with two fields,"
215
+ "`type` and `factor` or `type` and `train_seq_len`, "
216
+ f"got {self.alibi_scaling}")
217
+ alibi_scaling_type = self.alibi_scaling.get("type", None)
218
+ alibi_scaling_factor = self.alibi_scaling.get("factor", None)
219
+ alibi_dynamic_scaling = self.alibi_scaling.get("train_seq_len", None)
220
+ if alibi_scaling_type is None or alibi_scaling_type != "linear":
221
+ raise ValueError(f"`alibi_scaling`'s type field must be 'linear',"
222
+ f"got {alibi_scaling_type}")
223
+ if (alibi_scaling_factor is not None
224
+ and not isinstance(alibi_scaling_factor, float)
225
+ or (alibi_scaling_factor is not None
226
+ and alibi_scaling_factor <= 1.0)):
227
+ raise ValueError(
228
+ f"`alibi_scaling`'s factor field must be a float > 1.0,"
229
+ f"got {alibi_scaling_factor}")
230
+ if (alibi_dynamic_scaling is not None
231
+ and not isinstance(alibi_dynamic_scaling, int)
232
+ or (alibi_dynamic_scaling is not None
233
+ and alibi_dynamic_scaling <= 1)):
234
+ raise ValueError(
235
+ f"`alibi_scaling`'s `train_seq_len` field must be an"
236
+ f"integer > 1, got {alibi_dynamic_scaling}")
@@ -0,0 +1,178 @@
1
+ # coding=utf-8
2
+ # Copied from
3
+ # https://huggingface.co/mosaicml/mpt-7b/blob/main/configuration_mpt.py
4
+ """A HuggingFace-style model configuration."""
5
+ import warnings
6
+ from typing import Any, Dict, Optional, Union
7
+
8
+ from transformers import PretrainedConfig
9
+
10
+ attn_config_defaults: Dict = {
11
+ 'attn_type': 'multihead_attention',
12
+ 'attn_pdrop': 0.0,
13
+ 'attn_impl': 'triton',
14
+ 'qk_ln': False,
15
+ 'clip_qkv': None,
16
+ 'softmax_scale': None,
17
+ 'prefix_lm': False,
18
+ 'attn_uses_sequence_id': False,
19
+ 'alibi': False,
20
+ 'alibi_bias_max': 8
21
+ }
22
+ ffn_config_defaults: Dict = {'ffn_type': 'mptmlp'}
23
+ init_config_defaults: Dict = {
24
+ 'name': 'kaiming_normal_',
25
+ 'fan_mode': 'fan_in',
26
+ 'init_nonlinearity': 'relu',
27
+ 'init_div_is_residual': True,
28
+ 'emb_init_std': None,
29
+ 'emb_init_uniform_lim': None,
30
+ 'init_std': None,
31
+ 'init_gain': 0.0
32
+ }
33
+
34
+
35
+ class MPTConfig(PretrainedConfig):
36
+ model_type = 'mpt'
37
+ attribute_map = {
38
+ 'num_attention_heads': 'n_heads',
39
+ 'hidden_size': 'd_model',
40
+ 'num_hidden_layers': 'n_layers',
41
+ }
42
+
43
+ # pylint: disable=dangerous-default-value
44
+ def __init__(self,
45
+ d_model: int = 2048,
46
+ n_heads: int = 16,
47
+ n_layers: int = 24,
48
+ expansion_ratio: int = 4,
49
+ max_seq_len: int = 2048,
50
+ vocab_size: int = 50368,
51
+ resid_pdrop: float = 0.0,
52
+ emb_pdrop: float = 0.0,
53
+ learned_pos_emb: bool = True,
54
+ attn_config: Dict = attn_config_defaults,
55
+ ffn_config: Dict = ffn_config_defaults,
56
+ init_device: str = 'cpu',
57
+ logit_scale: Optional[Union[float, str]] = None,
58
+ no_bias: bool = False,
59
+ embedding_fraction: float = 1.0,
60
+ norm_type: str = 'low_precision_layernorm',
61
+ use_cache: bool = False,
62
+ init_config: Dict = init_config_defaults,
63
+ fc_type: str = 'torch',
64
+ verbose: Optional[int] = None,
65
+ **kwargs: Any):
66
+ self.d_model = d_model
67
+ self.n_heads = n_heads
68
+ self.n_layers = n_layers
69
+ self.expansion_ratio = expansion_ratio
70
+ self.max_seq_len = max_seq_len
71
+ self.vocab_size = vocab_size
72
+ self.resid_pdrop = resid_pdrop
73
+ self.emb_pdrop = emb_pdrop
74
+ self.learned_pos_emb = learned_pos_emb
75
+ self.attn_config = attn_config
76
+ self.ffn_config = ffn_config
77
+ self.init_device = init_device
78
+ self.logit_scale = logit_scale
79
+ self.no_bias = no_bias
80
+ self.embedding_fraction = embedding_fraction
81
+ self.norm_type = norm_type
82
+ self.use_cache = use_cache
83
+ self.init_config = init_config
84
+ self.fc_type = fc_type
85
+ if verbose is not None:
86
+ warnings.warn(DeprecationWarning(
87
+ 'verbose argument for MPTConfig is now ignored and '
88
+ 'will be removed. Use python_log_level instead.'),
89
+ stacklevel=2)
90
+ if 'name' in kwargs:
91
+ del kwargs['name']
92
+ if 'loss_fn' in kwargs:
93
+ del kwargs['loss_fn']
94
+ if self.attn_config.get('alibi', False):
95
+ self.learned_pos_emb = False
96
+ warnings.warn(
97
+ f'alibi is turned on, setting `learned_pos_emb` '
98
+ f'to {self.learned_pos_emb}`',
99
+ stacklevel=2)
100
+ super().__init__(**kwargs)
101
+ self._validate_config()
102
+
103
+ def _set_config_defaults(
104
+ self, config: Dict[str, Any],
105
+ config_defaults: Dict[str, Any]) -> Dict[str, Any]:
106
+ for (k, v) in config_defaults.items():
107
+ if k not in config:
108
+ config[k] = v
109
+ return config
110
+
111
+ def _validate_config(self) -> None:
112
+ self.attn_config = self._set_config_defaults(self.attn_config,
113
+ attn_config_defaults)
114
+ self.ffn_config = self._set_config_defaults(self.ffn_config,
115
+ ffn_config_defaults)
116
+ self.init_config = self._set_config_defaults(self.init_config,
117
+ init_config_defaults)
118
+ if self.d_model % self.n_heads != 0:
119
+ raise ValueError('d_model must be divisible by n_heads')
120
+ if any((
121
+ prob < 0 or prob > 1 for prob in
122
+ [self.attn_config['attn_pdrop'], self.resid_pdrop, self.emb_pdrop]
123
+ )):
124
+ raise ValueError(
125
+ "self.attn_config['attn_pdrop'], resid_pdrop, emb_pdrop are "
126
+ "probabilities and must be between 0 and 1")
127
+ if self.attn_config['attn_impl'] not in ['torch', 'flash', 'triton']:
128
+ raise ValueError(
129
+ f"Unknown attn_impl={self.attn_config['attn_impl']}")
130
+ if self.attn_config['prefix_lm'] and self.attn_config[
131
+ 'attn_impl'] not in ['torch', 'triton']:
132
+ raise NotImplementedError(
133
+ 'prefix_lm only implemented with torch and triton attention.')
134
+ if self.attn_config['alibi'] and self.attn_config['attn_impl'] not in [
135
+ 'torch', 'triton'
136
+ ]:
137
+ raise NotImplementedError(
138
+ 'alibi only implemented with torch and triton attention.')
139
+ if self.attn_config['attn_uses_sequence_id'] and self.attn_config[
140
+ 'attn_impl'] not in ['torch', 'triton']:
141
+ raise NotImplementedError(
142
+ 'attn_uses_sequence_id only implemented with torch '
143
+ 'and triton attention.')
144
+ if self.embedding_fraction > 1 or self.embedding_fraction <= 0:
145
+ raise ValueError(
146
+ 'model.embedding_fraction must be between 0 (exclusive) '
147
+ 'and 1 (inclusive)!')
148
+ if isinstance(self.logit_scale,
149
+ str) and self.logit_scale != 'inv_sqrt_d_model':
150
+ raise ValueError(
151
+ f"self.logit_scale={self.logit_scale!r} is not recognized as "
152
+ "an option; use numeric value or 'inv_sqrt_d_model'.")
153
+ if self.init_config.get('name', None) is None:
154
+ raise ValueError(
155
+ f"self.init_config={self.init_config!r} 'name' needs to be set."
156
+ )
157
+ if not self.learned_pos_emb and (not self.attn_config['alibi']):
158
+ warnings.warn(
159
+ 'Positional information not being provided to the model.',
160
+ stacklevel=2)
161
+ if self.fc_type == 'te' or self.ffn_config['ffn_type'] == 'te_ln_mlp':
162
+ try:
163
+ # pylint: disable=import-outside-toplevel
164
+ import transformer_engine.pytorch as te
165
+ del te
166
+ except Exception as exc:
167
+ raise ImportError(
168
+ 'TransformerEngine import fail. `fc_type: te` requires '
169
+ 'TransformerEngine be installed. '
170
+ 'The required version of transformer_engine also requires '
171
+ 'FlashAttention v1.0.6 is installed:\n'
172
+ 'pip install flash-attn==1.0.6 --no-build-isolation \n'
173
+ 'pip install git+https://github.com/NVIDIA/TransformerEngine.git@144e4888b2cdd60bd52e706d5b7a79cb9c1a7156'
174
+ ) from exc
175
+ if self.ffn_config['ffn_type'] == 'mptmlp':
176
+ self.ffn_config['fc_type'] = self.fc_type
177
+ elif self.ffn_config['ffn_type'] == 'te_ln_mlp':
178
+ self.ffn_config['bias'] = not self.no_bias