huggingface/text-generation-inference
5
0
Fork 0
mirror of https://github.com/huggingface/text-generation-inference.git synced 2025-09-09 11:24:53 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix santacoder sharded

Browse Source
This commit is contained in:
OlivierDehaene 2023-04-19 12:32:25 +02:00
parent b47edb28af
commit 0fc4f99379
2 changed files with 4 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
server/text_generation_server/models/__init__.py
View File

@ -90,10 +90,10 @@ def get_model(
raise NotImplementedError(
FLASH_ATT_ERROR_MESSAGE.format(f"Sharded Santacoder")
)
return FlashSantacoderSharded(model_id, revision=revision)
return FlashSantacoderSharded(model_id, revision, quantize=quantize)
else:
santacoder_cls = FlashSantacoder if FLASH_ATTENTION else SantaCoder
return santacoder_cls(model_id, revision, quantize)
return santacoder_cls(model_id, revision, quantize=quantize)
config = AutoConfig.from_pretrained(model_id, revision=revision)
model_type = config.model_type

4
server/text_generation_server/models/flash_llama.py
View File

@ -75,9 +75,9 @@ class FlashLlama(FlashCausalLM):
dtype: torch.dtype,
):
for filename in filenames:
state_dict = torch.load(filename, "cpu")
state_dict = torch.load(filename, map_location="cpu")
for key, value in state_dict.items():
value = value.to(dtype).to(device if not quantize else "cpu")
value = value.to(device if not quantize else "cpu").to(dtype)
layer_name = ".".join(key.split(".")[:4])