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,327 @@
1
+ # coding=utf-8
2
+ # Adapted from
3
+ # https://github.com/huggingface/transformers/blob/v4.28.0/src/transformers/models/bloom/modeling_bloom.py
4
+ # Copyright 2023 The vLLM team.
5
+ # Copyright 2022 HuggingFace Inc. team and BigScience workshop.
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
+ """Inference-only BLOOM model compatible with HuggingFace weights."""
19
+ import math
20
+ from typing import Iterable, List, Optional, Tuple
21
+
22
+ import torch
23
+ from torch import nn
24
+ from transformers import BloomConfig
25
+
26
+ from vllm.attention import Attention, AttentionMetadata
27
+ from vllm.distributed import (get_tensor_model_parallel_rank,
28
+ get_tensor_model_parallel_world_size)
29
+ from vllm.model_executor.layers.activation import get_act_fn
30
+ from vllm.model_executor.layers.linear import (ColumnParallelLinear,
31
+ QKVParallelLinear,
32
+ RowParallelLinear)
33
+ from vllm.model_executor.layers.logits_processor import LogitsProcessor
34
+ from vllm.model_executor.layers.quantization.base_config import (
35
+ QuantizationConfig)
36
+ from vllm.model_executor.layers.sampler import Sampler
37
+ from vllm.model_executor.layers.vocab_parallel_embedding import (
38
+ VocabParallelEmbedding)
39
+ from vllm.model_executor.model_loader.weight_utils import default_weight_loader
40
+ from vllm.model_executor.sampling_metadata import SamplingMetadata
41
+ from vllm.sequence import SamplerOutput
42
+
43
+
44
+ def _get_alibi_slopes(total_num_heads: int) -> torch.Tensor:
45
+ closest_power_of_2 = 2**math.floor(math.log2(total_num_heads))
46
+ base = torch.tensor(
47
+ 2**(-(2**-(math.log2(closest_power_of_2) - 3))),
48
+ dtype=torch.float32,
49
+ )
50
+ powers = torch.arange(1, 1 + closest_power_of_2, dtype=torch.int32)
51
+ slopes = torch.pow(base, powers)
52
+
53
+ if closest_power_of_2 != total_num_heads:
54
+ extra_base = torch.tensor(
55
+ 2**(-(2**-(math.log2(2 * closest_power_of_2) - 3))),
56
+ dtype=torch.float32,
57
+ )
58
+ num_remaining_heads = min(closest_power_of_2,
59
+ total_num_heads - closest_power_of_2)
60
+ extra_powers = torch.arange(start=1,
61
+ end=1 + 2 * num_remaining_heads,
62
+ step=2,
63
+ dtype=torch.int32)
64
+ slopes = torch.cat(
65
+ [slopes, torch.pow(extra_base, extra_powers)], dim=0)
66
+ return slopes
67
+
68
+
69
+ class BloomAttention(nn.Module):
70
+
71
+ def __init__(
72
+ self,
73
+ config: BloomConfig,
74
+ quant_config: Optional[QuantizationConfig] = None,
75
+ ):
76
+ super().__init__()
77
+ self.hidden_size = config.hidden_size
78
+ self.total_num_heads = config.n_head
79
+ self.head_dim = self.hidden_size // self.total_num_heads
80
+ assert self.head_dim * self.total_num_heads == self.hidden_size
81
+
82
+ tp_world_size = get_tensor_model_parallel_world_size()
83
+ assert self.total_num_heads % tp_world_size == 0
84
+ self.num_heads = self.total_num_heads // tp_world_size
85
+
86
+ self.query_key_value = QKVParallelLinear(
87
+ self.hidden_size,
88
+ self.head_dim,
89
+ self.total_num_heads,
90
+ bias=True,
91
+ quant_config=quant_config,
92
+ )
93
+ self.dense = RowParallelLinear(
94
+ self.hidden_size,
95
+ self.hidden_size,
96
+ bias=True,
97
+ quant_config=quant_config,
98
+ )
99
+
100
+ # Create the alibi slopes and slice them.
101
+ tp_rank = get_tensor_model_parallel_rank()
102
+ head_start = tp_rank * self.num_heads
103
+ head_end = (tp_rank + 1) * self.num_heads
104
+ alibi_slopes = _get_alibi_slopes(self.total_num_heads)
105
+ alibi_slopes = alibi_slopes[head_start:head_end].tolist()
106
+
107
+ scaling = self.head_dim**-0.5
108
+ self.attn = Attention(self.num_heads,
109
+ self.head_dim,
110
+ scaling,
111
+ alibi_slopes=alibi_slopes)
112
+
113
+ def forward(
114
+ self,
115
+ position_ids: torch.Tensor,
116
+ hidden_states: torch.Tensor,
117
+ kv_cache: torch.Tensor,
118
+ attn_metadata: AttentionMetadata,
119
+ ) -> torch.Tensor:
120
+ del position_ids # Unused.
121
+ qkv, _ = self.query_key_value(hidden_states)
122
+ q, k, v = qkv.chunk(chunks=3, dim=-1)
123
+ attn_output = self.attn(q, k, v, kv_cache, attn_metadata)
124
+ output, _ = self.dense(attn_output)
125
+ return output
126
+
127
+
128
+ class BloomMLP(nn.Module):
129
+
130
+ def __init__(
131
+ self,
132
+ config: BloomConfig,
133
+ quant_config: Optional[QuantizationConfig] = None,
134
+ ):
135
+ super().__init__()
136
+ hidden_size = config.hidden_size
137
+ self.dense_h_to_4h = ColumnParallelLinear(
138
+ hidden_size,
139
+ 4 * hidden_size,
140
+ quant_config=quant_config,
141
+ )
142
+ self.gelu_impl = get_act_fn("gelu", quant_config, 4 * hidden_size)
143
+ self.dense_4h_to_h = RowParallelLinear(
144
+ 4 * hidden_size,
145
+ hidden_size,
146
+ quant_config=quant_config,
147
+ )
148
+
149
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
150
+ x, _ = self.dense_h_to_4h(x)
151
+ x = self.gelu_impl(x)
152
+ x, _ = self.dense_4h_to_h(x)
153
+ return x
154
+
155
+
156
+ class BloomBlock(nn.Module):
157
+
158
+ def __init__(
159
+ self,
160
+ config: BloomConfig,
161
+ quant_config: Optional[QuantizationConfig] = None,
162
+ ):
163
+ super().__init__()
164
+ hidden_size = config.hidden_size
165
+
166
+ self.input_layernorm = nn.LayerNorm(hidden_size,
167
+ eps=config.layer_norm_epsilon)
168
+ self.self_attention = BloomAttention(config, quant_config)
169
+ self.post_attention_layernorm = nn.LayerNorm(
170
+ hidden_size, eps=config.layer_norm_epsilon)
171
+ self.mlp = BloomMLP(config, quant_config)
172
+ self.apply_residual_connection_post_layernorm = (
173
+ config.apply_residual_connection_post_layernorm)
174
+
175
+ def forward(
176
+ self,
177
+ position_ids: torch.Tensor,
178
+ hidden_states: torch.Tensor,
179
+ kv_cache: torch.Tensor,
180
+ attn_metadata: AttentionMetadata,
181
+ ) -> torch.Tensor:
182
+ # Layer norm at the beginning of the transformer layer.
183
+ layernorm_output = self.input_layernorm(hidden_states)
184
+
185
+ # Layer norm post the self attention.
186
+ if self.apply_residual_connection_post_layernorm:
187
+ residual = layernorm_output
188
+ else:
189
+ residual = hidden_states
190
+
191
+ # Self attention.
192
+ attention_output = self.self_attention(
193
+ position_ids=position_ids,
194
+ hidden_states=layernorm_output,
195
+ kv_cache=kv_cache,
196
+ attn_metadata=attn_metadata,
197
+ )
198
+ attention_output = attention_output + residual
199
+ layernorm_output = self.post_attention_layernorm(attention_output)
200
+
201
+ # Get residual
202
+ if self.apply_residual_connection_post_layernorm:
203
+ residual = layernorm_output
204
+ else:
205
+ residual = attention_output
206
+
207
+ # MLP.
208
+ output = self.mlp(layernorm_output) + residual
209
+ return output
210
+
211
+
212
+ class BloomModel(nn.Module):
213
+
214
+ def __init__(
215
+ self,
216
+ config: BloomConfig,
217
+ quant_config: Optional[QuantizationConfig] = None,
218
+ ):
219
+ super().__init__()
220
+ self.embed_dim = config.hidden_size
221
+
222
+ # Embedding + LN Embedding
223
+ self.word_embeddings = VocabParallelEmbedding(
224
+ config.vocab_size,
225
+ self.embed_dim,
226
+ )
227
+ self.word_embeddings_layernorm = nn.LayerNorm(
228
+ self.embed_dim, eps=config.layer_norm_epsilon)
229
+
230
+ # Transformer blocks
231
+ self.h = nn.ModuleList([
232
+ BloomBlock(config, quant_config)
233
+ for _ in range(config.num_hidden_layers)
234
+ ])
235
+
236
+ # Final Layer Norm
237
+ self.ln_f = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_epsilon)
238
+
239
+ def forward(
240
+ self,
241
+ input_ids: torch.Tensor,
242
+ position_ids: torch.Tensor,
243
+ kv_caches: List[torch.Tensor],
244
+ attn_metadata: AttentionMetadata,
245
+ ) -> torch.Tensor:
246
+ hidden_states = self.word_embeddings(input_ids)
247
+ hidden_states = self.word_embeddings_layernorm(hidden_states)
248
+ for i in range(len(self.h)):
249
+ layer = self.h[i]
250
+ hidden_states = layer(
251
+ position_ids,
252
+ hidden_states,
253
+ kv_caches[i],
254
+ attn_metadata,
255
+ )
256
+ hidden_states = self.ln_f(hidden_states)
257
+ return hidden_states
258
+
259
+
260
+ class BloomForCausalLM(nn.Module):
261
+
262
+ def __init__(
263
+ self,
264
+ config: BloomConfig,
265
+ quant_config: Optional[QuantizationConfig] = None,
266
+ ):
267
+ super().__init__()
268
+ self.config = config
269
+ self.quant_config = quant_config
270
+ self.transformer = BloomModel(config, quant_config)
271
+ self.lm_head_weight = self.transformer.word_embeddings.weight
272
+ self.logits_processor = LogitsProcessor(config.vocab_size)
273
+ self.sampler = Sampler()
274
+
275
+ def forward(
276
+ self,
277
+ input_ids: torch.Tensor,
278
+ positions: torch.Tensor,
279
+ kv_caches: List[torch.Tensor],
280
+ attn_metadata: AttentionMetadata,
281
+ ) -> torch.Tensor:
282
+ hidden_states = self.transformer(input_ids, positions, kv_caches,
283
+ attn_metadata)
284
+ return hidden_states
285
+
286
+ def compute_logits(self, hidden_states: torch.Tensor,
287
+ sampling_metadata: SamplingMetadata) -> torch.Tensor:
288
+ logits = self.logits_processor(self.lm_head_weight, hidden_states,
289
+ sampling_metadata)
290
+ return logits
291
+
292
+ def sample(
293
+ self,
294
+ logits: torch.Tensor,
295
+ sampling_metadata: SamplingMetadata,
296
+ ) -> Optional[SamplerOutput]:
297
+ next_tokens = self.sampler(logits, sampling_metadata)
298
+ return next_tokens
299
+
300
+ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
301
+ params_dict = dict(self.named_parameters(remove_duplicate=False))
302
+ for name, loaded_weight in weights:
303
+ if name == "lm_head.weight":
304
+ continue
305
+ if not name.startswith("transformer."):
306
+ name = "transformer." + name
307
+ param = params_dict[name]
308
+
309
+ if "query_key_value" in name:
310
+ # NOTE: BLOOM's fused QKV's output_dim has the shape of
311
+ # (num_heads * 3 * head_size), while the
312
+ # required shape is (3 * num_heads * head_size).
313
+ # Thus, we need weight conversion.
314
+ output_dim = getattr(param, "output_dim", None)
315
+ num_heads = self.config.num_attention_heads
316
+ if output_dim is not None:
317
+ loaded_weight_shape = loaded_weight.shape
318
+ loaded_weight = loaded_weight.view(
319
+ loaded_weight_shape[:output_dim] + (num_heads, 3, -1) +
320
+ loaded_weight_shape[output_dim + 1:])
321
+ loaded_weight = loaded_weight.transpose(
322
+ output_dim, output_dim + 1)
323
+ loaded_weight = loaded_weight.reshape(loaded_weight_shape)
324
+
325
+ weight_loader = getattr(param, "weight_loader",
326
+ default_weight_loader)
327
+ weight_loader(param, loaded_weight)
@@ -0,0 +1,386 @@
1
+ # coding=utf-8
2
+ # Adapted from
3
+ # https://github.com/THUDM/ChatGLM2-6B
4
+ """Inference-only ChatGLM model compatible with THUDM weights."""
5
+ from typing import Iterable, List, Optional, Tuple
6
+
7
+ import torch
8
+ from torch import nn
9
+ from torch.nn import LayerNorm
10
+
11
+ from vllm.attention import Attention, AttentionMetadata
12
+ from vllm.config import LoRAConfig
13
+ from vllm.distributed import get_tensor_model_parallel_world_size
14
+ from vllm.model_executor.layers.activation import SiluAndMul
15
+ from vllm.model_executor.layers.layernorm import RMSNorm
16
+ from vllm.model_executor.layers.linear import (MergedColumnParallelLinear,
17
+ QKVParallelLinear,
18
+ RowParallelLinear)
19
+ from vllm.model_executor.layers.logits_processor import LogitsProcessor
20
+ from vllm.model_executor.layers.quantization.base_config import (
21
+ QuantizationConfig)
22
+ from vllm.model_executor.layers.rotary_embedding import get_rope
23
+ from vllm.model_executor.layers.sampler import Sampler
24
+ from vllm.model_executor.layers.vocab_parallel_embedding import (
25
+ ParallelLMHead, VocabParallelEmbedding)
26
+ from vllm.model_executor.model_loader.weight_utils import default_weight_loader
27
+ from vllm.model_executor.sampling_metadata import SamplingMetadata
28
+ from vllm.sequence import SamplerOutput
29
+ from vllm.transformers_utils.configs import ChatGLMConfig
30
+
31
+
32
+ class GLMAttention(nn.Module):
33
+
34
+ def __init__(
35
+ self,
36
+ config,
37
+ quant_config: Optional[QuantizationConfig] = None,
38
+ ):
39
+ super().__init__()
40
+ self.hidden_size = config.hidden_size
41
+ tp_size = get_tensor_model_parallel_world_size()
42
+ self.total_num_heads = config.num_attention_heads
43
+ assert self.total_num_heads % tp_size == 0
44
+ self.num_heads = self.total_num_heads // tp_size
45
+ self.multi_query_attention = config.multi_query_attention
46
+ self.total_num_kv_heads = (config.multi_query_group_num
47
+ if config.multi_query_attention else
48
+ config.num_attention_heads)
49
+ if self.total_num_kv_heads >= tp_size:
50
+ # Number of KV heads is greater than TP size, so we partition
51
+ # the KV heads across multiple tensor parallel GPUs.
52
+ assert self.total_num_kv_heads % tp_size == 0
53
+ else:
54
+ # Number of KV heads is less than TP size, so we replicate
55
+ # the KV heads across multiple tensor parallel GPUs.
56
+ assert tp_size % self.total_num_kv_heads == 0
57
+ self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
58
+ self.head_dim = config.hidden_size // self.total_num_heads
59
+ self.q_size = self.num_heads * self.head_dim
60
+ self.kv_size = self.num_kv_heads * self.head_dim
61
+ self.scaling = self.head_dim**-0.5
62
+
63
+ self.query_key_value = QKVParallelLinear(
64
+ self.hidden_size,
65
+ self.head_dim,
66
+ self.total_num_heads,
67
+ self.total_num_kv_heads,
68
+ bias=config.add_bias_linear or config.add_qkv_bias,
69
+ quant_config=quant_config,
70
+ )
71
+ self.dense = RowParallelLinear(
72
+ self.total_num_heads * self.head_dim,
73
+ config.hidden_size,
74
+ bias=config.add_bias_linear,
75
+ quant_config=quant_config,
76
+ )
77
+
78
+ # https://huggingface.co/THUDM/chatglm3-6b-32k/blob/e210410255278dd9d74463cf396ba559c0ef801c/modeling_chatglm.py#L141
79
+ rope_ratio = getattr(config, "rope_ratio", 1.0)
80
+ max_positions = getattr(config, "seq_length", 8192)
81
+ self.rotary_emb = get_rope(
82
+ self.head_dim,
83
+ rotary_dim=self.head_dim // 2,
84
+ max_position=max_positions,
85
+ base=10000 * rope_ratio,
86
+ is_neox_style=False,
87
+ )
88
+ self.attn = Attention(
89
+ self.num_heads,
90
+ self.head_dim,
91
+ self.scaling,
92
+ num_kv_heads=self.num_kv_heads,
93
+ )
94
+
95
+ def forward(
96
+ self,
97
+ hidden_states: torch.Tensor,
98
+ position_ids: torch.Tensor,
99
+ kv_cache: torch.Tensor,
100
+ attn_metadata: AttentionMetadata,
101
+ ) -> torch.Tensor:
102
+ qkv, _ = self.query_key_value(hidden_states)
103
+ q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
104
+ q, k = self.rotary_emb(position_ids, q, k)
105
+ context_layer = self.attn(
106
+ q,
107
+ k,
108
+ v,
109
+ kv_cache,
110
+ attn_metadata,
111
+ )
112
+ attn_output, _ = self.dense(context_layer)
113
+ return attn_output
114
+
115
+
116
+ class GLMMLP(nn.Module):
117
+ """MLP.
118
+
119
+ MLP will take the input with h hidden state, project it to 4*h
120
+ hidden dimension, perform nonlinear transformation, and project the
121
+ state back into h hidden dimension.
122
+ """
123
+
124
+ def __init__(
125
+ self,
126
+ config,
127
+ quant_config: Optional[QuantizationConfig] = None,
128
+ ):
129
+ super().__init__()
130
+
131
+ self.add_bias = config.add_bias_linear
132
+
133
+ # Project to 4h.
134
+ self.dense_h_to_4h = MergedColumnParallelLinear(
135
+ config.hidden_size,
136
+ [config.ffn_hidden_size] * 2,
137
+ bias=config.add_bias_linear,
138
+ quant_config=quant_config,
139
+ )
140
+
141
+ self.activation_func = SiluAndMul()
142
+
143
+ # Project back to h.
144
+ self.dense_4h_to_h = RowParallelLinear(
145
+ config.ffn_hidden_size,
146
+ config.hidden_size,
147
+ bias=config.add_bias_linear,
148
+ quant_config=quant_config,
149
+ )
150
+
151
+ def forward(self, hidden_states):
152
+ # [s, b, 4hp]
153
+ intermediate_parallel, _ = self.dense_h_to_4h(hidden_states)
154
+ intermediate_parallel = self.activation_func(intermediate_parallel)
155
+ # [s, b, h]
156
+ output, _ = self.dense_4h_to_h(intermediate_parallel)
157
+ return output
158
+
159
+
160
+ class GLMBlock(nn.Module):
161
+ """A single transformer layer.
162
+
163
+ Transformer layer takes input with size [s, b, h] and returns an
164
+ output of the same size.
165
+ """
166
+
167
+ def __init__(
168
+ self,
169
+ config,
170
+ quant_config: Optional[QuantizationConfig] = None,
171
+ ):
172
+ super().__init__()
173
+ self.apply_residual_connection_post_layernorm = (
174
+ config.apply_residual_connection_post_layernorm)
175
+
176
+ self.fp32_residual_connection = config.fp32_residual_connection
177
+
178
+ layer_norm_func = RMSNorm if config.rmsnorm else LayerNorm
179
+ # Layernorm on the input data.
180
+ self.input_layernorm = layer_norm_func(config.hidden_size,
181
+ eps=config.layernorm_epsilon)
182
+
183
+ # Self attention.
184
+ self.self_attention = GLMAttention(config, quant_config)
185
+ self.hidden_dropout = config.hidden_dropout
186
+
187
+ # Layernorm on the attention output
188
+ self.post_attention_layernorm = layer_norm_func(
189
+ config.hidden_size, eps=config.layernorm_epsilon)
190
+
191
+ # MLP
192
+ self.mlp = GLMMLP(config, quant_config)
193
+
194
+ def forward(
195
+ self,
196
+ hidden_states: torch.Tensor,
197
+ position_ids: torch.Tensor,
198
+ kv_cache: torch.Tensor,
199
+ attn_metadata: AttentionMetadata,
200
+ ) -> torch.Tensor:
201
+ # hidden_states: [num_tokens, h]
202
+ # Layer norm at the beginning of the transformer layer.
203
+ layernorm_output = self.input_layernorm(hidden_states)
204
+ # Self attention.
205
+ attention_output = self.self_attention(
206
+ hidden_states=layernorm_output,
207
+ position_ids=position_ids,
208
+ kv_cache=kv_cache,
209
+ attn_metadata=attn_metadata,
210
+ )
211
+
212
+ # Residual connection.
213
+ if self.apply_residual_connection_post_layernorm:
214
+ residual = layernorm_output
215
+ else:
216
+ residual = hidden_states
217
+
218
+ layernorm_input = residual + attention_output
219
+
220
+ # Layer norm post the self attention.
221
+ layernorm_output = self.post_attention_layernorm(layernorm_input)
222
+
223
+ # Second residual connection.
224
+ if self.apply_residual_connection_post_layernorm:
225
+ residual = layernorm_output
226
+ else:
227
+ residual = layernorm_input
228
+
229
+ output = self.mlp(layernorm_output) + residual
230
+
231
+ return output
232
+
233
+
234
+ class GLMTransformer(nn.Module):
235
+ """Transformer class."""
236
+
237
+ def __init__(
238
+ self,
239
+ config,
240
+ quant_config: Optional[QuantizationConfig] = None,
241
+ ):
242
+ super().__init__()
243
+ self.post_layer_norm = config.post_layer_norm
244
+
245
+ # Number of layers.
246
+ self.num_layers = config.num_layers
247
+
248
+ # Transformer layers.
249
+ self.layers = nn.ModuleList(
250
+ [GLMBlock(config, quant_config) for i in range(self.num_layers)])
251
+
252
+ if self.post_layer_norm:
253
+ layer_norm_func = RMSNorm if config.rmsnorm else LayerNorm
254
+ # Final layer norm before output.
255
+ self.final_layernorm = layer_norm_func(
256
+ config.hidden_size, eps=config.layernorm_epsilon)
257
+
258
+ def forward(
259
+ self,
260
+ hidden_states: torch.Tensor,
261
+ position_ids: torch.Tensor,
262
+ kv_caches: List[torch.Tensor],
263
+ attn_metadata: AttentionMetadata,
264
+ ) -> torch.Tensor:
265
+ for i in range(self.num_layers):
266
+ layer = self.layers[i]
267
+ hidden_states = layer(
268
+ hidden_states=hidden_states,
269
+ position_ids=position_ids,
270
+ kv_cache=kv_caches[i],
271
+ attn_metadata=attn_metadata,
272
+ )
273
+ # Final layer norm.
274
+ if self.post_layer_norm:
275
+ hidden_states = self.final_layernorm(hidden_states)
276
+
277
+ return hidden_states
278
+
279
+
280
+ class ChatGLMModel(nn.Module):
281
+
282
+ def __init__(
283
+ self,
284
+ config,
285
+ quant_config: Optional[QuantizationConfig] = None,
286
+ ):
287
+ super().__init__()
288
+
289
+ self.embedding = VocabParallelEmbedding(config.padded_vocab_size,
290
+ config.hidden_size)
291
+
292
+ self.num_layers = config.num_layers
293
+ self.multi_query_group_num = config.multi_query_group_num
294
+ self.kv_channels = config.kv_channels
295
+ self.encoder = GLMTransformer(config, quant_config)
296
+
297
+ self.output_layer = ParallelLMHead(config.padded_vocab_size,
298
+ config.hidden_size)
299
+
300
+ def forward(
301
+ self,
302
+ input_ids: torch.Tensor,
303
+ position_ids: torch.Tensor,
304
+ kv_caches: List[torch.Tensor],
305
+ attn_metadata: AttentionMetadata,
306
+ ) -> torch.Tensor:
307
+ inputs_embeds = self.embedding(input_ids)
308
+
309
+ # Run encoder.
310
+ hidden_states = self.encoder(
311
+ hidden_states=inputs_embeds,
312
+ position_ids=position_ids,
313
+ kv_caches=kv_caches,
314
+ attn_metadata=attn_metadata,
315
+ )
316
+ return hidden_states
317
+
318
+
319
+ class ChatGLMForCausalLM(nn.Module):
320
+ packed_modules_mapping = {
321
+ "query_key_value": ["query_key_value"],
322
+ "dense_h_to_4h": ["dense_h_to_4h"]
323
+ }
324
+ # LoRA specific attributes
325
+ supported_lora_modules = [
326
+ "query_key_value",
327
+ "dense",
328
+ "dense_h_to_4h",
329
+ "dense_4h_to_h",
330
+ ]
331
+ embedding_modules = {}
332
+ embedding_padding_modules = []
333
+
334
+ def __init__(
335
+ self,
336
+ config: ChatGLMConfig,
337
+ quant_config: Optional[QuantizationConfig] = None,
338
+ lora_config: Optional[LoRAConfig] = None,
339
+ ):
340
+ super().__init__()
341
+ self.config: ChatGLMConfig = config
342
+ self.quant_config = quant_config
343
+ self.transformer = ChatGLMModel(config, quant_config)
344
+ self.lm_head_weight = self.transformer.output_layer.weight
345
+ self.logits_processor = LogitsProcessor(config.padded_vocab_size)
346
+ self.sampler = Sampler()
347
+
348
+ def forward(
349
+ self,
350
+ input_ids: torch.Tensor,
351
+ positions: torch.Tensor,
352
+ kv_caches: List[torch.Tensor],
353
+ attn_metadata: AttentionMetadata,
354
+ ) -> torch.Tensor:
355
+ hidden_states = self.transformer(input_ids, positions, kv_caches,
356
+ attn_metadata)
357
+ return hidden_states
358
+
359
+ def compute_logits(self, hidden_states: torch.Tensor,
360
+ sampling_metadata: SamplingMetadata) -> torch.Tensor:
361
+ logits = self.logits_processor(self.lm_head_weight, hidden_states,
362
+ sampling_metadata)
363
+ return logits
364
+
365
+ def sample(
366
+ self,
367
+ logits: torch.Tensor,
368
+ sampling_metadata: SamplingMetadata,
369
+ ) -> Optional[SamplerOutput]:
370
+ next_tokens = self.sampler(logits, sampling_metadata)
371
+ return next_tokens
372
+
373
+ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
374
+ params_dict = dict(self.named_parameters(remove_duplicate=False))
375
+ for name, loaded_weight in weights:
376
+ if "rotary_pos_emb.inv_freq" in name:
377
+ continue
378
+ if "word_embeddings" in name:
379
+ name = name.replace(".word_embeddings", "")
380
+ # Skip loading extra bias for GPTQ models.
381
+ if name.endswith(".bias") and name not in params_dict:
382
+ continue
383
+ param = params_dict[name]
384
+ weight_loader = getattr(param, "weight_loader",
385
+ default_weight_loader)
386
+ weight_loader(param, loaded_weight)