"mmdet3d/datasets/transforms/formating.py" did not exist on "260b5b4ad853074988ba700e5eae258d408afeb9"
mpt.py 3.13 KB
Newer Older
Casper's avatar
Casper committed
1
from .base import BaseAWQForCausalLM
Vik Paruchuri's avatar
Vik Paruchuri committed
2
from typing import Dict
3
from transformers.models.mpt.modeling_mpt import MptBlock as OldMptBlock, MptForCausalLM
Casper's avatar
Casper committed
4
5
6

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

9
    @staticmethod
Vik Paruchuri's avatar
Vik Paruchuri committed
10
    def fuse_layers(model: MptForCausalLM, quant_config: Dict):
Casper Hansen's avatar
Casper Hansen committed
11
        fuser = MptFuser(model)
12
        fuser.fuse_transformer()
13

14
    @staticmethod
Casper Hansen's avatar
Casper Hansen committed
15
    def get_model_layers(model: MptForCausalLM):
Casper's avatar
Casper committed
16
17
        return model.transformer.blocks
    
18
    @staticmethod
Casper Hansen's avatar
Casper Hansen committed
19
    def get_act_for_scaling(module: OldMptBlock):
20
21
22
23
24
25
26
27
        return dict(
            is_scalable=True,
            scale_name="ffn.act",
            scale_layer=module.ffn.act,
            scale_shape=module.ffn.up_proj.out_features
        )
    
    @staticmethod
Casper Hansen's avatar
Casper Hansen committed
28
    def move_embed(model: MptForCausalLM, device: str):
29
30
31
        model.transformer.wte = model.transformer.wte.to(device)
        model.transformer.emb_drop = model.transformer.emb_drop.to(device)
    
32
    @staticmethod
Casper Hansen's avatar
Casper Hansen committed
33
    def get_layers_for_scaling(module: OldMptBlock, input_feat, module_kwargs):
Casper's avatar
Casper committed
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
        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
54
            prev_op=module.norm_2,
Casper's avatar
Casper committed
55
56
57
58
59
60
61
62
63
64
65
66
            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']
        ))

Casper Hansen's avatar
Casper Hansen committed
67
68
69
70
        return layers

from typing import List, Tuple
from awq.utils.utils import set_module_name
71
72
from awq.modules.fused.block import MPTBlock
from awq.modules.fused.model import MPTModel
Casper Hansen's avatar
Casper Hansen committed
73
74

class MptFuser:
75
    def __init__(self, model: MptForCausalLM):
Casper Hansen's avatar
Casper Hansen committed
76
77
        self.model = model

Casper Hansen's avatar
Casper Hansen committed
78
        self.mpt_blocks: List[Tuple[str, OldMptBlock]] = [
Casper Hansen's avatar
Casper Hansen committed
79
            (name, module) for name, module in self.model.named_modules()
Casper Hansen's avatar
Casper Hansen committed
80
            if 'mptblock' in module.__class__.__name__.lower()
Casper Hansen's avatar
Casper Hansen committed
81
82
        ]

83
84
85
86
87
88
    def fuse_transformer(self):
        blocks = []

        module: OldMptBlock
        for module in self.model.transformer.blocks:
            blocks.append(MPTBlock(
Casper Hansen's avatar
Casper Hansen committed
89
90
91
92
                self.model.config.d_model,
                self.model.config.n_heads,
                module.attn.Wqkv,
                module.attn.out_proj,
Casper Hansen's avatar
Casper Hansen committed
93
                module.ffn,
94
95
96
97
98
                module.norm_1,
                module.norm_2,
                next(iter(module.state_dict().values())).device, 
                self.model.config.max_new_tokens
            ))
Casper Hansen's avatar
Casper Hansen committed
99

100
101
102
103
104
105
        self.model.transformer = MPTModel(
            self.model.config.vocab_size,
            blocks,
            self.model.transformer.wte,
            self.model.transformer.norm_f,
        )