[AutoConfig] add blacklist mechanism and scheduling in specified order (#60248)

Caozhou1995 · Difers · web-flow · commit f15a0cf072af · 2024-01-31T21:05:26.000+08:00
* add blacklist mechanism and scheduling in specified order

* change best cfg

* fix first few bug

* fix resume from history

---------

Co-authored-by: Difers &lt;707065510@qq.com&gt;
diff --git a/python/paddle/distributed/auto_tuner/prune.py b/python/paddle/distributed/auto_tuner/prune.py
@@ -615,3 +615,99 @@ def prune_by_sharding_overlap(tuner_cfg, cur_cfg, history_cfgs=[]):
         if not result[tuner_cfg['metric_cfg']['name']]:
             return True
     return False
+
+
+def is_invalid(cur_cfg, invalid_strategy):
+    mapping = {
+        "dp_degree": "dp",
+        "mp_degree": "mp",
+        "pp_degree": "pp",
+        "vpp_degree": "vpp",
+        "micro_batch_size": "mbs",
+        "sharding_degree": "sharding",
+        "sharding_stage": "stage",
+        "use_recompute": "recompute",
+        "recompute_granularity": "granularity",
+    }
+    granularity_mapping = {0: "full", 1: "full_attn", 2: "core_attn"}
+    reversed_mapping = {}
+    for key in mapping:
+        reversed_mapping[mapping[key]] = key
+
+    for strategy in invalid_strategy:
+        assert isinstance(strategy, str)
+        dims = strategy.split("_")
+        has_matched = 0
+        for dim in dims:
+            matched = None
+            for key in reversed_mapping:
+                if dim.startswith(key):
+                    matched = key
+                    break
+            if matched:
+                value = dim[len(matched)]
+                # * means this strategy turned on
+                if matched in ["dp", "mp", "pp", "vpp", "sharding"]:
+                    if value == "*":
+                        if cur_cfg[reversed_mapping[matched]] != 1:
+                            has_matched += 1
+                            continue
+                    else:
+                        value = int(value)
+                        if cur_cfg[reversed_mapping[matched]] == value:
+                            has_matched += 1
+                            continue
+                elif matched == "recompute":
+                    if value == "*":
+                        if cur_cfg[reversed_mapping[matched]]:
+                            has_matched += 1
+                            continue
+                    else:
+                        value = bool(int(value))
+                        if cur_cfg[reversed_mapping[matched]] == value:
+                            has_matched += 1
+                            continue
+                elif matched == "stage":
+                    if value == "*":
+                        if cur_cfg[reversed_mapping["sharding"]] != 1:
+                            has_matched += 1
+                            continue
+                    else:
+                        value = int(value)
+                        if cur_cfg[reversed_mapping[matched]] == value:
+                            has_matched += 1
+                            continue
+                elif matched == "mbs":
+                    if value == "*":
+                        has_matched += 1
+                        continue
+                    else:
+                        value = int(value)
+                        if cur_cfg[reversed_mapping[matched]] == value:
+                            has_matched += 1
+                            continue
+                elif matched == "granularity":
+                    if value == "*":
+                        if cur_cfg[reversed_mapping["use_recompute"]]:
+                            has_matched += 1
+                            continue
+                    else:
+                        value = int(value)
+                        granularity = granularity_mapping[value]
+                        if cur_cfg[reversed_mapping[matched]] == granularity:
+                            has_matched += 1
+                            continue
+    if has_matched == len(dims):
+        return True
+    return False
+
+
+@register_prune
+def prune_by_invalid_strategy(tuner_cfg, cur_cfg, history_cfgs=[]):
+    if tuner_cfg.get("invalid_strategy", None):
+        invalid_strategy = tuner_cfg["invalid_strategy"]
+        assert isinstance(invalid_strategy, list)
+        if is_invalid(cur_cfg, invalid_strategy):
+            return True
+
+    return False
diff --git a/python/paddle/distributed/auto_tuner/recorder.py b/python/paddle/distributed/auto_tuner/recorder.py
@@ -55,14 +55,14 @@ def get_best(self, metric, direction, mode=None) -> 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":
+        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 = 1
+            first_few = 0
             for cfg in self.history:
                 if (
                     not isinstance(cfg["max_mem_usage"], str)
@@ -71,7 +71,7 @@ def get_best(self, metric, direction, mode=None) -> Tuple[dict, bool]:
                 ):
                     best_cfg = cfg
                 first_few += 1
-                if first_few >= 5:
+                if first_few >= 3:
                     break
             return (best_cfg, False)
         if isinstance(self.history[0]["max_mem_usage"], str) or (
diff --git a/python/paddle/distributed/auto_tuner/search.py b/python/paddle/distributed/auto_tuner/search.py
@@ -49,17 +49,41 @@ def __init__(self, tuner_cfg):
         super().__init__(tuner_cfg)
         self.idx = 0
         self.all_tasks = search_all(tuner_cfg)
+        need_baseline = self.tuner_cfg.get("need_baseline", False)
+        self.baseline = None
+        if need_baseline:
+            from .utils import memory_sort
+
+            self.all_tasks.sort(key=memory_sort)
+        self.previous_cfg = None
 
     def search_once(self, history_cfgs):
         new_cfg = None
         stop = False
+        if history_cfgs:
+            if history_cfgs[-1].get("time", -1) > 0:
+                if self.baseline is None:
+                    from .utils import performance_sort
+
+                    self.baseline = history_cfgs[-1]
+                    self.all_tasks[self.idx :] = sorted(
+                        self.all_tasks[self.idx : len(self.all_tasks)],
+                        key=performance_sort,
+                    )
+                    if self.tuner_cfg.get("schedule_prior", False):
+                        from .utils import sort_by_sepecial
+
+                        self.all_tasks[self.idx :] = sort_by_sepecial(
+                            self.all_tasks[self.idx :], self.tuner_cfg
+                        )
         while not stop:
             if self.idx < len(self.all_tasks):
                 new_cfg = self.all_tasks[self.idx]
                 self.idx += 1
                 stop = not self.prune(self.tuner_cfg, new_cfg, history_cfgs)
             else:
                 return None
+
         return new_cfg
 
 
diff --git a/python/paddle/distributed/auto_tuner/utils.py b/python/paddle/distributed/auto_tuner/utils.py
@@ -416,9 +416,132 @@ def search_all(tuner_cfg):
     logger.info(
         f"{search_space_size_before_prune - search_space_size_after_prune} tasks are pruned before launching."
     )
+    if tuner_cfg.get("schedule_prior", False):
+        pruned_all_cfgs = sort_by_sepecial(pruned_all_cfgs, tuner_cfg)
     return pruned_all_cfgs
 
 
+def sort_by_sepecial(cfgs, tuner_cfg):
+    assert tuner_cfg.get("schedule_prior", False)
+    prior_strategy = tuner_cfg["schedule_prior"]
+    prior_strategy.sort(reverse=True)
+    for strategy in prior_strategy:
+        idx = 0
+        matched_count = 0
+        while idx < len(cfgs):
+            cfg = cfgs[idx]
+            if _matched(cfg, strategy):
+                cfgs.pop(idx)
+                cfgs.insert(0, cfg)
+                matched_count += 1
+            idx += 1
+        tmp = cfgs[:matched_count]
+        tmp.reverse()
+        cfgs[:matched_count] = tmp
+    return cfgs
+
+
+def memory_sort(cfg):
+    # ascending order in default
+    return (
+        -cfg['mp_degree'],
+        -cfg['pp_degree'],
+        -cfg['vpp_degree'],
+        -cfg["sharding_degree"],
+        -cfg["sharding_stage"],
+        cfg["micro_batch_size"],
+        -cfg["use_recompute"],
+    )
+
+
+def performance_sort(cfg):
+    return -cfg["micro_batch_size"]
+
+
+def _matched(cur_cfg, strategy):
+    mapping = {
+        "dp_degree": "dp",
+        "mp_degree": "mp",
+        "pp_degree": "pp",
+        "vpp_degree": "vpp",
+        "micro_batch_size": "mbs",
+        "sharding_degree": "sharding",
+        "sharding_stage": "stage",
+        "use_recompute": "recompute",
+        "recompute_granularity": "granularity",
+    }
+    granularity_mapping = {0: "full", 1: "full_attn", 2: "core_attn"}
+    reversed_mapping = {}
+    for key in mapping:
+        reversed_mapping[mapping[key]] = key
+
+    assert isinstance(strategy, str)
+    dims = strategy.split("_")
+    has_matched = 0
+    for dim in dims:
+        matched = None
+        for key in reversed_mapping:
+            if dim.startswith(key):
+                matched = key
+                break
+        if matched:
+            value = dim[len(matched)]
+            # * means this strategy turned on
+            if matched in ["dp", "mp", "pp", "vpp", "sharding"]:
+                if value == "*":
+                    if cur_cfg[reversed_mapping[matched]] > 1:
+                        has_matched += 1
+                        continue
+                else:
+                    value = int(value)
+                    if cur_cfg[reversed_mapping[matched]] == value:
+                        has_matched += 1
+                        continue
+            elif matched == "recompute":
+                if value == "*":
+                    if cur_cfg[reversed_mapping[matched]]:
+                        has_matched += 1
+                        continue
+                else:
+                    value = bool(int(value))
+                    if cur_cfg[reversed_mapping[matched]] == value:
+                        has_matched += 1
+                        continue
+            elif matched == "stage":
+                if value == "*":
+                    if cur_cfg[reversed_mapping["sharding"]] > 1:
+                        has_matched += 1
+                        continue
+                else:
+                    value = int(value)
+                    if cur_cfg[reversed_mapping[matched]] == value:
+                        has_matched += 1
+                        continue
+            elif matched == "mbs":
+                if value == "*":
+                    has_matched += 1
+                    continue
+                else:
+                    value = int(value)
+                    if cur_cfg[reversed_mapping[matched]] == value:
+                        has_matched += 1
+                        continue
+            elif matched == "granularity":
+                if value == "*":
+                    if cur_cfg[reversed_mapping["use_recompute"]]:
+                        has_matched += 1
+                        continue
+                else:
+                    value = int(value)
+                    granularity = granularity_mapping[value]
+                    if cur_cfg[reversed_mapping[matched]] == granularity:
+                        has_matched += 1
+                        continue
+    if has_matched == len(dims):
+        return True
+    return False
+
+
 def _param2range(param_from_json_file, max_value, param_key):
     """Convert a param from json file to candidates range."""
     selected_range = None
diff --git a/python/paddle/distributed/launch/main.py b/python/paddle/distributed/launch/main.py
@@ -1157,6 +1157,7 @@ def launch():
             cur_best_cfgs, err = recorder.get_best(
                 metric=tuner_cfg['metric_cfg']['name'],
                 direction=tuner_cfg['metric_cfg']['OptimizationDirection'],
+                mode=mode,
             )
             if not err:
                 ctx.logger.info(f"Current best config: {cur_best_cfgs}")

Original file line number	Diff line number	Diff line change
`@@ -1157,6 +1157,7 @@ def launch():`
`1157`	`1157`	`cur_best_cfgs, err = recorder.get_best(`
`1158`	`1158`	`metric=tuner_cfg['metric_cfg']['name'],`
`1159`	`1159`	`direction=tuner_cfg['metric_cfg']['OptimizationDirection'],`
	`1160`	`+ mode=mode,`
`1160`	`1161`	`)`
`1161`	`1162`	`if not err:`
`1162`	`1163`	`ctx.logger.info(f"Current best config: {cur_best_cfgs}")`