[AutoConfig] Add DP Estimation (#59748)

Caozhou1995 · web-flow · commit 942a9660c93c · 2023-12-08T10:46:19.000+08:00
* add dp estimation

* update prune to common prune and history prune

* fix prune bug

* fix sharding overlap bug in grid search

* remove has_error tmp info
diff --git a/python/paddle/distributed/auto_tuner/cost_model.py b/python/paddle/distributed/auto_tuner/cost_model.py
@@ -86,7 +86,7 @@ def divisor(num, reverse=False):
 def get_not_oom_cfgs(cfgs, tuner_cfg):
     """Get not OOM parallel strategies."""
     total_cards, l, h, a, V, s, gbs, per_card_memory = (
-        tuner_cfg["estimated_num_gpus"],
+        tuner_cfg["search_algo"]["estimated_num_gpus"],
         tuner_cfg["model_cfg"]["num_layers"],
         tuner_cfg["model_cfg"]["hidden_size"],
         tuner_cfg["model_cfg"]["num_attention_heads"],
diff --git a/python/paddle/distributed/auto_tuner/prune.py b/python/paddle/distributed/auto_tuner/prune.py
@@ -13,9 +13,10 @@
 # limitations under the License.
 
 _PRUNE_FUNC = []
+_PRUNE_HISTORY_FUNC = []
 
 
-def same_cfgs_beside(attr, cur_cfg, history_cfgs):
+def same_cfgs_beside(attr, cur_cfg, history_cfgs=[]):
     """
     Compare the current configuration with the history configuration,
     and obtain the same configurations as the current configuration except for the given attr.
@@ -36,6 +37,30 @@ def same_cfgs_beside(attr, cur_cfg, history_cfgs):
     return results
 
 
+def same_cfgs_beside_sharding_overlap(tuner_cfg, cur_cfg, history_cfgs=[]):
+    result = None
+    for cfg in history_cfgs:
+        keys = [
+            "dp_degree",
+            "mp_degree",
+            "pp_degree",
+            "vpp_degree",
+            "micro_batch_size",
+            "use_recompute",
+            "recompute_granularity",
+            "sharding_stage",
+        ]
+        same = True
+        for key in keys:
+            if cfg[key] != cur_cfg[key]:
+                same = False
+                break
+        if same:
+            result = cfg
+            break
+    return result
+
+
 def register_prune(func):
     def wrapper(*args, **kwargs):
         return func(*args, **kwargs)
@@ -44,8 +69,16 @@ def wrapper(*args, **kwargs):
     return wrapper
 
 
+def register_prune_history(func):
+    def wrapper(*args, **kwargs):
+        return func(*args, **kwargs)
+
+    _PRUNE_HISTORY_FUNC.append(wrapper)
+    return wrapper
+
+
 @register_prune
-def prune_by_mp(tuner_cfg, cur_cfg, history_cfgs=None):
+def prune_by_mp(tuner_cfg, cur_cfg, history_cfgs=[]):
     """
     Prune by mp, the rules are:
     1. MP degree should be evenly divided by hidden size and vocab size
@@ -89,7 +122,7 @@ def prune_by_mp(tuner_cfg, cur_cfg, history_cfgs=None):
 
 
 @register_prune
-def prune_by_pp(tuner_cfg, cur_cfg, history_cfgs=None):
+def prune_by_pp(tuner_cfg, cur_cfg, history_cfgs=[]):
     """
     Prune by pp (pipeline-parallelism), the rules are:
     1. PP degree should be evenly divided by number of layers.
@@ -98,7 +131,9 @@ def prune_by_pp(tuner_cfg, cur_cfg, history_cfgs=None):
     """
     pp_degree = cur_cfg.get("pp_degree", None)
     num_layers = tuner_cfg["model_cfg"].get("num_layers", None)
-    num_nodes = tuner_cfg.get("nodes", 1)
+    num_nodes = (
+        cur_cfg["nodes"] if "nodes" in cur_cfg else tuner_cfg.get("nodes", 1)
+    )
 
     if pp_degree is None:
         return False
@@ -120,7 +155,7 @@ def prune_by_pp(tuner_cfg, cur_cfg, history_cfgs=None):
 
 
 @register_prune
-def prune_by_vpp(tuner_cfg, cur_cfg, history_cfgs=None):
+def prune_by_vpp(tuner_cfg, cur_cfg, history_cfgs=[]):
     """
     Prune by vpp (virtual pipeline parallelism), the rules are:
     1. VPP degree should be evenly divided by number of layers.
@@ -137,6 +172,19 @@ def prune_by_vpp(tuner_cfg, cur_cfg, history_cfgs=None):
         return False
 
     if num_layers:
+        global_batch_size = (
+            cur_cfg["global_batch_size"]
+            if "global_batch_size" in cur_cfg
+            else tuner_cfg["model_cfg"].get("global_batch_size", None)
+        )
+        acc_steps = (
+            global_batch_size
+            // cur_cfg["dp_degree"]
+            // cur_cfg["sharding_degree"]
+            // cur_cfg["micro_batch_size"]
+        )
+        if vpp_degree > 1 and acc_steps % pp_degree != 0:
+            return True
         if num_layers % (pp_degree * vpp_degree) != 0:
             return True
         if pp_degree == 1 and vpp_degree != 1:
@@ -151,6 +199,14 @@ def prune_by_vpp(tuner_cfg, cur_cfg, history_cfgs=None):
         if vpp_degree not in vpp_degree_candidates:
             return True
 
+    return False
+
+
+@register_prune_history
+def prune_by_vpp_history(tuner_cfg, cur_cfg, history_cfgs=[]):
+    vpp_degree = cur_cfg.get("vpp_degree", None)
+    if vpp_degree is None:
+        return False
     cfgs = same_cfgs_beside("vpp_degree", cur_cfg, history_cfgs)
     if cfgs:
         for cfg in cfgs:
@@ -164,15 +220,19 @@ def prune_by_vpp(tuner_cfg, cur_cfg, history_cfgs=None):
 
 
 @register_prune
-def prune_by_mbs(tuner_cfg, cur_cfg, history_cfgs=None):
+def prune_by_mbs(tuner_cfg, cur_cfg, history_cfgs=[]):
     """
     Prune by mbs (micro batch size), the rules are:
     1. Micro batch size should be evenly divided by the local batch size.
     2. Micro batch size should be in the candidates of user defined.
     3. Prune if a similar configuration with a larger micro batch size resulted in a valid run.
     """
     micro_batch_size = cur_cfg.get("micro_batch_size", None)
-    global_batch_size = tuner_cfg["model_cfg"].get("global_batch_size", None)
+    global_batch_size = (
+        cur_cfg["global_batch_size"]
+        if "global_batch_size" in cur_cfg
+        else tuner_cfg["model_cfg"].get("global_batch_size", None)
+    )
     if global_batch_size == "auto":
         global_batch_size = cur_cfg["global_batch_size"]
     if global_batch_size:
@@ -207,6 +267,14 @@ def prune_by_mbs(tuner_cfg, cur_cfg, history_cfgs=None):
         if micro_batch_size not in mbs_candidates:
             return True
 
+    return False
+
+
+@register_prune_history
+def prune_by_mbs_history(tuner_cfg, cur_cfg, history_cfgs=[]):
+    micro_batch_size = cur_cfg.get("micro_batch_size", None)
+    if micro_batch_size is None:
+        return False
     cfgs = same_cfgs_beside("micro_batch_size", cur_cfg, history_cfgs)
     if cfgs:
         for cfg in cfgs:
@@ -222,12 +290,11 @@ def prune_by_mbs(tuner_cfg, cur_cfg, history_cfgs=None):
                 and cfg.get("max_mem_usage") == "OOM"
             ):
                 return True
-
     return False
 
 
 @register_prune
-def prune_by_sharding(tuner_cfg, cur_cfg, history_cfgs):
+def prune_by_sharding(tuner_cfg, cur_cfg, history_cfgs=[]):
     """
     Prune by sharding parameters, the rules are:
     1. Sharding stage and sharding degree should be specified.
@@ -264,6 +331,18 @@ def prune_by_sharding(tuner_cfg, cur_cfg, history_cfgs):
     if pp_degree and pp_degree != 1 and sharding_stage != 1:
         return True
 
+    return False
+
+
+@register_prune_history
+def prune_by_sharding_history(tuner_cfg, cur_cfg, history_cfgs=[]):
+    sharding_degree = cur_cfg.get("sharding_degree", None)
+    if sharding_degree is None:
+        return False
+
+    sharding_stage = cur_cfg.get("sharding_stage", None)
+    if sharding_stage is None:
+        return False
     cfgs = same_cfgs_beside("sharding_stage", cur_cfg, history_cfgs)
     if cfgs:
         for cfg in cfgs:
@@ -289,7 +368,7 @@ def prune_by_sharding(tuner_cfg, cur_cfg, history_cfgs):
 
 
 @register_prune
-def prune_by_recompute(tuner_cfg, cur_cfg, history_cfgs):
+def prune_by_recompute(tuner_cfg, cur_cfg, history_cfgs=[]):
     """
     Prune by recompute parameters, the rules are:
     1. If recompute is not used, return False directly.
@@ -317,6 +396,14 @@ def prune_by_recompute(tuner_cfg, cur_cfg, history_cfgs):
         if recompute_granularity not in recompute_granularity_candidates:
             return True
 
+    return False
+
+
+@register_prune_history
+def prune_by_recompute_history(tuner_cfg, cur_cfg, history_cfgs=[]):
+    use_recompute = cur_cfg.get("use_recompute", None)
+    if use_recompute is None:
+        return False
     cfgs = same_cfgs_beside("use_recompute", cur_cfg, history_cfgs)
     if cfgs:
         for cfg in cfgs:
@@ -343,8 +430,12 @@ def prune_by_recompute(tuner_cfg, cur_cfg, history_cfgs):
 
 
 @register_prune
-def prune_by_num_gpus(tuner_cfg, cur_cfg, history_cfgs):
-    num_gpus = tuner_cfg.get("num_gpus")
+def prune_by_num_gpus(tuner_cfg, cur_cfg, history_cfgs=[]):
+    num_gpus = (
+        cur_cfg["num_gpus"]
+        if "num_gpus" in cur_cfg
+        else tuner_cfg.get("num_gpus")
+    )
     dp_degree = cur_cfg.get("dp_degree", 1)
     mp_degree = cur_cfg.get("mp_degree", 1)
     pp_degree = cur_cfg.get("pp_degree", 1)
@@ -353,3 +444,17 @@ def prune_by_num_gpus(tuner_cfg, cur_cfg, history_cfgs):
         return True
 
     return False
+
+
+@register_prune
+def prune_by_sharding_overlap(tuner_cfg, cur_cfg, history_cfgs=[]):
+    """Prune by sharding overlap for single dp estimation"""
+    if "sharding_overlap" in cur_cfg:
+        result = same_cfgs_beside_sharding_overlap(
+            tuner_cfg, cur_cfg, history_cfgs
+        )
+        if not result:
+            return True
+        if not result[tuner_cfg['metric_cfg']['name']]:
+            return True
+    return False
diff --git a/python/paddle/distributed/auto_tuner/recorder.py b/python/paddle/distributed/auto_tuner/recorder.py
@@ -88,6 +88,8 @@ def store_history(self, path="./history.csv"):
         # check if 'time' exists
         if 'time' in df.columns:
             df = df.drop(columns=['time'])
+        if 'has_error' in df.columns:
+            df = df.drop(columns=['has_error'])
         # write to csv
         df.to_csv(self.store_path, index=False)
 
diff --git a/python/paddle/distributed/auto_tuner/search.py b/python/paddle/distributed/auto_tuner/search.py
@@ -15,7 +15,7 @@
 
 from abc import ABC, abstractmethod
 
-from .prune import _PRUNE_FUNC
+from .prune import _PRUNE_HISTORY_FUNC
 from .utils import gbs_search_all, search_all, search_by_dp_estimation
 
 
@@ -28,7 +28,7 @@ def search_once(self, history_cfgs):
         pass
 
     def prune(self, tuner_cfg, cur_cfg, history_cfgs):
-        for func in _PRUNE_FUNC:
+        for func in _PRUNE_HISTORY_FUNC:
             result = func(tuner_cfg, cur_cfg, history_cfgs)
             if result:
                 return True
@@ -59,15 +59,9 @@ def __init__(self, tuner_cfg):
         super().__init__(tuner_cfg)
         self.idx = 0
         self.all_tasks = search_by_dp_estimation(tuner_cfg)
-        assert len(self.all_tasks) > 0, "Unable to perform this search."
-        # change global_batch_size and dp_degree
-        tuner_cfg["model_cfg"]["global_batch_size"] = (
-            tuner_cfg["model_cfg"]["global_batch_size"]
-            // self.all_tasks[0]["dp_degree"]
-        )
-        for task in self.all_tasks:
-            task["estimated_dp_degree"] = task["dp_degree"]
-            task["dp_degree"] = 1
+        assert (
+            len(self.all_tasks) > 0
+        ), "Unable to perform single dp estimation search."
 
     def search_once(self, history_cfgs):
         new_cfg = None
diff --git a/python/paddle/distributed/auto_tuner/utils.py b/python/paddle/distributed/auto_tuner/utils.py
diff --git a/python/paddle/distributed/launch/main.py b/python/paddle/distributed/launch/main.py
