We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Can't load local flan-small models due to weight conversion failure 
### System Info

OS Version: 
Distributor ID: Ubuntu
Description:    Ubuntu 20.04.3 LTS
Release:        20.04
Codename:       focal

8 A-100 GPUS

Using latest text-generation-inference docker version. 

I've run fine-tuning on a [Flan-T5-Small](https://huggingface.co/google/flan-t5-small) model and saved the checkpoint in my local directory. I've stored this local model checkpoint in my data2 volume and run the command as follows:
docker run --gpus all --shm-size 1g -p 8080:80 -v $volume:/data2 ghcr.io/huggingface/text-generation-inference:0.9 --model-id /data2/checkpoint-20 --num-shard $num_shard

But I run into errors with the converting weights as mentioned below. 

### Information

- [X] Docker
- [ ] The CLI directly

### Tasks

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

### Reproduction

Run docker command above. 

I get this error now: 

2023-07-12T05:45:31.707548Z  INFO text_generation_launcher: Args { model_id: "/data2/checkpoint-20", revision: None, sharded: None, num_shard: Some(2), quantize: None, dtype: None, trust_remote_code: false, max_concurrent_requests: 128, max_best_of: 2, max_stop_sequences: 4, max_input_length: 1024, max_total_tokens: 2048, waiting_served_ratio: 1.2, max_batch_prefill_tokens: 4096, max_batch_total_tokens: 16000, max_waiting_tokens: 20, hostname: "0341f92fe465", port: 80, shard_uds_path: "/tmp/text-generation-server", master_addr: "localhost", master_port: 29500, huggingface_hub_cache: Some("/data"), weights_cache_override: None, disable_custom_kernels: false, json_output: false, otlp_endpoint: None, cors_allow_origin: [], watermark_gamma: None, watermark_delta: None, ngrok: false, ngrok_authtoken: None, ngrok_domain: None, ngrok_username: None, ngrok_password: None, env: false }
2023-07-12T05:45:31.707602Z  INFO text_generation_launcher: Sharding model on 2 processes
2023-07-12T05:45:31.707781Z  INFO text_generation_launcher: Starting download process.
2023-07-12T05:45:33.261253Z  WARN download: text_generation_launcher: No safetensors weights found for model /data2/checkpoint-20 at revision None. Converting PyTorch weights to safetensors.

2023-07-12T05:45:33.711218Z ERROR text_generation_launcher: Download encountered an error: Traceback (most recent call last):

  File "/opt/conda/bin/text-generation-server", line 8, in <module>
    sys.exit(app())

  File "/opt/conda/lib/python3.9/site-packages/text_generation_server/cli.py", line 164, in download_weights
    utils.convert_files(local_pt_files, local_st_files)

  File "/opt/conda/lib/python3.9/site-packages/text_generation_server/utils/convert.py", line 53, in convert_files
    convert_file(pt_file, sf_file)

  File "/opt/conda/lib/python3.9/site-packages/text_generation_server/utils/convert.py", line 21, in convert_file
    if "state_dict" in loaded:

TypeError: argument of type 'Seq2SeqTrainingArguments' is not iterable


Error: DownloadError


### Expected behavior

I would expect the local model to load as do the models from the hugging-face library. Appreciate any help!

--- 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/utils/convert.py`
Content:
```
1 import datetime
2 import torch
3 import os
4 
5 from loguru import logger
6 from pathlib import Path
7 from safetensors.torch import save_file, load_file, _find_shared_tensors, _is_complete
8 from typing import List, Dict
9 from collections import defaultdict
10 
11 
12 def _remove_duplicate_names(
13     state_dict: Dict[str, torch.Tensor],
14     *,
15     preferred_names: List[str] = None,
16     discard_names: List[str] = None,
17 ) -> Dict[str, List[str]]:
18     if preferred_names is None:
19         preferred_names = []
20     preferred_names = set(preferred_names)
21     if discard_names is None:
22         discard_names = []
23     discard_names = set(discard_names)
24 
25     shareds = _find_shared_tensors(state_dict)
26     to_remove = defaultdict(list)
27     for shared in shareds:
28         complete_names = set(
29             [name for name in shared if _is_complete(state_dict[name])]
30         )
31         if not complete_names:
32             raise RuntimeError(
33                 f"Error while trying to find names to remove to save state dict, but found no suitable name to keep for saving amongst: {shared}. None is covering the entire storage.Refusing to save/load the model since you could be storing much more memory than needed. Please refer to https://huggingface.co/docs/safetensors/torch_shared_tensors for more information. Or open an issue."
34             )
35 
36         keep_name = sorted(list(complete_names))[0]
37 
38         # Mecanism to preferentially select keys to keep
39         # coming from the on-disk file to allow
40         # loading models saved with a different choice
41         # of keep_name
42         preferred = complete_names.difference(discard_names)
43         if preferred:
44             keep_name = sorted(list(preferred))[0]
45 
46         if preferred_names:
47             preferred = preferred_names.intersection(complete_names)
48             if preferred:
49                 keep_name = sorted(list(preferred))[0]
50         for name in sorted(shared):
51             if name != keep_name:
52                 to_remove[keep_name].append(name)
53     return to_remove
54 
55 
56 def convert_file(pt_file: Path, sf_file: Path, discard_names: List[str]):
57     """
58     Convert a pytorch file to a safetensors file
59     This will remove duplicate tensors from the file.
60 
61     Unfortunately, this might not respect *transformers* convention.
62     Forcing us to check for potentially different keys during load when looking
63     for specific tensors (making tensor sharing explicit).
64     """
65     loaded = torch.load(pt_file, map_location="cpu")
66     if "state_dict" in loaded:
67         loaded = loaded["state_dict"]
68     to_removes = _remove_duplicate_names(loaded, discard_names=discard_names)
69 
70     metadata = {"format": "pt"}
71     for kept_name, to_remove_group in to_removes.items():
72         for to_remove in to_remove_group:
73             if to_remove not in metadata:
74                 metadata[to_remove] = kept_name
75             del loaded[to_remove]
76     # Force tensors to be contiguous
77     loaded = {k: v.contiguous() for k, v in loaded.items()}
78 
79     dirname = os.path.dirname(sf_file)
80     os.makedirs(dirname, exist_ok=True)
81     save_file(loaded, sf_file, metadata=metadata)
82     reloaded = load_file(sf_file)
83     for k in loaded:
84         pt_tensor = loaded[k]
85         sf_tensor = reloaded[k]
86         if not torch.equal(pt_tensor, sf_tensor):
87             raise RuntimeError(f"The output tensors do not match for key {k}")
88 
89 
90 def convert_files(pt_files: List[Path], sf_files: List[Path], discard_names: List[str]):
91     assert len(pt_files) == len(sf_files)
92 
93     N = len(pt_files)
94     # We do this instead of using tqdm because we want to parse the logs with the launcher
95 
96     for i, (pt_file, sf_file) in enumerate(zip(pt_files, sf_files)):
97         start = datetime.datetime.now()
98         convert_file(pt_file, sf_file, discard_names)
99         elapsed = datetime.datetime.now() - start
100         logger.info(f"Convert: [{i + 1}/{N}] -- Took: {elapsed}")
101
```
--- 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/utils/convert.py b/server/text_generation_server/utils/convert.py
--- a/server/text_generation_server/utils/convert.py
+++ b/server/text_generation_server/utils/convert.py
@@ -94,6 +94,14 @@
     # We do this instead of using tqdm because we want to parse the logs with the launcher
 
     for i, (pt_file, sf_file) in enumerate(zip(pt_files, sf_files)):
+        # Skip blacklisted files
+        if (
+            "arguments" in pt_file.name
+            or "args" in pt_file.name
+            or "training" in pt_file.name
+        ):
+            continue
+
         start = datetime.datetime.now()
         convert_file(pt_file, sf_file, discard_names)
         elapsed = datetime.datetime.now() - start

{"golden_diff": "diff --git a/server/text_generation_server/utils/convert.py b/server/text_generation_server/utils/convert.py\n--- a/server/text_generation_server/utils/convert.py\n+++ b/server/text_generation_server/utils/convert.py\n@@ -94,6 +94,14 @@\n     # We do this instead of using tqdm because we want to parse the logs with the launcher\n \n     for i, (pt_file, sf_file) in enumerate(zip(pt_files, sf_files)):\n+        # Skip blacklisted files\n+        if (\n+            \"arguments\" in pt_file.name\n+            or \"args\" in pt_file.name\n+            or \"training\" in pt_file.name\n+        ):\n+            continue\n+\n         start = datetime.datetime.now()\n         convert_file(pt_file, sf_file, discard_names)\n         elapsed = datetime.datetime.now() - start\n", "issue": "Can't load local flan-small models due to weight conversion failure \n### System Info\n\nOS Version: \r\nDistributor ID: Ubuntu\r\nDescription:    Ubuntu 20.04.3 LTS\r\nRelease:        20.04\r\nCodename:       focal\r\n\r\n8 A-100 GPUS\r\n\r\nUsing latest text-generation-inference docker version. \r\n\r\nI've run fine-tuning on a [Flan-T5-Small](https://huggingface.co/google/flan-t5-small) model and saved the checkpoint in my local directory. I've stored this local model checkpoint in my data2 volume and run the command as follows:\r\ndocker run --gpus all --shm-size 1g -p 8080:80 -v $volume:/data2 ghcr.io/huggingface/text-generation-inference:0.9 --model-id /data2/checkpoint-20 --num-shard $num_shard\r\n\r\nBut I run into errors with the converting weights as mentioned below. \n\n### Information\n\n- [X] Docker\n- [ ] The CLI directly\n\n### Tasks\n\n- [X] An officially supported command\n- [ ] My own modifications\n\n### Reproduction\n\nRun docker command above. \r\n\r\nI get this error now: \r\n\r\n2023-07-12T05:45:31.707548Z  INFO text_generation_launcher: Args { model_id: \"/data2/checkpoint-20\", revision: None, sharded: None, num_shard: Some(2), quantize: None, dtype: None, trust_remote_code: false, max_concurrent_requests: 128, max_best_of: 2, max_stop_sequences: 4, max_input_length: 1024, max_total_tokens: 2048, waiting_served_ratio: 1.2, max_batch_prefill_tokens: 4096, max_batch_total_tokens: 16000, max_waiting_tokens: 20, hostname: \"0341f92fe465\", port: 80, shard_uds_path: \"/tmp/text-generation-server\", master_addr: \"localhost\", master_port: 29500, huggingface_hub_cache: Some(\"/data\"), weights_cache_override: None, disable_custom_kernels: false, json_output: false, otlp_endpoint: None, cors_allow_origin: [], watermark_gamma: None, watermark_delta: None, ngrok: false, ngrok_authtoken: None, ngrok_domain: None, ngrok_username: None, ngrok_password: None, env: false }\r\n2023-07-12T05:45:31.707602Z  INFO text_generation_launcher: Sharding model on 2 processes\r\n2023-07-12T05:45:31.707781Z  INFO text_generation_launcher: Starting download process.\r\n2023-07-12T05:45:33.261253Z  WARN download: text_generation_launcher: No safetensors weights found for model /data2/checkpoint-20 at revision None. Converting PyTorch weights to safetensors.\r\n\r\n2023-07-12T05:45:33.711218Z ERROR text_generation_launcher: Download encountered an error: Traceback (most recent call last):\r\n\r\n  File \"/opt/conda/bin/text-generation-server\", line 8, in <module>\r\n    sys.exit(app())\r\n\r\n  File \"/opt/conda/lib/python3.9/site-packages/text_generation_server/cli.py\", line 164, in download_weights\r\n    utils.convert_files(local_pt_files, local_st_files)\r\n\r\n  File \"/opt/conda/lib/python3.9/site-packages/text_generation_server/utils/convert.py\", line 53, in convert_files\r\n    convert_file(pt_file, sf_file)\r\n\r\n  File \"/opt/conda/lib/python3.9/site-packages/text_generation_server/utils/convert.py\", line 21, in convert_file\r\n    if \"state_dict\" in loaded:\r\n\r\nTypeError: argument of type 'Seq2SeqTrainingArguments' is not iterable\r\n\r\n\r\nError: DownloadError\r\n\n\n### Expected behavior\n\nI would expect the local model to load as do the models from the hugging-face library. Appreciate any help!\n", "before_files": [{"content": "import datetime\nimport torch\nimport os\n\nfrom loguru import logger\nfrom pathlib import Path\nfrom safetensors.torch import save_file, load_file, _find_shared_tensors, _is_complete\nfrom typing import List, Dict\nfrom collections import defaultdict\n\n\ndef _remove_duplicate_names(\n    state_dict: Dict[str, torch.Tensor],\n    *,\n    preferred_names: List[str] = None,\n    discard_names: List[str] = None,\n) -> Dict[str, List[str]]:\n    if preferred_names is None:\n        preferred_names = []\n    preferred_names = set(preferred_names)\n    if discard_names is None:\n        discard_names = []\n    discard_names = set(discard_names)\n\n    shareds = _find_shared_tensors(state_dict)\n    to_remove = defaultdict(list)\n    for shared in shareds:\n        complete_names = set(\n            [name for name in shared if _is_complete(state_dict[name])]\n        )\n        if not complete_names:\n            raise RuntimeError(\n                f\"Error while trying to find names to remove to save state dict, but found no suitable name to keep for saving amongst: {shared}. None is covering the entire storage.Refusing to save/load the model since you could be storing much more memory than needed. Please refer to https://huggingface.co/docs/safetensors/torch_shared_tensors for more information. Or open an issue.\"\n            )\n\n        keep_name = sorted(list(complete_names))[0]\n\n        # Mecanism to preferentially select keys to keep\n        # coming from the on-disk file to allow\n        # loading models saved with a different choice\n        # of keep_name\n        preferred = complete_names.difference(discard_names)\n        if preferred:\n            keep_name = sorted(list(preferred))[0]\n\n        if preferred_names:\n            preferred = preferred_names.intersection(complete_names)\n            if preferred:\n                keep_name = sorted(list(preferred))[0]\n        for name in sorted(shared):\n            if name != keep_name:\n                to_remove[keep_name].append(name)\n    return to_remove\n\n\ndef convert_file(pt_file: Path, sf_file: Path, discard_names: List[str]):\n    \"\"\"\n    Convert a pytorch file to a safetensors file\n    This will remove duplicate tensors from the file.\n\n    Unfortunately, this might not respect *transformers* convention.\n    Forcing us to check for potentially different keys during load when looking\n    for specific tensors (making tensor sharing explicit).\n    \"\"\"\n    loaded = torch.load(pt_file, map_location=\"cpu\")\n    if \"state_dict\" in loaded:\n        loaded = loaded[\"state_dict\"]\n    to_removes = _remove_duplicate_names(loaded, discard_names=discard_names)\n\n    metadata = {\"format\": \"pt\"}\n    for kept_name, to_remove_group in to_removes.items():\n        for to_remove in to_remove_group:\n            if to_remove not in metadata:\n                metadata[to_remove] = kept_name\n            del loaded[to_remove]\n    # Force tensors to be contiguous\n    loaded = {k: v.contiguous() for k, v in loaded.items()}\n\n    dirname = os.path.dirname(sf_file)\n    os.makedirs(dirname, exist_ok=True)\n    save_file(loaded, sf_file, metadata=metadata)\n    reloaded = load_file(sf_file)\n    for k in loaded:\n        pt_tensor = loaded[k]\n        sf_tensor = reloaded[k]\n        if not torch.equal(pt_tensor, sf_tensor):\n            raise RuntimeError(f\"The output tensors do not match for key {k}\")\n\n\ndef convert_files(pt_files: List[Path], sf_files: List[Path], discard_names: List[str]):\n    assert len(pt_files) == len(sf_files)\n\n    N = len(pt_files)\n    # We do this instead of using tqdm because we want to parse the logs with the launcher\n\n    for i, (pt_file, sf_file) in enumerate(zip(pt_files, sf_files)):\n        start = datetime.datetime.now()\n        convert_file(pt_file, sf_file, discard_names)\n        elapsed = datetime.datetime.now() - start\n        logger.info(f\"Convert: [{i + 1}/{N}] -- Took: {elapsed}\")\n", "path": "server/text_generation_server/utils/convert.py"}], "after_files": [{"content": "import datetime\nimport torch\nimport os\n\nfrom loguru import logger\nfrom pathlib import Path\nfrom safetensors.torch import save_file, load_file, _find_shared_tensors, _is_complete\nfrom typing import List, Dict\nfrom collections import defaultdict\n\n\ndef _remove_duplicate_names(\n    state_dict: Dict[str, torch.Tensor],\n    *,\n    preferred_names: List[str] = None,\n    discard_names: List[str] = None,\n) -> Dict[str, List[str]]:\n    if preferred_names is None:\n        preferred_names = []\n    preferred_names = set(preferred_names)\n    if discard_names is None:\n        discard_names = []\n    discard_names = set(discard_names)\n\n    shareds = _find_shared_tensors(state_dict)\n    to_remove = defaultdict(list)\n    for shared in shareds:\n        complete_names = set(\n            [name for name in shared if _is_complete(state_dict[name])]\n        )\n        if not complete_names:\n            raise RuntimeError(\n                f\"Error while trying to find names to remove to save state dict, but found no suitable name to keep for saving amongst: {shared}. None is covering the entire storage.Refusing to save/load the model since you could be storing much more memory than needed. Please refer to https://huggingface.co/docs/safetensors/torch_shared_tensors for more information. Or open an issue.\"\n            )\n\n        keep_name = sorted(list(complete_names))[0]\n\n        # Mecanism to preferentially select keys to keep\n        # coming from the on-disk file to allow\n        # loading models saved with a different choice\n        # of keep_name\n        preferred = complete_names.difference(discard_names)\n        if preferred:\n            keep_name = sorted(list(preferred))[0]\n\n        if preferred_names:\n            preferred = preferred_names.intersection(complete_names)\n            if preferred:\n                keep_name = sorted(list(preferred))[0]\n        for name in sorted(shared):\n            if name != keep_name:\n                to_remove[keep_name].append(name)\n    return to_remove\n\n\ndef convert_file(pt_file: Path, sf_file: Path, discard_names: List[str]):\n    \"\"\"\n    Convert a pytorch file to a safetensors file\n    This will remove duplicate tensors from the file.\n\n    Unfortunately, this might not respect *transformers* convention.\n    Forcing us to check for potentially different keys during load when looking\n    for specific tensors (making tensor sharing explicit).\n    \"\"\"\n    loaded = torch.load(pt_file, map_location=\"cpu\")\n    if \"state_dict\" in loaded:\n        loaded = loaded[\"state_dict\"]\n    to_removes = _remove_duplicate_names(loaded, discard_names=discard_names)\n\n    metadata = {\"format\": \"pt\"}\n    for kept_name, to_remove_group in to_removes.items():\n        for to_remove in to_remove_group:\n            if to_remove not in metadata:\n                metadata[to_remove] = kept_name\n            del loaded[to_remove]\n    # Force tensors to be contiguous\n    loaded = {k: v.contiguous() for k, v in loaded.items()}\n\n    dirname = os.path.dirname(sf_file)\n    os.makedirs(dirname, exist_ok=True)\n    save_file(loaded, sf_file, metadata=metadata)\n    reloaded = load_file(sf_file)\n    for k in loaded:\n        pt_tensor = loaded[k]\n        sf_tensor = reloaded[k]\n        if not torch.equal(pt_tensor, sf_tensor):\n            raise RuntimeError(f\"The output tensors do not match for key {k}\")\n\n\ndef convert_files(pt_files: List[Path], sf_files: List[Path], discard_names: List[str]):\n    assert len(pt_files) == len(sf_files)\n\n    N = len(pt_files)\n    # We do this instead of using tqdm because we want to parse the logs with the launcher\n\n    for i, (pt_file, sf_file) in enumerate(zip(pt_files, sf_files)):\n        # Skip blacklisted files\n        if (\n            \"arguments\" in pt_file.name\n            or \"args\" in pt_file.name\n            or \"training\" in pt_file.name\n        ):\n            continue\n\n        start = datetime.datetime.now()\n        convert_file(pt_file, sf_file, discard_names)\n        elapsed = datetime.datetime.now() - start\n        logger.info(f\"Convert: [{i + 1}/{N}] -- Took: {elapsed}\")\n", "path": "server/text_generation_server/utils/convert.py"}]}