prompt: 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()
  
  ```
  
golden_diff: 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
  
verification_info: {"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"}]}