From f6ba71f60ffafd2ba87d4df98f244c678e065c64 Mon Sep 17 00:00:00 2001 From: hyunwoongko Date: Sat, 3 Jun 2023 12:00:07 +0900 Subject: [PATCH] [python] fix: Fix embedding mapping for deepspeed chat --- server/text_generation_server/models/bloom.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/server/text_generation_server/models/bloom.py b/server/text_generation_server/models/bloom.py index 45d7cd4c..c5889532 100644 --- a/server/text_generation_server/models/bloom.py +++ b/server/text_generation_server/models/bloom.py @@ -255,7 +255,7 @@ class BLOOMSharded(BLOOM): raise ValueError(f"Unexpected quantize `{quantize}`") module._parameters[param_name] = tensor - if name == "word_embeddings.weight": + if "word_embeddings.weight" in name: model.lm_head._parameters["weight"] = tensor def forward(