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

Fix shared weights load bug and T5 loading

Browse Source
This commit is contained in:
zhangsibo1129 2023-09-26 17:57:59 +08:00
parent c5de7cd886
commit 57433201b2
2 changed files with 4 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
server/text_generation_server/models/t5.py
View File

@ -63,6 +63,7 @@ class T5Sharded(Seq2SeqLM):
"shared.weight": [
"encoder.embed_tokens.weight",
"decoder.embed_tokens.weight",
"lm_head.weight",
]
},
)

6
server/text_generation_server/utils/weights.py
View File

@ -45,9 +45,9 @@ class Weights:
def get_filename(self, tensor_name: str) -> (str, str):
filename = self.routing.get(tensor_name, None)
if filename is None:
aliases = self.aliases.get(tensor_name, [])
for alias in aliases:
filename = self.routing.get(alias, None)
for alias, tensor_list in self.aliases.items():
if tensor_name in tensor_list:
filename = self.routing.get(alias, None)
if filename is not None:
return str(filename), alias
raise RuntimeError(f"weight {tensor_name} does not exist")