Skip to content

GitLab

Switch branch/tag
Find file Normal viewHistoryPermalink
setter.py 2.57 KB
Newer Older
wangsen's avatar
Initial commit  
wangsen committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
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
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113 
# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.


class ModelSetter:
    '''Model parameter setter.

    See convert.py for a full list of supported parameters and their names.
    '''

    @classmethod
    def set_tensor(cls, dst, src):
        '''Copy (in-place) src tensor to dst tensor.'''
        if src is not None:
            dst.data.copy_(src)

    @classmethod
    def has_position_embeddings(cls, model):
        '''
        Return True if learned parameters exist for position embeddings (e.g.,
        learned absolute), and False otherwise (e.g., RoPE).
        '''
        raise NotImplementedError

    @classmethod
    def set_embeddings(
        cls,
        model,
        word=None,
        pos=None,
    ):
        '''Set word and position embeddings.'''
        raise NotImplementedError

    @classmethod
    def set_output_word_embeddings(
        cls,
        model,
        emb=None,
    ):
        '''Set output word embeddings for final pipeline stage.'''
        raise NotImplementedError

    @classmethod
    def set_layer(
        cls,
        model,
        layer_idx,
        self_attn_norm_weight=None,
        self_attn_norm_bias=None,
        self_attn_qkv_weight=None,
        self_attn_qkv_bias=None,
        self_attn_proj_weight=None,
        self_attn_proj_bias=None,
        mlp_norm_weight=None,
        mlp_norm_bias=None,
        mlp_fc1_weight=None,
        mlp_fc1_bias=None,
        mlp_fc2_weight=None,
        mlp_fc2_bias=None,
    ):
        '''Set layer parameters.'''
        raise NotImplementedError

    @classmethod
    def set_final_norm(
        cls,
        model,
        weight=None,
        bias=None,
    ):
        '''Set final norm parameters (i.e., after last transformer layer).'''
        raise NotImplementedError

    @classmethod
    def set_output_layer(
        cls,
        model,
        weight=None,
    ):
        '''Set output (i.e., 'dense') weights.'''
        raise NotImplementedError

    @classmethod
    def set_pooler(
        cls,
        model,
        weight=None,
        bias=None,
    ):
        '''Set pooler parameters (e.g., for Bert).'''
        raise NotImplementedError

    @classmethod
    def set_lm_head(
        cls,
        model,
        dense_weight=None,
        dense_bias=None,
        norm_weight=None,
        norm_bias=None,
    ):
        '''Set LM head parameters.'''
        raise NotImplementedError

    @classmethod
    def set_binary_head(
        cls,
        model,
        weight=None,
        bias=None,
    ):
        '''Set binary head parameters.'''
        raise NotImplementedError