__init__.py 29.1 KB
Newer Older
1
import torch
2
import enum
Nicolas Patry's avatar
Nicolas Patry committed
3
import os
4

5
from loguru import logger
6
from transformers.configuration_utils import PretrainedConfig
7
from transformers.models.auto import modeling_auto
Nicolas Patry's avatar
Nicolas Patry committed
8
from huggingface_hub import hf_hub_download, HfApi
9
from typing import Optional
10
from pathlib import Path
11

Nicolas Patry's avatar
Nicolas Patry committed
12
from text_generation_server.utils.speculate import get_speculate, set_speculate
13
14
from text_generation_server.models.model import Model
from text_generation_server.models.causal_lm import CausalLM
15
from text_generation_server.models.flash_causal_lm import FlashCausalLM
16
from text_generation_server.models.bloom import BLOOMSharded
17
from text_generation_server.models.mpt import MPTSharded
18
from text_generation_server.models.seq2seq_lm import Seq2SeqLM
19
from text_generation_server.models.rw import RW
20
21
from text_generation_server.models.opt import OPTSharded
from text_generation_server.models.galactica import GalacticaSharded
22
23
from text_generation_server.models.santacoder import SantaCoder
from text_generation_server.models.t5 import T5Sharded
24
from text_generation_server.models.gpt_neox import GPTNeoxSharded
drbh's avatar
drbh committed
25
from text_generation_server.models.phi import Phi
26

27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
# 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",
    "GalacticaSharded",
    "Seq2SeqLM",
    "SantaCoder",
    "OPTSharded",
    "T5Sharded",
    "get_model",
]

49
FLASH_ATT_ERROR_MESSAGE = "{} requires Flash Attention enabled models."
50

51
FLASH_ATTENTION = True
52

53
try:
54
    from text_generation_server.models.flash_rw import FlashRWSharded
55
    from text_generation_server.models.flash_gpt2 import FlashGPT2
56
57
58
59
    from text_generation_server.models.flash_neox import FlashNeoXSharded
    from text_generation_server.models.flash_llama import (
        FlashLlama,
    )
OlivierDehaene's avatar
OlivierDehaene committed
60
61
62
    from text_generation_server.models.flash_qwen2 import (
        FlashQwen2,
    )
OlivierDehaene's avatar
OlivierDehaene committed
63
64
65
    from text_generation_server.models.flash_cohere import (
        FlashCohere,
    )
66
67
68
    from text_generation_server.models.flash_gemma import (
        FlashGemma,
    )
drbh's avatar
drbh committed
69
70
71
    from text_generation_server.models.pali_gemma import (
        PaliGemma,
    )
72
73
    from text_generation_server.models.flash_santacoder import (
        FlashSantacoderSharded,
74
    )
75
    from text_generation_server.models.idefics import IDEFICSSharded
76
    from text_generation_server.models.llava_next import LlavaNext
Nicolas Patry's avatar
Nicolas Patry committed
77
    from text_generation_server.models.idefics2 import Idefics2
78
79
    from text_generation_server.models.flash_mistral import FlashMistral
    from text_generation_server.models.flash_mixtral import FlashMixtral
drbh's avatar
drbh committed
80
    from text_generation_server.models.flash_phi import FlashPhi
OlivierDehaene's avatar
OlivierDehaene committed
81
    from text_generation_server.models.flash_starcoder2 import FlashStarcoder2
82
    from text_generation_server.models.flash_dbrx import FlashDbrx
fxmarty's avatar
fxmarty committed
83
84
85
86
    from text_generation_server.utils.flash_attn import (
        HAS_FLASH_ATTN_V2_CUDA,
        HAS_FLASH_ATTN_V2_ROCM,
    )
87
88
except ImportError as e:
    logger.warning(f"Could not import Flash Attention enabled models: {e}")
89
    FLASH_ATTENTION = False
90
    HAS_FLASH_ATTN_V2_CUDA = False
fxmarty's avatar
fxmarty committed
91
    HAS_FLASH_ATTN_V2_ROCM = False
92

93
if FLASH_ATTENTION:
94
    __all__.append(FlashGPT2)
95
    __all__.append(FlashNeoXSharded)
96
    __all__.append(FlashRWSharded)
97
    __all__.append(FlashSantacoderSharded)
98
    __all__.append(FlashLlama)
99
    __all__.append(IDEFICSSharded)
100
    __all__.append(FlashMistral)
OlivierDehaene's avatar
OlivierDehaene committed
101
    __all__.append(FlashMixtral)
102
    __all__.append(FlashDbrx)
drbh's avatar
drbh committed
103
    __all__.append(FlashPhi)
OlivierDehaene's avatar
OlivierDehaene committed
104
    __all__.append(FlashQwen2)
OlivierDehaene's avatar
OlivierDehaene committed
105
    __all__.append(FlashStarcoder2)
OlivierDehaene's avatar
OlivierDehaene committed
106
107
    __all__.append(FlashGemma)
    __all__.append(FlashCohere)
OlivierDehaene's avatar
OlivierDehaene committed
108

drbh's avatar
drbh committed
109
110
111
112
113
114
115
116
117
MAMBA_AVAILABLE = True
try:
    from text_generation_server.models.mamba import Mamba
except ImportError as e:
    logger.warning(f"Could not import Mamba: {e}")
    MAMBA_AVAILABLE = False

if MAMBA_AVAILABLE:
    __all__.append(Mamba)
OlivierDehaene's avatar
OlivierDehaene committed
118

119

120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
class ModelType(enum.Enum):
    IDEFICS2 = {
        "type": "idefics2",
        "name": "Idefics 2",
        "url": "https://huggingface.co/HuggingFaceM4/idefics2-8b",
        "multimodal": True,
    }
    LLAVA_NEXT = {
        "type": "llava_next",
        "name": "Llava Next (1.6)",
        "url": "https://huggingface.co/llava-hf/llava-v1.6-vicuna-13b-hf",
        "multimodal": True,
    }
    LLAMA = {
        "type": "llama",
        "name": "Llama",
        "url": "https://huggingface.co/meta-llama/Meta-Llama-3-8B-Instruct",
    }
    PHI3 = {
        "type": "phi3",
        "name": "Phi 3",
        "url": "https://huggingface.co/microsoft/Phi-3-mini-4k-instruct",
    }
    GEMMA = {
        "type": "gemma",
        "name": "Gemma",
        "url": "https://huggingface.co/google/gemma-7b",
    }
    COHERE = {
        "type": "cohere",
        "name": "Cohere",
        "url": "https://huggingface.co/CohereForAI/c4ai-command-r-plus",
    }
    DBRX = {
        "type": "dbrx",
        "name": "Dbrx",
        "url": "https://huggingface.co/databricks/dbrx-instruct",
    }
    MAMBA = {
        "type": "ssm",
        "name": "Mamba",
        "url": "https://huggingface.co/state-spaces/mamba-2.8b-slimpj",
    }
    MISTRAL = {
        "type": "mistral",
        "name": "Mistral",
        "url": "https://huggingface.co/mistralai/Mistral-7B-Instruct-v0.2",
    }
    MIXTRAL = {
        "type": "mixtral",
        "name": "Mixtral",
        "url": "https://huggingface.co/mistralai/Mixtral-8x22B-Instruct-v0.1",
    }
    GPT_BIGCODE = {
        "type": "gpt_bigcode",
        "name": "Gpt Bigcode",
        "url": "https://huggingface.co/bigcode/gpt_bigcode-santacoder",
    }
    PHI = {
        "type": "phi",
        "name": "Phi",
        "url": "https://huggingface.co/microsoft/phi-1_5",
    }
    BAICHUAN = {
        "type": "baichuan",
        "name": "Baichuan",
        "url": "https://huggingface.co/baichuan-inc/Baichuan2-7B-Chat",
    }
    FALCON = {
        "type": "falcon",
        "name": "Falcon",
        "url": "https://huggingface.co/tiiuae/falcon-7b-instruct",
    }
    STARCODER2 = {
        "type": "starcoder2",
        "name": "StarCoder 2",
        "url": "https://huggingface.co/bigcode/starcoder2-15b-instruct-v0.1",
    }
    QWEN2 = {
        "type": "qwen2",
        "name": "Qwen 2",
        "url": "https://huggingface.co/bigcode/starcoder2-15b-instruct-v0.1",
    }
    OPT = {
        "type": "opt",
        "name": "Opt",
        "url": "https://huggingface.co/facebook/opt-6.7b",
    }
    T5 = {
        "type": "t5",
        "name": "T5",
        "url": "https://huggingface.co/google/flan-t5-xxl",
    }
    GALACTICA = {
        "type": "galactica",
        "name": "Galactica",
        "url": "https://huggingface.co/facebook/galactica-120b",
    }
    SANTACODER = {
        "type": "santacoder",
        "name": "SantaCoder",
        "url": "https://huggingface.co/bigcode/santacoder",
    }
    BLOOM = {
        "type": "bloom",
        "name": "Bloom",
        "url": "https://huggingface.co/bigscience/bloom-560m",
    }
    MPT = {
        "type": "mpt",
        "name": "Mpt",
        "url": "https://huggingface.co/mosaicml/mpt-7b-instruct",
    }
    GPT2 = {
        "type": "gpt2",
        "name": "Gpt2",
        "url": "https://huggingface.co/openai-community/gpt2",
    }
    GPT_NEOX = {
        "type": "gpt_neox",
        "name": "Gpt Neox",
        "url": "https://huggingface.co/EleutherAI/gpt-neox-20b",
    }
    IDEFICS = {
        "type": "idefics",
        "name": "Idefics",
        "url": "https://huggingface.co/HuggingFaceM4/idefics-9b",
        "multimodal": True,
    }


__GLOBALS = locals()
for data in ModelType:
    __GLOBALS[data.name] = data.value["type"]


256
def get_model(
257
258
259
260
    model_id: str,
    revision: Optional[str],
    sharded: bool,
    quantize: Optional[str],
Nicolas Patry's avatar
Nicolas Patry committed
261
    speculate: Optional[int],
262
    dtype: Optional[str],
263
    trust_remote_code: bool,
264
) -> Model:
265
    if dtype is None:
266
267
268
        # Keep it as default for now and let
        # every model resolve their own default dtype.
        dtype = None
269
270
271
272
273
274
275
    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
276
277
278
279
280
    if speculate is not None:
        set_speculate(speculate)
    else:
        set_speculate(0)

OlivierDehaene's avatar
v0.8.2  
OlivierDehaene committed
281
282
283
    config_dict, _ = PretrainedConfig.get_config_dict(
        model_id, revision=revision, trust_remote_code=trust_remote_code
    )
Nicolas Patry's avatar
Nicolas Patry committed
284
    model_type = config_dict.get("model_type", None)
Nicolas Patry's avatar
Nicolas Patry committed
285

Nicolas Patry's avatar
Nicolas Patry committed
286
    speculator = None
Nicolas Patry's avatar
Nicolas Patry committed
287
    if "medusa_num_heads" in config_dict:
288
289
        medusa_model_id = model_id
        medusa_revision = revision
Nicolas Patry's avatar
Nicolas Patry committed
290
291
292
293
294
        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
295
                raise RuntimeError(
OlivierDehaene's avatar
OlivierDehaene committed
296
                    f"Speculate is set to `{speculate}` but this medusa models only has `{speculate_medusa}` heads, please make them match"
OlivierDehaene's avatar
OlivierDehaene committed
297
                )
Nicolas Patry's avatar
Nicolas Patry committed
298
299
300
301
302
303
304
305
            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
        )
Nicolas Patry's avatar
Nicolas Patry committed
306
307
        # Reload model type from parent.
        model_type = config_dict.get("model_type", None)
308
309
310
311
312
313
314
315
316
317
        is_local = Path(medusa_model_id).exists()
        if not is_local:
            medusa_config = hf_hub_download(
                medusa_model_id, revision=medusa_revision, filename="config.json"
            )
            hf_hub_download(
                medusa_model_id,
                revision=medusa_revision,
                filename="medusa_lm_head.safetensors",
            )
Nicolas Patry's avatar
Nicolas Patry committed
318
319
320
321
            speculator = {
                "path": Path(medusa_config).parent,
                "model_paths": ["medusa_lm_head.safetensors"],
            }
322
        else:
Nicolas Patry's avatar
Nicolas Patry committed
323
324
325
326
            speculator = {
                "path": Path(medusa_model_id),
                "model_paths": ["medusa_lm_head.safetensors"],
            }
327

Nicolas Patry's avatar
Nicolas Patry committed
328
        method = "medusa"
Nicolas Patry's avatar
Nicolas Patry committed
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
    elif model_type == "mlp_speculator":
        mlp_model_id = model_id
        mlp_revision = revision
        model_id = config_dict["base_model_name_or_path"]
        revision = "main"
        speculate_mlp = config_dict["n_predict"]
        if speculate is not None:
            if speculate > speculate_mlp:
                raise RuntimeError(
                    f"Speculate is set to `{speculate}` but this mlp_speculator models only has `{speculate_mlp}` heads, please make them match"
                )
            else:
                set_speculate(speculate)
        else:
            set_speculate(speculate_mlp)

        config_dict, _ = PretrainedConfig.get_config_dict(
            model_id, revision=revision, trust_remote_code=trust_remote_code
        )
        # Reload model type from parent.
        model_type = config_dict.get("model_type", None)
        is_local = Path(mlp_model_id).exists()
        extension = ".safetensors"
        if not is_local:
            mlp_speculator_config = hf_hub_download(
                mlp_model_id, revision=mlp_revision, filename="config.json"
            )
            api = HfApi()
            info = api.model_info(mlp_model_id, revision=mlp_revision)
            filenames = [
                s.rfilename
                for s in info.siblings
                if s.rfilename.endswith(extension)
                and len(s.rfilename.split("/")) == 1
                and "arguments" not in s.rfilename
                and "args" not in s.rfilename
                and "training" not in s.rfilename
            ]
            for filename in filenames:
                hf_hub_download(
                    mlp_model_id,
                    revision=mlp_revision,
                    filename=filename,
                )
            speculator = {
                "path": Path(mlp_speculator_config).parent,
                "model_paths": filenames,
            }
        else:
            speculator = Path(mlp_model_id)
            filenames = [p for p in os.listdir(speculator) if p.endswith(extension)]
            speculator = {"path": speculator, "model_paths": filenames}
        method = "mlp_speculator"
Nicolas Patry's avatar
Nicolas Patry committed
382
383
384
385
386
387
388
    else:
        method = "n-gram"

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

drbh's avatar
drbh committed
389
390
391
392
393
394
395
396
397
    if model_type is None:
        # TODO: fix how we determine model type for Mamba
        if "ssm_cfg" in config_dict:
            # *only happens in Mamba case
            model_type = "ssm"
        else:
            raise RuntimeError(
                f"Could not determine model type for {model_id} revision {revision}"
            )
398
399
400
401
402
403
404
405
    quantization_config = config_dict.get("quantization_config", None)
    if quantization_config is not None and quantize is None:
        method = quantization_config.get("quant_method", None)
        if method in {"gptq", "awq"}:
            logger.info(f"Auto selecting quantization method {method}")
            quantize = method
        else:
            logger.info(f"Unknown quantization method {method}")
drbh's avatar
drbh committed
406

407
    if model_type == MAMBA:
drbh's avatar
drbh committed
408
409
410
411
        return Mamba(
            model_id,
            revision,
            quantize=quantize,
Nicolas Patry's avatar
Nicolas Patry committed
412
            speculator=speculator,
drbh's avatar
drbh committed
413
414
415
            dtype=dtype,
            trust_remote_code=trust_remote_code,
        )
416

OlivierDehaene's avatar
OlivierDehaene committed
417
418
419
420
421
    if model_id.startswith("facebook/galactica"):
        return GalacticaSharded(
            model_id,
            revision,
            quantize=quantize,
Nicolas Patry's avatar
Nicolas Patry committed
422
            speculator=speculator,
OlivierDehaene's avatar
OlivierDehaene committed
423
424
425
426
            dtype=dtype,
            trust_remote_code=trust_remote_code,
        )

427
    if (
428
429
        model_type == GPT_BIGCODE
        or model_type == GPT2
430
431
        and model_id.startswith("bigcode/")
    ):
432
        if FLASH_ATTENTION:
433
434
435
436
            return FlashSantacoderSharded(
                model_id,
                revision,
                quantize=quantize,
Nicolas Patry's avatar
Nicolas Patry committed
437
                speculator=speculator,
438
                dtype=dtype,
439
440
                trust_remote_code=trust_remote_code,
            )
441
442
443
444
        elif sharded:
            raise NotImplementedError(
                FLASH_ATT_ERROR_MESSAGE.format("Sharded Santacoder")
            )
445
        else:
446
            return SantaCoder(
447
448
449
                model_id,
                revision,
                quantize=quantize,
Nicolas Patry's avatar
Nicolas Patry committed
450
                speculator=speculator,
451
                dtype=dtype,
452
453
                trust_remote_code=trust_remote_code,
            )
454

455
    if model_type == BLOOM:
456
        return BLOOMSharded(
457
458
459
            model_id,
            revision,
            quantize=quantize,
Nicolas Patry's avatar
Nicolas Patry committed
460
            speculator=speculator,
461
462
            dtype=dtype,
            trust_remote_code=trust_remote_code,
463
        )
464
    elif model_type == MPT:
465
        return MPTSharded(
OlivierDehaene's avatar
OlivierDehaene committed
466
467
468
            model_id,
            revision,
            quantize=quantize,
Nicolas Patry's avatar
Nicolas Patry committed
469
            speculator=speculator,
OlivierDehaene's avatar
OlivierDehaene committed
470
471
            dtype=dtype,
            trust_remote_code=trust_remote_code,
472
        )
473
    elif model_type == GPT2:
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
        if FLASH_ATTENTION:
            return FlashGPT2(
                model_id,
                revision,
                quantize=quantize,
                speculator=speculator,
                dtype=dtype,
                trust_remote_code=trust_remote_code,
            )
        elif sharded:
            raise NotImplementedError(FLASH_ATT_ERROR_MESSAGE.format("Sharded GPT-2"))
        else:
            return CausalLM(
                model_id,
                revision,
                quantize=quantize,
                speculator=speculator,
                dtype=dtype,
                trust_remote_code=trust_remote_code,
            )
494
    elif model_type == GPT_NEOX:
495
496
497
498
499
        if FLASH_ATTENTION:
            return FlashNeoXSharded(
                model_id,
                revision,
                quantize=quantize,
Nicolas Patry's avatar
Nicolas Patry committed
500
                speculator=speculator,
501
                dtype=dtype,
502
503
504
505
                trust_remote_code=trust_remote_code,
            )
        elif sharded:
            return GPTNeoxSharded(
506
507
508
                model_id,
                revision,
                quantize=quantize,
Nicolas Patry's avatar
Nicolas Patry committed
509
                speculator=speculator,
510
                dtype=dtype,
511
512
                trust_remote_code=trust_remote_code,
            )
513
        else:
514
            return CausalLM(
515
516
517
                model_id,
                revision,
                quantize=quantize,
Nicolas Patry's avatar
Nicolas Patry committed
518
                speculator=speculator,
519
                dtype=dtype,
520
521
                trust_remote_code=trust_remote_code,
            )
OlivierDehaene's avatar
OlivierDehaene committed
522

523
    elif model_type == PHI:
drbh's avatar
drbh committed
524
525
526
527
528
        if FLASH_ATTENTION:
            return FlashPhi(
                model_id,
                revision,
                quantize=quantize,
Nicolas Patry's avatar
Nicolas Patry committed
529
                speculator=speculator,
drbh's avatar
drbh committed
530
531
532
533
534
535
536
537
                dtype=dtype,
                trust_remote_code=trust_remote_code,
            )
        else:
            return CausalLM(
                model_id,
                revision,
                quantize=quantize,
Nicolas Patry's avatar
Nicolas Patry committed
538
                speculator=speculator,
drbh's avatar
drbh committed
539
540
541
542
543
544
                dtype=dtype,
                trust_remote_code=trust_remote_code,
            )

    elif model_type == "phi-msft":
        if FLASH_ATTENTION:
OlivierDehaene's avatar
OlivierDehaene committed
545
546
547
            raise NotImplementedError(
                "Legacy phi-msft is not supported with Flash Attention"
            )
drbh's avatar
drbh committed
548
549
550
551
552
        else:
            return Phi(
                model_id,
                revision,
                quantize=quantize,
Nicolas Patry's avatar
Nicolas Patry committed
553
                speculator=speculator,
drbh's avatar
drbh committed
554
555
556
                dtype=dtype,
                trust_remote_code=trust_remote_code,
            )
557

558
    elif model_type == LLAMA or model_type == BAICHUAN or model_type == PHI3:
559
560
        if FLASH_ATTENTION:
            return FlashLlama(
561
562
563
                model_id,
                revision,
                quantize=quantize,
Nicolas Patry's avatar
Nicolas Patry committed
564
                speculator=speculator,
565
                dtype=dtype,
566
567
                trust_remote_code=trust_remote_code,
            )
568
569
        elif sharded:
            raise NotImplementedError(FLASH_ATT_ERROR_MESSAGE.format("Sharded Llama"))
570
        else:
571
            return CausalLM(
572
573
574
                model_id,
                revision,
                quantize=quantize,
Nicolas Patry's avatar
Nicolas Patry committed
575
                speculator=speculator,
576
                dtype=dtype,
577
578
                trust_remote_code=trust_remote_code,
            )
579
    if model_type == GEMMA:
580
581
582
583
584
        if FLASH_ATTENTION:
            return FlashGemma(
                model_id,
                revision,
                quantize=quantize,
Nicolas Patry's avatar
Nicolas Patry committed
585
                speculator=speculator,
586
587
588
589
                dtype=dtype,
                trust_remote_code=trust_remote_code,
            )
        elif sharded:
OlivierDehaene's avatar
OlivierDehaene committed
590
            raise NotImplementedError(FLASH_ATT_ERROR_MESSAGE.format("Sharded Gemma"))
591
592
593
594
595
        else:
            return CausalLM(
                model_id,
                revision,
                quantize=quantize,
Nicolas Patry's avatar
Nicolas Patry committed
596
                speculator=speculator,
597
598
599
                dtype=dtype,
                trust_remote_code=trust_remote_code,
            )
600

601
    if model_type == COHERE:
OlivierDehaene's avatar
OlivierDehaene committed
602
603
604
605
606
        if FLASH_ATTENTION:
            return FlashCohere(
                model_id,
                revision,
                quantize=quantize,
Nicolas Patry's avatar
Nicolas Patry committed
607
                speculator=speculator,
OlivierDehaene's avatar
OlivierDehaene committed
608
609
610
611
612
613
614
615
616
617
                dtype=dtype,
                trust_remote_code=trust_remote_code,
            )
        elif sharded:
            raise NotImplementedError(FLASH_ATT_ERROR_MESSAGE.format("Sharded Cohere"))
        else:
            return CausalLM(
                model_id,
                revision,
                quantize=quantize,
Nicolas Patry's avatar
Nicolas Patry committed
618
                speculator=speculator,
OlivierDehaene's avatar
OlivierDehaene committed
619
620
621
622
                dtype=dtype,
                trust_remote_code=trust_remote_code,
            )

623
    if model_type == DBRX:
624
625
626
627
628
        if FLASH_ATTENTION:
            return FlashDbrx(
                model_id,
                revision,
                quantize=quantize,
Nicolas Patry's avatar
Nicolas Patry committed
629
                speculator=speculator,
630
631
632
633
634
635
636
637
638
639
                dtype=dtype,
                trust_remote_code=trust_remote_code,
            )
        elif sharded:
            raise NotImplementedError(FLASH_ATT_ERROR_MESSAGE.format("Sharded DBRX"))
        else:
            return CausalLM(
                model_id,
                revision,
                quantize=quantize,
Nicolas Patry's avatar
Nicolas Patry committed
640
                speculator=speculator,
641
642
643
644
                dtype=dtype,
                trust_remote_code=trust_remote_code,
            )

645
    if model_type in ["RefinedWeb", "RefinedWebModel", FALCON]:
646
647
        if sharded:
            if FLASH_ATTENTION:
648
                if config_dict.get("alibi", False):
649
650
651
652
653
                    raise NotImplementedError("sharded is not supported for this model")
                return FlashRWSharded(
                    model_id,
                    revision,
                    quantize=quantize,
Nicolas Patry's avatar
Nicolas Patry committed
654
                    speculator=speculator,
655
                    dtype=dtype,
656
657
                    trust_remote_code=trust_remote_code,
                )
658
            raise NotImplementedError(FLASH_ATT_ERROR_MESSAGE.format(f"Sharded Falcon"))
659
        else:
660
            if FLASH_ATTENTION and not config_dict.get("alibi", False):
661
                return FlashRWSharded(
662
663
664
                    model_id,
                    revision,
                    quantize=quantize,
Nicolas Patry's avatar
Nicolas Patry committed
665
                    speculator=speculator,
666
                    dtype=dtype,
667
668
669
670
671
672
673
                    trust_remote_code=trust_remote_code,
                )
            else:
                return RW(
                    model_id,
                    revision,
                    quantize=quantize,
Nicolas Patry's avatar
Nicolas Patry committed
674
                    speculator=speculator,
675
                    dtype=dtype,
676
677
678
                    trust_remote_code=trust_remote_code,
                )

679
    if model_type == MISTRAL:
680
681
        sliding_window = config_dict.get("sliding_window", -1)
        if (
fxmarty's avatar
fxmarty committed
682
683
684
685
            ((sliding_window is None or sliding_window == -1) and FLASH_ATTENTION)
            or HAS_FLASH_ATTN_V2_CUDA
            or HAS_FLASH_ATTN_V2_ROCM
        ):
686
687
688
689
            return FlashMistral(
                model_id,
                revision,
                quantize=quantize,
Nicolas Patry's avatar
Nicolas Patry committed
690
                speculator=speculator,
691
692
693
                dtype=dtype,
                trust_remote_code=trust_remote_code,
            )
OlivierDehaene's avatar
OlivierDehaene committed
694
695
696
697
698
699
700
        elif sharded:
            raise NotImplementedError(FLASH_ATT_ERROR_MESSAGE.format("Sharded Mistral"))
        else:
            return CausalLM(
                model_id,
                revision,
                quantize=quantize,
Nicolas Patry's avatar
Nicolas Patry committed
701
                speculator=speculator,
OlivierDehaene's avatar
OlivierDehaene committed
702
703
704
                dtype=dtype,
                trust_remote_code=trust_remote_code,
            )
OlivierDehaene's avatar
OlivierDehaene committed
705

706
    if model_type == MIXTRAL:
707
708
        sliding_window = config_dict.get("sliding_window", -1)
        if (
fxmarty's avatar
fxmarty committed
709
710
711
712
            ((sliding_window is None or sliding_window == -1) and FLASH_ATTENTION)
            or HAS_FLASH_ATTN_V2_CUDA
            or HAS_FLASH_ATTN_V2_ROCM
        ):
OlivierDehaene's avatar
OlivierDehaene committed
713
714
715
716
            return FlashMixtral(
                model_id,
                revision,
                quantize=quantize,
Nicolas Patry's avatar
Nicolas Patry committed
717
                speculator=speculator,
OlivierDehaene's avatar
OlivierDehaene committed
718
                dtype=dtype,
OlivierDehaene's avatar
OlivierDehaene committed
719
720
                trust_remote_code=trust_remote_code,
            )
OlivierDehaene's avatar
OlivierDehaene committed
721
722
723
724
725
726
727
        elif sharded:
            raise NotImplementedError(FLASH_ATT_ERROR_MESSAGE.format("Sharded Mixtral"))
        else:
            return CausalLM(
                model_id,
                revision,
                quantize=quantize,
Nicolas Patry's avatar
Nicolas Patry committed
728
                speculator=speculator,
OlivierDehaene's avatar
OlivierDehaene committed
729
730
731
732
                dtype=dtype,
                trust_remote_code=trust_remote_code,
            )

733
    if model_type == STARCODER2:
OlivierDehaene's avatar
OlivierDehaene committed
734
735
        sliding_window = config_dict.get("sliding_window", -1)
        if (
fxmarty's avatar
fxmarty committed
736
737
738
739
            ((sliding_window is None or sliding_window == -1) and FLASH_ATTENTION)
            or HAS_FLASH_ATTN_V2_CUDA
            or HAS_FLASH_ATTN_V2_ROCM
        ):
OlivierDehaene's avatar
OlivierDehaene committed
740
741
742
743
744
            return FlashStarcoder2(
                model_id,
                revision,
                quantize=quantize,
                dtype=dtype,
OlivierDehaene's avatar
OlivierDehaene committed
745
746
747
748
749
750
751
752
753
754
755
                trust_remote_code=trust_remote_code,
            )
        elif sharded:
            raise NotImplementedError(
                FLASH_ATT_ERROR_MESSAGE.format("Sharded Starcoder2")
            )
        else:
            return CausalLM(
                model_id,
                revision,
                quantize=quantize,
Nicolas Patry's avatar
Nicolas Patry committed
756
                speculator=speculator,
OlivierDehaene's avatar
OlivierDehaene committed
757
758
759
760
                dtype=dtype,
                trust_remote_code=trust_remote_code,
            )

761
    if model_type == QWEN2:
OlivierDehaene's avatar
OlivierDehaene committed
762
763
        sliding_window = config_dict.get("sliding_window", -1)
        if (
fxmarty's avatar
fxmarty committed
764
765
766
767
            ((sliding_window is None or sliding_window == -1) and FLASH_ATTENTION)
            or HAS_FLASH_ATTN_V2_CUDA
            or HAS_FLASH_ATTN_V2_ROCM
        ):
OlivierDehaene's avatar
OlivierDehaene committed
768
769
770
771
772
773
774
775
776
777
778
779
780
781
            return FlashQwen2(
                model_id,
                revision,
                quantize=quantize,
                dtype=dtype,
                trust_remote_code=trust_remote_code,
            )
        elif sharded:
            raise NotImplementedError(FLASH_ATT_ERROR_MESSAGE.format("Sharded Qwen2"))
        else:
            return CausalLM(
                model_id,
                revision,
                quantize=quantize,
Nicolas Patry's avatar
Nicolas Patry committed
782
                speculator=speculator,
OlivierDehaene's avatar
OlivierDehaene committed
783
                dtype=dtype,
OlivierDehaene's avatar
OlivierDehaene committed
784
785
                trust_remote_code=trust_remote_code,
            )
786

787
    if model_type == OPT:
788
        return OPTSharded(
789
790
791
            model_id,
            revision,
            quantize=quantize,
Nicolas Patry's avatar
Nicolas Patry committed
792
            speculator=speculator,
793
794
            dtype=dtype,
            trust_remote_code=trust_remote_code,
795
        )
796

797
    if model_type == T5:
798
799
800
801
        return T5Sharded(
            model_id,
            revision,
            quantize=quantize,
Nicolas Patry's avatar
Nicolas Patry committed
802
            speculator=speculator,
803
            dtype=dtype,
804
805
            trust_remote_code=trust_remote_code,
        )
806
    if model_type == IDEFICS:
807
        if FLASH_ATTENTION:
OlivierDehaene's avatar
OlivierDehaene committed
808
809
810
811
            return IDEFICSSharded(
                model_id,
                revision,
                quantize=quantize,
Nicolas Patry's avatar
Nicolas Patry committed
812
                speculator=speculator,
OlivierDehaene's avatar
OlivierDehaene committed
813
814
815
                dtype=dtype,
                trust_remote_code=trust_remote_code,
            )
816
817
        else:
            raise NotImplementedError(FLASH_ATT_ERROR_MESSAGE.format("Idefics"))
818
    if model_type == IDEFICS2:
Nicolas Patry's avatar
Nicolas Patry committed
819
820
821
822
823
        if FLASH_ATTENTION:
            return Idefics2(
                model_id,
                revision,
                quantize=quantize,
Nicolas Patry's avatar
Nicolas Patry committed
824
                speculator=speculator,
Nicolas Patry's avatar
Nicolas Patry committed
825
826
827
828
829
                dtype=dtype,
                trust_remote_code=trust_remote_code,
            )
        else:
            raise NotImplementedError(FLASH_ATT_ERROR_MESSAGE.format("Idefics"))
drbh's avatar
drbh committed
830
831
832
833
834
835
836
837
838
839
840
841
    if model_type == "paligemma":
        if FLASH_ATTENTION:
            return PaliGemma(
                model_id,
                revision,
                quantize=quantize,
                speculator=speculator,
                dtype=dtype,
                trust_remote_code=trust_remote_code,
            )
        else:
            raise NotImplementedError(FLASH_ATT_ERROR_MESSAGE.format("Idefics"))
842

843
    if model_type == LLAVA_NEXT:
844
845
846
847
848
        if FLASH_ATTENTION:
            return LlavaNext(
                model_id,
                revision,
                quantize=quantize,
Nicolas Patry's avatar
Nicolas Patry committed
849
                speculator=speculator,
850
851
852
853
854
855
                dtype=dtype,
                trust_remote_code=trust_remote_code,
            )
        else:
            raise NotImplementedError(FLASH_ATT_ERROR_MESSAGE.format("LlavaNext"))

856
    if sharded:
857
        raise NotImplementedError("sharded is not supported for AutoModel")
858
    if quantize == "gptq":
859
        raise NotImplementedError(
860
861
            "gptq quantization is not supported for AutoModel, you can try to quantize it with `text-generation-server quantize ORIGINAL_MODEL_ID NEW_MODEL_ID`"
        )
862
    if quantize == "awq":
863
        raise NotImplementedError("awq quantization is not supported for AutoModel")
Nicolas Patry's avatar
Nicolas Patry committed
864
    elif (quantize == "bitsandbytes-fp4") or (quantize == "bitsandbytes-nf4"):
865
        raise NotImplementedError("4bit quantization is not supported for AutoModel")
OlivierDehaene's avatar
OlivierDehaene committed
866
    elif quantize == "eetq":
867
        raise NotImplementedError("Eetq quantization is not supported for AutoModel")
868
    if model_type in modeling_auto.MODEL_FOR_CAUSAL_LM_MAPPING_NAMES:
869
        return CausalLM(
870
871
872
            model_id,
            revision,
            quantize=quantize,
Nicolas Patry's avatar
Nicolas Patry committed
873
            speculator=speculator,
874
875
            dtype=dtype,
            trust_remote_code=trust_remote_code,
876
        )
877
    if model_type in modeling_auto.MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING_NAMES:
878
        return Seq2SeqLM(
879
880
881
            model_id,
            revision,
            quantize=quantize,
Nicolas Patry's avatar
Nicolas Patry committed
882
            speculator=speculator,
883
884
            dtype=dtype,
            trust_remote_code=trust_remote_code,
885
886
        )

887
    auto_map = config_dict.get("auto_map", None)
888
889
890
891
892
893
    if trust_remote_code and auto_map is not None:
        if "AutoModelForCausalLM" in auto_map.keys():
            return CausalLM(
                model_id,
                revision,
                quantize=quantize,
Nicolas Patry's avatar
Nicolas Patry committed
894
                speculator=speculator,
895
                dtype=dtype,
896
897
                trust_remote_code=trust_remote_code,
            )
898
        if "AutoModelForSeq2SeqLM" in auto_map.keys():
899
900
901
902
            return Seq2SeqLM(
                model_id,
                revision,
                quantize=quantize,
Nicolas Patry's avatar
Nicolas Patry committed
903
                speculator=speculator,
904
                dtype=dtype,
905
906
                trust_remote_code=trust_remote_code,
            )
907
908

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