feat(server): Support Galactica (IBM#4)

Author: OlivierDehaene

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:
 

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:

server/text_generation/models/__init__.py

@@ -2,6 +2,7 @@
 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")

server/text_generation/models/bloom.py

@@ -63,6 +63,8 @@ def __init__(self, model_name: str, quantize: bool = False):
 
         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)

server/text_generation/models/causal_lm.py

@@ -156,31 +156,29 @@ def concatenate(cls, batches: List["CausalLMBatch"]) -> "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 @@ def concatenate(cls, batches: List["CausalLMBatch"]) -> "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 @@ def __init__(self, model_name: str, quantize=False):
             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,