Repository: huggingface/text-generation-inference

Issue:
endless Waiting for shard 0 to be ready... when using only 1 GPU (google/flan-t5-xl)
### System Info

ubuntu 20.04, CUDA 11.7, pytorch 2.0.1
text-generation-launcher docker latest


### Information

- [X] Docker
- [ ] The CLI directly

### Tasks

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

### Reproduction

1. Download the huggingface google/flan-t5-xl model to your local drive (for example, _$PWD/models/flan-t5-xl_)
2.  export CUDA_VISIBLE_DEVICES=0
3. docker run --gpus 0 -p 8080:80 -v $PWD/models/flan-t5-xl:/data ghcr.io/huggingface/text-generation-inference:latest  --model-id google/flan-t5-xl --sharded false  --weights-cache-override /data

I see  Waiting for shard 0 to be ready... for ever. Both GPUs are idle.

**However, if I change to use 2 gpus, with --num-shard 2, the model is loaded and distributed to 2 GPUs.**

### Expected behavior

I suppose that the docker image should work also with 1 GPU without sharding.


Discussion:

Resolving patch:
diff --git a/server/text_generation_server/models/__init__.py b/server/text_generation_server/models/__init__.py
--- a/server/text_generation_server/models/__init__.py
+++ b/server/text_generation_server/models/__init__.py
@@ -230,20 +230,12 @@ def get_model(
         )
 
     elif model_type == "t5":
-        if sharded:
-            return T5Sharded(
-                model_id,
-                revision,
-                quantize=quantize,
-                trust_remote_code=trust_remote_code,
-            )
-        else:
-            return Seq2SeqLM(
-                model_id,
-                revision,
-                quantize=quantize,
-                trust_remote_code=trust_remote_code,
-            )
+        return T5Sharded(
+            model_id,
+            revision,
+            quantize=quantize,
+            trust_remote_code=trust_remote_code,
+        )
 
     if sharded:
         raise ValueError("sharded is not supported for AutoModel")
diff --git a/server/text_generation_server/utils/logits_process.py b/server/text_generation_server/utils/logits_process.py
--- a/server/text_generation_server/utils/logits_process.py
+++ b/server/text_generation_server/utils/logits_process.py
@@ -42,25 +42,31 @@ def __init__(
         self.static_next_logprob = None
 
     def __call__(self, scores):
-        if self.cuda_graph is None:
-            self.static_scores = scores
-            self.cuda_graph = torch.cuda.CUDAGraph()
-
-            with torch.cuda.graph(self.cuda_graph, pool=mempool):
-                local_scores = self.static_scores
-                for warper in self.warpers:
-                    local_scores = warper(None, local_scores)
-
-                self.static_warped_scores = local_scores
-                # Compute logprobs
-                self.static_next_logprob = torch.log_softmax(
-                    self.static_warped_scores, -1
-                )
-
-        self.static_scores.copy_(scores)
-        self.cuda_graph.replay()
-
-        return self.static_warped_scores, self.static_next_logprob
+        if torch.cuda.is_available():
+            if self.cuda_graph is None:
+                self.static_scores = scores
+                self.cuda_graph = torch.cuda.CUDAGraph()
+
+                with torch.cuda.graph(self.cuda_graph, pool=mempool):
+                    local_scores = self.static_scores
+                    for warper in self.warpers:
+                        local_scores = warper(None, local_scores)
+
+                    self.static_warped_scores = local_scores
+                    # Compute logprobs
+                    self.static_next_logprob = torch.log_softmax(
+                        self.static_warped_scores, -1
+                    )
+
+            self.static_scores.copy_(scores)
+            self.cuda_graph.replay()
+
+            return self.static_warped_scores, self.static_next_logprob
+
+        # CPU branch
+        for warper in self.warpers:
+            scores = warper(None, scores)
+        return scores, torch.log_softmax(scores, -1)
 
 
 @lru_cache(10)


Test patch: