flash_rw.py 2.98 KB
Newer Older
1
2
3
4
import torch
import torch.distributed

from opentelemetry import trace
5
6
from transformers import AutoTokenizer
from typing import Optional
7
8
9
10
11
12
13
14
15

from text_generation_server.models import FlashCausalLM
from text_generation_server.models.custom_modeling.flash_rw_modeling import (
    RWConfig,
    FlashRWForCausalLM,
)
from text_generation_server.utils import (
    initialize_torch_distributed,
    weight_files,
16
    Weights,
17
)
xuxzh1's avatar
last  
xuxzh1 committed
18
from text_generation_server.utils.import_utils import SYSTEM
Nicolas Patry's avatar
Nicolas Patry committed
19

20
21
22
tracer = trace.get_tracer(__name__)


23
class FlashRWSharded(FlashCausalLM):
24
25
26
27
28
    def __init__(
        self,
        model_id: str,
        revision: Optional[str] = None,
        quantize: Optional[str] = None,
xuxzh1's avatar
last  
xuxzh1 committed
29
        speculator: Optional[str] = None,
30
        dtype: Optional[torch.dtype] = None,
31
32
33
34
35
        trust_remote_code: bool = False,
    ):
        self.process_group, rank, world_size = initialize_torch_distributed()
        if torch.cuda.is_available():
            device = torch.device(f"cuda:{rank}")
36
            dtype = torch.float16 if dtype is None else dtype
xuxzh1's avatar
last  
xuxzh1 committed
37
38
39
40
41
42
43
        elif SYSTEM == "ipex":
            if hasattr(torch, "xpu") and torch.xpu.is_available():
                device = torch.device(f"xpu:{rank}")
                dtype = torch.float16 if dtype is None else dtype
            else:
                device = torch.device("cpu")
                dtype = torch.bfloat16 if dtype is None else dtype
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
        else:
            raise NotImplementedError("FlashRW is only available on GPU")

        tokenizer = AutoTokenizer.from_pretrained(
            model_id,
            revision=revision,
            padding_side="left",
            truncation_side="left",
            trust_remote_code=trust_remote_code,
        )

        config = RWConfig.from_pretrained(
            model_id, revision=revision, trust_remote_code=trust_remote_code
        )

        torch.distributed.barrier(group=self.process_group)
        filenames = weight_files(model_id, revision=revision, extension=".safetensors")
61
62
63
64
65
        weights = Weights(
            filenames,
            device,
            dtype,
            process_group=self.process_group,
66
67
68
69
            aliases={
                "lm_head.weight": ["transformer.word_embeddings.weight"],
                "transformer.word_embeddings.weight": ["lm_head.weight"],
            },
70
        )
71

72
        config.quantize = quantize
xuxzh1's avatar
last  
xuxzh1 committed
73
74
        config.speculator = speculator
        if config.quantize in ["gptq", "marlin"]:
OlivierDehaene's avatar
OlivierDehaene committed
75
            weights._set_gptq_params(model_id, revision)
76

77
        model = FlashRWForCausalLM(config, weights)
78
79

        torch.distributed.barrier(group=self.process_group)
80
        super(FlashRWSharded, self).__init__(
xuxzh1's avatar
last  
xuxzh1 committed
81
            model_id=model_id,
82
83
            model=model.to(device),
            tokenizer=tokenizer,
84
85
86
            num_layers=len(model.transformer.h),
            num_kv_heads=model.transformer.cache_size,
            head_size=model.transformer.head_size,
87
88
89
90
91
            dtype=dtype,
            device=device,
            rank=rank,
            world_size=world_size,
        )