[ci] Fix gpt-j timeout issues in inf2 integration (#1535)

deepjavalibrary · Feb 9, 2024 · 8bf54f7 · 8bf54f7
1 parent 3e91fb6
commit 8bf54f7
Showing 1 changed file with 8 additions and 15 deletions.
diff --git a/tests/integration/llm/prepare.py b/tests/integration/llm/prepare.py
@@ -381,8 +381,7 @@
         "option.tensor_parallel_degree": 2,
         "option.n_positions": 512,
         "option.dtype": "fp16",
-        "option.model_loading_timeout": 600,
-        "option.enable_streaming": False
+        "option.model_loading_timeout": 600
     },
     "gpt2-quantize": {
         "option.model_id": "gpt2",
@@ -391,35 +390,31 @@
         "option.n_positions": 512,
         "option.dtype": "fp16",
         "option.model_loading_timeout": 600,
-        "option.quantize": "static_int8",
-        "option.enable_streaming": False
+        "option.quantize": "static_int8"
     },
     "opt-1.3b": {
         "option.model_id": "s3://djl-llm/opt-1.3b/",
         "batch_size": 4,
         "option.tensor_parallel_degree": 4,
         "option.n_positions": 512,
         "option.dtype": "fp16",
-        "option.model_loading_timeout": 600,
-        "option.enable_streaming": False
+        "option.model_loading_timeout": 600
     },
     "gpt-j-6b": {
         "option.model_id": "s3://djl-llm/gpt-j-6b/",
         "batch_size": 4,
         "option.tensor_parallel_degree": 8,
-        "option.n_positions": 1024,
+        "option.n_positions": 512,
         "option.dtype": "fp32",
-        "option.model_loading_timeout": 900,
-        "option.enable_streaming": False
+        "option.model_loading_timeout": 1200
     },
     "pythia-2.8b": {
         "option.model_id": "s3://djl-llm/pythia-2.8b/",
         "batch_size": 4,
         "option.tensor_parallel_degree": 2,
         "option.n_positions": 512,
         "option.dtype": "fp16",
-        "option.model_loading_timeout": 900,
-        "option.enable_streaming": False
+        "option.model_loading_timeout": 900
     },
     "open-llama-7b": {
         "option.model_id": "s3://djl-llm/open-llama-7b/",
@@ -428,17 +423,15 @@
         "option.n_positions": 512,
         "option.dtype": "fp16",
         "option.neuron_optimize_level": 1,
-        "option.model_loading_timeout": 1200,
-        "option.enable_streaming": False
+        "option.model_loading_timeout": 1200
     },
     "bloom-7b1": {
         "option.model_id": "s3://djl-llm/bloom-7b1/",
         "batch_size": 4,
         "option.tensor_parallel_degree": 4,
         "option.n_positions": 256,
         "option.dtype": "fp16",
-        "option.model_loading_timeout": 720,
-        "option.enable_streaming": False
+        "option.model_loading_timeout": 720
     },
     "llama-7b-split": {
         "option.model_id": "s3://djl-llm/llama-2-7b-split-inf2/split-model/",