[AutoTuner] Add custom search dim (#62724)

Difers · web-flow · commit 5c87fe1733ca · 2024-03-15T11:36:17.000+08:00
* add custom search dim

* fix sharding prune and custom pruned rules
diff --git a/python/paddle/distributed/auto_tuner/prune.py b/python/paddle/distributed/auto_tuner/prune.py
@@ -40,6 +40,12 @@ def log_pruned_info(cur_cfg, pruned_reason, tuner_cfg):
             strategy += str(cur_cfg[key])
             pruned_strategy = pruned_strategy + "_" + strategy
 
+    if "custom_search_dim" in tuner_cfg:
+        for key in tuner_cfg["custom_search_dim"]:
+            strategy = "".join(i.capitalize() for i in key.split("_"))
+            strategy += str(cur_cfg[key])
+            pruned_strategy = pruned_strategy + "_" + strategy
+
     try:
         from paddle.distributed.launch.main import ctx
 
@@ -204,13 +210,9 @@ def prune_by_mp_pp_history(tuner_cfg, cur_cfg, history_cfgs, pruned_cfgs):
 
     if mp_degree is None or pp_degree is None or use_recompute is None:
         return False
-
+    history_cfgs = copy.deepcopy(history_cfgs)
     history_cfgs.extend(pruned_cfgs)
     cfgs = same_cfgs_beside(["mp_degree", "pp_degree"], cur_cfg, history_cfgs)
-    if cur_cfg.get("sharding_degree") == 1:
-        cfgs = same_cfgs_beside(
-            ["mp_degree", "pp_degree", "sharding_satge"], cur_cfg, history_cfgs
-        )
 
     if cfgs:
         for cfg in cfgs:
@@ -281,14 +283,10 @@ def prune_by_vpp_history(tuner_cfg, cur_cfg, history_cfgs=[], pruned_cfgs=[]):
     vpp_degree = cur_cfg.get("vpp_degree", None)
     if vpp_degree is None:
         return False
-
+    history_cfgs = copy.deepcopy(history_cfgs)
     history_cfgs.extend(pruned_cfgs)
 
     cfgs = same_cfgs_beside("vpp_degree", cur_cfg, history_cfgs)
-    if cur_cfg.get("sharding_degree") == 1:
-        cfgs = same_cfgs_beside(
-            ["vpp_degree", "sharding_satge"], cur_cfg, history_cfgs
-        )
 
     if cfgs:
         for cfg in cfgs:
@@ -364,18 +362,12 @@ def prune_by_mbs_history(tuner_cfg, cur_cfg, history_cfgs=[], pruned_cfgs=[]):
     micro_batch_size = cur_cfg.get("micro_batch_size", None)
     if micro_batch_size is None:
         return False
-
+    history_cfgs = copy.deepcopy(history_cfgs)
     history_cfgs.extend(pruned_cfgs)
 
     cfgs = same_cfgs_beside(
         ["micro_batch_size", "acc_steps"], cur_cfg, history_cfgs
     )
-    if cur_cfg.get("sharding_degree") == 1:
-        cfgs = same_cfgs_beside(
-            ["micro_batch_size", "sharding_satge", "acc_steps"],
-            cur_cfg,
-            history_cfgs,
-        )
 
     if cfgs:
         for cfg in cfgs:
@@ -435,7 +427,12 @@ def prune_by_sharding(tuner_cfg, cur_cfg, history_cfgs=[]):
         if sharding_degree not in sharding_degree_candidates:
             return True
 
-    if pp_degree and pp_degree != 1 and sharding_stage != 1:
+    if (
+        pp_degree
+        and pp_degree != 1
+        and sharding_stage != 1
+        and sharding_degree != 1
+    ):
         return True
 
     if sharding_degree == 1:
@@ -457,7 +454,7 @@ def prune_by_sharding_history(
     sharding_stage = cur_cfg.get("sharding_stage", None)
     if sharding_stage is None:
         return False
-
+    history_cfgs = copy.deepcopy(history_cfgs)
     history_cfgs.extend(pruned_cfgs)
 
     cfgs = same_cfgs_beside("sharding_stage", cur_cfg, history_cfgs)
@@ -555,17 +552,12 @@ def prune_by_recompute_history(
     if recompute_level is None:
         return False
 
+    history_cfgs = copy.deepcopy(history_cfgs)
     history_cfgs.extend(pruned_cfgs)
 
     cfgs = same_cfgs_beside(
         ["use_recompute", "recompute_granularity"], cur_cfg, history_cfgs
     )
-    if cur_cfg.get("sharding_degree") == 1:
-        cfgs = same_cfgs_beside(
-            ["use_recompute", "recompute_granularity", "sharding_satge"],
-            cur_cfg,
-            history_cfgs,
-        )
 
     if cfgs:
         for cfg in cfgs:
@@ -862,6 +854,7 @@ def prune_by_refined_recompute_history(
     tuner_cfg, cur_cfg, history_cfgs=[], pruned_cfgs=[]
 ):
     if tuner_cfg.get("refined_recompute", None):
+        history_cfgs = copy.deepcopy(history_cfgs)
         history_cfgs.extend(pruned_cfgs)
         rr = tuner_cfg.get("refined_recompute")
         compare = copy.deepcopy(rr)
@@ -898,3 +891,44 @@ def prune_by_refined_recompute_history(
                         return True
 
     return False
+
+
+@register_prune_history
+def prune_by_custom_search_dim_history(
+    tuner_cfg, cur_cfg, history_cfgs=[], pruned_cfgs=[]
+):
+    history_cfgs = copy.deepcopy(history_cfgs)
+    custom_search_dim = tuner_cfg.get("custom_search_dim", None)
+    prune_custom_search_dim = []
+    custom_dim_level = {}
+    if custom_search_dim is not None:
+        for key, value in custom_search_dim.items():
+            if value["prune"]:
+                prune_custom_search_dim.append(key)
+                # In the custom_search_dim, the values are ordered according to the sequence specified in its custom configuration.
+                custom_dim_level[key] = {
+                    key: value for value, key in enumerate(value["value"])
+                }
+
+    for key in prune_custom_search_dim:
+        history_cfgs.extend(pruned_cfgs)
+        cfgs = same_cfgs_beside(key, cur_cfg, history_cfgs)
+        cur_value = cur_cfg.get(key, None)
+        if cur_value is None:
+            return False
+
+        # In the custom_search_dim, based on the order of values provided in its custom configuration, if a configuration is found to be executable, the subsequent configurations will be pruned.
+        if cfgs:
+            for cfg in cfgs:
+                cfg_value = cfg[key]
+                if (
+                    custom_dim_level[key][cfg_value]
+                    < custom_dim_level[key][cur_value]
+                    and cfg.get("time", -1) > 0
+                ):
+                    pruned_reason = f"{key}{cfg_value} may be slower because {key}{cur_value} has been already runnable."
+                    log_pruned_info(cur_cfg, pruned_reason, tuner_cfg)
+                    cur_cfg["time"] = cfg["time"]
+                    return True
+
+    return False
diff --git a/python/paddle/distributed/auto_tuner/tuner.py b/python/paddle/distributed/auto_tuner/tuner.py
@@ -138,6 +138,10 @@ def get_cfg_from_resume(self, cur_cfg):
             for rr in self.tuner_cfg["refined_recompute"]:
                 keys_to_compare.append(rr)
 
+        if self.tuner_cfg.get("custom_search_dim", None):
+            for key in self.tuner_cfg["custom_search_dim"]:
+                keys_to_compare.append(key)
+
         for cfg in self.resume_cfgs:
             ret_is_same = True
             for key in keys_to_compare:
diff --git a/python/paddle/distributed/auto_tuner/utils.py b/python/paddle/distributed/auto_tuner/utils.py
@@ -295,6 +295,11 @@ def default_candidates(tuner_cfg):
         raise ValueError(
             f"recompute_granularity only supports auto/{'/'.join(__SUPPORTED_RECOMPUTE_GRANULARITY__)}, but got {recompute_granularity}"
         )
+    custom_search_dim = tuner_cfg.get("custom_search_dim", None)
+    if custom_search_dim is not None:
+        candidates["custom_search_dim"] = []
+        for key, value in custom_search_dim.items():
+            candidates["custom_search_dim"].append(value["value"])
     return candidates
 
 
@@ -359,6 +364,17 @@ def search_all(tuner_cfg):
         )
     )
 
+    custom_search_dim = tuner_cfg.get("custom_search_dim", None)
+    if custom_search_dim is not None:
+        custom_search_dim_candidates = candidates["custom_search_dim"]
+        custom_dim_cfgs = list(itertools.product(*custom_search_dim_candidates))
+        other_cfgs_without_cumtom = other_dim_cfgs
+        other_dim_cfgs = []
+        for cfg_without_cumtom in other_cfgs_without_cumtom:
+            for custom_cfg in custom_dim_cfgs:
+                cfg = list(cfg_without_cumtom) + list(custom_cfg)
+                other_dim_cfgs.append(cfg)
+
     all_cfgs = []
     refined_recompute = tuner_cfg.get("refined_recompute", None)
     for valid_degree in valid_degrees:
@@ -370,7 +386,7 @@ def search_all(tuner_cfg):
                 vpp,
                 use_recompute,
                 recompute_granularity,
-            ) = list(other_dim_cfg)
+            ) = list(other_dim_cfg[:5])
             if (
                 tuner_cfg["model_cfg"]["global_batch_size"]
                 % (mbs * sharding_degree * dp_degree)
@@ -456,6 +472,10 @@ def search_all(tuner_cfg):
         8: "recompute_granularity",
     }
 
+    if custom_search_dim is not None:
+        for key, _ in custom_search_dim.items():
+            mapping[len(mapping)] = key
+
     if refined_recompute is not None:
         for dim in refined_recompute:
             mapping[len(mapping)] = dim
@@ -1467,20 +1487,29 @@ def _gen_new_arg(arg, cmd, cfg, res_args, tuner_cfg):
     cmd = copy.deepcopy(tuner_cfg["run_cmd"])
     res_args = copy.deepcopy(raw_args)
 
-    _gen_new_arg("dp_degree", cmd, cfg, res_args, tuner_cfg)
-    _gen_new_arg("mp_degree", cmd, cfg, res_args, tuner_cfg)
-    _gen_new_arg("pp_degree", cmd, cfg, res_args, tuner_cfg)
-    _gen_new_arg("vpp_degree", cmd, cfg, res_args, tuner_cfg)
-    _gen_new_arg("micro_batch_size", cmd, cfg, res_args, tuner_cfg)
-    _gen_new_arg("sharding_degree", cmd, cfg, res_args, tuner_cfg)
-    _gen_new_arg("sharding_stage", cmd, cfg, res_args, tuner_cfg)
-    _gen_new_arg("use_recompute", cmd, cfg, res_args, tuner_cfg)
-    _gen_new_arg("recompute_granularity", cmd, cfg, res_args, tuner_cfg)
-    _gen_new_arg("local_batch_size", cmd, cfg, res_args, tuner_cfg)
-    _gen_new_arg("gradient_accumulation_steps", cmd, cfg, res_args, tuner_cfg)
-    _gen_new_arg("global_batch_size", cmd, cfg, res_args, tuner_cfg)
-    _gen_new_arg("sequence_parallel", cmd, cfg, res_args, tuner_cfg)
-    _gen_new_arg("refined_recompute", cmd, cfg, res_args, tuner_cfg)
+    new_args = [
+        "dp_degree",
+        "mp_degree",
+        "pp_degree",
+        "vpp_degree",
+        "micro_batch_size",
+        "sharding_degree",
+        "sharding_stage",
+        "use_recompute",
+        "recompute_granularity",
+        "local_batch_size",
+        "gradient_accumulation_steps",
+        "global_batch_size",
+        "sequence_parallel",
+        "refined_recompute",
+    ]
+
+    if "custom_search_dim" in tuner_cfg:
+        for key in tuner_cfg["custom_search_dim"]:
+            new_args.append(key)
+
+    for arg in new_args:
+        _gen_new_arg(arg, cmd, cfg, res_args, tuner_cfg)
 
     if tuner_cfg["run_cmd"].get("search_stage", None) and not run_best:
         cmd = copy.deepcopy(tuner_cfg["run_cmd"]["search_stage"])
diff --git a/python/paddle/distributed/launch/main.py b/python/paddle/distributed/launch/main.py
@@ -649,6 +649,12 @@ def launch():
                     dir_name += str(cur_cfg[key])
                     log_dir = log_dir + "_" + dir_name
 
+            if "custom_search_dim" in tuner_cfg:
+                for key in tuner_cfg["custom_search_dim"]:
+                    dir_name = "".join(i.capitalize() for i in key.split("_"))
+                    dir_name += str(cur_cfg[key])
+                    log_dir = log_dir + "_" + dir_name
+
             ctx.args.log_dir = os.path.join(
                 os.path.dirname(ctx.args.auto_tuner_json), log_dir
             )
