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

4
from loguru import logger
5
6
7
from transformers import AutoConfig
from typing import Optional

8
9
10
11
12
13
from text_generation_server.models.model import Model
from text_generation_server.models.causal_lm import CausalLM
from text_generation_server.models.bloom import BLOOM, BLOOMSharded
from text_generation_server.models.seq2seq_lm import Seq2SeqLM
from text_generation_server.models.galactica import Galactica, GalacticaSharded
from text_generation_server.models.santacoder import SantaCoder
14
from text_generation_server.models.gpt_neox import GPTNeoxSharded
15
from text_generation_server.models.t5 import T5Sharded
16

17
18
try:
    from text_generation_server.models.flash_neox import FlashNeoX, FlashNeoXSharded
19

20
21
22
23
24
25
    FLASH_NEOX = torch.cuda.is_available() and int(os.environ.get("FLASH_NEOX", 0)) == 1
except ImportError:
    if int(os.environ.get("FLASH_NEOX", 0)) == 1:
        logger.exception("Could not import FlashNeoX")
    FLASH_NEOX = False

26
27
28
29
30
__all__ = [
    "Model",
    "BLOOM",
    "BLOOMSharded",
    "CausalLM",
31
32
33
    "Galactica",
    "GalacticaSharded",
    "GPTNeoxSharded",
34
35
    "Seq2SeqLM",
    "SantaCoder",
36
    "T5Sharded",
37
38
39
    "get_model",
]

40
41
42
43
if FLASH_NEOX:
    __all__.append(FlashNeoX)
    __all__.append(FlashNeoXSharded)

44
45
46
# 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
47

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

51
52
53
# Disable gradients
torch.set_grad_enabled(False)

54

55
def get_model(
56
    model_id: str, revision: Optional[str], sharded: bool, quantize: bool
57
) -> Model:
58
    if "facebook/galactica" in model_id:
59
60
61
62
63
64
65
66
        if sharded:
            return GalacticaSharded(model_id, revision, quantize=quantize)
        else:
            return Galactica(model_id, revision, quantize=quantize)

    if "santacoder" in model_id:
        return SantaCoder(model_id, revision, quantize)

67
    config = AutoConfig.from_pretrained(model_id, revision=revision)
68
69
70

    if config.model_type == "bloom":
        if sharded:
71
            return BLOOMSharded(model_id, revision, quantize=quantize)
72
        else:
73
            return BLOOM(model_id, revision, quantize=quantize)
74
75

    if config.model_type == "gpt_neox":
76
        if sharded:
77
78
            neox_cls = FlashNeoXSharded if FLASH_NEOX else GPTNeoxSharded
            return neox_cls(model_id, revision, quantize=quantize)
79
        else:
80
81
            neox_cls = FlashNeoX if FLASH_NEOX else CausalLM
            return neox_cls(model_id, revision, quantize=quantize)
82
83

    if config.model_type == "t5":
84
85
86
87
        if sharded:
            return T5Sharded(model_id, revision, quantize=quantize)
        else:
            return Seq2SeqLM(model_id, revision, quantize=quantize)
88
89
90
91
92
93
94

    if sharded:
        raise ValueError("sharded is not supported for AutoModel")
    try:
        return CausalLM(model_id, revision, quantize=quantize)
    except Exception:
        return Seq2SeqLM(model_id, revision, quantize=quantize)