Embeddings weights alias fix for latest falcon models

Author: njhill

server/text_generation_server/inference_engine/hf_custom_tp.py

@@ -85,7 +85,10 @@ def __init__(
         elif model_type in ["RefinedWeb", "RefinedWebModel", "falcon"]:
             if sharded and self._config.alibi:
                 raise NotImplementedError("TP is not supported for Falcon models using alibi")
-            aliases = {"transformer.word_embeddings.weight": ["lm_head.weight"]}
+            aliases = {
+                "transformer.word_embeddings.weight": ["lm_head.weight"],
+                "lm_head.weight": ["transformer.word_embeddings.weight"],
+            }
             from text_generation_server.models.custom_modeling.flash_rw_modeling import FlashRWForCausalLM
             model_class = FlashRWForCausalLM