From 7e11c5d92ba93e9cb0be488d26ed50b22ccc38dd Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Mon, 8 May 2023 09:45:27 +0000 Subject: [PATCH] Hotfixes for santacoder/bigcode. --- server/text_generation_server/models/__init__.py | 8 ++++---- server/text_generation_server/models/flash_santacoder.py | 3 +++ 2 files changed, 7 insertions(+), 4 deletions(-) diff --git a/server/text_generation_server/models/__init__.py b/server/text_generation_server/models/__init__.py index 221c9139..5cb66382 100644 --- a/server/text_generation_server/models/__init__.py +++ b/server/text_generation_server/models/__init__.py @@ -99,7 +99,10 @@ def get_model( else: return Galactica(model_id, revision, quantize=quantize) - if "bigcode" in model_id: + config = AutoConfig.from_pretrained(model_id, revision=revision) + model_type = config.model_type + + if model_type == "gpt_bigcode": if sharded: if not FLASH_ATTENTION: raise NotImplementedError( @@ -110,9 +113,6 @@ def get_model( santacoder_cls = FlashSantacoder if FLASH_ATTENTION else SantaCoder return santacoder_cls(model_id, revision, quantize=quantize) - config = AutoConfig.from_pretrained(model_id, revision=revision) - model_type = config.model_type - if model_type == "bloom": if sharded: return BLOOMSharded(model_id, revision, quantize=quantize) diff --git a/server/text_generation_server/models/flash_santacoder.py b/server/text_generation_server/models/flash_santacoder.py index 333180e8..de858a99 100644 --- a/server/text_generation_server/models/flash_santacoder.py +++ b/server/text_generation_server/models/flash_santacoder.py @@ -373,6 +373,9 @@ class FlashSantacoderSharded(FlashSantacoder): else: module._buffers[param_name] = tensor + + model.lm_head.weight = torch.nn.Parameter(model.transformer.wte.weight) + uninitialized_parameters = [] for n, p in model.named_parameters(): if p.data.device == torch.device("meta"):