We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Tied weight optimization for checkpoints doesn't work with text-generation-inference.
### System Info
Ubuntu 20.04
4 A10 NVIDIA GPU's

I think checkpoints saved after this feature was merged don't work with text-generation-inference.
https://github.com/huggingface/transformers/issues/23868

With falcon models getting "`lm_head` not found"
I'll add more details once I find minimal steps to reproduce.

### Information

- [X] Docker
- [ ] The CLI directly

### Tasks

- [X] An officially supported command
- [ ] My own modifications

### Reproduction

Save tiiuae/falcon-40b checkpoint using transformers==4.30.2
launch text-generation-inference server

(using transformers==4.27.4 works without issue)

### Expected behavior

Expect the text-generation-inference weight loader to be able to find the `lm_head` weight in the checkpoint. Note this may be a safetensor issue.

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `server/text_generation_server/models/flash_rw.py`
Content:
```
1 import torch
2 import torch.distributed
3 
4 from opentelemetry import trace
5 from transformers import AutoTokenizer
6 from typing import Optional
7 
8 from text_generation_server.models import FlashCausalLM
9 from text_generation_server.models.custom_modeling.flash_rw_modeling import (
10     RWConfig,
11     FlashRWForCausalLM,
12 )
13 from text_generation_server.utils import (
14     initialize_torch_distributed,
15     weight_files,
16     Weights,
17 )
18 
19 tracer = trace.get_tracer(__name__)
20 
21 
22 class FlashRWSharded(FlashCausalLM):
23     def __init__(
24         self,
25         model_id: str,
26         revision: Optional[str] = None,
27         quantize: Optional[str] = None,
28         dtype: Optional[torch.dtype] = None,
29         trust_remote_code: bool = False,
30     ):
31         self.process_group, rank, world_size = initialize_torch_distributed()
32         if torch.cuda.is_available():
33             device = torch.device(f"cuda:{rank}")
34             dtype = torch.float16 if dtype is None else dtype
35         else:
36             raise NotImplementedError("FlashRW is only available on GPU")
37 
38         tokenizer = AutoTokenizer.from_pretrained(
39             model_id,
40             revision=revision,
41             padding_side="left",
42             truncation_side="left",
43             trust_remote_code=trust_remote_code,
44         )
45 
46         config = RWConfig.from_pretrained(
47             model_id, revision=revision, trust_remote_code=trust_remote_code
48         )
49 
50         torch.distributed.barrier(group=self.process_group)
51         filenames = weight_files(model_id, revision=revision, extension=".safetensors")
52         weights = Weights(filenames, device, dtype, process_group=self.process_group)
53 
54         config.quantize = quantize
55 
56         model = FlashRWForCausalLM(config, weights)
57 
58         torch.distributed.barrier(group=self.process_group)
59         super(FlashRWSharded, self).__init__(
60             model=model.to(device),
61             tokenizer=tokenizer,
62             num_layers=len(model.transformer.h),
63             num_kv_heads=model.transformer.cache_size,
64             head_size=model.transformer.head_size,
65             dtype=dtype,
66             device=device,
67             rank=rank,
68             world_size=world_size,
69         )
70
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/server/text_generation_server/models/flash_rw.py b/server/text_generation_server/models/flash_rw.py
--- a/server/text_generation_server/models/flash_rw.py
+++ b/server/text_generation_server/models/flash_rw.py
@@ -49,7 +49,13 @@
 
         torch.distributed.barrier(group=self.process_group)
         filenames = weight_files(model_id, revision=revision, extension=".safetensors")
-        weights = Weights(filenames, device, dtype, process_group=self.process_group)
+        weights = Weights(
+            filenames,
+            device,
+            dtype,
+            process_group=self.process_group,
+            aliases={"transformer.word_embeddings.weight": ["lm_head.weight"]},
+        )
 
         config.quantize = quantize

{"golden_diff": "diff --git a/server/text_generation_server/models/flash_rw.py b/server/text_generation_server/models/flash_rw.py\n--- a/server/text_generation_server/models/flash_rw.py\n+++ b/server/text_generation_server/models/flash_rw.py\n@@ -49,7 +49,13 @@\n \n         torch.distributed.barrier(group=self.process_group)\n         filenames = weight_files(model_id, revision=revision, extension=\".safetensors\")\n-        weights = Weights(filenames, device, dtype, process_group=self.process_group)\n+        weights = Weights(\n+            filenames,\n+            device,\n+            dtype,\n+            process_group=self.process_group,\n+            aliases={\"transformer.word_embeddings.weight\": [\"lm_head.weight\"]},\n+        )\n \n         config.quantize = quantize\n", "issue": "Tied weight optimization for checkpoints doesn't work with text-generation-inference.\n### System Info\r\nUbuntu 20.04\r\n4 A10 NVIDIA GPU's\r\n\r\nI think checkpoints saved after this feature was merged don't work with text-generation-inference.\r\nhttps://github.com/huggingface/transformers/issues/23868\r\n\r\nWith falcon models getting \"`lm_head` not found\"\r\nI'll add more details once I find minimal steps to reproduce.\r\n\r\n### Information\r\n\r\n- [X] Docker\r\n- [ ] The CLI directly\r\n\r\n### Tasks\r\n\r\n- [X] An officially supported command\r\n- [ ] My own modifications\r\n\r\n### Reproduction\r\n\r\nSave tiiuae/falcon-40b checkpoint using transformers==4.30.2\r\nlaunch text-generation-inference server\r\n\r\n(using transformers==4.27.4 works without issue)\r\n\r\n### Expected behavior\r\n\r\nExpect the text-generation-inference weight loader to be able to find the `lm_head` weight in the checkpoint. Note this may be a safetensor issue.\n", "before_files": [{"content": "import torch\nimport torch.distributed\n\nfrom opentelemetry import trace\nfrom transformers import AutoTokenizer\nfrom typing import Optional\n\nfrom text_generation_server.models import FlashCausalLM\nfrom text_generation_server.models.custom_modeling.flash_rw_modeling import (\n    RWConfig,\n    FlashRWForCausalLM,\n)\nfrom text_generation_server.utils import (\n    initialize_torch_distributed,\n    weight_files,\n    Weights,\n)\n\ntracer = trace.get_tracer(__name__)\n\n\nclass FlashRWSharded(FlashCausalLM):\n    def __init__(\n        self,\n        model_id: str,\n        revision: Optional[str] = None,\n        quantize: Optional[str] = None,\n        dtype: Optional[torch.dtype] = None,\n        trust_remote_code: bool = False,\n    ):\n        self.process_group, rank, world_size = initialize_torch_distributed()\n        if torch.cuda.is_available():\n            device = torch.device(f\"cuda:{rank}\")\n            dtype = torch.float16 if dtype is None else dtype\n        else:\n            raise NotImplementedError(\"FlashRW is only available on GPU\")\n\n        tokenizer = AutoTokenizer.from_pretrained(\n            model_id,\n            revision=revision,\n            padding_side=\"left\",\n            truncation_side=\"left\",\n            trust_remote_code=trust_remote_code,\n        )\n\n        config = RWConfig.from_pretrained(\n            model_id, revision=revision, trust_remote_code=trust_remote_code\n        )\n\n        torch.distributed.barrier(group=self.process_group)\n        filenames = weight_files(model_id, revision=revision, extension=\".safetensors\")\n        weights = Weights(filenames, device, dtype, process_group=self.process_group)\n\n        config.quantize = quantize\n\n        model = FlashRWForCausalLM(config, weights)\n\n        torch.distributed.barrier(group=self.process_group)\n        super(FlashRWSharded, self).__init__(\n            model=model.to(device),\n            tokenizer=tokenizer,\n            num_layers=len(model.transformer.h),\n            num_kv_heads=model.transformer.cache_size,\n            head_size=model.transformer.head_size,\n            dtype=dtype,\n            device=device,\n            rank=rank,\n            world_size=world_size,\n        )\n", "path": "server/text_generation_server/models/flash_rw.py"}], "after_files": [{"content": "import torch\nimport torch.distributed\n\nfrom opentelemetry import trace\nfrom transformers import AutoTokenizer\nfrom typing import Optional\n\nfrom text_generation_server.models import FlashCausalLM\nfrom text_generation_server.models.custom_modeling.flash_rw_modeling import (\n    RWConfig,\n    FlashRWForCausalLM,\n)\nfrom text_generation_server.utils import (\n    initialize_torch_distributed,\n    weight_files,\n    Weights,\n)\n\ntracer = trace.get_tracer(__name__)\n\n\nclass FlashRWSharded(FlashCausalLM):\n    def __init__(\n        self,\n        model_id: str,\n        revision: Optional[str] = None,\n        quantize: Optional[str] = None,\n        dtype: Optional[torch.dtype] = None,\n        trust_remote_code: bool = False,\n    ):\n        self.process_group, rank, world_size = initialize_torch_distributed()\n        if torch.cuda.is_available():\n            device = torch.device(f\"cuda:{rank}\")\n            dtype = torch.float16 if dtype is None else dtype\n        else:\n            raise NotImplementedError(\"FlashRW is only available on GPU\")\n\n        tokenizer = AutoTokenizer.from_pretrained(\n            model_id,\n            revision=revision,\n            padding_side=\"left\",\n            truncation_side=\"left\",\n            trust_remote_code=trust_remote_code,\n        )\n\n        config = RWConfig.from_pretrained(\n            model_id, revision=revision, trust_remote_code=trust_remote_code\n        )\n\n        torch.distributed.barrier(group=self.process_group)\n        filenames = weight_files(model_id, revision=revision, extension=\".safetensors\")\n        weights = Weights(\n            filenames,\n            device,\n            dtype,\n            process_group=self.process_group,\n            aliases={\"transformer.word_embeddings.weight\": [\"lm_head.weight\"]},\n        )\n\n        config.quantize = quantize\n\n        model = FlashRWForCausalLM(config, weights)\n\n        torch.distributed.barrier(group=self.process_group)\n        super(FlashRWSharded, self).__init__(\n            model=model.to(device),\n            tokenizer=tokenizer,\n            num_layers=len(model.transformer.h),\n            num_kv_heads=model.transformer.cache_size,\n            head_size=model.transformer.head_size,\n            dtype=dtype,\n            device=device,\n            rank=rank,\n            world_size=world_size,\n        )\n", "path": "server/text_generation_server/models/flash_rw.py"}]}