feat: support multiple models on autotune service (#107)

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
BaguaSys · Jul 7, 2021 · 0aec789 · 0aec789
1 parent 328141c
commit 0aec789
Show file tree

Hide file tree

Showing 25 changed files with 430 additions and 665 deletions.
diff --git a/tests/nccl/__init__.py → .test_to_migrate/nccl/__init__.py b/tests/nccl/__init__.py → .test_to_migrate/nccl/__init__.py
diff --git a/tests/nccl/test_gather.py → .test_to_migrate/nccl/test_gather.py b/tests/nccl/test_gather.py → .test_to_migrate/nccl/test_gather.py
diff --git a/tests/nccl/test_p2p.py → .test_to_migrate/nccl/test_p2p.py b/tests/nccl/test_p2p.py → .test_to_migrate/nccl/test_p2p.py
diff --git a/tests/qsgd/test_norm_qsgd.py → .test_to_migrate/qsgd/test_norm_qsgd.py b/tests/qsgd/test_norm_qsgd.py → .test_to_migrate/qsgd/test_norm_qsgd.py
diff --git a/tests/qsgd/test_qsgd.py → .test_to_migrate/qsgd/test_qsgd.py b/tests/qsgd/test_qsgd.py → .test_to_migrate/qsgd/test_qsgd.py
diff --git a/tests/qsgd/utils.py → .test_to_migrate/qsgd/utils.py b/tests/qsgd/utils.py → .test_to_migrate/qsgd/utils.py
diff --git a/tests/test_communication.py → .test_to_migrate/test_communication.py b/tests/test_communication.py → .test_to_migrate/test_communication.py
diff --git a/tests/test_flatten.py → .test_to_migrate/test_flatten.py b/tests/test_flatten.py → .test_to_migrate/test_flatten.py
diff --git a/tests/test_fused_optimizer.py → .test_to_migrate/test_fused_optimizer.py b/tests/test_fused_optimizer.py → .test_to_migrate/test_fused_optimizer.py
diff --git a/tests/tools/test_loadbalance.py → .test_to_migrate/tools/test_loadbalance.py b/tests/tools/test_loadbalance.py → .test_to_migrate/tools/test_loadbalance.py
diff --git a/bagua/bagua_define.py b/bagua/bagua_define.py
@@ -37,6 +37,7 @@ class BaguaHyperparameter(BaseModel):
     """
 
     buckets: List[List[TensorDeclaration]] = []
+    bucket_size: int = 0
     is_hierarchical_reduce: bool = False
 
     def update(self, param_dict: dict):

diff --git a/bagua/distributed/launch.py b/bagua/distributed/launch.py
@@ -111,7 +111,10 @@ def parse_args():
         help="Bagua automatic super parameter search level. The higher the level, the better the theoretical effect, and the longer it takes",
     )
     parser.add_argument(
-        "--autotune_logfile", type=str, default="/tmp/bagua_autotune.log"
+        "--is_output_autotune_log",
+        action="store_true",
+        default=False,
+        help="Whether autotune output log or not. default is False",
     )
     parser.add_argument(
         "--report_metrics",
@@ -152,7 +155,7 @@ def set_bagua_env(args, current_env):
         args.autotune_sampling_confidence_time
     )
     current_env["BAGUA_AUTOTUNE_WARMUP_TIME_S"] = str(args.autotune_warmup_time)
-    current_env["BAGUA_AUTOTUNE_LOGFILE_PATH"] = args.autotune_logfile
+    current_env["BAGUA_IS_OUTPUT_AUTOTUNE_LOG"] = str(int(args.is_output_autotune_log))
 
 
 def main():

diff --git a/bagua/distributed/run.py b/bagua/distributed/run.py
@@ -383,7 +383,11 @@ def get_args_parser() -> ArgumentParser:
     parser.add_argument("--autotune_sampling_confidence_time", type=float, default=5.0)
     parser.add_argument("--autotune_warmup_time", type=float, default=30.0)
     parser.add_argument(
-        "--autotune_logfile", type=str, default="/tmp/bagua_autotune.log"
+        "--is_output_autotune_log",
+        type=bool,
+        action="store_true",
+        default=False,
+        help="Whether autotune output log or not. default is False",
     )
 
     #
@@ -573,7 +577,7 @@ def set_bagua_env(args, current_env):
         args.autotune_sampling_confidence_time
     )
     current_env["BAGUA_AUTOTUNE_WARMUP_TIME_S"] = str(args.autotune_warmup_time)
-    current_env["BAGUA_AUTOTUNE_LOGFILE_PATH"] = args.autotune_logfile
+    current_env["BAGUA_IS_OUTPUT_AUTOTUNE_LOG"] = args.is_output_autotune_log
 
 
 def run(args):

diff --git a/bagua/service/__init__.py b/bagua/service/__init__.py
@@ -1,5 +1,2 @@
-from .autotune_service import AutotuneService, BaguaHyperparameter
-from .service_discovery import (
-    generate_and_broadcast_server_addr,
-    pick_n_free_ports,
-)
+from .autotune_service import AutotuneService, AutotuneClient  # noqa: F401
+from . import autotune  # noqa: F401
diff --git a/bagua/service/autotune.py b/bagua/service/autotune.py
@@ -43,16 +43,21 @@ class BayesianOptimizer:
     Simple package of beyasian optimizer
     """
 
-    def __init__(self, param_declaration: dict, n_initial_points: int = 10):
+    def __init__(
+        self,
+        param_declaration: dict,
+        n_initial_points: int = 20,
+        initial_point_generator: str = "halton",
+    ):
         self.param_declaration = collections.OrderedDict(param_declaration)
         search_space = [
             declar.space_dimension for _, declar in self.param_declaration.items()
         ]
 
         self.bayesian_optimizer = skopt.Optimizer(
             dimensions=search_space,
-            base_estimator="GP",
             n_initial_points=n_initial_points,
+            initial_point_generator=initial_point_generator,
             n_jobs=-1,
         )