add buffer and best_cfg_log for autotuner

Author: Difers

python/paddle/distributed/auto_tuner/recorder.py

@@ -51,33 +51,26 @@ def sort_metric(self, direction, metric_name) -> None:
                 reverse=False,
             )
 
-    def get_best(self, metric, direction, mode=None) -> Tuple[dict, bool]:
+    def get_best(self, metric, direction, buffer, memory) -> Tuple[dict, bool]:
         self.sort_metric(direction=direction, metric_name=metric)
         if len(self.history) == 0:
             return (self.history[0], True)
-        if mode == "SFT" or mode == "LoRA" or mode == "Pretrain":
-            best_cfg = self.history[0]
-            if (
-                isinstance(best_cfg["max_mem_usage"], str)
-                or best_cfg["time"] == -1
-            ):
-                return (best_cfg, True)
-            first_few = 0
+
+        best_cfg = self.history[0]
+        if isinstance(best_cfg["max_mem_usage"], str) or best_cfg["time"] == -1:
+            return (best_cfg, True)
+
+        if buffer > 0:
             for cfg in self.history:
+                best_cfg = cfg
                 if (
                     not isinstance(cfg["max_mem_usage"], str)
-                    and cfg["max_mem_usage"] < best_cfg["max_mem_usage"]
+                    and cfg["max_mem_usage"] < memory - buffer
                     and cfg["time"] != -1
                 ):
-                    best_cfg = cfg
-                first_few += 1
-                if first_few >= 3:
                     break
             return (best_cfg, False)
-        if isinstance(self.history[0]["max_mem_usage"], str) or (
-            "time" in self.history[0] and self.history[0]["time"] == -1
-        ):
-            return (self.history[0], True)
+
         return (self.history[0], False)
 
     def _store_history_impl(self, data, path="./history.csv"):

python/paddle/distributed/launch/main.py

@@ -436,6 +436,12 @@ def launch():
         # max_search_time
         max_search_time = tuner_cfg.get("max_search_time", None)
 
+        # buffer and memory
+        buffer = tuner_cfg.get("buffer", 0)
+        tuner_cfg['buffer'] = buffer
+        memory = tuner_cfg.get("memory", None)
+        tuner_cfg['buffer'] = buffer
+
         is_first_task = True
         # build history recorder
         recorder = HistoryRecorder(tuner_cfg)
@@ -681,6 +687,8 @@ def launch():
                 cur_best_cfgs, err = recorder.get_best(
                     metric=tuner_cfg['metric_cfg']['name'],
                     direction=tuner_cfg['metric_cfg']['OptimizationDirection'],
+                    buffer=tuner_cfg['buffer'],
+                    memory=tuner_cfg['memory'],
                 )
                 if not err:
                     ctx.logger.info(f"Current best config: {cur_best_cfgs}")
@@ -781,6 +789,8 @@ def launch():
                         direction=tuner_cfg['metric_cfg'][
                             'OptimizationDirection'
                         ],
+                        buffer=tuner_cfg['buffer'],
+                        memory=tuner_cfg['memory'],
                     )
                     if not err:
                         ctx.logger.info(f"Current best config: {cur_best_cfgs}")
@@ -1158,7 +1168,8 @@ def launch():
             cur_best_cfgs, err = recorder.get_best(
                 metric=tuner_cfg['metric_cfg']['name'],
                 direction=tuner_cfg['metric_cfg']['OptimizationDirection'],
-                mode=mode,
+                buffer=tuner_cfg['buffer'],
+                memory=tuner_cfg['memory'],
             )
             if not err:
                 ctx.logger.info(f"Current best config: {cur_best_cfgs}")
@@ -1206,7 +1217,8 @@ def launch():
                 best_cfg, err = recorder.get_best(
                     metric=tuner_cfg['metric_cfg']['name'],
                     direction=tuner_cfg['metric_cfg']['OptimizationDirection'],
-                    mode=mode,
+                    buffer=tuner_cfg['buffer'],
+                    memory=tuner_cfg['memory'],
                 )
                 if err:
                     raise ValueError(
@@ -1232,7 +1244,8 @@ def launch():
             best_cfg, err = recorder.get_best(
                 metric=tuner_cfg['metric_cfg']['name'],
                 direction=tuner_cfg['metric_cfg']['OptimizationDirection'],
-                mode=mode,
+                buffer=tuner_cfg['buffer'],
+                memory=tuner_cfg['memory'],
             )
             if err:
                 raise ValueError(
@@ -1255,9 +1268,13 @@ def launch():
         ctx.args.job_id = "best_cfg"
         ctx.logger.info(f"Launch best cfg: {best_cfg}")
         logger.info(f"Launch best cfg: {best_cfg}")
-        ctx.args.log_dir = ctx.args.log_dir = os.path.join(
-            os.path.dirname(ctx.args.auto_tuner_json), "best_cfg"
-        )
+
+        if tuner_cfg.get("best_cfg_dir", None):
+            ctx.args.log_dir = tuner_cfg["best_cfg_dir"]
+        else:
+            ctx.args.log_dir = os.path.join(
+                os.path.dirname(ctx.args.auto_tuner_json), "best_cfg"
+            )
         # run best cfg
         c = controllers.init(ctx)
         c.run()