Repository: huggingface/text-generation-inference

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.


Discussion:
> @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?

Yes everything should be very close.

Very close is hard to quantify though. Pytorch version can screw things up, dtype can mess up (some models are more susceptible to `bfloat16` to `float16`  change.

https://github.com/huggingface/transformers/blob/main/src/transformers/models/t5/modeling_t5.py#LL316C15-L316C71
 https://github.com/huggingface/transformers/issues/20287
 
 Those might be relevant to your issue too. (Maybe the sharding accentuates some precision issue)
Thanks @Narsil!
@OlivierDehaene thanks for https://github.com/huggingface/text-generation-inference/pull/356, I had also just discovered that bf16 works better for flan-t5-xxl sharded, per @Narsil's suggestion.

However it would be good if we could get it to work in fp16 so that V100's could be used, and since the results are more stable generally for variable length input.

It's interesting that:
- Non-sharded flan-t5-xxl works ok in fp16 (I know changes were made to the impl in the past to accommodate this)
- flan-ul2 actually appears to work ok when sharded in fp16. The output does begin to diverge slightly from a non-TP deployment but the differences are fairly minor
💡 : Imagine a world, where we could change a model range values arbitrarily without worrying about dtypes :) 

Since the model contains all of the weights, it should be doable to find invariants to allow some form of rescaling..

Resolving patch:
diff --git a/server/text_generation_server/models/bloom.py b/server/text_generation_server/models/bloom.py
--- a/server/text_generation_server/models/bloom.py
+++ b/server/text_generation_server/models/bloom.py
@@ -160,7 +160,10 @@ def load_weights(
                             # XXX: Hack for Rowlinear to add the bias only once.
                             if rank != 0:
                                 tensor = torch.zeros_like(tensor)
-                    elif isinstance(module, TensorParallelEmbedding) or name == "lm_head.weight":
+                    elif (
+                        isinstance(module, TensorParallelEmbedding)
+                        or name == "lm_head.weight"
+                    ):
                         size = slice_.get_shape()[0]
                         block_size = size // world_size
                         start = rank * block_size
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 @@ def __init__(
         self.process_group, rank, world_size = initialize_torch_distributed()
         if torch.cuda.is_available():
             device = torch.device(f"cuda:{rank}")
-            dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float32
+            dtype = torch.float16
         else:
             device = torch.device("cpu")
             dtype = torch.float32
@@ -154,9 +154,15 @@ def load_weights(
                             f"Name {name} -- Current {current_parameter_tensor.shape} and got {tensor.shape}"
                         )
 
-                    tensor = tensor.contiguous().to(dtype)
+                    tensor = tensor.contiguous()
 
-                    if quantize == "bitsandbytes":
+                    # See: https://github.com/huggingface/transformers/blob/1fe1e3caa44617047f149bcc0c0b566343b714a7/src/transformers/models/t5/modeling_t5.py#LL316C15-L316C71
+                    if module_name.endswith("wo"):
+                        tensor = tensor.to(torch.float32)
+                    else:
+                        tensor = tensor.to(dtype)
+
+                    if quantize == "bitsandbytes" and not module_name.endswith("wo"):
                         if not HAS_BITS_AND_BYTES:
                             raise ImportError(
                                 "bitsandbytes is not available on your machine either because it is not installed "
@@ -207,7 +213,7 @@ def linear(input, weight, bias):
 
                             module.linear = replace_linear(state)
 
-                        elif quantize == "gptq":
+                        elif quantize == "gptq" and not module_name.endswith("wo"):
                             raise NotImplementedError(
                                 "`gptq` is not implemented for now"
                             )


Test patch:
diff --git a/integration-tests/models/__snapshots__/test_t5_sharded/test_t5_sharded.json b/integration-tests/models/__snapshots__/test_t5_sharded/test_t5_sharded.json
new file mode 100644
--- /dev/null
+++ b/integration-tests/models/__snapshots__/test_t5_sharded/test_t5_sharded.json
@@ -0,0 +1,60 @@
+{
+  "details": {
+    "best_of_sequences": null,
+    "finish_reason": "eos_token",
+    "generated_tokens": 7,
+    "prefill": [
+      {
+        "id": 0,
+        "logprob": null,
+        "text": "<pad>"
+      }
+    ],
+    "seed": null,
+    "tokens": [
+      {
+        "id": 3,
+        "logprob": -0.7001953,
+        "special": false,
+        "text": " "
+      },
+      {
+        "id": 18,
+        "logprob": -1.1943359,
+        "special": false,
+        "text": "-"
+      },
+      {
+        "id": 26937,
+        "logprob": -1.2099609,
+        "special": false,
+        "text": "196"
+      },
+      {
+        "id": 3,
+        "logprob": -1.2451172,
+        "special": false,
+        "text": " "
+      },
+      {
+        "id": 1956,
+        "logprob": -0.3322754,
+        "special": false,
+        "text": "°"
+      },
+      {
+        "id": 254,
+        "logprob": -0.19213867,
+        "special": false,
+        "text": "C"
+      },
+      {
+        "id": 1,
+        "logprob": -0.030151367,
+        "special": true,
+        "text": "
"
+      }
+    ]
+  },
+  "generated_text": "-196 °C"
+}
diff --git a/integration-tests/models/__snapshots__/test_t5_sharded/test_t5_sharded_load.json b/integration-tests/models/__snapshots__/test_t5_sharded/test_t5_sharded_load.json
new file mode 100644
--- /dev/null
+++ b/integration-tests/models/__snapshots__/test_t5_sharded/test_t5_sharded_load.json
@@ -0,0 +1,242 @@
+[
+  {
+    "details": {
+      "best_of_sequences": null,
+      "finish_reason": "eos_token",
+      "generated_tokens": 7,
+      "prefill": [
+        {
+          "id": 0,
+          "logprob": null,
+          "text": "<pad>"
+        }
+      ],
+      "seed": null,
+      "tokens": [
+        {
+          "id": 3,
+          "logprob": -0.7001953,
+          "special": false,
+          "text": " "
+        },
+        {
+          "id": 18,
+          "logprob": -1.1943359,
+          "special": false,
+          "text": "-"
+        },
+        {
+          "id": 26937,
+          "logprob": -1.2119141,
+          "special": false,
+          "text": "196"
+        },
+        {
+          "id": 3,
+          "logprob": -1.2480469,
+          "special": false,
+          "text": " "
+        },
+        {
+          "id": 1956,
+          "logprob": -0.33203125,
+          "special": false,
+          "text": "°"
+        },
+        {
+          "id": 254,
+          "logprob": -0.19250488,
+          "special": false,
+          "text": "C"
+        },
+        {
+          "id": 1,
+          "logprob": -0.030166626,
+          "special": true,
+          "text": "
"
+        }
+      ]
+    },
+    "generated_text": "-196 °C"
+  },
+  {
+    "details": {
+      "best_of_sequences": null,
+      "finish_reason": "eos_token",
+      "generated_tokens": 7,
+      "prefill": [
+        {
+          "id": 0,
+          "logprob": null,
+          "text": "<pad>"
+        }
+      ],
+      "seed": null,
+      "tokens": [
+        {
+          "id": 3,
+          "logprob": -0.7001953,
+          "special": false,
+          "text": " "
+        },
+        {
+          "id": 18,
+          "logprob": -1.1943359,
+          "special": false,
+          "text": "-"
+        },
+        {
+          "id": 26937,
+          "logprob": -1.2119141,
+          "special": false,
+          "text": "196"
+        },
+        {
+          "id": 3,
+          "logprob": -1.2480469,
+          "special": false,
+          "text": " "
+        },
+        {
+          "id": 1956,
+          "logprob": -0.33203125,
+          "special": false,
+          "text": "°"
+        },
+        {
+          "id": 254,
+          "logprob": -0.19250488,
+          "special": false,
+          "text": "C"
+        },
+        {
+          "id": 1,
+          "logprob": -0.030166626,
+          "special": true,
+          "text": "
"
+        }
+      ]
+    },
+    "generated_text": "-196 °C"
+  },
+  {
+    "details": {
+      "best_of_sequences": null,
+      "finish_reason": "eos_token",
+      "generated_tokens": 7,
+      "prefill": [
+        {
+          "id": 0,
+          "logprob": null,
+          "text": "<pad>"
+        }
+      ],
+      "seed": null,
+      "tokens": [
+        {
+          "id": 3,
+          "logprob": -0.7001953,
+          "special": false,
+          "text": " "
+        },
+        {
+          "id": 18,
+          "logprob": -1.1943359,
+          "special": false,
+          "text": "-"
+        },
+        {
+          "id": 26937,
+          "logprob": -1.2119141,
+          "special": false,
+          "text": "196"
+        },
+        {
+          "id": 3,
+          "logprob": -1.2480469,
+          "special": false,
+          "text": " "
+        },
+        {
+          "id": 1956,
+          "logprob": -0.33203125,
+          "special": false,
+          "text": "°"
+        },
+        {
+          "id": 254,
+          "logprob": -0.19250488,
+          "special": false,
+          "text": "C"
+        },
+        {
+          "id": 1,
+          "logprob": -0.030166626,
+          "special": true,
+          "text": "
"
+        }
+      ]
+    },
+    "generated_text": "-196 °C"
+  },
+  {
+    "details": {
+      "best_of_sequences": null,
+      "finish_reason": "eos_token",
+      "generated_tokens": 7,
+      "prefill": [
+        {
+          "id": 0,
+          "logprob": null,
+          "text": "<pad>"
+        }
+      ],
+      "seed": null,
+      "tokens": [
+        {
+          "id": 3,
+          "logprob": -0.7001953,
+          "special": false,
+          "text": " "
+        },
+        {
+          "id": 18,
+          "logprob": -1.1943359,
+          "special": false,
+          "text": "-"
+        },
+        {
+          "id": 26937,
+          "logprob": -1.2099609,
+          "special": false,
+          "text": "196"
+        },
+        {
+          "id": 3,
+          "logprob": -1.2451172,
+          "special": false,
+          "text": " "
+        },
+        {
+          "id": 1956,
+          "logprob": -0.3322754,
+          "special": false,
+          "text": "°"
+        },
+        {
+          "id": 254,
+          "logprob": -0.19213867,
+          "special": false,
+          "text": "C"
+        },
+        {
+          "id": 1,
+          "logprob": -0.030151367,
+          "special": true,
+          "text": "
"
+        }
+      ]
+    },
+    "generated_text": "-196 °C"
+  }
+]
diff --git a/integration-tests/models/test_flash_neox.py b/integration-tests/models/test_flash_neox.py
--- a/integration-tests/models/test_flash_neox.py
+++ b/integration-tests/models/test_flash_neox.py
@@ -36,6 +36,8 @@ async def test_flash_neox_load(flash_neox, generate_load, response_snapshot):
     generated_texts = [r.generated_text for r in responses]
 
     assert len(generated_texts) == 4
-    assert generated_texts, all([text == generated_texts[0] for text in generated_texts])
+    assert generated_texts, all(
+        [text == generated_texts[0] for text in generated_texts]
+    )
 
     assert responses == response_snapshot
diff --git a/integration-tests/models/test_t5_sharded.py b/integration-tests/models/test_t5_sharded.py
new file mode 100644
--- /dev/null
+++ b/integration-tests/models/test_t5_sharded.py
@@ -0,0 +1,38 @@
+import pytest
+
+
+@pytest.fixture(scope="module")
+def t5_sharded_handle(launcher):
+    with launcher("google/flan-t5-xxl", num_shard=2) as handle:
+        yield handle
+
+
+@pytest.fixture(scope="module")
+async def t5_sharded(t5_sharded_handle):
+    await t5_sharded_handle.health(240)
+    return t5_sharded_handle.client
+
+
+@pytest.mark.asyncio
+async def test_t5_sharded(t5_sharded, response_snapshot):
+    response = await t5_sharded.generate(
+        "Please answer the following question. What is the boiling point of Nitrogen?",
+        max_new_tokens=10,
+    )
+
+    assert response == response_snapshot
+
+
+@pytest.mark.asyncio
+async def test_t5_sharded_load(t5_sharded, generate_load, response_snapshot):
+    responses = await generate_load(
+        t5_sharded,
+        "Please answer the following question. What is the boiling point of Nitrogen?",
+        max_new_tokens=10,
+        n=4,
+    )
+
+    assert len(responses) == 4
+    assert all([r.generated_text == responses[0].generated_text for r in responses])
+
+    assert responses == response_snapshot