We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Question about sharding / TP
@OlivierDehaene @Narsil is it expected that the output should be the same (or very close) when using the TP implementation for a given model vs non-sharded/single GPU?

Am seeing quite different output, this is for example with flan-ul2 or flan-t5-xxl with 2 GPUs, using float16 for both single and double GPU cases.

This is using a different fork of the code - I'm still investigating and will also try with the latest from the main branch of this repo as-is, but would be very helpful to know generally what you observe / what's expected.

--- 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/t5.py`
Content:
```
1 import torch
2 import torch.distributed
3 
4 from typing import List, Optional, Tuple
5 
6 from accelerate import init_empty_weights
7 from safetensors import safe_open
8 from transformers import (
9     AutoTokenizer,
10     AutoModelForSeq2SeqLM,
11     AutoConfig,
12 )
13 
14 from text_generation_server.models import Seq2SeqLM
15 from text_generation_server.utils import (
16     initialize_torch_distributed,
17     weight_files,
18 )
19 from transformers.models.t5.parallel_layers import (
20     TensorParallelRowLinear,
21     TensorParallelColumnLinear,
22     TensorParallelEmbedding,
23 )
24 
25 HAS_BITS_AND_BYTES = True
26 try:
27     import bitsandbytes as bnb
28     from bitsandbytes.nn import Int8Params
29 except ImportError as e:
30     HAS_BITS_AND_BYTES = False
31 
32 
33 class T5Sharded(Seq2SeqLM):
34     def __init__(
35         self,
36         model_id: str,
37         revision: Optional[str] = None,
38         quantize: Optional[str] = None,
39     ):
40         self.process_group, rank, world_size = initialize_torch_distributed()
41         if torch.cuda.is_available():
42             device = torch.device(f"cuda:{rank}")
43             dtype = torch.float16
44         else:
45             device = torch.device("cpu")
46             dtype = torch.float32
47 
48         tokenizer = AutoTokenizer.from_pretrained(
49             model_id, revision=revision, padding_side="left", truncation_side="left"
50         )
51 
52         config = AutoConfig.from_pretrained(
53             model_id, revision=revision, tp_parallel=True
54         )
55         tokenizer.bos_token_id = config.decoder_start_token_id
56 
57         torch.distributed.barrier(group=self.process_group)
58         filenames = weight_files(model_id, revision=revision, extension=".safetensors")
59 
60         with init_empty_weights():
61             model = AutoModelForSeq2SeqLM.from_config(config)
62 
63         torch.distributed.barrier(group=self.process_group)
64         self.load_weights(
65             model,
66             filenames,
67             quantize=quantize,
68             device=device,
69             dtype=dtype,
70             rank=rank,
71             world_size=world_size,
72         )
73         torch.distributed.barrier(group=self.process_group)
74         super(Seq2SeqLM, self).__init__(
75             model=model,
76             tokenizer=tokenizer,
77             requires_padding=True,
78             dtype=dtype,
79             device=device,
80             rank=rank,
81             world_size=world_size,
82         )
83 
84     @staticmethod
85     def load_weights(
86         model,
87         filenames: List[str],
88         quantize: Optional[str],
89         device: torch.device,
90         dtype: torch.dtype,
91         rank: int,
92         world_size: int,
93     ):
94         parameters = dict(model.named_parameters())
95         for file in filenames:
96             with safe_open(
97                 file, framework="pt", device=str(device) if quantize is None else "cpu"
98             ) as f:
99                 for name in f.keys():
100                     module_name, param_name = name.rsplit(".", 1)
101                     module = model.get_submodule(module_name)
102 
103                     current_parameter_tensor = parameters.get(name, None)
104 
105                     slice_ = f.get_slice(name)
106 
107                     if isinstance(module, TensorParallelColumnLinear):
108                         size = slice_.get_shape()[0]
109                         block_size = size // world_size
110                         start = rank * block_size
111                         stop = (rank + 1) * block_size
112                         tensor = slice_[start:stop]
113                     elif isinstance(module, TensorParallelRowLinear):
114                         if param_name == "weight":
115                             size = slice_.get_shape()[1]
116                             block_size = size // world_size
117                             start = rank * block_size
118                             stop = (rank + 1) * block_size
119                             tensor = slice_[:, start:stop]
120                         else:
121                             tensor = slice_[:]
122                             # XXX: Hack for Rowlinear to add the bias only once.
123                             if rank != 0:
124                                 tensor = torch.zeros_like(tensor)
125                     elif isinstance(module, TensorParallelEmbedding):
126                         size = slice_.get_shape()[0]
127                         block_size = size // world_size
128                         start = rank * block_size
129                         stop = (rank + 1) * block_size
130                         tensor = slice_[start:stop]
131                     elif name == "lm_head.weight":
132                         size = slice_.get_shape()[0]
133                         block_size = size // world_size
134                         start = rank * block_size
135                         stop = (rank + 1) * block_size
136                         tensor = slice_[start:stop]
137                     elif "relative_attention_bias.weight" in name:
138                         size = slice_.get_shape()[1]
139                         block_size = size // world_size
140                         start = rank * block_size
141                         stop = (rank + 1) * block_size
142                         tensor = slice_[:, start:stop]
143                     else:
144                         try:
145                             tensor = slice_[:]
146                         except:
147                             tensor = f.get_tensor(name)
148 
149                     if (
150                         current_parameter_tensor is not None
151                         and current_parameter_tensor.shape != tensor.shape
152                     ):
153                         raise ValueError(
154                             f"Name {name} -- Current {current_parameter_tensor.shape} and got {tensor.shape}"
155                         )
156 
157                     tensor = tensor.contiguous().to(dtype)
158 
159                     if quantize == "bitsandbytes":
160                         if not HAS_BITS_AND_BYTES:
161                             raise ImportError(
162                                 "bitsandbytes is not available on your machine either because it is not installed "
163                                 "or you don't have a GPU.\n"
164                                 "You can install it with `pip install bitsandbytes`."
165                             )
166 
167                         if (
168                             type(module)
169                             in [TensorParallelRowLinear, TensorParallelColumnLinear]
170                             and param_name == "weight"
171                         ):
172                             tensor = Int8Params(
173                                 tensor,
174                                 has_fp16_weights=False,
175                                 requires_grad=False,
176                             ).to(device)
177                             state = bnb.MatmulLtState()
178                             state.threshold = 6.0
179                             state.has_fp16_weights = False
180                             state.memory_efficient_backward = False
181                             state.use_pool = True
182                             state.CB = tensor.CB
183                             state.SCB = tensor.SCB
184                             tensor.CB = None
185                             tensor.SCB = None
186 
187                             def replace_linear(state):
188                                 def linear(input, weight, bias):
189                                     out = bnb.matmul(
190                                         input,
191                                         weight,
192                                         state=state,
193                                         threshold=state.threshold,
194                                         bias=bias,
195                                     )
196 
197                                     if state.CB is not None:
198                                         # we converted 8-bit row major to turing/ampere format
199                                         # in the first inference pass
200                                         # we no longer need the row-major weight
201                                         del state.CB
202                                         weight.data = state.CxB
203 
204                                     return out
205 
206                                 return linear
207 
208                             module.linear = replace_linear(state)
209 
210                         elif quantize == "gptq":
211                             raise NotImplementedError(
212                                 "`gptq` is not implemented for now"
213                             )
214                         elif quantize is None:
215                             tensor = tensor.to(device)
216                         else:
217                             raise ValueError(f"Unexpected quantize `{quantize}`")
218 
219                     if current_parameter_tensor is not None:
220                         module._parameters[param_name] = tensor
221                     else:
222                         module._buffers[param_name] = tensor
223 
224     def forward(
225         self,
226         input_ids,
227         attention_mask,
228         decoder_input_ids,
229         decoder_attention_mask: Optional,
230         encoder_last_hidden_state: Optional,
231         past_key_values: Optional = None,
232     ) -> Tuple[
233         torch.Tensor,
234         torch.Tensor,
235         List[Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]],
236     ]:
237         # Model Forward
238         outputs = self.model.forward(
239             input_ids=input_ids,
240             attention_mask=attention_mask,
241             decoder_input_ids=decoder_input_ids,
242             decoder_attention_mask=decoder_attention_mask,
243             encoder_outputs=encoder_last_hidden_state,
244             past_key_values=past_key_values,
245             use_cache=True,
246         )
247 
248         # Logits are sharded, so we need to gather them
249         logits = [torch.empty_like(outputs.logits) for _ in range(self.world_size)]
250         torch.distributed.all_gather(logits, outputs.logits, group=self.process_group)
251         logits = torch.cat(logits, dim=2)
252 
253         return (
254             logits,
255             outputs.encoder_last_hidden_state,
256             outputs.past_key_values,
257         )
258
```
--- 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/models/t5.py b/server/text_generation_server/models/t5.py
--- a/server/text_generation_server/models/t5.py
+++ b/server/text_generation_server/models/t5.py
@@ -40,7 +40,7 @@
         self.process_group, rank, world_size = initialize_torch_distributed()
         if torch.cuda.is_available():
             device = torch.device(f"cuda:{rank}")
-            dtype = torch.float16
+            dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float32
         else:
             device = torch.device("cpu")
             dtype = torch.float32

{"golden_diff": "diff --git a/server/text_generation_server/models/t5.py b/server/text_generation_server/models/t5.py\n--- a/server/text_generation_server/models/t5.py\n+++ b/server/text_generation_server/models/t5.py\n@@ -40,7 +40,7 @@\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\n+            dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float32\n         else:\n             device = torch.device(\"cpu\")\n             dtype = torch.float32\n", "issue": "Question about sharding / TP\n@OlivierDehaene @Narsil is it expected that the output should be the same (or very close) when using the TP implementation for a given model vs non-sharded/single GPU?\r\n\r\nAm seeing quite different output, this is for example with flan-ul2 or flan-t5-xxl with 2 GPUs, using float16 for both single and double GPU cases.\r\n\r\nThis is using a different fork of the code - I'm still investigating and will also try with the latest from the main branch of this repo as-is, but would be very helpful to know generally what you observe / what's expected.\n", "before_files": [{"content": "import torch\nimport torch.distributed\n\nfrom typing import List, Optional, Tuple\n\nfrom accelerate import init_empty_weights\nfrom safetensors import safe_open\nfrom transformers import (\n    AutoTokenizer,\n    AutoModelForSeq2SeqLM,\n    AutoConfig,\n)\n\nfrom text_generation_server.models import Seq2SeqLM\nfrom text_generation_server.utils import (\n    initialize_torch_distributed,\n    weight_files,\n)\nfrom transformers.models.t5.parallel_layers import (\n    TensorParallelRowLinear,\n    TensorParallelColumnLinear,\n    TensorParallelEmbedding,\n)\n\nHAS_BITS_AND_BYTES = True\ntry:\n    import bitsandbytes as bnb\n    from bitsandbytes.nn import Int8Params\nexcept ImportError as e:\n    HAS_BITS_AND_BYTES = False\n\n\nclass T5Sharded(Seq2SeqLM):\n    def __init__(\n        self,\n        model_id: str,\n        revision: Optional[str] = None,\n        quantize: Optional[str] = None,\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\n        else:\n            device = torch.device(\"cpu\")\n            dtype = torch.float32\n\n        tokenizer = AutoTokenizer.from_pretrained(\n            model_id, revision=revision, padding_side=\"left\", truncation_side=\"left\"\n        )\n\n        config = AutoConfig.from_pretrained(\n            model_id, revision=revision, tp_parallel=True\n        )\n        tokenizer.bos_token_id = config.decoder_start_token_id\n\n        torch.distributed.barrier(group=self.process_group)\n        filenames = weight_files(model_id, revision=revision, extension=\".safetensors\")\n\n        with init_empty_weights():\n            model = AutoModelForSeq2SeqLM.from_config(config)\n\n        torch.distributed.barrier(group=self.process_group)\n        self.load_weights(\n            model,\n            filenames,\n            quantize=quantize,\n            device=device,\n            dtype=dtype,\n            rank=rank,\n            world_size=world_size,\n        )\n        torch.distributed.barrier(group=self.process_group)\n        super(Seq2SeqLM, self).__init__(\n            model=model,\n            tokenizer=tokenizer,\n            requires_padding=True,\n            dtype=dtype,\n            device=device,\n            rank=rank,\n            world_size=world_size,\n        )\n\n    @staticmethod\n    def load_weights(\n        model,\n        filenames: List[str],\n        quantize: Optional[str],\n        device: torch.device,\n        dtype: torch.dtype,\n        rank: int,\n        world_size: int,\n    ):\n        parameters = dict(model.named_parameters())\n        for file in filenames:\n            with safe_open(\n                file, framework=\"pt\", device=str(device) if quantize is None else \"cpu\"\n            ) as f:\n                for name in f.keys():\n                    module_name, param_name = name.rsplit(\".\", 1)\n                    module = model.get_submodule(module_name)\n\n                    current_parameter_tensor = parameters.get(name, None)\n\n                    slice_ = f.get_slice(name)\n\n                    if isinstance(module, TensorParallelColumnLinear):\n                        size = slice_.get_shape()[0]\n                        block_size = size // world_size\n                        start = rank * block_size\n                        stop = (rank + 1) * block_size\n                        tensor = slice_[start:stop]\n                    elif isinstance(module, TensorParallelRowLinear):\n                        if param_name == \"weight\":\n                            size = slice_.get_shape()[1]\n                            block_size = size // world_size\n                            start = rank * block_size\n                            stop = (rank + 1) * block_size\n                            tensor = slice_[:, start:stop]\n                        else:\n                            tensor = slice_[:]\n                            # XXX: Hack for Rowlinear to add the bias only once.\n                            if rank != 0:\n                                tensor = torch.zeros_like(tensor)\n                    elif isinstance(module, TensorParallelEmbedding):\n                        size = slice_.get_shape()[0]\n                        block_size = size // world_size\n                        start = rank * block_size\n                        stop = (rank + 1) * block_size\n                        tensor = slice_[start:stop]\n                    elif name == \"lm_head.weight\":\n                        size = slice_.get_shape()[0]\n                        block_size = size // world_size\n                        start = rank * block_size\n                        stop = (rank + 1) * block_size\n                        tensor = slice_[start:stop]\n                    elif \"relative_attention_bias.weight\" in name:\n                        size = slice_.get_shape()[1]\n                        block_size = size // world_size\n                        start = rank * block_size\n                        stop = (rank + 1) * block_size\n                        tensor = slice_[:, start:stop]\n                    else:\n                        try:\n                            tensor = slice_[:]\n                        except:\n                            tensor = f.get_tensor(name)\n\n                    if (\n                        current_parameter_tensor is not None\n                        and current_parameter_tensor.shape != tensor.shape\n                    ):\n                        raise ValueError(\n                            f\"Name {name} -- Current {current_parameter_tensor.shape} and got {tensor.shape}\"\n                        )\n\n                    tensor = tensor.contiguous().to(dtype)\n\n                    if quantize == \"bitsandbytes\":\n                        if not HAS_BITS_AND_BYTES:\n                            raise ImportError(\n                                \"bitsandbytes is not available on your machine either because it is not installed \"\n                                \"or you don't have a GPU.\\n\"\n                                \"You can install it with `pip install bitsandbytes`.\"\n                            )\n\n                        if (\n                            type(module)\n                            in [TensorParallelRowLinear, TensorParallelColumnLinear]\n                            and param_name == \"weight\"\n                        ):\n                            tensor = Int8Params(\n                                tensor,\n                                has_fp16_weights=False,\n                                requires_grad=False,\n                            ).to(device)\n                            state = bnb.MatmulLtState()\n                            state.threshold = 6.0\n                            state.has_fp16_weights = False\n                            state.memory_efficient_backward = False\n                            state.use_pool = True\n                            state.CB = tensor.CB\n                            state.SCB = tensor.SCB\n                            tensor.CB = None\n                            tensor.SCB = None\n\n                            def replace_linear(state):\n                                def linear(input, weight, bias):\n                                    out = bnb.matmul(\n                                        input,\n                                        weight,\n                                        state=state,\n                                        threshold=state.threshold,\n                                        bias=bias,\n                                    )\n\n                                    if state.CB is not None:\n                                        # we converted 8-bit row major to turing/ampere format\n                                        # in the first inference pass\n                                        # we no longer need the row-major weight\n                                        del state.CB\n                                        weight.data = state.CxB\n\n                                    return out\n\n                                return linear\n\n                            module.linear = replace_linear(state)\n\n                        elif quantize == \"gptq\":\n                            raise NotImplementedError(\n                                \"`gptq` is not implemented for now\"\n                            )\n                        elif quantize is None:\n                            tensor = tensor.to(device)\n                        else:\n                            raise ValueError(f\"Unexpected quantize `{quantize}`\")\n\n                    if current_parameter_tensor is not None:\n                        module._parameters[param_name] = tensor\n                    else:\n                        module._buffers[param_name] = tensor\n\n    def forward(\n        self,\n        input_ids,\n        attention_mask,\n        decoder_input_ids,\n        decoder_attention_mask: Optional,\n        encoder_last_hidden_state: Optional,\n        past_key_values: Optional = None,\n    ) -> Tuple[\n        torch.Tensor,\n        torch.Tensor,\n        List[Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]],\n    ]:\n        # Model Forward\n        outputs = self.model.forward(\n            input_ids=input_ids,\n            attention_mask=attention_mask,\n            decoder_input_ids=decoder_input_ids,\n            decoder_attention_mask=decoder_attention_mask,\n            encoder_outputs=encoder_last_hidden_state,\n            past_key_values=past_key_values,\n            use_cache=True,\n        )\n\n        # Logits are sharded, so we need to gather them\n        logits = [torch.empty_like(outputs.logits) for _ in range(self.world_size)]\n        torch.distributed.all_gather(logits, outputs.logits, group=self.process_group)\n        logits = torch.cat(logits, dim=2)\n\n        return (\n            logits,\n            outputs.encoder_last_hidden_state,\n            outputs.past_key_values,\n        )\n", "path": "server/text_generation_server/models/t5.py"}], "after_files": [{"content": "import torch\nimport torch.distributed\n\nfrom typing import List, Optional, Tuple\n\nfrom accelerate import init_empty_weights\nfrom safetensors import safe_open\nfrom transformers import (\n    AutoTokenizer,\n    AutoModelForSeq2SeqLM,\n    AutoConfig,\n)\n\nfrom text_generation_server.models import Seq2SeqLM\nfrom text_generation_server.utils import (\n    initialize_torch_distributed,\n    weight_files,\n)\nfrom transformers.models.t5.parallel_layers import (\n    TensorParallelRowLinear,\n    TensorParallelColumnLinear,\n    TensorParallelEmbedding,\n)\n\nHAS_BITS_AND_BYTES = True\ntry:\n    import bitsandbytes as bnb\n    from bitsandbytes.nn import Int8Params\nexcept ImportError as e:\n    HAS_BITS_AND_BYTES = False\n\n\nclass T5Sharded(Seq2SeqLM):\n    def __init__(\n        self,\n        model_id: str,\n        revision: Optional[str] = None,\n        quantize: Optional[str] = None,\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.bfloat16 if torch.cuda.is_bf16_supported() else torch.float32\n        else:\n            device = torch.device(\"cpu\")\n            dtype = torch.float32\n\n        tokenizer = AutoTokenizer.from_pretrained(\n            model_id, revision=revision, padding_side=\"left\", truncation_side=\"left\"\n        )\n\n        config = AutoConfig.from_pretrained(\n            model_id, revision=revision, tp_parallel=True\n        )\n        tokenizer.bos_token_id = config.decoder_start_token_id\n\n        torch.distributed.barrier(group=self.process_group)\n        filenames = weight_files(model_id, revision=revision, extension=\".safetensors\")\n\n        with init_empty_weights():\n            model = AutoModelForSeq2SeqLM.from_config(config)\n\n        torch.distributed.barrier(group=self.process_group)\n        self.load_weights(\n            model,\n            filenames,\n            quantize=quantize,\n            device=device,\n            dtype=dtype,\n            rank=rank,\n            world_size=world_size,\n        )\n        torch.distributed.barrier(group=self.process_group)\n        super(Seq2SeqLM, self).__init__(\n            model=model,\n            tokenizer=tokenizer,\n            requires_padding=True,\n            dtype=dtype,\n            device=device,\n            rank=rank,\n            world_size=world_size,\n        )\n\n    @staticmethod\n    def load_weights(\n        model,\n        filenames: List[str],\n        quantize: Optional[str],\n        device: torch.device,\n        dtype: torch.dtype,\n        rank: int,\n        world_size: int,\n    ):\n        parameters = dict(model.named_parameters())\n        for file in filenames:\n            with safe_open(\n                file, framework=\"pt\", device=str(device) if quantize is None else \"cpu\"\n            ) as f:\n                for name in f.keys():\n                    module_name, param_name = name.rsplit(\".\", 1)\n                    module = model.get_submodule(module_name)\n\n                    current_parameter_tensor = parameters.get(name, None)\n\n                    slice_ = f.get_slice(name)\n\n                    if isinstance(module, TensorParallelColumnLinear):\n                        size = slice_.get_shape()[0]\n                        block_size = size // world_size\n                        start = rank * block_size\n                        stop = (rank + 1) * block_size\n                        tensor = slice_[start:stop]\n                    elif isinstance(module, TensorParallelRowLinear):\n                        if param_name == \"weight\":\n                            size = slice_.get_shape()[1]\n                            block_size = size // world_size\n                            start = rank * block_size\n                            stop = (rank + 1) * block_size\n                            tensor = slice_[:, start:stop]\n                        else:\n                            tensor = slice_[:]\n                            # XXX: Hack for Rowlinear to add the bias only once.\n                            if rank != 0:\n                                tensor = torch.zeros_like(tensor)\n                    elif isinstance(module, TensorParallelEmbedding):\n                        size = slice_.get_shape()[0]\n                        block_size = size // world_size\n                        start = rank * block_size\n                        stop = (rank + 1) * block_size\n                        tensor = slice_[start:stop]\n                    elif name == \"lm_head.weight\":\n                        size = slice_.get_shape()[0]\n                        block_size = size // world_size\n                        start = rank * block_size\n                        stop = (rank + 1) * block_size\n                        tensor = slice_[start:stop]\n                    elif \"relative_attention_bias.weight\" in name:\n                        size = slice_.get_shape()[1]\n                        block_size = size // world_size\n                        start = rank * block_size\n                        stop = (rank + 1) * block_size\n                        tensor = slice_[:, start:stop]\n                    else:\n                        try:\n                            tensor = slice_[:]\n                        except:\n                            tensor = f.get_tensor(name)\n\n                    if (\n                        current_parameter_tensor is not None\n                        and current_parameter_tensor.shape != tensor.shape\n                    ):\n                        raise ValueError(\n                            f\"Name {name} -- Current {current_parameter_tensor.shape} and got {tensor.shape}\"\n                        )\n\n                    tensor = tensor.contiguous().to(dtype)\n\n                    if quantize == \"bitsandbytes\":\n                        if not HAS_BITS_AND_BYTES:\n                            raise ImportError(\n                                \"bitsandbytes is not available on your machine either because it is not installed \"\n                                \"or you don't have a GPU.\\n\"\n                                \"You can install it with `pip install bitsandbytes`.\"\n                            )\n\n                        if (\n                            type(module)\n                            in [TensorParallelRowLinear, TensorParallelColumnLinear]\n                            and param_name == \"weight\"\n                        ):\n                            tensor = Int8Params(\n                                tensor,\n                                has_fp16_weights=False,\n                                requires_grad=False,\n                            ).to(device)\n                            state = bnb.MatmulLtState()\n                            state.threshold = 6.0\n                            state.has_fp16_weights = False\n                            state.memory_efficient_backward = False\n                            state.use_pool = True\n                            state.CB = tensor.CB\n                            state.SCB = tensor.SCB\n                            tensor.CB = None\n                            tensor.SCB = None\n\n                            def replace_linear(state):\n                                def linear(input, weight, bias):\n                                    out = bnb.matmul(\n                                        input,\n                                        weight,\n                                        state=state,\n                                        threshold=state.threshold,\n                                        bias=bias,\n                                    )\n\n                                    if state.CB is not None:\n                                        # we converted 8-bit row major to turing/ampere format\n                                        # in the first inference pass\n                                        # we no longer need the row-major weight\n                                        del state.CB\n                                        weight.data = state.CxB\n\n                                    return out\n\n                                return linear\n\n                            module.linear = replace_linear(state)\n\n                        elif quantize == \"gptq\":\n                            raise NotImplementedError(\n                                \"`gptq` is not implemented for now\"\n                            )\n                        elif quantize is None:\n                            tensor = tensor.to(device)\n                        else:\n                            raise ValueError(f\"Unexpected quantize `{quantize}`\")\n\n                    if current_parameter_tensor is not None:\n                        module._parameters[param_name] = tensor\n                    else:\n                        module._buffers[param_name] = tensor\n\n    def forward(\n        self,\n        input_ids,\n        attention_mask,\n        decoder_input_ids,\n        decoder_attention_mask: Optional,\n        encoder_last_hidden_state: Optional,\n        past_key_values: Optional = None,\n    ) -> Tuple[\n        torch.Tensor,\n        torch.Tensor,\n        List[Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]],\n    ]:\n        # Model Forward\n        outputs = self.model.forward(\n            input_ids=input_ids,\n            attention_mask=attention_mask,\n            decoder_input_ids=decoder_input_ids,\n            decoder_attention_mask=decoder_attention_mask,\n            encoder_outputs=encoder_last_hidden_state,\n            past_key_values=past_key_values,\n            use_cache=True,\n        )\n\n        # Logits are sharded, so we need to gather them\n        logits = [torch.empty_like(outputs.logits) for _ in range(self.world_size)]\n        torch.distributed.all_gather(logits, outputs.logits, group=self.process_group)\n        logits = torch.cat(logits, dim=2)\n\n        return (\n            logits,\n            outputs.encoder_last_hidden_state,\n            outputs.past_key_values,\n        )\n", "path": "server/text_generation_server/models/t5.py"}]}