From daa1d81d5ec4ef9bc59a4d6e850687b788732c90 Mon Sep 17 00:00:00 2001 From: OlivierDehaene Date: Thu, 1 Dec 2022 19:31:54 +0100 Subject: [PATCH] feat(server): Support Galactica (#4) --- README.md | 1 + server/Makefile | 10 +- server/text_generation/models/__init__.py | 6 + server/text_generation/models/bloom.py | 2 + server/text_generation/models/causal_lm.py | 46 ++- server/text_generation/models/galactica.py | 346 +++++++++++++++++++++ server/text_generation/utils.py | 2 +- 7 files changed, 383 insertions(+), 30 deletions(-) create mode 100644 server/text_generation/models/galactica.py diff --git a/README.md b/README.md index d081ad7f..e76e38d0 100644 --- a/README.md +++ b/README.md @@ -21,6 +21,7 @@ to power Bloom, BloomZ and MT0-XXL api-inference widgets. - [BLOOM](https://huggingface.co/bigscience/bloom) - [BLOOMZ](https://huggingface.co/bigscience/bloomz) - [MT0-XXL](https://huggingface.co/bigscience/mt0-xxl) +- ~~[Galactica](https://huggingface.co/facebook/galactica-120b)~~ (deactivated) Other models are supported on a best effort basis using: diff --git a/server/Makefile b/server/Makefile index 57dea48f..10028dad 100644 --- a/server/Makefile +++ b/server/Makefile @@ -9,11 +9,11 @@ gen-server: install-transformers: # Install specific version of transformers with custom cuda kernels rm transformers || true - rm transformers-b55f16c5b71aeef47a66a4270e19c154f050a7a7 || true - curl -L -O https://github.com/OlivierDehaene/transformers/archive/b55f16c5b71aeef47a66a4270e19c154f050a7a7.zip - unzip b55f16c5b71aeef47a66a4270e19c154f050a7a7.zip - rm b55f16c5b71aeef47a66a4270e19c154f050a7a7.zip - mv transformers-b55f16c5b71aeef47a66a4270e19c154f050a7a7 transformers + rm transformers-text_generation_inference || true + curl -L -O https://github.com/OlivierDehaene/transformers/archive/refs/heads/text_generation_inference.zip + unzip text_generation_inference.zip + rm text_generation_inference.zip + mv transformers-text_generation_inference transformers cd transformers && python setup.py install install-torch: diff --git a/server/text_generation/models/__init__.py b/server/text_generation/models/__init__.py index bf44115c..b364309a 100644 --- a/server/text_generation/models/__init__.py +++ b/server/text_generation/models/__init__.py @@ -2,6 +2,7 @@ from text_generation.models.model import Model from text_generation.models.causal_lm import CausalLM from text_generation.models.bloom import BLOOMSharded from text_generation.models.seq2seq_lm import Seq2SeqLM +from text_generation.models.galactica import Galactica, GalacticaSharded __all__ = ["Model", "BLOOMSharded", "CausalLM", "Seq2SeqLM"] @@ -12,6 +13,11 @@ def get_model(model_name: str, sharded: bool, quantize: bool) -> Model: return BLOOMSharded(model_name, quantize=quantize) else: return CausalLM(model_name, quantize=quantize) + elif model_name.startswith("facebook/galactica"): + if sharded: + return GalacticaSharded(model_name, quantize=quantize) + else: + return Galactica(model_name, quantize=quantize) else: if sharded: raise ValueError("sharded is not supported for AutoModel") diff --git a/server/text_generation/models/bloom.py b/server/text_generation/models/bloom.py index 38ef8ef7..008288f8 100644 --- a/server/text_generation/models/bloom.py +++ b/server/text_generation/models/bloom.py @@ -63,6 +63,8 @@ class BLOOMSharded(CausalLM): torch.distributed.barrier(group=self.process_group) filenames = weight_files(model_name, extension=".safetensors") + if not filenames: + raise ValueError("No safetensors weights found") with init_empty_weights(): model = AutoModelForCausalLM.from_config(config) diff --git a/server/text_generation/models/causal_lm.py b/server/text_generation/models/causal_lm.py index c1057635..ca8ea575 100644 --- a/server/text_generation/models/causal_lm.py +++ b/server/text_generation/models/causal_lm.py @@ -156,31 +156,29 @@ class CausalLMBatch: past_keys = past_keys.view(batch.size, -1, *past_keys.shape[-2:]) past_values = past_values.view(batch.size, -1, *past_values.shape[-2:]) - _, num_heads, head_dim, padded_sequence_length = past_keys.shape + _, num_heads, padded_sequence_length, head_dim = past_values.shape - padded_past_keys_shape = ( + padded_past_values_shape = ( total_batch_size, num_heads, - head_dim, max_sequence_length - 1, + head_dim, ) - # head_dim is last for BLOOM - if past_values.shape[-1] == head_dim: - past_values_head_dim_last = True - padded_past_values_shape = ( + # seq_length is last for BLOOM + if past_keys.shape[-2] == head_dim: + past_keys_head_dim_last = False + padded_past_keys_shape = ( total_batch_size, num_heads, - max_sequence_length - 1, head_dim, + max_sequence_length - 1, ) - elif past_values.shape[-2] == head_dim: - past_values_head_dim_last = False - padded_past_values_shape = padded_past_keys_shape + elif past_keys.shape[-1] == head_dim: + past_keys_head_dim_last = True + padded_past_keys_shape = padded_past_values_shape else: - raise ValueError( - f"past_values shape {past_values.shape} is not valid" - ) + raise ValueError(f"past_keys shape {past_keys.shape} is not valid") # This will run only once per layer if j == len(past_key_values): @@ -197,24 +195,24 @@ class CausalLMBatch: past_key_values.append((padded_past_keys, padded_past_values)) # We slice the past keys and values to remove the padding from previous batches - past_key_values[j][0][ - start_index:end_index, :, :, -(batch.max_sequence_length - 1) : - ] = past_keys[:, :, :, -(batch.max_sequence_length - 1) :] - - if past_values_head_dim_last: - past_key_values[j][1][ + if past_keys_head_dim_last: + past_key_values[j][0][ start_index:end_index, :, -(batch.max_sequence_length - 1) :, :, - ] = past_values[:, :, -(batch.max_sequence_length - 1) :, :] + ] = past_keys[:, :, -(batch.max_sequence_length - 1) :, :] else: - past_key_values[j][1][ + past_key_values[j][0][ start_index:end_index, :, :, -(batch.max_sequence_length - 1) :, - ] = past_values[:, :, :, -(batch.max_sequence_length - 1) :] + ] = past_keys[:, :, :, -(batch.max_sequence_length - 1) :] + + past_key_values[j][1][ + start_index:end_index, :, -(batch.max_sequence_length - 1) :, : + ] = past_values[:, :, -(batch.max_sequence_length - 1) :, :] start_index += batch.size @@ -243,13 +241,13 @@ class CausalLM(Model): dtype = torch.float32 tokenizer = AutoTokenizer.from_pretrained(model_name, padding_side="left") - tokenizer.add_special_tokens({"pad_token": "[PAD]"}) self.model = AutoModelForCausalLM.from_pretrained( model_name, torch_dtype=dtype, device_map="auto" if torch.cuda.is_available() else None, load_in_8bit=quantize, ).eval() + tokenizer.pad_token_id = self.model.config.pad_token_id super(CausalLM, self).__init__( tokenizer=tokenizer, diff --git a/server/text_generation/models/galactica.py b/server/text_generation/models/galactica.py new file mode 100644 index 00000000..abc3c36c --- /dev/null +++ b/server/text_generation/models/galactica.py @@ -0,0 +1,346 @@ +import re +import torch +import torch.distributed + +from typing import List, Optional, Type + +from accelerate import init_empty_weights +from safetensors import safe_open +from transformers import AutoTokenizer, AutoModelForCausalLM, AutoConfig +from transformers.models.opt.parallel_layers import ( + TensorParallelColumnLinear, + TensorParallelEmbedding, + TensorParallelRowLinear, +) + +from text_generation.models import CausalLM +from text_generation.pb import generate_pb2 +from text_generation.models.causal_lm import CausalLMBatch +from text_generation.utils import ( + NextTokenChooser, + StoppingCriteria, + initialize_torch_distributed, + weight_files, + download_weights, +) + +HAS_BITS_AND_BYTES = True +try: + import bitsandbytes as bnb + from bitsandbytes.nn import Int8Params +except Exception as e: + HAS_BITS_AND_BYTES = False + +torch.manual_seed(0) + +# CREDIT: Papers with code => https://github.com/paperswithcode/galai/blob/main/galai/utils.py + +# we split individual characters inside special tokens like [START_DNA] +CUSTOM_SEQ_RE = re.compile(r"(\[START_(DNA|SMILES|I_SMILES|AMINO)])(.*?)(\[END_\2])") + +# token added to implement a custom sequence tokenization. This token is added at +# corpus cleaning step and removed in pretokenization. The digits are added to increase the chance +# that they do not occur in the corpus. The digits are escaped so that the token does not appear +# literally in the source code in case we ever include it in the training data. +SPLIT_MARKER = f"SPL{1}T-TH{1}S-Pl3A5E" + + +def _insert_split_marker(m: re.Match): + """ + Applies split marker based on a regex match of special tokens such as + [START_DNA]. + Parameters + ---------- + n : str + Input text to split + Returns + ---------- + str - the text with the split token added + """ + start_token, _, sequence, end_token = m.groups() + sequence = re.sub(r"(.)", rf"{SPLIT_MARKER}\1", sequence, flags=re.DOTALL) + return f"{start_token}{sequence}{SPLIT_MARKER}{end_token}" + + +def escape_custom_split_sequence(text): + """ + Applies custom splitting to the text for GALILEO's tokenization + Parameters + ---------- + text : str + Input text to split + Returns + ---------- + str - the text with the split token added + """ + return CUSTOM_SEQ_RE.sub(_insert_split_marker, text) + + +# END CREDIT + + +class GalacticaCausalLMBatch(CausalLMBatch): + @classmethod + def from_pb( + cls, pb: generate_pb2.Batch, tokenizer: AutoTokenizer, device: torch.device + ) -> "CausalLMBatch": + inputs = [] + next_token_choosers = [] + stopping_criterias = [] + input_lengths = [] + + # Parse batch + for r in pb.requests: + # Add escape_custom_split_sequence to the CausalLMBatch logic + inputs.append(escape_custom_split_sequence(r.inputs)) + input_lengths.append(r.input_length) + next_token_choosers.append( + NextTokenChooser( + temperature=r.parameters.temperature, + top_k=r.parameters.top_k, + top_p=r.parameters.top_p, + do_sample=r.parameters.do_sample, + ) + ) + stopping_criterias.append( + StoppingCriteria( + eos_token_id=tokenizer.eos_token_id, max_new_tokens=r.max_new_tokens + ) + ) + + tokenized_inputs = tokenizer( + inputs, return_tensors="pt", padding=True, pad_to_multiple_of=8 + ).to(device) + all_input_ids = tokenized_inputs["input_ids"].unsqueeze(-1) + + return cls( + batch_id=pb.id, + requests=pb.requests, + input_ids=tokenized_inputs["input_ids"], + attention_mask=tokenized_inputs["attention_mask"], + past_key_values=None, + all_input_ids=all_input_ids, + input_lengths=input_lengths, + next_token_choosers=next_token_choosers, + stopping_criterias=stopping_criterias, + size=pb.size, + max_sequence_length=max(input_lengths), + ) + + +class Galactica(CausalLM): + @property + def batch_type(self) -> Type[CausalLMBatch]: + return GalacticaCausalLMBatch + + +class GalacticaSharded(Galactica): + def __init__(self, model_name: str, quantize: bool = False): + if not model_name.startswith("facebook/galactica"): + raise ValueError(f"Model {model_name} is not supported") + + self.process_group, self.rank, self.world_size = initialize_torch_distributed() + self.master = self.rank == 0 + if torch.cuda.is_available(): + device = torch.device(f"cuda:{self.rank}") + dtype = torch.bfloat16 + else: + device = torch.device("cpu") + dtype = torch.float32 + + tokenizer = AutoTokenizer.from_pretrained(model_name, padding_side="left") + + config = AutoConfig.from_pretrained(model_name, tp_parallel=True) + tokenizer.pad_token_id = config.pad_token_id + + # 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 + + # Only download weights for small models + if self.master and model_name == "facebook/galactica-125m": + download_weights(model_name, extension=".safetensors") + + torch.distributed.barrier(group=self.process_group) + filenames = weight_files(model_name, extension=".safetensors") + if not filenames: + raise ValueError("No safetensors weights found") + + with init_empty_weights(): + model = AutoModelForCausalLM.from_config(config) + + torch.distributed.barrier(group=self.process_group) + self.load_weights( + model, + filenames, + quantize=quantize, + device=device, + rank=self.rank, + world_size=self.world_size, + ) + self.model = model.eval().to(dtype) + torch.distributed.barrier(group=self.process_group) + super(CausalLM, self).__init__( + tokenizer=tokenizer, + num_heads=config.num_attention_heads // self.process_group.size(), + device=device, + ) + + @staticmethod + def load_weights( + model, + filenames: List[str], + quantize: bool, + device: torch.device, + rank: int, + world_size: int, + ): + parameters = dict(model.named_parameters()) + for file in filenames: + with safe_open( + file, framework="pt", device=str(device) if not quantize else "cpu" + ) as f: + for name in f.keys(): + if name == "lm_head.weight": + continue + + module_name, param_name = name.rsplit(".", 1) + try: + module = model.get_submodule(module_name) + except Exception as e: + print(type(model), name, module_name, param_name) + raise e + current_tensor = parameters[name] + + slice_ = f.get_slice(name) + + if isinstance(module, TensorParallelColumnLinear): + if param_name == "weight": + size = slice_.get_shape()[0] + block_size = size // world_size + start = rank * block_size + stop = (rank + 1) * block_size + tensor = slice_[start:stop] + tensor = tensor.transpose(1, 0) + else: + size = slice_.get_shape()[0] + block_size = size // world_size + start = rank * block_size + stop = (rank + 1) * block_size + tensor = slice_[start:stop] + elif isinstance(module, TensorParallelRowLinear): + if param_name == "weight": + size = slice_.get_shape()[1] + block_size = size // world_size + start = rank * block_size + stop = (rank + 1) * block_size + tensor = slice_[:, start:stop] + tensor = tensor.transpose(1, 0) + else: + tensor = slice_[:] + # XXX: Hack for Rowlinear to add the bias only once. + if rank != 0: + tensor = torch.zeros_like(tensor) + elif isinstance(module, TensorParallelEmbedding): + size = slice_.get_shape()[0] + block_size = size // world_size + start = rank * block_size + stop = (rank + 1) * block_size + tensor = slice_[start:stop] + else: + tensor = slice_[:] + + if current_tensor.shape != tensor.shape: + raise ValueError( + f"Name {name} -- Current {current_tensor.shape} and got {tensor.shape}" + ) + + tensor = tensor.contiguous() + + if quantize: + if not HAS_BITS_AND_BYTES: + raise ImportError( + "bitsandbytes is not available on your machine either because it is not installed " + "or you don't have a GPU.\n" + "You can install it with `pip install bitsandbytes`." + ) + + if ( + type(module) + in [TensorParallelRowLinear, TensorParallelColumnLinear] + and param_name == "weight" + ): + tensor = Int8Params( + tensor.transpose(1, 0), + has_fp16_weights=False, + requires_grad=False, + ).to(device) + state = bnb.MatmulLtState() + state.threshold = 6.0 + state.has_fp16_weights = False + state.memory_efficient_backward = False + state.use_pool = True + state.CB = tensor.CB + state.SCB = tensor.SCB + tensor.CB = None + tensor.SCB = None + + def replace_linear(state, in_features, out_features): + def linear(input, weight, bias): + size_out = input.size()[:-1] + (out_features,) + input = input.view(-1, in_features) + out = torch.empty( + size_out, device=input.device, dtype=input.dtype + ) + out = bnb.matmul( + input, + weight, + out=out.view(-1, out_features), + state=state, + threshold=state.threshold, + bias=bias, + ) + + if state.CB is not None: + # we converted 8-bit row major to turing/ampere format + # in the first inference pass + # we no longer need the row-major weight + del state.CB + weight.data = state.CxB + + return out.view(size_out) + + return linear + + module.linear = replace_linear( + state, module.in_features, module.out_features + ) + + else: + tensor = tensor.to(device) + + module._parameters[param_name] = tensor + if name == "model.decoder.embed_tokens.weight": + model.lm_head._parameters["weight"] = tensor + + def forward(self, input_ids, attention_mask, past_key_values: Optional = None): + outputs = self.model.forward( + input_ids=input_ids, + attention_mask=attention_mask, + past_key_values=past_key_values, + use_cache=True, + ) + + # Logits are sharded, so we need to gather them + logits_shard = outputs.logits[:, -1, :].contiguous() + + batch_size, vocab_shard_size = logits_shard.shape + vocab_size = self.world_size * vocab_shard_size + logits = [torch.empty_like(logits_shard) for _ in range(self.world_size)] + torch.distributed.all_gather(logits, logits_shard, group=self.process_group) + logits = torch.cat(logits, dim=1).view(batch_size, 1, vocab_size) + + return logits, outputs.past_key_values diff --git a/server/text_generation/utils.py b/server/text_generation/utils.py index 0e2d9ae4..e6bfc391 100644 --- a/server/text_generation/utils.py +++ b/server/text_generation/utils.py @@ -11,7 +11,7 @@ from functools import partial from huggingface_hub import HfApi, hf_hub_download, try_to_load_from_cache from huggingface_hub.utils import LocalEntryNotFoundError from tqdm import tqdm -from transformers.generation_logits_process import ( +from transformers.generation.logits_process import ( LogitsProcessorList, TemperatureLogitsWarper, TopPLogitsWarper,