"configs/datasets/PJExam/PJExam_gen_8cd97c.py" did not exist on "c94cc943485e275897ad95cfa5192ff8e066378a"
mpt.py 1.63 KB
Newer Older
Casper's avatar
Casper committed
1
2
3
4
from .base import BaseAWQForCausalLM

class MptAWQForCausalLM(BaseAWQForCausalLM):
    layer_type = "MPTBlock"
5
    max_new_tokens_key = "max_seq_len"
Casper's avatar
Casper committed
6

7
8
    @staticmethod
    def get_model_layers(model):
Casper's avatar
Casper committed
9
10
        return model.transformer.blocks
    
11
12
13
14
15
16
17
18
19
20
21
22
23
24
    @staticmethod
    def get_act_for_scaling(module):
        return dict(
            is_scalable=True,
            scale_name="ffn.act",
            scale_layer=module.ffn.act,
            scale_shape=module.ffn.up_proj.out_features
        )
    
    @staticmethod
    def move_embed(model, device):
        model.transformer.wte = model.transformer.wte.to(device)
        model.transformer.emb_drop = model.transformer.emb_drop.to(device)
    
25
26
    @staticmethod
    def get_layers_for_scaling(module, input_feat, module_kwargs):
Casper's avatar
Casper committed
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
        layers = []

        # attention input
        layers.append(dict(
            prev_op=module.norm_1,
            layers=[module.attn.Wqkv],
            inp=input_feat['attn.Wqkv'],
            module2inspect=module.attn,
            kwargs=module_kwargs
        ))

        # attention output
        layers.append(dict(
            prev_op=module.attn.Wqkv,
            layers=[module.attn.out_proj],
            inp=input_feat['attn.out_proj']
        ))

        # linear 1
        layers.append(dict(
Casper Hansen's avatar
Casper Hansen committed
47
            prev_op=module.norm_2,
Casper's avatar
Casper committed
48
49
50
51
52
53
54
55
56
57
58
59
            layers=[module.ffn.up_proj],
            inp=input_feat['ffn.up_proj'],
            module2inspect=module.ffn
        ))

        # linear 2
        layers.append(dict(
            prev_op=module.ffn.act,
            layers=[module.ffn.down_proj],
            inp=input_feat['ffn.down_proj']
        ))

60
        return layers