[INF2] give better room for more tokens (deepjavalibrary#710)

* give better room for more tokens * revert the tests
Lokiiiiii · May 12, 2023 · 14d0e67 · 14d0e67
1 parent a583888
commit 14d0e67
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 5 deletions.
diff --git a/tests/integration/llm/client.py b/tests/integration/llm/client.py
@@ -94,7 +94,7 @@
         "seq_length": [64, 128]
     },
     "bigscience/bloom-3b": {
-        "max_memory_per_gpu": [0.1, 5.0, 6.0, 0.1],
+        "max_memory_per_gpu": [5.0, 6.0],
         "batch_size": [1, 4],
         "seq_length": [16, 32],
         "worker": 1,
@@ -139,7 +139,7 @@
         "worker": 2
     },
     "gpt-neo-1.3b": {
-        "max_memory_per_gpu": [0.1, 4.0, 0.1, 5.0],
+        "max_memory_per_gpu": [4.0, 5.0],
         "batch_size": [1, 4],
         "seq_length": [16],
         "worker": 1,

diff --git a/tests/integration/llm/prepare.py b/tests/integration/llm/prepare.py
@@ -238,15 +238,15 @@
         "option.model_id": "gpt2",
         "option.batch_size": 4,
         "option.tensor_parallel_degree": 2,
-        "option.n_positions": 256,
+        "option.n_positions": 512,
         "option.dtype": "fp16",
         "option.model_loading_timeout": 600
     },
     "opt-1.3b": {
         "option.model_id": "s3://djl-llm/opt-1.3b/",
         "option.batch_size": 4,
         "option.tensor_parallel_degree": 4,
-        "option.n_positions": 256,
+        "option.n_positions": 512,
         "option.dtype": "fp16",
         "option.model_loading_timeout": 600
     },
@@ -262,7 +262,7 @@
         "option.model_id": "s3://djl-llm/opt-1.3b/",
         "option.batch_size": 1,
         "option.tensor_parallel_degree": 4,
-        "option.n_positions": 256,
+        "option.n_positions": 512,
         "option.dtype": "fp16",
         "option.model_loading_timeout": 600,
         "option.enable_streaming": True