__init__.py 13.4 KB
Newer Older
1
2
import torch

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

Nicolas Patry's avatar
Nicolas Patry committed
8
from text_generation_server.utils.speculate import get_speculate, set_speculate
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
drbh's avatar
drbh committed
21
from text_generation_server.models.phi import Phi
22

23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
# 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",
]

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

48
FLASH_ATTENTION = True
49
try:
50
51
52
53
54
55
56
    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,
57
    )
58
    from text_generation_server.models.idefics import IDEFICSSharded
59
60
    from text_generation_server.models.flash_mistral import FlashMistral
    from text_generation_server.models.flash_mixtral import FlashMixtral
drbh's avatar
drbh committed
61
    from text_generation_server.models.flash_phi import FlashPhi
62
    from text_generation_server.utils.flash_attn import HAS_FLASH_ATTN_V2_CUDA
63
64
65

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

69
if FLASH_ATTENTION:
70
    __all__.append(FlashNeoXSharded)
71
    __all__.append(FlashRWSharded)
72
    __all__.append(FlashSantacoderSharded)
73
    __all__.append(FlashLlama)
74
    __all__.append(IDEFICSSharded)
75
    __all__.append(FlashMistral)
OlivierDehaene's avatar
OlivierDehaene committed
76
    __all__.append(FlashMixtral)
drbh's avatar
drbh committed
77
    __all__.append(FlashPhi)
OlivierDehaene's avatar
OlivierDehaene committed
78
79


80
def get_model(
81
82
83
84
    model_id: str,
    revision: Optional[str],
    sharded: bool,
    quantize: Optional[str],
Nicolas Patry's avatar
Nicolas Patry committed
85
    speculate: Optional[int],
86
    dtype: Optional[str],
87
    trust_remote_code: bool,
88
) -> Model:
89
    if dtype is None:
90
91
92
        # Keep it as default for now and let
        # every model resolve their own default dtype.
        dtype = None
93
94
95
96
97
98
99
    elif dtype == "float16":
        dtype = torch.float16
    elif dtype == "bfloat16":
        dtype = torch.bfloat16
    else:
        raise RuntimeError(f"Unknown dtype {dtype}")

Nicolas Patry's avatar
Nicolas Patry committed
100
101
102
103
104
    if speculate is not None:
        set_speculate(speculate)
    else:
        set_speculate(0)

105
    if "facebook/galactica" in model_id:
106
        return GalacticaSharded(
107
108
109
110
            model_id,
            revision,
            quantize=quantize,
            dtype=dtype,
Nicolas Patry's avatar
Nicolas Patry committed
111
            trust_remote_code=trust_remote_code,
112
        )
113

114
    if model_id.startswith("bigcode/"):
115
        if FLASH_ATTENTION:
116
117
118
119
            return FlashSantacoderSharded(
                model_id,
                revision,
                quantize=quantize,
120
                dtype=dtype,
121
122
                trust_remote_code=trust_remote_code,
            )
123
124
125
126
        elif sharded:
            raise NotImplementedError(
                FLASH_ATT_ERROR_MESSAGE.format("Sharded Santacoder")
            )
127
        else:
128
            return SantaCoder(
129
130
131
                model_id,
                revision,
                quantize=quantize,
132
                dtype=dtype,
133
134
                trust_remote_code=trust_remote_code,
            )
135

OlivierDehaene's avatar
v0.8.2  
OlivierDehaene committed
136
137
138
    config_dict, _ = PretrainedConfig.get_config_dict(
        model_id, revision=revision, trust_remote_code=trust_remote_code
    )
Nicolas Patry's avatar
Nicolas Patry committed
139
140
141
142
143
144
145
146
147

    use_medusa = None
    if "medusa_num_heads" in config_dict:
        use_medusa = model_id
        model_id = config_dict["base_model_name_or_path"]
        revision = "main"
        speculate_medusa = config_dict["medusa_num_heads"]
        if speculate is not None:
            if speculate > speculate_medusa:
OlivierDehaene's avatar
OlivierDehaene committed
148
149
150
                raise RuntimeError(
                    "Speculate is set to `{speculate}` but this medusa models only has `{speculate_medusa}` heads, please make them match"
                )
Nicolas Patry's avatar
Nicolas Patry committed
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
            else:
                set_speculate(speculate)
        else:
            set_speculate(speculate_medusa)

        config_dict, _ = PretrainedConfig.get_config_dict(
            model_id, revision=revision, trust_remote_code=trust_remote_code
        )
        method = "medusa"
    else:
        method = "n-gram"

    speculate = get_speculate()
    if speculate > 0:
        logger.info(f"Using speculation {method} with {speculate} input ids.")

167
    model_type = config_dict["model_type"]
168

169
    if model_type == "gpt_bigcode":
170
        if FLASH_ATTENTION:
171
172
173
174
            return FlashSantacoderSharded(
                model_id,
                revision,
                quantize=quantize,
175
                dtype=dtype,
176
177
                trust_remote_code=trust_remote_code,
            )
178
179
180
181
        elif sharded:
            raise NotImplementedError(
                FLASH_ATT_ERROR_MESSAGE.format("Sharded Santacoder")
            )
182
        else:
183
            return SantaCoder(
184
185
186
                model_id,
                revision,
                quantize=quantize,
187
                dtype=dtype,
188
189
                trust_remote_code=trust_remote_code,
            )
190

191
    if model_type == "bloom":
192
        return BLOOMSharded(
193
194
195
196
197
            model_id,
            revision,
            quantize=quantize,
            dtype=dtype,
            trust_remote_code=trust_remote_code,
198
        )
199
200
    elif model_type == "mpt":
        return MPTSharded(
OlivierDehaene's avatar
OlivierDehaene committed
201
202
203
204
205
            model_id,
            revision,
            quantize=quantize,
            dtype=dtype,
            trust_remote_code=trust_remote_code,
206
        )
207
208
209
210
211
212
213

    elif model_type == "gpt_neox":
        if FLASH_ATTENTION:
            return FlashNeoXSharded(
                model_id,
                revision,
                quantize=quantize,
214
                dtype=dtype,
215
216
217
218
                trust_remote_code=trust_remote_code,
            )
        elif sharded:
            return GPTNeoxSharded(
219
220
221
                model_id,
                revision,
                quantize=quantize,
222
                dtype=dtype,
223
224
                trust_remote_code=trust_remote_code,
            )
225
        else:
226
            return CausalLM(
227
228
229
                model_id,
                revision,
                quantize=quantize,
230
                dtype=dtype,
231
232
                trust_remote_code=trust_remote_code,
            )
drbh's avatar
drbh committed
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
        
    elif model_type == "phi":
        if FLASH_ATTENTION:
            return FlashPhi(
                model_id,
                revision,
                quantize=quantize,
                dtype=dtype,
                trust_remote_code=trust_remote_code,
                use_medusa=use_medusa,
            )
        else:
            return CausalLM(
                model_id,
                revision,
                quantize=quantize,
                dtype=dtype,
                trust_remote_code=trust_remote_code,
            )

    elif model_type == "phi-msft":
        if FLASH_ATTENTION:
            raise NotImplementedError("Legacy phi-msft is not supported with Flash Attention")
        else:
            return Phi(
                model_id,
                revision,
                quantize=quantize,
                dtype=dtype,
                trust_remote_code=trust_remote_code,
            )
264

xiaobin's avatar
xiaobin committed
265
    elif model_type == "llama" or model_type == "baichuan":
266
267
        if FLASH_ATTENTION:
            return FlashLlama(
268
269
270
                model_id,
                revision,
                quantize=quantize,
271
                dtype=dtype,
272
                trust_remote_code=trust_remote_code,
OlivierDehaene's avatar
OlivierDehaene committed
273
                use_medusa=use_medusa,
274
            )
275
276
        elif sharded:
            raise NotImplementedError(FLASH_ATT_ERROR_MESSAGE.format("Sharded Llama"))
277
        else:
278
            return CausalLM(
279
280
281
                model_id,
                revision,
                quantize=quantize,
282
                dtype=dtype,
283
284
                trust_remote_code=trust_remote_code,
            )
285

286
    if model_type in ["RefinedWeb", "RefinedWebModel", "falcon"]:
287
288
        if sharded:
            if FLASH_ATTENTION:
289
                if config_dict.get("alibi", False):
290
291
292
293
294
                    raise NotImplementedError("sharded is not supported for this model")
                return FlashRWSharded(
                    model_id,
                    revision,
                    quantize=quantize,
295
                    dtype=dtype,
296
297
                    trust_remote_code=trust_remote_code,
                )
298
            raise NotImplementedError(FLASH_ATT_ERROR_MESSAGE.format(f"Sharded Falcon"))
299
        else:
300
            if FLASH_ATTENTION and not config_dict.get("alibi", False):
301
                return FlashRWSharded(
302
303
304
                    model_id,
                    revision,
                    quantize=quantize,
305
                    dtype=dtype,
306
307
308
309
310
311
312
                    trust_remote_code=trust_remote_code,
                )
            else:
                return RW(
                    model_id,
                    revision,
                    quantize=quantize,
313
                    dtype=dtype,
314
315
316
                    trust_remote_code=trust_remote_code,
                )

317
    if model_type == "mistral":
318
319
320
321
        sliding_window = config_dict.get("sliding_window", -1)
        if (
            (sliding_window is None or sliding_window == -1) and FLASH_ATTENTION
        ) or HAS_FLASH_ATTN_V2_CUDA:
322
323
324
325
326
327
328
            return FlashMistral(
                model_id,
                revision,
                quantize=quantize,
                dtype=dtype,
                trust_remote_code=trust_remote_code,
            )
OlivierDehaene's avatar
OlivierDehaene committed
329
330

    if model_type == "mixtral":
331
332
333
334
        sliding_window = config_dict.get("sliding_window", -1)
        if (
            (sliding_window is None or sliding_window == -1) and FLASH_ATTENTION
        ) or HAS_FLASH_ATTN_V2_CUDA:
OlivierDehaene's avatar
OlivierDehaene committed
335
336
337
338
339
340
341
            return FlashMixtral(
                model_id,
                revision,
                quantize=quantize,
                dtype=dtype,
                trust_remote_code=trust_remote_code,
            )
342
343

    if model_type == "opt":
344
        return OPTSharded(
345
346
347
348
349
            model_id,
            revision,
            quantize=quantize,
            dtype=dtype,
            trust_remote_code=trust_remote_code,
350
        )
351

352
    if model_type == "t5":
353
354
355
356
        return T5Sharded(
            model_id,
            revision,
            quantize=quantize,
357
            dtype=dtype,
358
359
            trust_remote_code=trust_remote_code,
        )
360
    if model_type == "idefics":
361
        if FLASH_ATTENTION:
OlivierDehaene's avatar
OlivierDehaene committed
362
363
364
365
366
367
368
            return IDEFICSSharded(
                model_id,
                revision,
                quantize=quantize,
                dtype=dtype,
                trust_remote_code=trust_remote_code,
            )
369
370
        else:
            raise NotImplementedError(FLASH_ATT_ERROR_MESSAGE.format("Idefics"))
371
372

    if sharded:
373
        raise NotImplementedError("sharded is not supported for AutoModel")
374
    if quantize == "gptq":
375
        raise NotImplementedError(
376
377
            "gptq quantization is not supported for AutoModel, you can try to quantize it with `text-generation-server quantize ORIGINAL_MODEL_ID NEW_MODEL_ID`"
        )
378
    if quantize == "awq":
379
        raise NotImplementedError("awq quantization is not supported for AutoModel")
Nicolas Patry's avatar
Nicolas Patry committed
380
    elif (quantize == "bitsandbytes-fp4") or (quantize == "bitsandbytes-nf4"):
381
        raise NotImplementedError("4bit quantization is not supported for AutoModel")
OlivierDehaene's avatar
OlivierDehaene committed
382
    elif quantize == "eetq":
383
        raise NotImplementedError("Eetq quantization is not supported for AutoModel")
384
    if model_type in modeling_auto.MODEL_FOR_CAUSAL_LM_MAPPING_NAMES:
385
        return CausalLM(
386
387
388
389
390
            model_id,
            revision,
            quantize=quantize,
            dtype=dtype,
            trust_remote_code=trust_remote_code,
391
        )
392
    if model_type in modeling_auto.MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING_NAMES:
393
        return Seq2SeqLM(
394
395
396
397
398
            model_id,
            revision,
            quantize=quantize,
            dtype=dtype,
            trust_remote_code=trust_remote_code,
399
400
        )

401
    auto_map = config_dict.get("auto_map", None)
402
403
404
405
406
407
    if trust_remote_code and auto_map is not None:
        if "AutoModelForCausalLM" in auto_map.keys():
            return CausalLM(
                model_id,
                revision,
                quantize=quantize,
408
                dtype=dtype,
409
410
                trust_remote_code=trust_remote_code,
            )
411
        if "AutoModelForSeq2SeqLM" in auto_map.keys():
412
413
414
415
            return Seq2SeqLM(
                model_id,
                revision,
                quantize=quantize,
416
                dtype=dtype,
417
418
                trust_remote_code=trust_remote_code,
            )
419
420

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