From ce8eaaf2be2227aab88fd827ed5c82b275de1957 Mon Sep 17 00:00:00 2001 From: Nicolas Patry Date: Mon, 25 Sep 2023 09:50:42 +0000 Subject: [PATCH] Better fix. --- .../models/custom_modeling/t5_modeling.py | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/server/text_generation_server/models/custom_modeling/t5_modeling.py b/server/text_generation_server/models/custom_modeling/t5_modeling.py index a9dff6c6..d3e4f53a 100644 --- a/server/text_generation_server/models/custom_modeling/t5_modeling.py +++ b/server/text_generation_server/models/custom_modeling/t5_modeling.py @@ -1032,9 +1032,17 @@ class T5ForConditionalGeneration(T5PreTrainedModel): embed_tokens=self.shared, ) - self.lm_head = TensorParallelHead.load( - config, prefix="shared", weights=weights - ) + try: + self.lm_head = TensorParallelHead.load( + config, prefix="lm_head", weights=weights + ) + except RuntimeError: + # Some models like t5-small were saved with shared weights unlike flan + # Since they are declared as the same arch we have no choice but hope + # that this is OK instead of using a proper flag. + self.lm_head = TensorParallelHead.load( + config, prefix="shared", weights=weights + ) def forward( self,