Repository: huggingface/text-generation-inference

Issue:
exllama fails with sharded model
### System Info

image: sha-1da642b (released 2 days ago latest branch)

setup: a server with 2 RTX A6000 GPUs on it
256GB RAM
32 CPU cores

### Information

- [X] Docker
- [ ] The CLI directly

### Tasks

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

### Reproduction

```bash
#!/bin/bash

export GPTQ_BITS=4
export GPTQ_GROUPSIZE=128
export CUDA_VISIBLE_DEVICES=0,1

text-generation-launcher --model-id TheBloke/Wizard-Vicuna-30B-Uncensored-GPTQ \
                         --revision gptq-4bit-128g-actorder_True \
                         --quantize gptq \
                         --port 8080 \
                         --max-input-length 512 \
                         --max-best-of 4 \
                         --max-total-tokens 600 \
                         --max-batch-prefill-tokens 2048
```

OUTPUT:
2023-07-23T16:48:27.611908Z  INFO text_generation_launcher: Sharding model on 2 processes
2023-07-23T16:48:27.611982Z  INFO download: text_generation_launcher: Starting download process.
2023-07-23T16:48:28.856188Z  INFO text_generation_launcher: Files are already present on the host. Skipping download.

2023-07-23T16:48:29.113858Z  INFO download: text_generation_launcher: Successfully downloaded weights.
2023-07-23T16:48:29.114100Z  INFO shard-manager: text_generation_launcher: Starting shard rank=0
2023-07-23T16:48:29.114562Z  INFO shard-manager: text_generation_launcher: Starting shard rank=1
2023-07-23T16:48:39.123016Z  INFO shard-manager: text_generation_launcher: Waiting for shard to be ready... rank=1
2023-07-23T16:48:39.123228Z  INFO shard-manager: text_generation_launcher: Waiting for shard to be ready... rank=0
2023-07-23T16:48:47.629123Z ERROR shard-manager: text_generation_launcher: Shard complete standard error output:

You are using a model of type llama to instantiate a model of type . This is not supported for all configurations of models and can yield errors.
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 180, 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 150, in serve_inner
    create_exllama_buffers()

  File "/opt/conda/lib/python3.9/site-packages/text_generation_server/utils/gptq/exllama.py", line 52, in create_exllama_buffers
    prepare_buffers(DEVICE, temp_state, temp_dq)

TypeError: prepare_buffers(): incompatible function arguments. The following argument types are supported:
    1. (arg0: torch.device, arg1: torch.Tensor, arg2: torch.Tensor) -> None

Invoked with: None, tensor([[0.]], dtype=torch.float16), tensor([[0.]], dtype=torch.float16)
 rank=1
2023-07-23T16:48:47.727702Z ERROR text_generation_launcher: Shard 1 failed to start
2023-07-23T16:48:47.727724Z  INFO text_generation_launcher: Shutting down shards
2023-07-23T16:48:47.730006Z ERROR shard-manager: text_generation_launcher: Shard complete standard error output:

You are using a model of type llama to instantiate a model of type . This is not supported for all configurations of models and can yield errors.
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 180, 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 150, in serve_inner
    create_exllama_buffers()

  File "/opt/conda/lib/python3.9/site-packages/text_generation_server/utils/gptq/exllama.py", line 52, in create_exllama_buffers
    prepare_buffers(DEVICE, temp_state, temp_dq)

TypeError: prepare_buffers(): incompatible function arguments. The following argument types are supported:
    1. (arg0: torch.device, arg1: torch.Tensor, arg2: torch.Tensor) -> None

Invoked with: None, tensor([[0.]], dtype=torch.float16), tensor([[0.]], dtype=torch.float16)
 rank=0
Error: ShardCannotStart

### Expected behavior

It works fine when I remove the CUDA_VISIBLE_DEVICES=0,1 and when it then runs on a single process.


Discussion:
Can confirm this is happening on other models as well. Tried various models including FreeWilly2 GPTQ by Bloke and seems like only gptq-4bit-128g-actorder_False branch is able to load and rest all give the error above.

Resolving patch:
diff --git a/server/text_generation_server/server.py b/server/text_generation_server/server.py
--- a/server/text_generation_server/server.py
+++ b/server/text_generation_server/server.py
@@ -105,21 +105,21 @@ async def Decode(self, request, context):
 
 
 def serve(
-    model_id: str,
-    revision: Optional[str],
-    sharded: bool,
-    quantize: Optional[str],
-    dtype: Optional[str],
-    trust_remote_code: bool,
-    uds_path: Path,
-):
-    async def serve_inner(
         model_id: str,
         revision: Optional[str],
-        sharded: bool = False,
-        quantize: Optional[str] = None,
-        dtype: Optional[str] = None,
-        trust_remote_code: bool = False,
+        sharded: bool,
+        quantize: Optional[str],
+        dtype: Optional[str],
+        trust_remote_code: bool,
+        uds_path: Path,
+):
+    async def serve_inner(
+            model_id: str,
+            revision: Optional[str],
+            sharded: bool = False,
+            quantize: Optional[str] = None,
+            dtype: Optional[str] = None,
+            trust_remote_code: bool = False,
     ):
         unix_socket_template = "unix://{}-{}"
         if sharded:
@@ -147,8 +147,10 @@ async def serve_inner(
                 # This will allocate those buffers.
                 from text_generation_server.utils.gptq.exllama import (
                     create_exllama_buffers,
+                    set_device,
                 )
 
+                set_device(model.device)
                 create_exllama_buffers()
             except ImportError:
                 pass
diff --git a/server/text_generation_server/utils/gptq/exllama.py b/server/text_generation_server/utils/gptq/exllama.py
--- a/server/text_generation_server/utils/gptq/exllama.py
+++ b/server/text_generation_server/utils/gptq/exllama.py
@@ -32,9 +32,16 @@ def ext_q4_matmul(x, q4, q4_width):
 TEMP_DQ = None
 
 
+def set_device(device):
+    global DEVICE
+    DEVICE = device
+
+
 def create_exllama_buffers():
     global MAX_DQ, MAX_INNER, ACT_ORDER, DEVICE, TEMP_STATE, TEMP_DQ
 
+    assert DEVICE is not None, "call set_device first"
+
     if ACT_ORDER:
         # TODO: this should be set to rust side `max_total_tokens`, but TGI
         # does not offer an API to expose this variable to python, as this variable


Test patch: