__init__.py 9.16 KB
Newer Older
1
import os
2
3
import torch

4
from loguru import logger
5
from transformers.configuration_utils import PretrainedConfig
6
from transformers.models.auto import modeling_auto
7
8
from typing import Optional

9
10
from text_generation_server.models.model import Model
from text_generation_server.models.causal_lm import CausalLM
11
from text_generation_server.models.flash_causal_lm import FlashCausalLM
12
from text_generation_server.models.bloom import BLOOMSharded
13
from text_generation_server.models.mpt import MPTSharded
14
from text_generation_server.models.seq2seq_lm import Seq2SeqLM
15
from text_generation_server.models.rw import RW
16
17
from text_generation_server.models.opt import OPTSharded
from text_generation_server.models.galactica import GalacticaSharded
18
19
from text_generation_server.models.santacoder import SantaCoder
from text_generation_server.models.t5 import T5Sharded
20
from text_generation_server.models.gpt_neox import GPTNeoxSharded
21

22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
# The flag below controls whether to allow TF32 on matmul. This flag defaults to False
# in PyTorch 1.12 and later.
torch.backends.cuda.matmul.allow_tf32 = True

# The flag below controls whether to allow TF32 on cuDNN. This flag defaults to True.
torch.backends.cudnn.allow_tf32 = True

# Disable gradients
torch.set_grad_enabled(False)

__all__ = [
    "Model",
    "BLOOMSharded",
    "CausalLM",
    "FlashCausalLM",
    "GalacticaSharded",
    "Seq2SeqLM",
    "SantaCoder",
    "OPTSharded",
    "T5Sharded",
    "get_model",
]

45
FLASH_ATT_ERROR_MESSAGE = "{} requires Flash Attention enabled models."
46

47
FLASH_ATTENTION = True
48
try:
49
50
51
52
53
54
55
    from text_generation_server.models.flash_rw import FlashRWSharded
    from text_generation_server.models.flash_neox import FlashNeoXSharded
    from text_generation_server.models.flash_llama import (
        FlashLlama,
    )
    from text_generation_server.models.flash_santacoder import (
        FlashSantacoderSharded,
56
    )
57
58
59

except ImportError as e:
    logger.warning(f"Could not import Flash Attention enabled models: {e}")
60
    FLASH_ATTENTION = False
61

62
if FLASH_ATTENTION:
63
    __all__.append(FlashNeoXSharded)
64
    __all__.append(FlashRWSharded)
65
    __all__.append(FlashSantacoderSharded)
66
67
    __all__.append(FlashLlama)

68

69
def get_model(
70
71
72
73
    model_id: str,
    revision: Optional[str],
    sharded: bool,
    quantize: Optional[str],
74
    dtype: Optional[str],
75
    trust_remote_code: bool,
76
) -> Model:
77
78
79
80
81
82
83
84
85
    if dtype is None:
        dtype = torch.float16
    elif dtype == "float16":
        dtype = torch.float16
    elif dtype == "bfloat16":
        dtype = torch.bfloat16
    else:
        raise RuntimeError(f"Unknown dtype {dtype}")

86
    if "facebook/galactica" in model_id:
87
        return GalacticaSharded(
88
89
90
91
92
            model_id,
            revision,
            quantize=quantize,
            dtype=dtype,
            dtypetrust_remote_code=trust_remote_code,
93
        )
94

95
    if model_id.startswith("bigcode/"):
96
        if FLASH_ATTENTION:
97
98
99
100
            return FlashSantacoderSharded(
                model_id,
                revision,
                quantize=quantize,
101
                dtype=dtype,
102
103
                trust_remote_code=trust_remote_code,
            )
104
105
106
107
        elif sharded:
            raise NotImplementedError(
                FLASH_ATT_ERROR_MESSAGE.format("Sharded Santacoder")
            )
108
        else:
109
            return SantaCoder(
110
111
112
                model_id,
                revision,
                quantize=quantize,
113
                dtype=dtype,
114
115
                trust_remote_code=trust_remote_code,
            )
116

OlivierDehaene's avatar
v0.8.2  
OlivierDehaene committed
117
118
119
    config_dict, _ = PretrainedConfig.get_config_dict(
        model_id, revision=revision, trust_remote_code=trust_remote_code
    )
120
    model_type = config_dict["model_type"]
121

122
    if model_type == "gpt_bigcode":
123
        if FLASH_ATTENTION:
124
125
126
127
            return FlashSantacoderSharded(
                model_id,
                revision,
                quantize=quantize,
128
                dtype=dtype,
129
130
                trust_remote_code=trust_remote_code,
            )
131
132
133
134
        elif sharded:
            raise NotImplementedError(
                FLASH_ATT_ERROR_MESSAGE.format("Sharded Santacoder")
            )
135
        else:
136
            return SantaCoder(
137
138
139
                model_id,
                revision,
                quantize=quantize,
140
                dtype=dtype,
141
142
                trust_remote_code=trust_remote_code,
            )
143

144
    if model_type == "bloom":
145
        return BLOOMSharded(
146
147
148
149
150
            model_id,
            revision,
            quantize=quantize,
            dtype=dtype,
            trust_remote_code=trust_remote_code,
151
        )
152
153
154
155
    elif model_type == "mpt":
        return MPTSharded(
            model_id, revision, quantize=quantize, trust_remote_code=trust_remote_code
        )
156
157
158
159
160
161
162

    elif model_type == "gpt_neox":
        if FLASH_ATTENTION:
            return FlashNeoXSharded(
                model_id,
                revision,
                quantize=quantize,
163
                dtype=dtype,
164
165
166
167
                trust_remote_code=trust_remote_code,
            )
        elif sharded:
            return GPTNeoxSharded(
168
169
170
                model_id,
                revision,
                quantize=quantize,
171
                dtype=dtype,
172
173
                trust_remote_code=trust_remote_code,
            )
174
        else:
175
            return CausalLM(
176
177
178
                model_id,
                revision,
                quantize=quantize,
179
                dtype=dtype,
180
181
                trust_remote_code=trust_remote_code,
            )
182

183
184
185
    elif model_type == "llama":
        if FLASH_ATTENTION:
            return FlashLlama(
186
187
188
                model_id,
                revision,
                quantize=quantize,
189
                dtype=dtype,
190
191
                trust_remote_code=trust_remote_code,
            )
192
193
        elif sharded:
            raise NotImplementedError(FLASH_ATT_ERROR_MESSAGE.format("Sharded Llama"))
194
        else:
195
            return CausalLM(
196
197
198
                model_id,
                revision,
                quantize=quantize,
199
                dtype=dtype,
200
201
                trust_remote_code=trust_remote_code,
            )
202

203
204
205
    if model_type in ["RefinedWeb", "RefinedWebModel"]:
        if sharded:
            if FLASH_ATTENTION:
206
207
208
                if config_dict.get("alibi", False) or (
                    model_type == "RefinedWebModel"
                    and config_dict.get("multi_query", True)
209
210
211
212
213
214
                ):
                    raise NotImplementedError("sharded is not supported for this model")
                return FlashRWSharded(
                    model_id,
                    revision,
                    quantize=quantize,
215
                    dtype=dtype,
216
217
218
219
220
221
                    trust_remote_code=trust_remote_code,
                )
            raise NotImplementedError(
                FLASH_ATT_ERROR_MESSAGE.format(f"Sharded RefinedWeb")
            )
        else:
222
            if FLASH_ATTENTION and not config_dict.get("alibi", False):
223
                return FlashRWSharded(
224
225
226
                    model_id,
                    revision,
                    quantize=quantize,
227
                    dtype=dtype,
228
229
230
231
232
233
234
                    trust_remote_code=trust_remote_code,
                )
            else:
                return RW(
                    model_id,
                    revision,
                    quantize=quantize,
235
                    dtype=dtype,
236
237
238
                    trust_remote_code=trust_remote_code,
                )

239
240
    elif model_type == "opt":
        return OPTSharded(
241
242
243
244
245
            model_id,
            revision,
            quantize=quantize,
            dtype=dtype,
            trust_remote_code=trust_remote_code,
246
        )
247

248
    elif model_type == "t5":
249
250
251
252
        return T5Sharded(
            model_id,
            revision,
            quantize=quantize,
253
            dtype=dtype,
254
255
            trust_remote_code=trust_remote_code,
        )
256
257
258

    if sharded:
        raise ValueError("sharded is not supported for AutoModel")
259
260
261
262
    if quantize == "gptq":
        raise ValueError(
            "gptq quantization is not supported for AutoModel, you can try to quantize it with `text-generation-server quantize ORIGINAL_MODEL_ID NEW_MODEL_ID`"
        )
263
264

    if model_type in modeling_auto.MODEL_FOR_CAUSAL_LM_MAPPING_NAMES:
265
        return CausalLM(
266
267
268
269
270
            model_id,
            revision,
            quantize=quantize,
            dtype=dtype,
            trust_remote_code=trust_remote_code,
271
        )
272
    if model_type in modeling_auto.MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING_NAMES:
273
        return Seq2SeqLM(
274
275
276
277
278
            model_id,
            revision,
            quantize=quantize,
            dtype=dtype,
            trust_remote_code=trust_remote_code,
279
280
        )

281
    auto_map = config_dict.get("auto_map", None)
282
283
284
285
286
287
    if trust_remote_code and auto_map is not None:
        if "AutoModelForCausalLM" in auto_map.keys():
            return CausalLM(
                model_id,
                revision,
                quantize=quantize,
288
                dtype=dtype,
289
290
                trust_remote_code=trust_remote_code,
            )
291
        if "AutoModelForSeq2SeqLM" in auto_map.keys():
292
293
294
295
            return Seq2SeqLM(
                model_id,
                revision,
                quantize=quantize,
296
                dtype=dtype,
297
298
                trust_remote_code=trust_remote_code,
            )
299
300

    raise ValueError(f"Unsupported model type {model_type}")