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.
- vllm/__init__.py +23 -0
- vllm/_custom_ops.py +251 -0
- vllm/attention/__init__.py +13 -0
- vllm/attention/backends/__init__.py +0 -0
- vllm/attention/backends/abstract.py +127 -0
- vllm/attention/backends/flash_attn.py +271 -0
- vllm/attention/backends/flashinfer.py +220 -0
- vllm/attention/backends/rocm_flash_attn.py +374 -0
- vllm/attention/backends/torch_sdpa.py +250 -0
- vllm/attention/backends/xformers.py +393 -0
- vllm/attention/layer.py +56 -0
- vllm/attention/ops/__init__.py +0 -0
- vllm/attention/ops/paged_attn.py +216 -0
- vllm/attention/ops/prefix_prefill.py +792 -0
- vllm/attention/ops/triton_flash_attention.py +810 -0
- vllm/attention/selector.py +91 -0
- vllm/block.py +84 -0
- vllm/config.py +1225 -0
- vllm/core/__init__.py +0 -0
- vllm/core/block/__init__.py +0 -0
- vllm/core/block/block_table.py +295 -0
- vllm/core/block/common.py +199 -0
- vllm/core/block/cpu_gpu_block_allocator.py +228 -0
- vllm/core/block/interfaces.py +205 -0
- vllm/core/block/naive_block.py +318 -0
- vllm/core/block/prefix_caching_block.py +606 -0
- vllm/core/block_manager_v1.py +625 -0
- vllm/core/block_manager_v2.py +258 -0
- vllm/core/evictor_v1.py +105 -0
- vllm/core/evictor_v2.py +127 -0
- vllm/core/interfaces.py +113 -0
- vllm/core/policy.py +45 -0
- vllm/core/scheduler.py +1163 -0
- vllm/distributed/__init__.py +3 -0
- vllm/distributed/communication_op.py +237 -0
- vllm/distributed/device_communicators/__init__.py +0 -0
- vllm/distributed/device_communicators/custom_all_reduce.py +274 -0
- vllm/distributed/device_communicators/pynccl.py +287 -0
- vllm/distributed/device_communicators/pynccl_utils.py +66 -0
- vllm/distributed/parallel_state.py +339 -0
- vllm/distributed/utils.py +136 -0
- vllm/engine/__init__.py +0 -0
- vllm/engine/arg_utils.py +649 -0
- vllm/engine/async_llm_engine.py +737 -0
- vllm/engine/llm_engine.py +784 -0
- vllm/engine/metrics.py +368 -0
- vllm/engine/output_processor/__init__.py +0 -0
- vllm/engine/output_processor/interfaces.py +76 -0
- vllm/engine/output_processor/multi_step.py +142 -0
- vllm/engine/output_processor/single_step.py +284 -0
- vllm/engine/output_processor/stop_checker.py +101 -0
- vllm/engine/output_processor/util.py +19 -0
- vllm/entrypoints/__init__.py +0 -0
- vllm/entrypoints/api_server.py +119 -0
- vllm/entrypoints/llm.py +259 -0
- vllm/entrypoints/openai/__init__.py +0 -0
- vllm/entrypoints/openai/api_server.py +186 -0
- vllm/entrypoints/openai/cli_args.py +115 -0
- vllm/entrypoints/openai/protocol.py +460 -0
- vllm/entrypoints/openai/serving_chat.py +392 -0
- vllm/entrypoints/openai/serving_completion.py +347 -0
- vllm/entrypoints/openai/serving_engine.py +234 -0
- vllm/envs.py +217 -0
- vllm/executor/__init__.py +0 -0
- vllm/executor/cpu_executor.py +152 -0
- vllm/executor/distributed_gpu_executor.py +115 -0
- vllm/executor/executor_base.py +115 -0
- vllm/executor/gpu_executor.py +150 -0
- vllm/executor/multiproc_worker_utils.py +263 -0
- vllm/executor/neuron_executor.py +91 -0
- vllm/executor/ray_gpu_executor.py +327 -0
- vllm/executor/ray_utils.py +119 -0
- vllm/logger.py +153 -0
- vllm/logging/__init__.py +5 -0
- vllm/logging/formatter.py +15 -0
- vllm/lora/__init__.py +0 -0
- vllm/lora/fully_sharded_layers.py +262 -0
- vllm/lora/layers.py +1181 -0
- vllm/lora/lora.py +167 -0
- vllm/lora/models.py +645 -0
- vllm/lora/punica.py +213 -0
- vllm/lora/request.py +32 -0
- vllm/lora/utils.py +98 -0
- vllm/lora/worker_manager.py +251 -0
- vllm/model_executor/__init__.py +7 -0
- vllm/model_executor/guided_decoding/__init__.py +25 -0
- vllm/model_executor/guided_decoding/lm_format_enforcer_decoding.py +70 -0
- vllm/model_executor/guided_decoding/outlines_decoding.py +130 -0
- vllm/model_executor/guided_decoding/outlines_logits_processors.py +184 -0
- vllm/model_executor/layers/__init__.py +0 -0
- vllm/model_executor/layers/activation.py +173 -0
- vllm/model_executor/layers/fused_moe/__init__.py +7 -0
- vllm/model_executor/layers/fused_moe/configs/E=16,N=1344,device_name=NVIDIA_A100-SXM4-40GB.json +146 -0
- vllm/model_executor/layers/fused_moe/configs/E=16,N=1344,device_name=NVIDIA_A100-SXM4-80GB.json +146 -0
- vllm/model_executor/layers/fused_moe/configs/E=16,N=1344,device_name=NVIDIA_H100_80GB_HBM3.json +146 -0
- vllm/model_executor/layers/fused_moe/configs/E=16,N=2688,device_name=NVIDIA_A100-SXM4-80GB.json +146 -0
- vllm/model_executor/layers/fused_moe/configs/E=16,N=2688,device_name=NVIDIA_H100_80GB_HBM3.json +146 -0
- vllm/model_executor/layers/fused_moe/configs/E=8,N=1792,device_name=NVIDIA_A100-SXM4-40GB.json +146 -0
- vllm/model_executor/layers/fused_moe/configs/E=8,N=1792,device_name=NVIDIA_A100-SXM4-80GB.json +146 -0
- vllm/model_executor/layers/fused_moe/configs/E=8,N=1792,device_name=NVIDIA_H100_80GB_HBM3.json +146 -0
- vllm/model_executor/layers/fused_moe/configs/E=8,N=2048,device_name=NVIDIA_A100-SXM4-80GB.json +146 -0
- vllm/model_executor/layers/fused_moe/configs/E=8,N=2048,device_name=NVIDIA_H100_80GB_HBM3.json +146 -0
- vllm/model_executor/layers/fused_moe/configs/E=8,N=3584,device_name=NVIDIA_A100-SXM4-40GB.json +146 -0
- vllm/model_executor/layers/fused_moe/configs/E=8,N=3584,device_name=NVIDIA_A100-SXM4-80GB.json +146 -0
- vllm/model_executor/layers/fused_moe/configs/E=8,N=3584,device_name=NVIDIA_H100_80GB_HBM3,dtype=float8.json +140 -0
- vllm/model_executor/layers/fused_moe/configs/E=8,N=3584,device_name=NVIDIA_H100_80GB_HBM3.json +146 -0
- vllm/model_executor/layers/fused_moe/configs/E=8,N=4096,device_name=NVIDIA_A100-SXM4-80GB.json +146 -0
- vllm/model_executor/layers/fused_moe/configs/E=8,N=4096,device_name=NVIDIA_H100_80GB_HBM3.json +146 -0
- vllm/model_executor/layers/fused_moe/configs/E=8,N=7168,device_name=NVIDIA_A100-SXM4-80GB.json +146 -0
- vllm/model_executor/layers/fused_moe/configs/E=8,N=7168,device_name=NVIDIA_H100_80GB_HBM3,dtype=float8.json +146 -0
- vllm/model_executor/layers/fused_moe/configs/E=8,N=7168,device_name=NVIDIA_H100_80GB_HBM3.json +146 -0
- vllm/model_executor/layers/fused_moe/fused_moe.py +479 -0
- vllm/model_executor/layers/layernorm.py +71 -0
- vllm/model_executor/layers/linear.py +709 -0
- vllm/model_executor/layers/logits_processor.py +115 -0
- vllm/model_executor/layers/ops/__init__.py +0 -0
- vllm/model_executor/layers/ops/rand.py +157 -0
- vllm/model_executor/layers/ops/sample.py +406 -0
- vllm/model_executor/layers/quantization/__init__.py +35 -0
- vllm/model_executor/layers/quantization/aqlm.py +376 -0
- vllm/model_executor/layers/quantization/awq.py +175 -0
- vllm/model_executor/layers/quantization/base_config.py +97 -0
- vllm/model_executor/layers/quantization/fp8.py +265 -0
- vllm/model_executor/layers/quantization/gptq.py +224 -0
- vllm/model_executor/layers/quantization/gptq_marlin.py +438 -0
- vllm/model_executor/layers/quantization/marlin.py +227 -0
- vllm/model_executor/layers/quantization/schema.py +84 -0
- vllm/model_executor/layers/quantization/squeezellm.py +137 -0
- vllm/model_executor/layers/rejection_sampler.py +405 -0
- vllm/model_executor/layers/rotary_embedding.py +525 -0
- vllm/model_executor/layers/sampler.py +1051 -0
- vllm/model_executor/layers/vocab_parallel_embedding.py +155 -0
- vllm/model_executor/model_loader/__init__.py +30 -0
- vllm/model_executor/model_loader/loader.py +362 -0
- vllm/model_executor/model_loader/neuron.py +136 -0
- vllm/model_executor/model_loader/tensorizer.py +368 -0
- vllm/model_executor/model_loader/utils.py +41 -0
- vllm/model_executor/model_loader/weight_utils.py +372 -0
- vllm/model_executor/models/__init__.py +119 -0
- vllm/model_executor/models/baichuan.py +410 -0
- vllm/model_executor/models/bloom.py +327 -0
- vllm/model_executor/models/chatglm.py +386 -0
- vllm/model_executor/models/commandr.py +373 -0
- vllm/model_executor/models/dbrx.py +413 -0
- vllm/model_executor/models/decilm.py +122 -0
- vllm/model_executor/models/deepseek.py +438 -0
- vllm/model_executor/models/falcon.py +444 -0
- vllm/model_executor/models/gemma.py +393 -0
- vllm/model_executor/models/gpt2.py +266 -0
- vllm/model_executor/models/gpt_bigcode.py +274 -0
- vllm/model_executor/models/gpt_j.py +281 -0
- vllm/model_executor/models/gpt_neox.py +295 -0
- vllm/model_executor/models/internlm2.py +323 -0
- vllm/model_executor/models/jais.py +333 -0
- vllm/model_executor/models/llama.py +442 -0
- vllm/model_executor/models/llava.py +239 -0
- vllm/model_executor/models/minicpm.py +531 -0
- vllm/model_executor/models/mixtral.py +583 -0
- vllm/model_executor/models/mixtral_quant.py +404 -0
- vllm/model_executor/models/mpt.py +295 -0
- vllm/model_executor/models/olmo.py +356 -0
- vllm/model_executor/models/opt.py +349 -0
- vllm/model_executor/models/orion.py +319 -0
- vllm/model_executor/models/phi.py +300 -0
- vllm/model_executor/models/qwen.py +284 -0
- vllm/model_executor/models/qwen2.py +367 -0
- vllm/model_executor/models/qwen2_moe.py +447 -0
- vllm/model_executor/models/stablelm.py +301 -0
- vllm/model_executor/models/starcoder2.py +302 -0
- vllm/model_executor/models/xverse.py +366 -0
- vllm/model_executor/sampling_metadata.py +588 -0
- vllm/model_executor/utils.py +35 -0
- vllm/outputs.py +150 -0
- vllm/py.typed +2 -0
- vllm/sampling_params.py +340 -0
- vllm/sequence.py +766 -0
- vllm/spec_decode/__init__.py +0 -0
- vllm/spec_decode/batch_expansion.py +397 -0
- vllm/spec_decode/interfaces.py +73 -0
- vllm/spec_decode/metrics.py +191 -0
- vllm/spec_decode/multi_step_worker.py +203 -0
- vllm/spec_decode/ngram_worker.py +176 -0
- vllm/spec_decode/spec_decode_worker.py +472 -0
- vllm/spec_decode/top1_proposer.py +200 -0
- vllm/spec_decode/util.py +228 -0
- vllm/test_utils.py +41 -0
- vllm/transformers_utils/__init__.py +0 -0
- vllm/transformers_utils/config.py +58 -0
- vllm/transformers_utils/configs/__init__.py +16 -0
- vllm/transformers_utils/configs/chatglm.py +68 -0
- vllm/transformers_utils/configs/dbrx.py +278 -0
- vllm/transformers_utils/configs/falcon.py +87 -0
- vllm/transformers_utils/configs/jais.py +236 -0
- vllm/transformers_utils/configs/mpt.py +178 -0
- vllm/transformers_utils/detokenizer.py +313 -0
- vllm/transformers_utils/tokenizer.py +149 -0
- vllm/transformers_utils/tokenizer_group/__init__.py +33 -0
- vllm/transformers_utils/tokenizer_group/base_tokenizer_group.py +55 -0
- vllm/transformers_utils/tokenizer_group/ray_tokenizer_group.py +169 -0
- vllm/transformers_utils/tokenizer_group/tokenizer_group.py +78 -0
- vllm/transformers_utils/tokenizers/__init__.py +5 -0
- vllm/transformers_utils/tokenizers/baichuan.py +255 -0
- vllm/usage/__init__.py +0 -0
- vllm/usage/usage_lib.py +209 -0
- vllm/utils.py +677 -0
- vllm/worker/__init__.py +0 -0
- vllm/worker/cache_engine.py +105 -0
- vllm/worker/cpu_model_runner.py +346 -0
- vllm/worker/cpu_worker.py +321 -0
- vllm/worker/model_runner.py +1168 -0
- vllm/worker/neuron_model_runner.py +196 -0
- vllm/worker/neuron_worker.py +98 -0
- vllm/worker/worker.py +345 -0
- vllm/worker/worker_base.py +146 -0
- vllm_npu-0.4.2.dist-info/LICENSE +201 -0
- vllm_npu-0.4.2.dist-info/METADATA +173 -0
- vllm_npu-0.4.2.dist-info/RECORD +219 -0
- vllm_npu-0.4.2.dist-info/WHEEL +5 -0
- 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)
|