Repository: huggingface/text-generation-inference

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.


Discussion:
Could you share the name fo the affected model ?

It's simply a matter of weight naming, the conversion method here is a bit crude (but very efficient memory wise), we just need some weights renaming.
Hi @Narsil,

Just wanted to add some more detail as I have been dealing with this issue as well. If I load and save one of the falcon models:
```
model = AutoModelForCausalLM.from_pretrained("tiiuae/falcon-7b-instrusct", trust_remote_code=True)
model.save_pretrained("/path/to/model")
```

Then copy over the tokenizer and use that saved model to start the text-generation-inference server:
```
docker run --gpus '"device=2,3"' --shm-size 1g -p 8080:80 -v /data:/data ghcr.io/huggingface/text-generation-inference:0.9 --model-id /path/to/model --num-shard 1 --max-input-length 9000 --max-total-tokens 10000 --max-best-of 8 --trust-remote-code --max-batch-prefill-tokens 9000
```

When transformers version = `4.30.2` I get an error that looks something like this:
```
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 78, in serve
    server.serve(

  File "/opt/conda/lib/python3.9/site-packages/text_generation_server/server.py", line 166, in serve
    asyncio.run(

  File "/opt/conda/lib/python3.9/asyncio/runners.py", line 44, in run
    return loop.run_until_complete(main)

  File "/opt/conda/lib/python3.9/asyncio/base_events.py", line 647, in run_until_complete
    return future.result()

  File "/opt/conda/lib/python3.9/site-packages/text_generation_server/server.py", line 133, in serve_inner
    model = get_model(

  File "/opt/conda/lib/python3.9/site-packages/text_generation_server/models/__init__.py", line 253, in get_model
    return FlashRWSharded(

  File "/opt/conda/lib/python3.9/site-packages/text_generation_server/models/flash_rw.py", line 56, in __init__
    model = FlashRWForCausalLM(config, weights)

  File "/opt/conda/lib/python3.9/site-packages/text_generation_server/models/custom_modeling/flash_rw_modeling.py", line 628, in __init__
    self.transformer = FlashRWModel(config, weights)

  File "/opt/conda/lib/python3.9/site-packages/text_generation_server/models/custom_modeling/flash_rw_modeling.py", line 553, in __init__
    self.word_embeddings = TensorParallelEmbedding(

  File "/opt/conda/lib/python3.9/site-packages/text_generation_server/utils/layers.py", line 280, in __init__
    weight = weights.get_sharded(f"{prefix}.weight", dim=0)

  File "/opt/conda/lib/python3.9/site-packages/text_generation_server/utils/weights.py", line 73, in get_sharded
    filename, tensor_name = self.get_filename(tensor_name)

  File "/opt/conda/lib/python3.9/site-packages/text_generation_server/utils/weights.py", line 49, in get_filename
    raise RuntimeError(f"weight {tensor_name} does not exist")

RuntimeError: weight transformer.word_embeddings.weight does not exist
```

However using transformers version = `4.27.4` to load and save the model allows the tgi server to start up as expected from the locally saved weights.

I have also tried using the PR you've linked above but that does not solve the issue which I think is related to how the model weights are saved rather than how they are converted to safetensors. Maybe this issue belongs in the transformers repo?
Can you try with 

```
ghcr.io/huggingface/text-generation-inference:sha-b4024ed
```

This is what got modified 4 days ago (We changed a bit how we choose the actual tensors to copy).
Yeah just tried that now it runs into the same issue. Here is the full output when pointing to the model saved with `4.30.2`:

```
2023-07-10T18:06:11.828477Z  INFO text_generation_launcher: Args { model_id: "/data/test-models/falcon-7b-instruct", revision: None, sharded: None, num_shard: Some(1), quantize: None, dtype: None, trust_remote_code: true, max_concurrent_requests: 128, max_best_of: 8, max_stop_sequences: 4, max_input_length: 9000, max_total_tokens: 10000, waiting_served_ratio: 1.2, max_batch_prefill_tokens: 9000, max_batch_total_tokens: 16000, max_waiting_tokens: 20, hostname: "dc570373c710", 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-10T18:06:11.828614Z  INFO text_generation_launcher: Starting download process.
2023-07-10T18:06:13.292872Z  WARN download: text_generation_launcher: No safetensors weights found for model /data/test-models/falcon-7b-instruct at revision None. Converting PyTorch weights to safetensors.

2023-07-10T18:06:27.323366Z  INFO download: text_generation_launcher: Convert: [1/3] -- Took: 0:00:13.984885

2023-07-10T18:06:38.581181Z  INFO download: text_generation_launcher: Convert: [2/3] -- Took: 0:00:11.257167

2023-07-10T18:06:53.949943Z  INFO download: text_generation_launcher: Convert: [3/3] -- Took: 0:00:15.368202

2023-07-10T18:06:54.369737Z  INFO text_generation_launcher: Successfully downloaded weights.
2023-07-10T18:06:54.369814Z  WARN text_generation_launcher: `trust_remote_code` is set. Trusting that model `/data/test-models/falcon-7b-instruct` do not contain malicious code.
2023-07-10T18:06:54.369821Z  WARN text_generation_launcher: Explicitly passing a `revision` is encouraged when loading a model with custom code to ensure no malicious code has been contributed in a newer revision.
2023-07-10T18:06:54.370572Z  INFO text_generation_launcher: Starting shard 0
2023-07-10T18:06:56.617177Z  WARN shard-manager: text_generation_launcher: We're not using custom kernels.
 rank=0
2023-07-10T18:06:56.818961Z ERROR shard-manager: text_generation_launcher: Error when initializing model
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/typer/main.py", line 311, in __call__
    return get_command(self)(*args, **kwargs)
  File "/opt/conda/lib/python3.9/site-packages/click/core.py", line 1130, in __call__
    return self.main(*args, **kwargs)
  File "/opt/conda/lib/python3.9/site-packages/typer/core.py", line 778, in main
    return _main(
  File "/opt/conda/lib/python3.9/site-packages/typer/core.py", line 216, in _main
    rv = self.invoke(ctx)
  File "/opt/conda/lib/python3.9/site-packages/click/core.py", line 1657, in invoke
    return _process_result(sub_ctx.command.invoke(sub_ctx))
  File "/opt/conda/lib/python3.9/site-packages/click/core.py", line 1404, in invoke
    return ctx.invoke(self.callback, **ctx.params)
  File "/opt/conda/lib/python3.9/site-packages/click/core.py", line 760, in invoke
    return __callback(*args, **kwargs)
  File "/opt/conda/lib/python3.9/site-packages/typer/main.py", line 683, in wrapper
    return callback(**use_params)  # type: ignore
  File "/opt/conda/lib/python3.9/site-packages/text_generation_server/cli.py", line 78, in serve
    server.serve(
  File "/opt/conda/lib/python3.9/site-packages/text_generation_server/server.py", line 166, in serve
    asyncio.run(
  File "/opt/conda/lib/python3.9/asyncio/runners.py", line 44, in run
    return loop.run_until_complete(main)
  File "/opt/conda/lib/python3.9/asyncio/base_events.py", line 634, in run_until_complete
    self.run_forever()
  File "/opt/conda/lib/python3.9/asyncio/base_events.py", line 601, in run_forever
    self._run_once()
  File "/opt/conda/lib/python3.9/asyncio/base_events.py", line 1905, in _run_once
    handle._run()
  File "/opt/conda/lib/python3.9/asyncio/events.py", line 80, in _run
    self._context.run(self._callback, *self._args)
> File "/opt/conda/lib/python3.9/site-packages/text_generation_server/server.py", line 133, in serve_inner
    model = get_model(
  File "/opt/conda/lib/python3.9/site-packages/text_generation_server/models/__init__.py", line 253, in get_model
    return FlashRWSharded(
  File "/opt/conda/lib/python3.9/site-packages/text_generation_server/models/flash_rw.py", line 56, in __init__
    model = FlashRWForCausalLM(config, weights)
  File "/opt/conda/lib/python3.9/site-packages/text_generation_server/models/custom_modeling/flash_rw_modeling.py", line 634, in __init__
    self.transformer = FlashRWModel(config, weights)
  File "/opt/conda/lib/python3.9/site-packages/text_generation_server/models/custom_modeling/flash_rw_modeling.py", line 559, in __init__
    self.word_embeddings = TensorParallelEmbedding(
  File "/opt/conda/lib/python3.9/site-packages/text_generation_server/utils/layers.py", line 280, in __init__
    weight = weights.get_sharded(f"{prefix}.weight", dim=0)
  File "/opt/conda/lib/python3.9/site-packages/text_generation_server/utils/weights.py", line 73, in get_sharded
    filename, tensor_name = self.get_filename(tensor_name)
  File "/opt/conda/lib/python3.9/site-packages/text_generation_server/utils/weights.py", line 49, in get_filename
    raise RuntimeError(f"weight {tensor_name} does not exist")
RuntimeError: weight transformer.word_embeddings.weight does not exist
```

compared to pointing at the model saved with `4.27.4`:

```
2023-07-10T18:09:53.995176Z  INFO text_generation_launcher: Args { model_id: "/data/test-models/falcon-7b-instruct-oldhf", revision: None, sharded: None, num_shard: Some(1), quantize: None, dtype: None, trust_remote_code: true, max_concurrent_requests: 128, max_best_of: 8, max_stop_sequences: 4, max_input_length: 9000, max_total_tokens: 10000, waiting_served_ratio: 1.2, max_batch_prefill_tokens: 9000, max_batch_total_tokens: 16000, max_waiting_tokens: 20, hostname: "9ae087ae066c", 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-10T18:09:53.995278Z  INFO text_generation_launcher: Starting download process.
2023-07-10T18:09:55.481628Z  WARN download: text_generation_launcher: No safetensors weights found for model /data/test-models/falcon-7b-instruct-oldhf at revision None. Converting PyTorch weights to safetensors.

2023-07-10T18:10:09.814556Z  INFO download: text_generation_launcher: Convert: [1/3] -- Took: 0:00:14.287668

2023-07-10T18:10:22.541350Z  INFO download: text_generation_launcher: Convert: [2/3] -- Took: 0:00:12.726082

2023-07-10T18:10:37.206531Z  INFO download: text_generation_launcher: Convert: [3/3] -- Took: 0:00:14.664828

2023-07-10T18:10:37.940380Z  INFO text_generation_launcher: Successfully downloaded weights.
2023-07-10T18:10:37.940456Z  WARN text_generation_launcher: `trust_remote_code` is set. Trusting that model `/data/test-models/falcon-7b-instruct-oldhf` do not contain malicious code.
2023-07-10T18:10:37.940464Z  WARN text_generation_launcher: Explicitly passing a `revision` is encouraged when loading a model with custom code to ensure no malicious code has been contributed in a newer revision.
2023-07-10T18:10:37.940728Z  INFO text_generation_launcher: Starting shard 0
2023-07-10T18:10:40.184657Z  WARN shard-manager: text_generation_launcher: We're not using custom kernels.
 rank=0
2023-07-10T18:10:47.951716Z  INFO text_generation_launcher: Waiting for shard 0 to be ready...
2023-07-10T18:10:50.542364Z  INFO shard-manager: text_generation_launcher: Server started at unix:///tmp/text-generation-server-0
 rank=0
2023-07-10T18:10:50.554101Z  INFO text_generation_launcher: Shard 0 ready in 12.611513023s
2023-07-10T18:10:50.650702Z  INFO text_generation_launcher: Starting Webserver
2023-07-10T18:10:50.728175Z  WARN text_generation_router: router/src/main.rs:186: no pipeline tag found for model /data/test-models/falcon-7b-instruct-oldhf
2023-07-10T18:10:50.733065Z  INFO text_generation_router: router/src/main.rs:205: Warming up model
2023-07-10T18:10:53.408244Z  INFO text_generation_router: router/src/main.rs:214: Connected
2023-07-10T18:10:53.408264Z  WARN text_generation_router: router/src/main.rs:219: Invalid hostname, defaulting to 0.0.0.0
```
Another thing I have noticed (though not sure if it is at all related), the newer version of transformers does not save the configuration and modeling python files when `save_pretrained` is used which results in an error when trying to load the saved model with `from_pretrained`

Edit: I believe this issue is unrelated I can download those files and place them in the saved model and load it up as expected. I have loaded both models with from pretrained and compared `model.transformer.word_embeddings.weight` and `model.lm_head.weight` and they are identical for both so it would seem that none of the weights are being skipped over when using `save_pretrained`
I just think the previous weights are already saved.

The new PR doesn't fix it because we're still pointing to a `trust_remote_code` model, which doesn't have the heursitics to choose the weights better.

PR incoming.
Yeah they are I had just wanted to confirm that by loading the models directly. 

Great to hear let me know if there's anything I can do to help!
>  the newer version of transformers does not save the configuration and modeling python files when save_pretrained is used which results in an error when trying to load the saved model with from_pretrained

This is not new, you used only `AutoModelFor...from_pretrained(..).save_pretrained(...)`.
Which never saved anything else than the model.

You need to do the same with `AutoConfig` and `AutoTokenizer` if you want those files in. Or do `pipe = pipeline(...); pipe.save_pretrained(...")` 

Resolving patch:
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 @@ def __init__(
 
         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
 


Test patch: