[python] making rolling batch default as auto (#2607)

deepjavalibrary · Nov 27, 2024 · 05d07a9 · 05d07a9
1 parent 5a6562a
commit 05d07a9
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/engines/python/setup/djl_python/properties_manager/properties.py b/engines/python/setup/djl_python/properties_manager/properties.py
@@ -48,7 +48,7 @@ class Properties(BaseModel):
     # Optional configurations with default values
     model_dir: Optional[str] = None
     # Make the default to auto, after java front end changes and test cases are changed.
-    rolling_batch: RollingBatchEnum = RollingBatchEnum.disable
+    rolling_batch: RollingBatchEnum = RollingBatchEnum.auto
     tensor_parallel_degree: int = 1
     cluster_size: int = 1
     trust_remote_code: bool = False

diff --git a/engines/python/setup/djl_python/tests/test_properties_manager.py b/engines/python/setup/djl_python/tests/test_properties_manager.py
@@ -328,6 +328,7 @@ def test_hf_configs(self):
             "low_cpu_mem_usage": "true",
             "disable_flash_attn": "false",
             "mpi_mode": "true",
+            "rolling_batch": "disable"
         }
 
         hf_configs = HuggingFaceProperties(**properties)
@@ -395,7 +396,8 @@ def test_hf_device_map(self, mock_device_count):
         hf_configs = HuggingFaceProperties(**properties, rolling_batch="auto")
         self.assertEqual(hf_configs.kwargs.get("device_map"), "auto")
 
-        hf_configs = HuggingFaceProperties(**properties)
+        hf_configs = HuggingFaceProperties(**properties,
+                                           rolling_batch="disable")
         self.assertIsNone(hf_configs.kwargs.get("device_map"))
 
     def test_hf_quantize(self):