Auto max prefill (#2797)

* Attempt at automatic max batch prefill. * Taking into account number of shards. * Adding more cards. * Adding A100 + H100 * Adding a few more cards. * Logprobs cost too much. * h100 better name, and keep factor of 2 * Damn inflated sparse tflops. * Typo in h100. * Updated the flops calculation (checked with fvcore). * chunking by default. * Fix prefix caching for chat completion since we removed logprobs. * More tests. * Dropping all the prefill logprobs. * Add a flag that enables users to get logprobs back. * Repairing prompt token counting. * Fixing a few tests. * Remove some scaffolding. * Attempting to reduces the issues (workarounds for now).
huggingface · Dec 6, 2024 · 5df8059 · 5df8059
1 parent 8c3669b
commit 5df8059
Show file tree

Hide file tree

Showing 130 changed files with 1,154 additions and 74,942 deletions.
diff --git a/docs/source/reference/launcher.md b/docs/source/reference/launcher.md
@@ -467,6 +467,16 @@ Options:
           [env: PAYLOAD_LIMIT=]
           [default: 2000000]
 
+```
+## ENABLE_PREFILL_LOGPROBS
+```shell
+      --enable-prefill-logprobs
+          Enables prefill logprobs
+
+          Logprobs in the prompt are deactivated by default because they consume a large amount of VRAM (especially for long prompts). Using this flag reallows users to ask for them.
+
+          [env: ENABLE_PREFILL_LOGPROBS=]
+
 ```
 ## HELP
 ```shell

diff --git a/integration-tests/models/__snapshots__/test.py b/integration-tests/models/__snapshots__/test.py
@@ -0,0 +1,22 @@
+import os
+import json
+
+
+for root, dirs, files in os.walk("."):
+    for filename in files:
+        if filename.endswith(".json"):
+            with open(os.path.join(root, filename), "r") as f:
+                data = json.load(f)
+
+            print(os.path.join(root, filename))
+            try:
+                if filename.endswith("_load.json"):
+                    for i in range(len(data)):
+                        data[i]["details"]["prefill"] = []
+                else:
+                    data["details"]["prefill"] = []
+            except Exception:
+                pass
+
+            with open(os.path.join(root, filename), "w") as f:
+                json.dump(data, f, indent=2, ensure_ascii=False)
diff --git a/...dels/__snapshots__/test_compressed_tensors_w8a8_int/test_compressed_tensors_w8a8_int.json b/...dels/__snapshots__/test_compressed_tensors_w8a8_int/test_compressed_tensors_w8a8_int.json
@@ -3,38 +3,7 @@
     "best_of_sequences": null,
     "finish_reason": "length",
     "generated_tokens": 10,
-    "prefill": [
-      {
-        "id": 128000,
-        "logprob": null,
-        "text": "<|begin_of_text|>"
-      },
-      {
-        "id": 3923,
-        "logprob": -6.3867188,
-        "text": "What"
-      },
-      {
-        "id": 374,
-        "logprob": -1.1318359,
-        "text": " is"
-      },
-      {
-        "id": 5655,
-        "logprob": -9.6875,
-        "text": " deep"
-      },
-      {
-        "id": 6975,
-        "logprob": -1.3007812,
-        "text": " learning"
-      },
-      {
-        "id": 30,
-        "logprob": -2.4902344,
-        "text": "?"
-      }
-    ],
+    "prefill": [],
     "seed": null,
     "tokens": [
       {

diff --git a/...shots__/test_compressed_tensors_w8a8_int/test_compressed_tensors_w8a8_int_all_params.json b/...shots__/test_compressed_tensors_w8a8_int/test_compressed_tensors_w8a8_int_all_params.json
@@ -3,33 +3,7 @@
     "best_of_sequences": null,
     "finish_reason": "length",
     "generated_tokens": 10,
-    "prefill": [
-      {
-        "id": 128000,
-        "logprob": null,
-        "text": "<|begin_of_text|>"
-      },
-      {
-        "id": 3923,
-        "logprob": -6.3867188,
-        "text": "What"
-      },
-      {
-        "id": 374,
-        "logprob": -1.1318359,
-        "text": " is"
-      },
-      {
-        "id": 5655,
-        "logprob": -9.6875,
-        "text": " deep"
-      },
-      {
-        "id": 6975,
-        "logprob": -1.3007812,
-        "text": " learning"
-      }
-    ],
+    "prefill": [],
     "seed": 0,
     "tokens": [
       {

diff --git a/...__snapshots__/test_compressed_tensors_w8a8_int/test_compressed_tensors_w8a8_int_load.json b/...__snapshots__/test_compressed_tensors_w8a8_int/test_compressed_tensors_w8a8_int_load.json
@@ -4,38 +4,7 @@
       "best_of_sequences": null,
       "finish_reason": "length",
       "generated_tokens": 10,
-      "prefill": [
-        {
-          "id": 128000,
-          "logprob": null,
-          "text": "<|begin_of_text|>"
-        },
-        {
-          "id": 3923,
-          "logprob": -6.3867188,
-          "text": "What"
-        },
-        {
-          "id": 374,
-          "logprob": -1.1318359,
-          "text": " is"
-        },
-        {
-          "id": 5655,
-          "logprob": -9.6875,
-          "text": " deep"
-        },
-        {
-          "id": 6975,
-          "logprob": -1.3007812,
-          "text": " learning"
-        },
-        {
-          "id": 30,
-          "logprob": -2.4902344,
-          "text": "?"
-        }
-      ],
+      "prefill": [],
       "seed": null,
       "tokens": [
         {
@@ -108,38 +77,7 @@
       "best_of_sequences": null,
       "finish_reason": "length",
       "generated_tokens": 10,
-      "prefill": [
-        {
-          "id": 128000,
-          "logprob": null,
-          "text": "<|begin_of_text|>"
-        },
-        {
-          "id": 3923,
-          "logprob": -6.3867188,
-          "text": "What"
-        },
-        {
-          "id": 374,
-          "logprob": -1.1318359,
-          "text": " is"
-        },
-        {
-          "id": 5655,
-          "logprob": -9.6875,
-          "text": " deep"
-        },
-        {
-          "id": 6975,
-          "logprob": -1.3007812,
-          "text": " learning"
-        },
-        {
-          "id": 30,
-          "logprob": -2.4902344,
-          "text": "?"
-        }
-      ],
+      "prefill": [],
       "seed": null,
       "tokens": [
         {
@@ -212,38 +150,7 @@
       "best_of_sequences": null,
       "finish_reason": "length",
       "generated_tokens": 10,
-      "prefill": [
-        {
-          "id": 128000,
-          "logprob": null,
-          "text": "<|begin_of_text|>"
-        },
-        {
-          "id": 3923,
-          "logprob": -6.3867188,
-          "text": "What"
-        },
-        {
-          "id": 374,
-          "logprob": -1.1318359,
-          "text": " is"
-        },
-        {
-          "id": 5655,
-          "logprob": -9.6875,
-          "text": " deep"
-        },
-        {
-          "id": 6975,
-          "logprob": -1.3007812,
-          "text": " learning"
-        },
-        {
-          "id": 30,
-          "logprob": -2.4902344,
-          "text": "?"
-        }
-      ],
+      "prefill": [],
       "seed": null,
       "tokens": [
         {
@@ -316,38 +223,7 @@
       "best_of_sequences": null,
       "finish_reason": "length",
       "generated_tokens": 10,
-      "prefill": [
-        {
-          "id": 128000,
-          "logprob": null,
-          "text": "<|begin_of_text|>"
-        },
-        {
-          "id": 3923,
-          "logprob": -6.3867188,
-          "text": "What"
-        },
-        {
-          "id": 374,
-          "logprob": -1.1318359,
-          "text": " is"
-        },
-        {
-          "id": 5655,
-          "logprob": -9.6875,
-          "text": " deep"
-        },
-        {
-          "id": 6975,
-          "logprob": -1.3007812,
-          "text": " learning"
-        },
-        {
-          "id": 30,
-          "logprob": -2.4902344,
-          "text": "?"
-        }
-      ],
+      "prefill": [],
       "seed": null,
       "tokens": [
         {

diff --git a/...ssed_tensors_w8a8_int_dynamic_weight/test_compressed_tensors_w8a8_int_dynamic_weight.json b/...ssed_tensors_w8a8_int_dynamic_weight/test_compressed_tensors_w8a8_int_dynamic_weight.json
@@ -3,33 +3,7 @@
     "best_of_sequences": null,
     "finish_reason": "length",
     "generated_tokens": 10,
-    "prefill": [
-      {
-        "id": 3838,
-        "logprob": null,
-        "text": "What"
-      },
-      {
-        "id": 374,
-        "logprob": -8.59375,
-        "text": " is"
-      },
-      {
-        "id": 5538,
-        "logprob": -10.921875,
-        "text": " deep"
-      },
-      {
-        "id": 6832,
-        "logprob": -0.56347656,
-        "text": " learning"
-      },
-      {
-        "id": 30,
-        "logprob": -1.5,
-        "text": "?"
-      }
-    ],
+    "prefill": [],
     "seed": null,
     "tokens": [
       {

diff --git a/...s_w8a8_int_dynamic_weight/test_compressed_tensors_w8a8_int_dynamic_weight_all_params.json b/...s_w8a8_int_dynamic_weight/test_compressed_tensors_w8a8_int_dynamic_weight_all_params.json
@@ -3,28 +3,7 @@
     "best_of_sequences": null,
     "finish_reason": "length",
     "generated_tokens": 10,
-    "prefill": [
-      {
-        "id": 3838,
-        "logprob": null,
-        "text": "What"
-      },
-      {
-        "id": 374,
-        "logprob": -8.59375,
-        "text": " is"
-      },
-      {
-        "id": 5538,
-        "logprob": -10.921875,
-        "text": " deep"
-      },
-      {
-        "id": 6832,
-        "logprob": -0.56347656,
-        "text": " learning"
-      }
-    ],
+    "prefill": [],
     "seed": 0,
     "tokens": [
       {