flash_santacoder.py 2.39 KB
Newer Older
1
2
3
4
import torch
import torch.distributed

from opentelemetry import trace
5
from transformers import AutoTokenizer, AutoConfig
6
7
8
9
from typing import Optional, List

from text_generation_server.models import FlashCausalLM
from text_generation_server.models.custom_modeling.flash_santacoder_modeling import (
10
    FlashSantacoderForCausalLM,
11
12
)
from text_generation_server.utils import (
13
    initialize_torch_distributed,
14
    weight_files,
15
    Weights,
16
17
18
19
20
)

tracer = trace.get_tracer(__name__)


21
class FlashSantacoderSharded(FlashCausalLM):
22
    def __init__(
23
24
25
26
        self,
        model_id: str,
        revision: Optional[str] = None,
        quantize: Optional[str] = None,
27
        trust_remote_code: bool = False,
28
    ):
29
        self.process_group, rank, world_size = initialize_torch_distributed()
30
        if torch.cuda.is_available():
31
            device = torch.device(f"cuda:{rank}")
32
            dtype = torch.float16
33
34
35
36
        else:
            raise NotImplementedError("FlashSantacoderSharded is only available on GPU")

        tokenizer = AutoTokenizer.from_pretrained(
37
38
39
40
41
            model_id,
            revision=revision,
            padding_side="left",
            truncation_side="left",
            trust_remote_code=trust_remote_code,
42
43
        )

44
        config = AutoConfig.from_pretrained(
45
46
            model_id,
            revision=revision,
47
            trust_remote_code=True,
48
        )
49
50
        config.quantize = quantize
        config.transpose = config.architectures[0].startswith("GPT2")
51
52
53

        torch.distributed.barrier(group=self.process_group)
        filenames = weight_files(model_id, revision=revision, extension=".safetensors")
54
55
56
        weights = Weights(
            filenames, device=device, dtype=dtype, process_group=self.process_group
        )
57

58
        model = FlashSantacoderForCausalLM(config, weights)
59
60
61

        torch.distributed.barrier(group=self.process_group)
        super(FlashCausalLM, self).__init__(
62
            model=model.to(device),
63
            tokenizer=tokenizer,
64
65
            requires_padding=False,
            dtype=dtype,
66
            device=device,
67
68
            rank=rank,
            world_size=world_size,
69
70
        )

71
72
73
74
75
    def decode(self, generated_ids: List[int]) -> str:
        # Do not skip special tokens as they are used for custom parsing rules of the generated text
        return self.tokenizer.decode(
            generated_ids, skip_special_tokens=False, clean_up_tokenization_spaces=False
        )