chore: import upstream snapshot with attribution

python/paddle/distributed/auto_tuner/__init__.py

@@ -0,0 +1,15 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+__all__ = []

python/paddle/distributed/auto_tuner/cost_model.py

@@ -0,0 +1,143 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+def all_params(mp, pp, sharding, h, l, V):
+    # TODO: TBD - add some fixed structure models.
+    return 1
+
+
+def full_recompute_acts(mp, pp, s, b, h, l):
+    # TODO: TBD - add some fixed structure models.
+    return 1
+
+
+def all_acts(mp, pp, s, b, h, l, a):
+    # TODO: TBD - add some fixed structure models.
+    return 1
+
+
+def to_gb(p):
+    return p / (2**30)
+
+
+def get_mem(total_cards, parallel_cfg, l, h, a, V, s, gbs):
+    """Estimate the memory of model unset parallel strategy."""
+    sharding = parallel_cfg["sharding_degree"]
+    mp = parallel_cfg["mp_degree"]
+    b = parallel_cfg["micro_batch_size"]
+    pp = parallel_cfg["pp_degree"]
+    vpp = parallel_cfg["vpp_degree"]
+    use_recompute = parallel_cfg["use_recompute"]
+
+    sep = 1
+
+    lbs = int(gbs / sharding / s)
+    lbs = int(lbs / pp) * pp
+    assert s % sep == 0
+    s_sep = s // sep
+    assert a % (sep * mp) == 0, f'{a} vs {sep * mp}'
+
+    vpp_ratio = 1
+    if vpp > 1:
+        assert l % (pp * vpp) == 0
+        vpp_ratio = 1 + (pp - 1) / (pp * vpp)
+
+    params = to_gb(all_params(mp, pp, sharding, h, l, V))
+
+    acts = 0
+    assert l % pp == 0
+
+    if use_recompute:
+        acts = to_gb(full_recompute_acts(mp, pp, s_sep, b, h, l)) * vpp_ratio
+    else:
+        acts = to_gb(all_acts(mp, pp, s, b, h, l, a)) * vpp_ratio
+    assert acts > 0
+
+    peak_mem = params + acts
+    return peak_mem
+
+
+def divisor(num, reverse=False):
+    """Get the divisor of a given number."""
+    results = set()
+    i = 1
+    mid = num // 2 + 1
+    while i < mid:
+        if num % i == 0:
+            results.add(i)
+            results.add(num // i)
+        i += 1
+    results = list(results)
+    return sorted(results, reverse=reverse)
+
+
+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["search_algo"]["estimated_num_gpus"],
+        tuner_cfg["model_cfg"]["num_layers"],
+        tuner_cfg["model_cfg"]["hidden_size"],
+        tuner_cfg["model_cfg"]["num_attention_heads"],
+        tuner_cfg["model_cfg"]["vocab_size"],
+        tuner_cfg["model_cfg"]["seq_length"],
+        tuner_cfg["model_cfg"]["global_batch_size"],
+        tuner_cfg.get("per_card_memory", 80),
+    )
+    pruned_cfgs = []
+    for cfg in cfgs:
+        mp = cfg["mp_degree"]
+        sharding = cfg["sharding_degree"]
+        mbs = cfg["micro_batch_size"]
+        pp = cfg["pp_degree"]
+        vpp = cfg["vpp_degree"]
+        dp = cfg["dp_degree"]
+        use_recompute = cfg["use_recompute"]
+
+        if mp * sharding * pp * dp != total_cards:
+            continue
+        if gbs % sharding != 0:
+            continue
+        if gbs // sharding % dp != 0:
+            continue
+        if gbs // sharding // dp % mbs != 0:
+            continue
+        if l % pp != 0:
+            continue
+        if l // pp % vpp != 0:
+            continue
+        if vpp != 1 and pp <= 2:
+            continue
+        if a % mp != 0 or V % mp != 0 or h % mp != 0:
+            continue
+
+        pruned_cfgs.append(cfg)
+    valid_cfgs = []
+    for cfg in pruned_cfgs:
+        mem = get_mem(total_cards, cfg, l, h, a, V, s, gbs)
+        # TODO: Uncomment when it is actually implemented.
+        # if (
+        #     mem < per_card_memory
+        #     and mem
+        #     > tuner_cfg.get(
+        #         "search_algo", {"name": "dp_estimation", "threshold": 0.7}
+        #     ).get("threshold", 0.7)
+        #     * per_card_memory
+        # ):
+        # cfg["memory_cost"] = mem
+        # valid_cfgs.append(cfg)
+        cfg["memory_cost"] = mem
+        valid_cfgs.append(cfg)
+    assert valid_cfgs
+    return valid_cfgs

python/paddle/distributed/auto_tuner/memory_cost_model.py

@@ -0,0 +1,95 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from argparse import ArgumentParser
+
+
+def parse_arguments():
+    parser = ArgumentParser()
+
+    # for distributed strategy
+    parser.add_argument(
+        "--dp_degree", type=int, required=True, help="dp degree"
+    )
+    parser.add_argument(
+        "--mp_degree", type=int, required=True, help="mp degree"
+    )
+    parser.add_argument(
+        "--pp_degree", type=int, required=True, help="pp degree"
+    )
+    parser.add_argument(
+        "--vpp_degree", type=int, required=True, help="vpp degree"
+    )
+    parser.add_argument(
+        "--sharding_degree", type=int, required=True, help="sharding degree"
+    )
+    parser.add_argument(
+        "--sharding_stage", type=int, required=True, help="sharding stage"
+    )
+    parser.add_argument(
+        "--micro_batch_size", type=int, required=True, help="micro batch size"
+    )
+    parser.add_argument(
+        "--use_recompute", type=bool, required=True, help="use recompute"
+    )
+    parser.add_argument(
+        "--recompute_granularity",
+        type=str,
+        required=True,
+        choices=["None", "core_attn", "full_attn", "full"],
+        help="recompute granularity",
+    )
+
+    # for model config
+    parser.add_argument(
+        "--hidden_size", type=int, required=False, help="hidden size"
+    )
+    parser.add_argument(
+        "--num_attention_heads",
+        type=int,
+        required=False,
+        help="number of attention heads",
+    )
+    parser.add_argument(
+        "--num_layers", type=int, required=False, help="number of hidden layers"
+    )
+    parser.add_argument(
+        "--max_sequence_length",
+        type=int,
+        required=False,
+        help="maximum sequence length",
+    )
+    parser.add_argument(
+        "--vocab_size", type=int, required=False, help="vocabulary size"
+    )
+    parser.add_argument(
+        "--intermediate_size",
+        type=int,
+        required=False,
+        help="intermediate size",
+    )
+
+    return parser.parse_args()
+
+
+def get_model_memory_usage(args):
+    # evaluate model memory usage based on distributed strategy and model setting
+    raise NotImplementedError(
+        "Please implement this function for memory usage estimation based on distributed strategy and model setting."
+    )
+
+
+if __name__ == "__main__":
+    args = parse_arguments()
+    print(get_model_memory_usage(args))

python/paddle/distributed/auto_tuner/prune.py

@@ -0,0 +1,934 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import copy
+import logging
+import os
+import subprocess
+
+logger = logging.getLogger('auto_tuner')
+_PRUNE_FUNC = []
+_PRUNE_HISTORY_FUNC = []
+
+
+def log_pruned_info(cur_cfg, pruned_reason, tuner_cfg):
+    pruned_strategy = "DP{}_MP{}_PP{}_VPP{}_Sharding{}_Stage{}_MBS{}_Recompute_{}_Granularity_{}".format(
+        cur_cfg["dp_degree"],
+        cur_cfg["mp_degree"],
+        cur_cfg["pp_degree"],
+        cur_cfg["vpp_degree"],
+        cur_cfg["sharding_degree"],
+        cur_cfg["sharding_stage"],
+        cur_cfg["micro_batch_size"],
+        cur_cfg["use_recompute"],
+        cur_cfg["recompute_granularity"],
+    )
+    if "refined_recompute" in tuner_cfg:
+        for key in tuner_cfg["refined_recompute"]:
+            strategy = "".join(i.capitalize() for i in key.split("_"))
+            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
+
+        ctx.logger.info(
+            f"Strategy {pruned_strategy} has been pruned that {pruned_reason}"
+        )
+    except:
+        pass
+    logger.info(
+        f"Strategy {pruned_strategy} has been pruned that {pruned_reason}"
+    )
+
+
+def same_cfgs_beside(attrs, 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.
+    """
+    results = []
+    same = True
+
+    for cfg in history_cfgs:
+        for key in cur_cfg:
+            if key in attrs:
+                continue
+            if key not in cfg or (
+                cfg[key] != cur_cfg[key]
+                and key not in ["estimated_memory_usage"]
+            ):
+                same = False
+                break
+        if same:
+            results.append(cfg)
+        else:
+            same = True
+
+    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)
+
+    _PRUNE_FUNC.append(wrapper)
+    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=[]):
+    """
+    Prune by mp, the rules are:
+    1. MP degree should be evenly divided by hidden size and vocab size
+    2. MP degree should be in the candidates of user defined.
+    3. MP degree should be less than 8 if no candidates.
+    """
+    mp_degree = cur_cfg.get("mp_degree", None)
+    hidden_size = tuner_cfg["model_cfg"].get("hidden_size", None)
+    vocab_size = tuner_cfg["model_cfg"].get("vocab_size", None)
+    num_attention_heads = tuner_cfg["model_cfg"].get(
+        "num_attention_heads", None
+    )
+    seq_length = tuner_cfg["model_cfg"].get("seq_length", None)
+    use_sequence_parallel = tuner_cfg.get("use_sequence_parallel", False)
+
+    if mp_degree is None:
+        return False
+
+    if hidden_size and hidden_size % mp_degree != 0:
+        return True
+
+    if vocab_size and vocab_size % mp_degree != 0:
+        return True
+
+    if num_attention_heads and num_attention_heads % mp_degree != 0:
+        return True
+
+    if seq_length and seq_length % mp_degree != 0 and use_sequence_parallel:
+        return True
+
+    mp_degree_candidates = tuner_cfg.get("mp_degree", None)
+
+    if mp_degree_candidates == "auto":
+        mp_degree_candidates = tuner_cfg["candidates"]["mp_degree"]
+
+    if mp_degree_candidates:
+        if mp_degree not in mp_degree_candidates:
+            return True
+
+    return False
+
+
+@register_prune
+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.
+    2. PP degree should be in the candidates of user defined.
+    3. If no candidates, PP degree should be less than or equal to the number of nodes.
+    """
+    pp_degree = cur_cfg.get("pp_degree", None)
+    num_layers = tuner_cfg["model_cfg"].get("num_layers", None)
+    num_nodes = (
+        cur_cfg["nodes"] if "nodes" in cur_cfg else tuner_cfg.get("nodes", 1)
+    )
+
+    if pp_degree is None:
+        return False
+
+    if num_layers:
+        if num_layers % pp_degree != 0:
+            return True
+
+    pp_degree_candidates = tuner_cfg.get("pp_degree", None)
+    if pp_degree_candidates == "auto":
+        pp_degree_candidates = tuner_cfg["candidates"]["pp_degree"]
+    if pp_degree_candidates:
+        if pp_degree not in pp_degree_candidates:
+            return True
+    else:
+        if num_nodes != 1 and pp_degree > num_nodes:
+            return True
+    return False
+
+
+@register_prune_history
+def prune_by_mp_pp_history(tuner_cfg, cur_cfg, history_cfgs, pruned_cfgs):
+    mp_degree = cur_cfg.get("mp_degree", None)
+    pp_degree = cur_cfg.get("pp_degree", None)
+    use_recompute = cur_cfg.get("recompute", None)
+
+    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 cfgs:
+        for cfg in cfgs:
+            if (
+                not use_recompute
+                and cfg["mp_degree"] * cfg["pp_degree"] == mp_degree * pp_degree
+                and cfg["mp_degree"] > mp_degree
+                and cfg.get("max_mem_usage") == "OOM"
+            ):
+                pruned_reason = f"mp_degree {mp_degree}, pp_degree {pp_degree} may cause oom because {cfg['mp_degree']}, {cfg['pp_degree']} already oom."
+                log_pruned_info(cur_cfg, pruned_reason, tuner_cfg)
+                cur_cfg["max_mem_usage"] = "OOM"
+                return True
+
+    return False
+
+
+@register_prune
+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.
+    2. VPP degree should be in the candidates of user defined.
+    """
+    pp_degree = cur_cfg.get("pp_degree", None)
+    vpp_degree = cur_cfg.get("vpp_degree", None)
+    num_layers = tuner_cfg["model_cfg"].get("num_layers", None)
+
+    if pp_degree is None:
+        return False
+
+    if vpp_degree is 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:
+            return True
+        if pp_degree <= 2 and vpp_degree != 1:
+            return True
+
+    vpp_degree_candidates = tuner_cfg.get("vpp_degree", None)
+    if vpp_degree_candidates == "auto":
+        vpp_degree_candidates = tuner_cfg["candidates"]["vpp_degree"]
+    if vpp_degree_candidates:
+        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=[], 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 cfgs:
+        for cfg in cfgs:
+            # memory prune
+            if (
+                cfg["vpp_degree"] > vpp_degree
+                and cfg.get("max_mem_usage") == "OOM"
+            ):
+                pruned_reason = f"vpp_degree {vpp_degree} may cause oom because {cfg['vpp_degree']} already oom."
+                log_pruned_info(cur_cfg, pruned_reason, tuner_cfg)
+                cur_cfg["max_mem_usage"] = "OOM"
+                return True
+
+    return False
+
+
+@register_prune
+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 = (
+        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:
+        local_batch_size = (
+            global_batch_size
+            // cur_cfg["dp_degree"]
+            // cur_cfg["sharding_degree"]
+        )
+        if local_batch_size == 0:
+            return True
+
+    mbs_candidates = tuner_cfg.get("micro_batch_size", None)
+
+    if mbs_candidates == "auto":
+        mbs_candidates = tuner_cfg["candidates"]["micro_batch_size"]
+
+    if micro_batch_size is None:
+        return False
+
+    if local_batch_size:
+        if local_batch_size % micro_batch_size != 0:
+            return True
+        acc_steps = local_batch_size // micro_batch_size
+        pp_degree = cur_cfg.get("pp_degree", None)
+        if pp_degree is not None:
+            if acc_steps < pp_degree:
+                return True
+        vpp_degree = cur_cfg.get("vpp_degree", None)
+        if vpp_degree is not None and vpp_degree > 1:
+            if pp_degree is not None:
+                if acc_steps % pp_degree != 0:
+                    return True
+
+    if mbs_candidates:
+        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=[], 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 cfgs:
+        for cfg in cfgs:
+            if (
+                cfg["micro_batch_size"] > micro_batch_size
+                and cfg.get("time", -1) > 0
+            ):
+                pruned_reason = f"micro_batch_size {micro_batch_size} may be slower because {cfg['micro_batch_size']} has been already runnable."
+                log_pruned_info(cur_cfg, pruned_reason, tuner_cfg)
+                cur_cfg["time"] = cfg["time"]
+                return True
+            # memory prune
+            if (
+                cfg["micro_batch_size"] < micro_batch_size
+                and cfg.get("max_mem_usage") == "OOM"
+            ):
+                pruned_reason = f"micro_batch_size {micro_batch_size} may cause oom because {cfg['micro_batch_size']} already oom."
+                log_pruned_info(cur_cfg, pruned_reason, tuner_cfg)
+                cur_cfg["max_mem_usage"] = "OOM"
+                return True
+    return False
+
+
+@register_prune
+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.
+    2. Sharding stage and degree should be in the candidates of user defined.
+    3. If PP (pipeline-parallelism) degree is not 1, sharding stage must be 1.
+    4. Prune if a similar configuration with a lower sharding stage resulted in a valid run.
+    5. If sharding degree is 1, sharding stage is invalid.
+    """
+    sharding_stage = cur_cfg.get("sharding_stage", None)
+    sharding_degree = cur_cfg.get("sharding_degree", None)
+    pp_degree = cur_cfg.get("pp_degree", None)
+
+    if not sharding_stage:
+        return False
+
+    if not sharding_degree:
+        return False
+
+    sharding_stage_candidates = tuner_cfg.get("sharding_stage", None)
+    if sharding_stage_candidates == "auto":
+        sharding_stage_candidates = tuner_cfg["candidates"]["sharding_stage"]
+
+    sharding_degree_candidates = tuner_cfg.get("sharding_degree", None)
+    if sharding_degree_candidates == "auto":
+        sharding_degree_candidates = tuner_cfg["candidates"]["sharding_degree"]
+
+    if sharding_stage_candidates:
+        if sharding_stage not in sharding_stage_candidates:
+            return True
+
+    if sharding_degree_candidates:
+        if sharding_degree not in sharding_degree_candidates:
+            return True
+
+    if (
+        pp_degree
+        and pp_degree != 1
+        and sharding_stage != 1
+        and sharding_degree != 1
+    ):
+        return True
+
+    if sharding_degree == 1:
+        cfgs = same_cfgs_beside("sharding_stage", cur_cfg, history_cfgs)
+        if cfgs:
+            return True
+
+    return False
+
+
+@register_prune_history
+def prune_by_sharding_history(
+    tuner_cfg, cur_cfg, history_cfgs=[], pruned_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
+    history_cfgs = copy.deepcopy(history_cfgs)
+    history_cfgs.extend(pruned_cfgs)
+
+    cfgs = same_cfgs_beside("sharding_stage", cur_cfg, history_cfgs)
+    if cfgs:
+        for cfg in cfgs:
+            if (
+                cfg["sharding_stage"] < sharding_stage
+                and cfg.get("time", -1) > 0
+            ):
+                pruned_reason = f"sharding_stage {sharding_stage} may be slower because {cfg['sharding_stage']} has been already runnable."
+                log_pruned_info(cur_cfg, pruned_reason, tuner_cfg)
+                cur_cfg["time"] = cfg["time"]
+                return True
+
+            # memory prune
+            if (
+                cfg["sharding_stage"] > sharding_stage
+                and cfg.get("max_mem_usage") == "OOM"
+            ):
+                pruned_reason = f"sharding_stage {sharding_stage} may cause oom because {cfg['sharding_stage']} already oom."
+                log_pruned_info(cur_cfg, pruned_reason, tuner_cfg)
+                cur_cfg["max_mem_usage"] = "OOM"
+                return True
+
+    return False
+
+
+@register_prune
+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.
+    2. Usage of recompute and recompute granularity should be in the candidates of user defined.
+    3. If recompute is not used, but recompute granularity is set, return True for pruning.
+    4. Prune if a similar configuration without using recompute resulted in a valid run.
+    5. If recompute is false, prune redundant recompute granularity
+    """
+    recompute_granularity = cur_cfg.get("recompute_granularity", None)
+    use_recompute = cur_cfg.get("use_recompute", None)
+    recompute_level = get_config_recompute_level(cur_cfg)
+
+    if use_recompute is None:
+        return False
+
+    recompute_granularity_candidates = tuner_cfg["candidates"].get(
+        "recompute_granularity", None
+    )
+    use_recompute_candidates = tuner_cfg["candidates"].get(
+        "use_recompute", None
+    )
+
+    if use_recompute_candidates:
+        if use_recompute not in use_recompute_candidates:
+            return True
+
+    if recompute_granularity_candidates and recompute_granularity:
+        if recompute_granularity not in recompute_granularity_candidates:
+            return True
+
+    if not use_recompute:
+        if recompute_granularity != "full":
+            return True
+
+        cfgs = same_cfgs_beside(
+            ["use_recompute", "recompute_granularity"], cur_cfg, history_cfgs
+        )
+        if cfgs:
+            for cfg in cfgs:
+                if recompute_level == get_config_recompute_level(cfg):
+                    return True
+
+    return False
+
+
+def get_config_recompute_level(cfg):
+    recompute_granularity_level = {"full": 3, "full_attn": 2, "core_attn": 1}
+    use_recompute = cfg.get("use_recompute", None)
+    recompute_granularity = cfg.get("recompute_granularity", None)
+
+    if use_recompute is None:
+        return None
+
+    if not use_recompute:
+        return 0
+    else:
+        return recompute_granularity_level[recompute_granularity]
+
+
+@register_prune_history
+def prune_by_recompute_history(
+    tuner_cfg, cur_cfg, history_cfgs=[], pruned_cfgs=[]
+):
+    recompute_level = get_config_recompute_level(cur_cfg)
+
+    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 cfgs:
+        for cfg in cfgs:
+            cfg["recompute_level"] = get_config_recompute_level(cfg)
+
+            if (
+                cfg["recompute_level"] < recompute_level
+                and cfg.get("time", -1) > 0
+            ):
+                pruned_reason = f"use_recompute may be slower because {cfg['use_recompute']} has been already runnable."
+                log_pruned_info(cur_cfg, pruned_reason, tuner_cfg)
+                cur_cfg["time"] = cfg["time"]
+                return True
+
+            if (
+                cfg["recompute_level"] > recompute_level
+                and cfg.get("max_mem_usage") == "OOM"
+            ):
+                pruned_reason = f"use_recompute may cause oom because {cfg['use_recompute']} already oom."
+                log_pruned_info(cur_cfg, pruned_reason, tuner_cfg)
+                cur_cfg["max_mem_usage"] = "OOM"
+                return True
+
+    return False
+
+
+@register_prune
+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)
+    sharding_degree = cur_cfg.get("sharding_degree", 1)
+    if dp_degree * mp_degree * pp_degree * sharding_degree != num_gpus:
+        return True
+
+    return False
+
+
+@register_prune
+def prune_by_memory_estimation(tuner_cfg, cur_cfg, history_cfgs=[]):
+    memory_estimation_tool = tuner_cfg.get("memory_estimation_tool", None)
+    # TODO(@gexiao): get from system api
+    max_memory_usage = tuner_cfg.get("max_mem_usage", None)
+    model_cfg = tuner_cfg["model_cfg"]
+
+    if memory_estimation_tool is None:
+        return False
+
+    if not os.path.exists(memory_estimation_tool):
+        raise ValueError(
+            f"memory_estimation_tool should be a valid path, but got {memory_estimation_tool}"
+        )
+
+    if max_memory_usage is None:
+        raise ValueError(
+            "max_mem_usage should be set when using memory estimation tool"
+        )
+
+    # get distributed strategy
+    dp_degree = cur_cfg['dp_degree']
+    mp_degree = cur_cfg['mp_degree']
+    pp_degree = cur_cfg['pp_degree']
+    vpp_degree = cur_cfg['vpp_degree']
+    sharding_degree = cur_cfg['sharding_degree']
+    sharding_stage = cur_cfg['sharding_stage']
+    use_recompute = cur_cfg['use_recompute']
+    micro_batch_size = cur_cfg['micro_batch_size']
+    recompute_granularity = cur_cfg['recompute_granularity']
+
+    memory_estimation_cmd = [
+        "python",
+        memory_estimation_tool,
+        "--dp_degree",
+        str(dp_degree),
+        "--mp_degree",
+        str(mp_degree),
+        "--pp_degree",
+        str(pp_degree),
+        "--vpp_degree",
+        str(vpp_degree),
+        "--sharding_degree",
+        str(sharding_degree),
+        "--sharding_stage",
+        str(sharding_stage),
+        "--use_recompute",
+        str(use_recompute),
+        "--micro_batch_size",
+        str(micro_batch_size),
+        "--recompute_granularity",
+        str(recompute_granularity),
+    ]
+
+    # get model config
+    hidden_size = model_cfg.get('hidden_size', None)
+    if hidden_size is not None:
+        memory_estimation_cmd.extend(["--hidden_size", str(hidden_size)])
+
+    num_attention_heads = model_cfg.get('num_attention_heads', None)
+    if num_attention_heads is not None:
+        memory_estimation_cmd.extend(
+            ["--num_attention_heads", str(num_attention_heads)]
+        )
+
+    num_layers = model_cfg.get('num_layers', None)
+    if num_layers is not None:
+        memory_estimation_cmd.extend(["--num_layers", str(num_layers)])
+
+    max_sequence_length = model_cfg.get('max_sequence_length', None)
+    if max_sequence_length is not None:
+        memory_estimation_cmd.extend(
+            ["--max_sequence_length", str(max_sequence_length)]
+        )
+
+    vocab_size = model_cfg.get('vocab_size', None)
+    if vocab_size is not None:
+        memory_estimation_cmd.extend(["--vocab_size", str(vocab_size)])
+
+    intermediate_size = model_cfg.get('intermediate_size', None)
+    if intermediate_size is not None:
+        memory_estimation_cmd.extend(
+            ["--intermediate_size", str(intermediate_size)]
+        )
+
+    result = subprocess.run(
+        memory_estimation_cmd,
+        capture_output=True,
+        text=True,
+    )
+
+    if result.returncode == 0:
+        cur_memory_usage = int(round(float(result.stdout), 2))
+        cur_cfg["estimated_memory_usage"] = cur_memory_usage
+        msg = f"Estimated {cur_cfg} memory usage: {cur_memory_usage} MB"
+        memory_exceeded = cur_memory_usage > (max_memory_usage * 1024)
+        if memory_exceeded:
+            msg += ", it will be pruned!"
+        logger.info(msg)
+        return memory_exceeded
+    else:
+        raise ValueError(
+            f"memory_estimation_tool failed with error: {result.stderr}"
+        )
+
+
+@register_prune_history
+def prune_by_sharding_overlap(
+    tuner_cfg, cur_cfg, history_cfgs=[], pruned_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
+
+
+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
+
+
+@register_prune
+def prune_by_refined_recompute(tuner_cfg, cur_cfg, history_cfgs=[]):
+    if tuner_cfg.get("refined_recompute", None):
+        rr = tuner_cfg.get("refined_recompute")
+        pp_degree = cur_cfg["pp_degree"]
+        recompute = cur_cfg["use_recompute"]
+        recompute_granularity = cur_cfg["recompute_granularity"]
+        compare = [cur_cfg[item] for item in rr]
+        if recompute:
+            if recompute_granularity and recompute_granularity != "full":
+                if compare.count(0) != len(compare):
+                    return True
+        if pp_degree == 1 and compare.count(0) != len(compare):
+            return True
+        if tuner_cfg["model_cfg"]["num_layers"] % pp_degree != 0:
+            return True
+        max_value = tuner_cfg["model_cfg"]["num_layers"] / pp_degree
+        if cur_cfg[rr[0]] > max_value:
+            return True
+        i = 1
+        while i < len(rr):
+            if cur_cfg[rr[i]] > max_value or (
+                cur_cfg[rr[i - 1]] != max_value and cur_cfg[rr[i]] != 0
+            ):
+                return True
+            i += 1
+
+    return False
+
+
+@register_prune_history
+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)
+        compare.append("use_recompute")
+        cfgs = same_cfgs_beside(compare, cur_cfg, history_cfgs)
+        for item in rr:
+            if cfgs:
+                for cfg in cfgs:
+                    if not cfg["use_recompute"] and cfg.get("time", -1) > 0:
+                        pruned_reason = f"{item} {cur_cfg[item]} may be slower because not recompute has been already runnable."
+                        log_pruned_info(cur_cfg, pruned_reason, tuner_cfg)
+                        cur_cfg["time"] = cfg["time"]
+                        return True
+                    if (
+                        cfg[item] > cur_cfg[item]
+                        and cfg.get("time", -1) > 0
+                        and cfg["use_recompute"]
+                        and cur_cfg["use_recompute"]
+                    ):
+                        pruned_reason = f"{item} {cur_cfg[item]} may be slower because {cfg[item]} has been already runnable."
+                        log_pruned_info(cur_cfg, pruned_reason, tuner_cfg)
+                        cur_cfg["time"] = cfg["time"]
+                        return True
+                    # memory prune
+                    if (
+                        cfg[item] < cur_cfg[item]
+                        and cfg.get("max_mem_usage") == "OOM"
+                        and cfg["use_recompute"]
+                        and cur_cfg["use_recompute"]
+                    ):
+                        pruned_reason = f"{item} {cur_cfg[item]} may cause oom because {cfg[item]} already oom."
+                        log_pruned_info(cur_cfg, pruned_reason, tuner_cfg)
+                        cur_cfg["max_mem_usage"] = "OOM"
+                        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

python/paddle/distributed/auto_tuner/recorder.py

@@ -0,0 +1,166 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from __future__ import annotations
+
+import copy
+import csv
+import os
+
+import pandas as pd
+
+
+class HistoryRecorder:
+    # NOTE increase extenable ablitity
+    def __init__(self, tuner_cfg) -> None:
+        self.tuner_cfg = tuner_cfg
+        self.search_algo = self.tuner_cfg['search_algo']['name']
+        self.history = []
+        self.store_path = None
+        self.additional_metric_key = None
+
+    def add_cfg(self, **kwargs):
+        cur_configs = {}
+        for key, val in kwargs.items():
+            cur_configs[key] = val
+        self.history.append(cur_configs)
+
+    def sort_metric(self, direction, metric_name) -> None:
+        if direction == 'Maximize':
+            self.history.sort(
+                key=lambda x: (
+                    x[metric_name]
+                    if x[metric_name] is not None
+                    else float('-inf')
+                ),
+                reverse=True,
+            )
+        else:
+            self.history.sort(
+                key=lambda x: (
+                    x[metric_name]
+                    if x[metric_name] is not None
+                    else float('inf')
+                ),
+                reverse=False,
+            )
+
+    def get_best(
+        self, metric, direction, buffer=None, max_mem_usage=None
+    ) -> tuple[dict, bool]:
+        self.sort_metric(direction=direction, metric_name=metric)
+        if len(self.history) == 0:
+            return (None, True)
+
+        best_cfg = self.history[0]
+        if isinstance(best_cfg["max_mem_usage"], str) or best_cfg["time"] == -1:
+            return (best_cfg, True)
+
+        if buffer is not None:
+            if buffer < 0:
+                raise ValueError("The buffer should be not less than 0.")
+            assert max_mem_usage is not None, (
+                "max_mem_usage cannot be None when buffer is greater than 0."
+            )
+            if max_mem_usage <= 0:
+                raise ValueError("max_mem_usage should be greater than 0.")
+
+            for cfg in self.history:
+                if (
+                    not best_cfg["max_mem_usage"]
+                    and cfg["max_mem_usage"]
+                    and not isinstance(cfg["max_mem_usage"], str)
+                    and cfg["time"] != -1
+                ):
+                    best_cfg = cfg
+                    continue
+
+                if (
+                    not isinstance(cfg["max_mem_usage"], str)
+                    and cfg["max_mem_usage"]
+                    and cfg["max_mem_usage"] < best_cfg["max_mem_usage"]
+                    and cfg["time"] != -1
+                ):
+                    best_cfg = cfg
+
+                if (
+                    not isinstance(cfg["max_mem_usage"], str)
+                    and cfg["max_mem_usage"]
+                    and cfg["max_mem_usage"] < max_mem_usage - buffer
+                    and cfg["time"] != -1
+                ):
+                    break
+            return (best_cfg, False)
+
+        return (self.history[0], False)
+
+    def _store_history_impl(self, data, path="./history.csv"):
+        """Store history to csv file."""
+        # convert to pd dataframe
+        df = pd.DataFrame(data)
+        # move 'job_id' to the first column
+        cols = df.columns.tolist()
+        cols.insert(0, cols.pop(cols.index('job_id')))
+        df = df.reindex(columns=cols)
+        # 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(path, index=False)
+
+    def store_history(self, path="./history.csv"):
+        # get enhanced report in dp-estimation mode
+        if self.search_algo == "dp_estimation":
+            metric_name = self.tuner_cfg['metric_cfg']['name']
+            if self.additional_metric_key:
+                _history = []
+                for cfg in self.history:
+                    if (
+                        "sharding_overlap" not in cfg.keys()
+                        or cfg["sharding_overlap"] is None
+                    ) and cfg["error_info"] is None:
+                        _history.append(copy.deepcopy(cfg))
+                _history.sort(
+                    key=lambda x: (
+                        x[self.additional_metric_key]
+                        if x[self.additional_metric_key] is not None
+                        else float('-inf')
+                    ),
+                    reverse=True,
+                )
+                self._store_history_impl(
+                    data=_history, path=path.split('.csv')[0] + '_enhanced.csv'
+                )
+
+        """Store history to csv file."""
+        self.store_path = path
+        self._store_history_impl(data=self.history, path=path)
+
+    def load_history(self, path="./history.csv") -> tuple[list, bool]:
+        """Load history from csv file."""
+        err = False
+        if self.store_path is None:
+            self.store_path = path
+        if not os.path.exists(self.store_path):
+            err = True
+        else:
+            with open(self.store_path, "r") as f:
+                reader = csv.reader(f)
+                self.history = list(reader)
+        return (self.history, err)
+
+    def clean_history(self) -> None:
+        """Clean history."""
+        self.history = []

python/paddle/distributed/auto_tuner/search.py

@@ -0,0 +1,157 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import logging
+import os
+from abc import ABC, abstractmethod
+
+from .prune import _PRUNE_HISTORY_FUNC
+from .utils import (
+    gbs_search_all,
+    load_configs_from_csv,
+    search_all,
+    search_by_dp_estimation,
+)
+
+logger = logging.getLogger('auto_tuner')
+
+
+class SearchAlgo(ABC):
+    def __init__(self, tuner_cfg):
+        self.tuner_cfg = tuner_cfg
+        self.pruned_cfgs = []
+
+    @abstractmethod
+    def search_once(self, history_cfgs):
+        pass
+
+    def prune(self, tuner_cfg, cur_cfg, history_cfgs, pruned_cfgs):
+        for func in _PRUNE_HISTORY_FUNC:
+            result = func(tuner_cfg, cur_cfg, history_cfgs, pruned_cfgs)
+            if result:
+                return True
+        return False
+
+
+class GridSearch(SearchAlgo):
+    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 and self.tuner_cfg.get(
+                    "need_baseline", False
+                ):
+                    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_special
+
+                        self.all_tasks[self.idx :] = sort_by_special(
+                            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, self.pruned_cfgs
+                )
+                self.pruned_cfgs.append(new_cfg)
+            else:
+                return None
+
+        return new_cfg
+
+
+class DpEstimationSearch(SearchAlgo):
+    def __init__(self, tuner_cfg):
+        super().__init__(tuner_cfg)
+        self.idx = 0
+        if tuner_cfg["candidates"]["dp_degree"] != [1]:
+            logger.warning(
+                "dp_degree should be [1] in dp estimation search mode. Modify it to [1] automatically."
+            )
+            tuner_cfg["candidates"]["dp_degree"] = [1]
+        self.all_tasks = search_by_dp_estimation(tuner_cfg)
+        assert len(self.all_tasks) > 0, (
+            "Unable to perform single dp estimation search."
+        )
+
+    def search_once(self, history_cfgs):
+        new_cfg = None
+        stop = False
+        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
+
+
+class GBSSearch(SearchAlgo):
+    def __init__(self, tuner_cfg):
+        super().__init__(tuner_cfg)
+        self.idx = 0
+        self.all_tasks = gbs_search_all(tuner_cfg)
+
+    def search_once(self, history_cfgs):
+        new_cfg = None
+        stop = False
+        while not stop:
+            if self.idx < len(self.all_tasks):
+                new_cfg = self.all_tasks[self.idx]
+                self.idx += 1
+                glb = new_cfg.get("global_batch_size", None)
+                self.tuner_cfg["model_cfg"]["global_batch_size"] = glb
+                stop = not self.prune(self.tuner_cfg, new_cfg, history_cfgs)
+            else:
+                return None
+        return new_cfg
+
+
+class CustomizeSearch(SearchAlgo):
+    def __init__(self, tuner_cfg):
+        super().__init__(tuner_cfg)
+        self.idx = 0
+        self.configs_csv = tuner_cfg.get("configs_csv", None)
+        assert os.path.exists(self.configs_csv), (
+            "configs_csv file is necessary in CustomizeSearch mode."
+        )
+        self.all_tasks = load_configs_from_csv(self.configs_csv)
+
+    def search_once(self, history_cfgs):
+        new_cfg = self.all_tasks[self.idx]
+        self.idx += 1
+        return new_cfg

python/paddle/distributed/auto_tuner/tuner.py

@@ -0,0 +1,155 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import csv
+import os
+
+from .utils import default_candidates, gbs_default_candidates
+
+
+class AutoTuner:
+    """
+    The AutoTuner can automatically provide running task based on user-defined settings
+    and the task will be launched for execution.
+
+    Args:
+        tuner_cfg (dict): The configuration of auto tuner user defined.
+    """
+
+    def __init__(self, tuner_cfg):
+        self.cur_task_id = 1
+        self.task_limit = tuner_cfg.get("task_limit", 100)
+
+        search_algo = tuner_cfg.get("search_algo", {"name": "grid"})["name"]
+
+        if search_algo == "grid":
+            from .search import GridSearch
+
+            tuner_cfg["candidates"] = default_candidates(tuner_cfg)
+            self.algo = GridSearch(tuner_cfg)
+        elif search_algo == "dp_estimation":
+            from .search import DpEstimationSearch
+
+            tuner_cfg["candidates"] = default_candidates(tuner_cfg)
+            self.algo = DpEstimationSearch(tuner_cfg)
+        elif search_algo == "gbs":
+            from .search import GBSSearch
+
+            tuner_cfg["candidates"] = gbs_default_candidates(tuner_cfg)
+            self.algo = GBSSearch(tuner_cfg)
+        elif search_algo == "customize":
+            from .search import CustomizeSearch
+
+            self.algo = CustomizeSearch(tuner_cfg)
+        else:
+            raise NotImplementedError
+
+        self.history_cfgs = []
+        self.resume_cfgs = []
+        self.tuner_cfg = tuner_cfg
+
+    def search_once(self):
+        """Return a new task config."""
+        if self.cur_task_id > self.task_limit:
+            return None
+        new_cfg = self.algo.search_once(self.history_cfgs)
+        self.cur_task_id += 1
+
+        return new_cfg
+
+    def add_cfg(self, cfg):
+        """Add cfg into history cfgs"""
+        self.history_cfgs.append(cfg)
+
+    def resume_form_history(self, history_csv_path="./history.csv"):
+        """Resume form history csv file"""
+        # The breakpoint resume function does not start when the resume csv file does not exist.
+        if not os.path.exists(history_csv_path):
+            return
+        resume_csv_path = os.path.join(
+            os.path.dirname(history_csv_path),
+            f'{os.path.basename(history_csv_path).split(".")[0]}_copy.csv',
+        )
+        with open(history_csv_path, "r") as fread:
+            reader = csv.reader(fread)
+            data_list = list(reader)
+            with open(resume_csv_path, "w") as fwrite:
+                writer = csv.writer(fwrite)
+                for row in data_list:
+                    writer.writerow(row)
+        # chang str type to real type
+        for row in data_list:
+            for i, value in enumerate(row):
+                try:
+                    row[i] = int(value)
+                except ValueError:
+                    try:
+                        row[i] = float(value)
+                    except ValueError:
+                        pass
+
+        data_dict = []
+        keys = data_list[0]
+        values = data_list[1:]
+        for val in values:
+            val = [x if x != '' else None for x in val]
+            val = [True if x == 'True' else x for x in val]
+            val = [False if x == 'False' else x for x in val]
+            dictionary = dict(zip(keys, val))
+            time_val = -1
+            target_key = self.tuner_cfg["metric_cfg"]["name"]
+            if dictionary[target_key]:
+                time_val = dictionary[target_key]
+            dictionary["time"] = time_val
+            data_dict.append(dictionary)
+        self.resume_cfgs = data_dict
+
+    def get_cfg_from_resume(self, cur_cfg):
+        """Get cfg from resume cfgs"""
+        keys_to_compare = [
+            'mp_degree',
+            'sharding_degree',
+            'pp_degree',
+            'dp_degree',
+            'sharding_stage',
+            'micro_batch_size',
+            'vpp_degree',
+            'use_recompute',
+            'recompute_granularity',
+            'num_gpus',
+            'nodes',
+            'global_batch_size',
+            'sharding_overlap',
+            'acc_steps',
+        ]
+
+        if self.tuner_cfg.get("refined_recompute", None):
+            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:
+                if not cfg.get(key) and not cur_cfg.get(key):
+                    continue
+                else:
+                    is_same = str(cfg.get(key)) == str(cur_cfg.get(key))
+                ret_is_same = ret_is_same and is_same
+            if ret_is_same:
+                return cfg
+        return None