falcon.py 21.6 KB
Newer Older
1
2
# SPDX-License-Identifier: Apache-2.0

Zhuohan Li's avatar
Zhuohan Li committed
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
# Adapted from
# https://github.com/huggingface/transformers/blob/a5cc30d72ae2dc19af534e4b35c986cc28db1275/src/transformers/models/falcon/modeling_falcon.py
# Copyright 2023 The vLLM team.
# Copyright 2023 the Falcon authors and HuggingFace Inc. team.  All rights
# reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""PyTorch Falcon model."""

import math
23
from typing import Iterable, List, Optional, Set, Tuple, Union
Zhuohan Li's avatar
Zhuohan Li committed
24
25
26
27
28
29

import torch
from torch import nn
from torch.nn import LayerNorm
from transformers import FalconConfig as HF_FalconConfig

30
from vllm.attention import Attention, AttentionMetadata
31
from vllm.compilation.decorators import support_torch_compile
32
from vllm.config import CacheConfig, VllmConfig
33
from vllm.distributed import (get_pp_group, get_tensor_model_parallel_rank,
34
35
                              get_tensor_model_parallel_world_size,
                              tensor_model_parallel_all_reduce)
36
from vllm.model_executor.layers.activation import get_act_fn
37
38
39
from vllm.model_executor.layers.linear import (ColumnParallelLinear,
                                               QKVParallelLinear,
                                               RowParallelLinear)
40
from vllm.model_executor.layers.logits_processor import LogitsProcessor
41
from vllm.model_executor.layers.quantization import QuantizationConfig
42
from vllm.model_executor.layers.rotary_embedding import get_rope
Joe Runde's avatar
Joe Runde committed
43
from vllm.model_executor.layers.sampler import SamplerOutput, get_sampler
44
from vllm.model_executor.layers.vocab_parallel_embedding import (
45
    ParallelLMHead, VocabParallelEmbedding)
46
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
47
from vllm.model_executor.sampling_metadata import SamplingMetadata
48
from vllm.sequence import IntermediateTensors
Zhuohan Li's avatar
Zhuohan Li committed
49
50
from vllm.transformers_utils.configs import RWConfig

51
52
from .interfaces import SupportsPP
from .utils import (is_pp_missing_parameter,
53
54
                    make_empty_intermediate_tensors_factory, make_layers,
                    maybe_prefix)
55

Zhuohan Li's avatar
Zhuohan Li committed
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
FalconConfig = Union[HF_FalconConfig, RWConfig]


def _get_alibi_slopes(total_num_heads: int) -> torch.Tensor:
    closest_power_of_2 = 2**math.floor(math.log2(total_num_heads))
    base = torch.tensor(2**(-(2**-(math.log2(closest_power_of_2) - 3))),
                        dtype=torch.float32)
    powers = torch.arange(1, 1 + closest_power_of_2, dtype=torch.int32)
    slopes = torch.pow(base, powers)

    if closest_power_of_2 != total_num_heads:
        extra_base = torch.tensor(
            2**(-(2**-(math.log2(2 * closest_power_of_2) - 3))),
            dtype=torch.float32)
        num_remaining_heads = min(closest_power_of_2,
                                  total_num_heads - closest_power_of_2)
        extra_powers = torch.arange(1,
                                    1 + 2 * num_remaining_heads,
                                    2,
                                    dtype=torch.int32)
        slopes = torch.cat(
            [slopes, torch.pow(extra_base, extra_powers)], dim=0)

    return slopes


class FalconAttention(nn.Module):

84
85
86
    def __init__(
        self,
        config: FalconConfig,
87
        cache_config: Optional[CacheConfig] = None,
88
        quant_config: Optional[QuantizationConfig] = None,
89
        prefix: str = "",
90
    ):
Zhuohan Li's avatar
Zhuohan Li committed
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
        super().__init__()

        self.hidden_size = config.hidden_size
        tp_size = get_tensor_model_parallel_world_size()

        self.total_num_heads = config.num_attention_heads
        assert self.total_num_heads % tp_size == 0
        self.num_heads = self.total_num_heads // tp_size
        self.head_dim = self.hidden_size // self.total_num_heads
        assert self.head_dim * self.total_num_heads == self.hidden_size

        self.new_decoder_architecture = config.new_decoder_architecture
        self.multi_query = config.multi_query

        if self.new_decoder_architecture:
            self.total_num_kv_heads = config.num_kv_heads
        elif self.multi_query:
            self.total_num_kv_heads = 1
        else:
            self.total_num_kv_heads = self.total_num_heads
111
112
113
114
115
116
117
118
119
        if self.total_num_kv_heads >= tp_size:
            # Number of KV heads is greater than TP size, so we partition
            # the KV heads across multiple tensor parallel GPUs.
            assert self.total_num_kv_heads % tp_size == 0
        else:
            # Number of KV heads is less than TP size, so we replicate
            # the KV heads across multiple tensor parallel GPUs.
            assert tp_size % self.total_num_kv_heads == 0
        self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
Zhuohan Li's avatar
Zhuohan Li committed
120

121
122
123
124
125
126
127
        self.query_key_value = QKVParallelLinear(
            self.hidden_size,
            self.head_dim,
            self.total_num_heads,
            self.total_num_kv_heads,
            bias=config.bias,
            skip_bias_add=True,
128
            quant_config=quant_config,
129
        )
Zhuohan Li's avatar
Zhuohan Li committed
130
131
132
133
134
135
136
137
138
139
140
141
        self.q_size = self.num_heads * self.head_dim
        self.kv_size = self.num_kv_heads * self.head_dim

        # Layer-wise attention scaling
        self.inv_norm_factor = 1.0 / math.sqrt(self.head_dim)
        self.reduce_row_parallel_results = not (config.new_decoder_architecture
                                                or config.parallel_attn)
        self.dense = RowParallelLinear(
            self.hidden_size,
            self.hidden_size,
            bias=config.bias,
            skip_bias_add=True,
142
            quant_config=quant_config,
Zhuohan Li's avatar
Zhuohan Li committed
143
144
145
146
147
148
149
150
            reduce_results=self.reduce_row_parallel_results)

        self.use_rotary = config.rotary
        self.use_alibi = config.alibi
        assert not (self.use_rotary and self.use_alibi), (
            "Rotary and alibi are mutually exclusive.")

        if self.use_rotary:
151
152
153
            rope_theta = getattr(config, "rope_theta", 10000)
            max_position_embeddings = getattr(config,
                                              "max_position_embeddings", 8192)
Woosuk Kwon's avatar
Woosuk Kwon committed
154
            self.rotary_emb = get_rope(
155
156
                self.head_dim,
                rotary_dim=self.head_dim,
Woosuk Kwon's avatar
Woosuk Kwon committed
157
158
159
                max_position=max_position_embeddings,
                base=rope_theta,
            )
160
161
162
            self.attn = Attention(self.num_heads,
                                  self.head_dim,
                                  self.inv_norm_factor,
163
                                  num_kv_heads=self.num_kv_heads,
164
165
                                  quant_config=quant_config,
                                  prefix=f"{prefix}.attn")
Zhuohan Li's avatar
Zhuohan Li committed
166
167
168
169
170
171
172
        elif self.use_alibi:
            tp_rank = get_tensor_model_parallel_rank()
            head_start = tp_rank * self.num_heads
            head_end = (tp_rank + 1) * self.num_heads
            alibi_slopes = (_get_alibi_slopes(self.total_num_heads) *
                            self.inv_norm_factor)
            alibi_slopes = alibi_slopes[head_start:head_end].tolist()
173
174
175
176
            self.attn = Attention(self.num_heads,
                                  self.head_dim,
                                  self.inv_norm_factor,
                                  num_kv_heads=self.num_kv_heads,
177
                                  alibi_slopes=alibi_slopes,
178
179
                                  quant_config=quant_config,
                                  prefix=f"{prefix}.attn")
Zhuohan Li's avatar
Zhuohan Li committed
180
        else:
181
182
183
            self.attn = Attention(self.num_heads,
                                  self.head_dim,
                                  scale=self.inv_norm_factor,
184
                                  num_kv_heads=self.num_kv_heads,
185
                                  cache_config=cache_config,
186
187
                                  quant_config=quant_config,
                                  prefix=f"{prefix}.attn")
Zhuohan Li's avatar
Zhuohan Li committed
188
189
190
191
192

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
193
194
        kv_cache: torch.Tensor,
        attn_metadata: AttentionMetadata,
Zhuohan Li's avatar
Zhuohan Li committed
195
    ) -> torch.Tensor:
196
197
198
199
        qkv, bias = self.query_key_value(hidden_states)
        if bias is not None:
            qkv += bias
        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
Zhuohan Li's avatar
Zhuohan Li committed
200
        if self.use_rotary:
Woosuk Kwon's avatar
Woosuk Kwon committed
201
            q, k = self.rotary_emb(positions, q, k)
202
        attn_output = self.attn(q, k, v, kv_cache, attn_metadata)
Zhuohan Li's avatar
Zhuohan Li committed
203
204
205
206
207
208
        attn_output, bias = self.dense(attn_output)
        return attn_output, bias


class FalconMLP(nn.Module):

209
210
211
    def __init__(
        self,
        config: FalconConfig,
212
        quant_config: Optional[QuantizationConfig] = None,
213
    ):
Zhuohan Li's avatar
Zhuohan Li committed
214
215
216
217
218
219
        super().__init__()
        hidden_size = config.hidden_size

        self.dense_h_to_4h = ColumnParallelLinear(hidden_size,
                                                  4 * hidden_size,
                                                  bias=config.bias,
220
                                                  skip_bias_add=True,
221
                                                  quant_config=quant_config)
222
        self.act = get_act_fn("gelu")
Zhuohan Li's avatar
Zhuohan Li committed
223
224
225
226
227
228
229
        self.reduce_row_parallel_results = not (config.new_decoder_architecture
                                                or config.parallel_attn)
        self.dense_4h_to_h = RowParallelLinear(
            4 * hidden_size,
            hidden_size,
            bias=config.bias,
            skip_bias_add=True,
230
            reduce_results=self.reduce_row_parallel_results,
231
            quant_config=quant_config)
Zhuohan Li's avatar
Zhuohan Li committed
232
233
234
235
236
237
238
239
240
241
242
243
244

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # NOTE(zhuohan): Following huggingface, we do not fuse bias add here.
        x, bias = self.dense_h_to_4h(x)
        if bias is not None:
            x += bias
        x = self.act(x)
        x, bias = self.dense_4h_to_h(x)
        return x, bias


class FalconDecoderLayer(nn.Module):

245
246
247
    def __init__(
        self,
        config: FalconConfig,
248
        cache_config: Optional[CacheConfig] = None,
249
        quant_config: Optional[QuantizationConfig] = None,
250
        prefix: str = "",
251
    ):
Zhuohan Li's avatar
Zhuohan Li committed
252
253
254
        super().__init__()
        hidden_size = config.hidden_size
        self.num_heads = config.num_attention_heads
255
256
257
258
259
        self.self_attention = FalconAttention(
            config,
            cache_config,
            quant_config,
            prefix=f"{prefix}.self_attention")
260
        self.mlp = FalconMLP(config, quant_config)
Zhuohan Li's avatar
Zhuohan Li committed
261
262
        self.config = config

263
264
265
        if (not hasattr(config, "num_ln_in_parallel_attn")):
            config.num_ln_in_parallel_attn = None

266
267
268
269
270
271
272
        if (config.num_ln_in_parallel_attn is None
                and config.new_decoder_architecture):
            config.num_ln_in_parallel_attn = 2

        if not config.parallel_attn:
            self.post_attention_layernorm = LayerNorm(
                hidden_size, eps=config.layer_norm_epsilon)
Zhuohan Li's avatar
Zhuohan Li committed
273
274
            self.input_layernorm = LayerNorm(hidden_size,
                                             eps=config.layer_norm_epsilon)
275
276
277
278
279
280
281
282
283
284
285
        else:
            if config.num_ln_in_parallel_attn == 2:
                # The layer norm before self-attention
                self.ln_attn = LayerNorm(hidden_size,
                                         eps=config.layer_norm_epsilon)
                # The layer norm before the MLP
                self.ln_mlp = LayerNorm(hidden_size,
                                        eps=config.layer_norm_epsilon)
            else:
                self.input_layernorm = LayerNorm(hidden_size,
                                                 eps=config.layer_norm_epsilon)
Zhuohan Li's avatar
Zhuohan Li committed
286
287
288
289
290
291
292
293

        self.reduce_row_parallel_results = not (config.new_decoder_architecture
                                                or config.parallel_attn)

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
294
295
        kv_cache: torch.Tensor,
        attn_metadata: AttentionMetadata,
296
    ) -> torch.Tensor:
Zhuohan Li's avatar
Zhuohan Li committed
297
298
        residual = hidden_states

299
        if self.config.num_ln_in_parallel_attn == 2:
Zhuohan Li's avatar
Zhuohan Li committed
300
301
302
303
304
305
306
307
308
309
            attention_layernorm_out = self.ln_attn(hidden_states)
            mlp_layernorm_out = self.ln_mlp(hidden_states)
        else:
            attention_layernorm_out = self.input_layernorm(hidden_states)

        # Self attention.
        attention_output, attention_bias = self.self_attention(
            positions=positions,
            hidden_states=attention_layernorm_out,
            kv_cache=kv_cache,
310
            attn_metadata=attn_metadata,
Zhuohan Li's avatar
Zhuohan Li committed
311
312
313
314
315
316
317
318
319
320
321
        )
        if self.reduce_row_parallel_results and attention_bias is not None:
            attention_output += attention_bias

        if not self.config.new_decoder_architecture:
            if self.config.parallel_attn:
                mlp_layernorm_out = attention_layernorm_out
            else:
                residual += attention_output
                mlp_layernorm_out = self.post_attention_layernorm(residual)

322
323
324
325
        if (self.config.new_decoder_architecture and self.config.parallel_attn
                and self.config.num_ln_in_parallel_attn == 1):
            mlp_layernorm_out = attention_layernorm_out

Zhuohan Li's avatar
Zhuohan Li committed
326
327
328
329
330
331
332
333
334
335
        # MLP.
        mlp_output, mlp_bias = self.mlp(mlp_layernorm_out)
        if self.reduce_row_parallel_results and mlp_bias is not None:
            mlp_output += mlp_bias

        if not self.reduce_row_parallel_results:
            # When MLP and Attention layers are parallel, we can use
            # only one all-reduce operator to reduce the results from
            # both MLP and Attention layers.
            mlp_output += attention_output
336
            mlp_output = tensor_model_parallel_all_reduce(mlp_output)
Zhuohan Li's avatar
Zhuohan Li committed
337
338
339
340
341
342
343
344
345
            if attention_bias is not None:
                mlp_output += attention_bias
            if mlp_bias is not None:
                mlp_output += mlp_bias

        output = mlp_output + residual
        return output


346
@support_torch_compile
Zhuohan Li's avatar
Zhuohan Li committed
347
348
class FalconModel(nn.Module):

349
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
Zhuohan Li's avatar
Zhuohan Li committed
350
        super().__init__()
351
352
353
354
355

        config = vllm_config.model_config.hf_config
        cache_config = vllm_config.cache_config
        quant_config = vllm_config.quant_config

Zhuohan Li's avatar
Zhuohan Li committed
356
357
358
359
360
361
362
        self.config = config
        self.embed_dim = config.hidden_size
        self.num_heads = config.num_attention_heads
        self.use_alibi = config.alibi

        # Embedding + LN Embedding
        self.word_embeddings = VocabParallelEmbedding(
363
364
365
            config.vocab_size,
            self.embed_dim,
        )
Zhuohan Li's avatar
Zhuohan Li committed
366
367

        # Transformer blocks
368
369
        self.start_layer, self.end_layer, self.h = make_layers(
            config.num_hidden_layers,
370
371
            lambda prefix: FalconDecoderLayer(
                config, cache_config, quant_config, prefix=prefix),
372
            prefix=f"{prefix}.h")
Zhuohan Li's avatar
Zhuohan Li committed
373
374
375

        # Final Layer Norm
        self.ln_f = LayerNorm(self.embed_dim, eps=config.layer_norm_epsilon)
376
377
378
        self.make_empty_intermediate_tensors = (
            make_empty_intermediate_tensors_factory(["hidden_states"],
                                                    config.hidden_size))
Zhuohan Li's avatar
Zhuohan Li committed
379

380
381
382
    def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.word_embeddings(input_ids)

Zhuohan Li's avatar
Zhuohan Li committed
383
384
    def forward(
        self,
385
        input_ids: torch.Tensor,
Zhuohan Li's avatar
Zhuohan Li committed
386
        positions: torch.Tensor,
387
388
        kv_caches: List[torch.Tensor],
        attn_metadata: AttentionMetadata,
389
        intermediate_tensors: Optional[IntermediateTensors],
390
        inputs_embeds: Optional[torch.Tensor] = None,
391
392
    ) -> Union[torch.Tensor, IntermediateTensors]:
        if get_pp_group().is_first_rank:
393
394
395
396
            if inputs_embeds is not None:
                hidden_states = inputs_embeds
            else:
                hidden_states = self.get_input_embeddings(input_ids)
397
398
399
        else:
            hidden_states = intermediate_tensors["hidden_states"]
        for i in range(self.start_layer, self.end_layer):
Zhuohan Li's avatar
Zhuohan Li committed
400
401
402
403
            layer = self.h[i]
            hidden_states = layer(
                positions,
                hidden_states,
404
                kv_caches[i - self.start_layer],
405
                attn_metadata,
Zhuohan Li's avatar
Zhuohan Li committed
406
            )
407
408
        if not get_pp_group().is_last_rank:
            return IntermediateTensors({"hidden_states": hidden_states})
Zhuohan Li's avatar
Zhuohan Li committed
409
410
411
412
        hidden_states = self.ln_f(hidden_states)
        return hidden_states


413
class FalconForCausalLM(nn.Module, SupportsPP):
414
415
416
    packed_modules_mapping = {
        "query_key_value": ["query_key_value"],
    }
417

418
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
Zhuohan Li's avatar
Zhuohan Li committed
419
        super().__init__()
420
421
        config = vllm_config.model_config.hf_config
        quant_config = vllm_config.quant_config
Zhuohan Li's avatar
Zhuohan Li committed
422
        self.config = config
423
        self.quant_config = quant_config
424
425
426
        self.transformer = FalconModel(vllm_config=vllm_config,
                                       prefix=maybe_prefix(
                                           prefix, "transformer"))
427
428
429
430
431
432
433
        # only Falcon-11B doesn't share lm_head weight with word embeddings
        # and previous Falcon model doesn't have tie_word_embeddings config
        # so we set tie_word_embeddings to True by default
        self.tie_word_embeddings = (config.tie_word_embeddings
                                    if config.tie_word_embeddings is not None
                                    else True)
        if self.tie_word_embeddings:
434
            self.lm_head = self.transformer.word_embeddings
435
436
437
438
        else:
            self.lm_head = ParallelLMHead(
                config.vocab_size,
                config.hidden_size,
439
                quant_config=quant_config,
440
            )
441
        self.logits_processor = LogitsProcessor(config.vocab_size)
Joe Runde's avatar
Joe Runde committed
442
        self.sampler = get_sampler()
443
444
        self.make_empty_intermediate_tensors = (
            self.transformer.make_empty_intermediate_tensors)
Zhuohan Li's avatar
Zhuohan Li committed
445

446
447
448
    def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.transformer.get_input_embeddings(input_ids)

Zhuohan Li's avatar
Zhuohan Li committed
449
450
451
452
    def forward(
        self,
        input_ids: torch.LongTensor,
        positions: torch.Tensor,
453
454
        kv_caches: List[torch.Tensor],
        attn_metadata: AttentionMetadata,
455
        intermediate_tensors: Optional[IntermediateTensors] = None,
456
        inputs_embeds: Optional[torch.Tensor] = None,
457
    ) -> torch.Tensor:
458
        hidden_states = self.transformer(input_ids, positions, kv_caches,
459
460
                                         attn_metadata, intermediate_tensors,
                                         inputs_embeds)
461
        return hidden_states
Zhuohan Li's avatar
Zhuohan Li committed
462

463
464
465
466
467
    def compute_logits(
        self,
        hidden_states: torch.Tensor,
        sampling_metadata: SamplingMetadata,
    ) -> Optional[torch.Tensor]:
468
        logits = self.logits_processor(self.lm_head, hidden_states,
469
470
471
                                       sampling_metadata)
        return logits

472
473
    def sample(
        self,
474
        logits: torch.Tensor,
475
        sampling_metadata: SamplingMetadata,
476
    ) -> Optional[SamplerOutput]:
477
        next_tokens = self.sampler(logits, sampling_metadata)
Zhuohan Li's avatar
Zhuohan Li committed
478
479
        return next_tokens

480
481
    def load_weights(self, weights: Iterable[Tuple[str,
                                                   torch.Tensor]]) -> Set[str]:
Zhuohan Li's avatar
Zhuohan Li committed
482
483
484
485
486
487
488
489
        total_num_heads = self.config.num_attention_heads
        if self.config.new_decoder_architecture:
            total_num_kv_heads = self.config.num_kv_heads
        elif self.config.multi_query:
            total_num_kv_heads = 1
        else:
            total_num_kv_heads = total_num_heads
        num_query_heads_per_kv_head = total_num_heads // total_num_kv_heads
490
        params_dict = dict(self.named_parameters(remove_duplicate=False))
491
        loaded_params: Set[str] = set()
492
        for name, loaded_weight in weights:
493
494
            if name == "lm_head.weight" and self.tie_word_embeddings:
                # Falcon uses tied embeddings except Falcon-11b.
495
                continue
CHU Tianxiang's avatar
CHU Tianxiang committed
496
497
498
            # Skip loading extra bias for GPTQ models.
            if name.endswith(".bias") and name not in params_dict:
                continue
499
500
            if is_pp_missing_parameter(name, self):
                continue
501
            param = params_dict[name]
Zhuohan Li's avatar
Zhuohan Li committed
502
            if "query_key_value" in name:
503
504
                output_dim = getattr(param, "output_dim", None)
                loaded_weight_shape = loaded_weight.shape
CHU Tianxiang's avatar
CHU Tianxiang committed
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
                if output_dim is not None:
                    loaded_weight = loaded_weight.view(
                        loaded_weight_shape[:output_dim] +
                        (total_num_kv_heads, num_query_heads_per_kv_head + 2,
                         -1) + loaded_weight_shape[output_dim + 1:])
                    wq = loaded_weight.narrow(
                        output_dim + 1, 0,
                        num_query_heads_per_kv_head).reshape(
                            *loaded_weight_shape[:output_dim], -1,
                            *loaded_weight_shape[output_dim + 1:])
                    wk = loaded_weight.narrow(
                        output_dim + 1, num_query_heads_per_kv_head,
                        1).reshape(*loaded_weight_shape[:output_dim], -1,
                                   *loaded_weight_shape[output_dim + 1:])
                    wv = loaded_weight.narrow(
                        output_dim + 1, num_query_heads_per_kv_head + 1,
                        1).reshape(*loaded_weight_shape[:output_dim], -1,
                                   *loaded_weight_shape[output_dim + 1:])
                    loaded_weight = torch.cat([wq, wk, wv], dim=output_dim)
524
525
526
527

            weight_loader = getattr(param, "weight_loader",
                                    default_weight_loader)
            weight_loader(param, loaded_weight)
528
529
            loaded_params.add(name)
        return loaded_params