chore: import upstream snapshot with attribution

python/paddle/distributed/auto_parallel/static/tuner/__init__.py

@@ -0,0 +1,17 @@
+#   Copyright (c) 2022 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 .profiler import profiler  # noqa: F401
+
+__all__ = []

python/paddle/distributed/auto_parallel/static/tuner/algorithms.py

@@ -0,0 +1,241 @@
+#   Copyright (c) 2022 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
+from abc import ABC, abstractmethod
+
+from ..utils import get_logger, is_recompute_op
+from .trial import (
+    OptimizationTunerTrial as Trial,
+    TrialStatus,
+)
+
+
+class AlgorithmBase(ABC):
+    """
+    An Tuning algorithm is a class to find out an optimal configuration
+    given the selected tuning optimization pass(es) and the arguments to be tuned.
+    Different optimization pass(es) will correspond to a different algorithm,
+    where different search space **pruning rules** will applied.
+
+    In another word, the key "algorithm" for this class is the
+    search space pruning rules specific for the given optimization scenario.
+    """
+
+    _REGISTERED_ALGORITHMS = {}
+
+    name = None
+
+    @staticmethod
+    def _register(algo_name, algo_class):
+        assert issubclass(algo_class, AlgorithmBase)
+        AlgorithmBase._REGISTERED_ALGORITHMS[algo_name] = algo_class
+
+    def __init__(self, config):
+        self._config = config
+        self._init_spaces()
+        self._logger = get_logger(logging.INFO)
+        self._changed_configs = []
+
+    @property
+    def changed_configs(self):
+        return self._changed_configs[:]
+
+    def collect_model_info(self, main_prog, startup_prog):
+        """
+        Collect the model static info (from programs) that could be used to
+        pruning candidate trials and saving tuning time. For instance,
+        model info like number of model parameters and activation memory could be
+        used to prune candidate trial and decide the next trial.
+        """
+        pass
+
+    @abstractmethod
+    def _init_spaces(self):
+        pass
+
+    @abstractmethod
+    def next_trial(self):
+        pass
+
+    @abstractmethod
+    def update(self, results):
+        """
+        Update the algorithm with the results of last trial. Using this information is used to
+        pruning the search space of the future trial.
+        """
+        pass
+
+    def get_config_from_trial(self, trial):
+        """
+        Return a new fleet.DistributedStrategy with the configurations in trial.
+        """
+        assert len(self._changed_configs) > 0
+        new_strategy = copy.deepcopy(self._config.dist_strategy)
+        for name in self._changed_configs:
+            config = getattr(trial.space, name)
+            setattr(new_strategy, name, config)
+        return new_strategy
+
+
+def register_algor(name):
+    def impl(cls):
+        AlgorithmBase._register(name, cls)
+        cls.name = name
+        return cls
+
+    return impl
+
+
+def new_algorithm(name, config):
+    algor_class = AlgorithmBase._REGISTERED_ALGORITHMS.get(name)
+    assert algor_class is not None, f"Algorithm {name} is not defined."
+    algor_obj = algor_class(config)
+    return algor_obj
+
+
+@register_algor("sharding")
+class ShardingStageAlgorithm(AlgorithmBase):
+    # TODO import trial class & copy strategy
+    def __init__(self, config):
+        super().__init__(config)
+        self._changed_configs = ["sharding"]
+
+    def _init_spaces(self):
+        self._max_stage = 3
+        self._trial_idx = 0
+
+        stage_range = self._config.sharding.get("tuning_range", None)
+        if stage_range:
+            assert set(stage_range).issubset({0, 1, 2, 3}), (
+                f"Sharding Stage should belong into range within 0 - 3 but got {stage_range}."
+            )
+            stage_range.sort(reverse=True)
+        else:
+            stage_range = list(range(self._max_stage + 1))
+            stage_range.sort(reverse=True)
+        self._stage_range = stage_range[:]
+        self._total_num_trial = len(self._stage_range)
+
+    def next_trial(self):
+        if self._trial_idx < self._total_num_trial:
+            stage = self._stage_range[self._trial_idx]
+
+            new_strategy = copy.deepcopy(self._config.dist_strategy)
+            sharding = new_strategy.sharding
+            sharding.stage = stage
+
+            name = f"trial-sharding-stage{stage}"
+            trial = Trial(new_strategy, name, self.changed_configs)
+
+            return trial
+        else:
+            return Trial(None, None, None, status=TrialStatus.STOPPED)
+
+    def update(self, results):
+        et = results.get("ErrorType", None)
+        if et and et == "ResourceExhaustedError":
+            self._trial_idx = self._total_num_trial
+            self._logger.info(
+                "Last trial is failed with OOM, all remaining trials are pruned to save time !"
+            )
+        else:
+            self._trial_idx += 1
+
+
+@register_algor("recompute")
+class RecomputeCheckpointAlgorithm(AlgorithmBase):
+    def __init__(self, config):
+        super().__init__(config)
+        self._changed_configs = ["recompute"]
+
+    def collect_model_info(self, main_prog, startup_prog):
+        segments = []
+        for op in main_prog.global_block().ops:
+            if not is_recompute_op(op):
+                continue
+
+            seg_name = op.attr('op_namescope')
+            if seg_name not in segments:
+                segments.append(seg_name)
+
+        self._total_num_trial = len(segments)
+        self._tuning_segments = list(range(len(segments)))
+        self._trial_left = 0
+        self._trial_right = len(segments) - 1
+        self._trial_idx = int(0 + (len(segments) - 1) / 2)
+
+    def _init_spaces(self):
+        self._recompute_mode = "all"
+
+    def next_trial(self):
+        if self._trial_idx < self._total_num_trial:
+            if self._recompute_mode == "all":
+                self._recompute_flag = False
+                new_strategy = copy.deepcopy(self._config.dist_strategy)
+                name = "trial-recompute-all-segments"
+                return Trial(new_strategy, name, self.changed_configs)
+            elif self._recompute_mode == "none":
+                self._recompute_flag = False
+                new_strategy = copy.deepcopy(self._config.dist_strategy)
+                recompute = new_strategy.recompute
+                recompute.enable = False
+                name = "trial-recompute-none-segments"
+                return Trial(new_strategy, name, self.changed_configs)
+            elif self._recompute_mode == "part":
+                new_no_recompute = self._tuning_segments[: self._trial_idx]
+                new_strategy = copy.deepcopy(self._config.dist_strategy)
+                recompute = new_strategy.recompute
+                recompute.no_recompute_segments.extend(new_no_recompute)
+                name = f"trial-recompute-part-segments-idx{self._trial_idx}"
+                return Trial(new_strategy, name, self.changed_configs)
+        else:
+            return Trial(None, None, None, status=TrialStatus.STOPPED)
+
+    def update(self, results):
+        et = results.get("ErrorType", None)
+        if self._recompute_mode == "all":
+            if et and et == "ResourceExhaustedError":
+                self._trial_idx = self._total_num_trial
+                self._logger.info(
+                    "Recompute all candidate segments is failed with OOM, please reduce model size or batch size."
+                )
+            else:
+                self._recompute_mode = "none"
+        elif self._recompute_mode == "none":
+            if et and et == "ResourceExhaustedError":
+                self._recompute_mode = "part"
+            else:
+                self._trial_idx = self._total_num_trial
+                self._logger.info(
+                    "Recompute is unnecessary for this model size, which will reduce the Throughput."
+                )
+        else:
+            if self._trail_left >= self._trail_right:
+                self._trial_idx = self._total_num_trial
+            elif et and et == "ResourceExhaustedError":
+                self._trail_left = self._trail_left
+                self._trail_right = self._trial_idx - 1
+                self._trial_idx = int(
+                    self._trail_left
+                    + (self._trail_right - self._trail_left) / 2
+                )
+            else:
+                self._trail_left = self._trial_idx + 1
+                self._trail_right = self._trail_right
+                self._trial_idx = int(
+                    self._trail_left
+                    + (self._trail_right - self._trail_left) / 2
+                )

python/paddle/distributed/auto_parallel/static/tuner/config.py

@@ -0,0 +1,122 @@
+#   Copyright (c) 2022 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 os
+
+from ...strategy import Strategy
+
+_tuning_supported_passes = ["sharding", "recompute"]
+
+
+def _get_pass_config(strategy, pass_name):
+    config = getattr(strategy, pass_name)
+    return config
+
+
+class TuningConfig:
+    """
+    A uniform config wrap:
+    distributed strategy: the user defined configuration for optimization pass
+    tuning config: configuration for the tuning process: mode (profile or cost model), log dir, extra tuning config for optimization like search range for specific
+    """
+
+    def __init__(self, strategy):
+        if not isinstance(strategy, Strategy):
+            raise TypeError("'strategy' must be object of class `Strategy`.")
+
+        self._tuning_passes_name = set()
+        self._dist_strategy = copy.deepcopy(strategy)
+        self._mode = None
+        self._profile_start_step = None
+        self._profile_end_step = None
+        self._project_dir = None
+        self._max_num_trial = None
+        self._early_stop = None
+        self._debug = None
+
+        self._initialize()
+
+    @property
+    def mode(self):
+        return self._mode
+
+    @property
+    def profile_start_step(self):
+        return self._profile_start_step
+
+    @property
+    def profile_end_step(self):
+        return self._profile_end_step
+
+    @property
+    def project_dir(self):
+        return self._project_dir
+
+    @property
+    def tuning_passes_name(self):
+        return self._tuning_passes_name
+
+    @property
+    def max_num_trial(self):
+        return self._max_num_trial
+
+    @property
+    def early_stop(self):
+        return self._early_stop
+
+    @property
+    def debug(self):
+        return self._debug
+
+    @property
+    def dist_strategy(self):
+        return self._dist_strategy
+
+    # initialize config with user define value or default value
+    def _initialize(self):
+        tuning_strategy = self._dist_strategy.tuning
+
+        self._mode = tuning_strategy.get("mode", "PROFILE")
+        self._profile_start_step = tuning_strategy.get("profile_start_step", 10)
+        self._profile_end_step = tuning_strategy.get("profile_end_step", 30)
+        self._max_num_trial = tuning_strategy.get("max_num_trial", 50)
+        self._early_stop = tuning_strategy.get("early_stop", None)
+        self._debug = tuning_strategy.get("debug", False)
+
+        project_dir = tuning_strategy.get("project_dir", None)
+        if not project_dir:
+            project_dir = os.path.join(os.getcwd(), "OptimizationTuning")
+        self._project_dir = project_dir
+
+        for p in _tuning_supported_passes:
+            if (
+                getattr(self._dist_strategy, p)
+                and _get_pass_config(self._dist_strategy, p).enable_tuning
+            ):
+                # TODO distinguish different args of each passes
+                self._tuning_passes_name.add(p)
+
+                p_strategy = getattr(self._dist_strategy, p)
+                self.__dict__[p] = p_strategy
+
+                # # TODO verify the user defined configs
+                # tuning_config_for_pass = tuning_strategy.get(p, None)
+                # if tuning_config_for_pass:
+                #     for k, v in tuning_config_for_pass.items():
+                #         self.__dict__[p][k] = v
+
+    # (NOTE)tuning config ONLY wraps dist strategy for pass config which is to be tuned
+    def __getattr__(self, item):
+        return getattr(self._dist_strategy, item)

python/paddle/distributed/auto_parallel/static/tuner/optimization_tuner.py

@@ -0,0 +1,643 @@
+#   Copyright (c) 2022 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 json
+import logging
+
+# import yaml
+import os
+import pathlib
+import pickle
+import shlex
+import shutil
+import subprocess
+import sys
+import time
+
+import paddle
+from paddle.distributed.auto_parallel.static.completion import Completer
+from paddle.distributed.auto_parallel.static.dist_context import (
+    DistributedContext,
+)
+from paddle.distributed.auto_parallel.static.partitioner import Partitioner
+from paddle.distributed.auto_parallel.static.process_group import (
+    clear_all_process_groups,
+    get_all_process_groups,
+    new_process_group,
+)
+from paddle.distributed.auto_parallel.static.reshard import Resharder
+from paddle.distributed.auto_parallel.static.utils import debug_program
+from paddle.distributed.passes import PassContext, new_pass
+from paddle.static import append_backward, program_guard
+
+from ..utils import get_logger
+from .algorithms import new_algorithm
+from .config import TuningConfig
+from .trial import TrialStatus
+
+
+def _get_new_params_grads(target_program, ref_program, ref_params_grads):
+    ref_block = ref_program.global_block()
+    target_block = target_program.global_block()
+    target_params_grads = []
+
+    for p, g in ref_params_grads:
+        # NOTE grad var might not be generated
+        assert ref_block.has_var(p.name)
+        assert target_block.has_var(p.name)
+        new_p = target_block.var(p.name)
+        if g:
+            new_g = target_block.var(g.name)
+        else:
+            new_g = None
+
+        target_params_grads.append((new_p, new_g))
+
+    return target_params_grads
+
+
+def _get_new_loss(target_program, ref_program, loss):
+    ref_block = ref_program.global_block()
+    target_block = target_program.global_block()
+    assert ref_block.has_var(loss.name)
+
+    return target_block.var(loss.name)
+
+
+def parse_process_groups():
+    group_map = {}
+    all_process_groups = get_all_process_groups()
+    for process_group in all_process_groups:
+        group_map[process_group.id] = process_group.ranks
+    return group_map
+
+
+def get_metric(results):
+    assert isinstance(results, dict), (
+        f"results should be type of dictionary, but got {type(results)}."
+    )
+    if 'Throughput' in results and isinstance(results['Throughput'], float):
+        return float(results['Throughput'])
+    else:
+        return -1.0
+
+
+def parse_results(results):
+    if results['Throughput'] > 0:
+        return "Throughput: {} step / s.".format(results['Throughput'])
+    et = results.get("ErrorType", None)
+    if et == "ResourceExhaustedError":
+        return "Fail with OOM"
+    else:
+        return "Fail with UNKNOWN ERROR"
+
+
+# TODO only dependent on dist context
+# all env need to be start a new pass are member of dist context
+def _copy_context(ref_dist_context):
+    # clear all process groups and recover the world process group
+    clear_all_process_groups()
+    ranks = []
+    for process_mesh in ref_dist_context._process_meshes:
+        ranks.extend(process_mesh.process_ids)
+    new_process_group(list(set(ranks)))
+
+    new_dist_context = DistributedContext()
+    new_dist_context._serial_main_program = (
+        ref_dist_context.serial_main_program.clone(for_test=False)
+    )
+    new_dist_context._serial_startup_program = (
+        ref_dist_context.serial_startup_program.clone(for_test=False)
+    )
+
+    # mapping variable into new dist context
+    if getattr(ref_dist_context, '_params_grads', None):
+        new_dist_context._params_grads = _get_new_params_grads(
+            new_dist_context.serial_main_program,
+            ref_dist_context.serial_main_program,
+            ref_dist_context._params_grads,
+        )
+    new_dist_context._serial_loss = _get_new_loss(
+        new_dist_context.serial_main_program,
+        ref_dist_context.serial_main_program,
+        ref_dist_context.serial_loss,
+    )
+
+    for key, var_list in ref_dist_context._serial_feed_vars.items():
+        new_var_list = []
+        for var in var_list:
+            block_idx = var.block.idx
+            var_name = var.name
+            var = new_dist_context._serial_main_program.blocks[
+                block_idx
+            ]._var_recursive(var_name)
+            new_var_list.append(var)
+        new_dist_context._serial_feed_vars[key] = new_var_list
+
+    for key, var_list in ref_dist_context._serial_fetch_vars.items():
+        new_var_list = []
+        # metrics is a list of list
+        if key == "metrics":
+            for inner_var_list in var_list:
+                new_inner_var_list = []
+                for var in inner_var_list:
+                    block_idx = var.block.idx
+                    var_name = var.name
+                    var = new_dist_context._serial_main_program.blocks[
+                        block_idx
+                    ]._var_recursive(var_name)
+                    new_inner_var_list.append(var)
+                new_var_list.append(new_inner_var_list)
+        else:
+            for var in var_list:
+                block_idx = var.block.idx
+                var_name = var.name
+                var = new_dist_context._serial_main_program.blocks[
+                    block_idx
+                ]._var_recursive(var_name)
+                new_var_list.append(var)
+        new_dist_context._serial_fetch_vars[key] = new_var_list
+
+    # copy information in forward and backward
+    new_dist_context._serial_optimizer = copy.deepcopy(
+        ref_dist_context.serial_optimizer
+    )
+    new_dist_context._dist_tensors_for_program = copy.deepcopy(
+        ref_dist_context._dist_tensors_for_program
+    )
+    new_dist_context._dist_ops_for_program = copy.deepcopy(
+        ref_dist_context._dist_ops_for_program
+    )
+    for pm in ref_dist_context.process_meshes:
+        new_dist_context.add_process_mesh(pm)
+    new_dist_context._dist_op_context = copy.deepcopy(
+        ref_dist_context._dist_op_context
+    )
+    new_dist_context._block_state = copy.deepcopy(ref_dist_context.block_state)
+
+    return new_dist_context
+
+
+class OptimizationTuner:
+    """
+    OptimizationTuner is used to manage the tuning procedure of hyper-parameters (configs)
+    of Optimization Pass in AutoParallel.
+    """
+
+    def __init__(
+        self,
+        dist_context,
+        dataset,
+        inputs_spec,
+        labels_spec,
+        batch_size,
+        rank,
+    ):
+        self._config = TuningConfig(dist_context.strategy)
+        # should not modify dist context from calling function
+        self._baseline_dist_context = _copy_context(dist_context)
+        self._baseline_completer = Completer(self._baseline_dist_context)
+
+        self._rank = rank
+        self._inputs_spec = inputs_spec
+        self._labels_spec = labels_spec
+        self._dataset = dataset
+        self._batch_size = batch_size
+
+        self._finished_trials = []
+        self._best_metric = None
+        self._best_iter = float("-inf")
+
+        self._logger = get_logger(logging.INFO)
+
+        self._build_programs_without_optimization()
+        self._select_tuning_algorithm()
+
+    @property
+    def project_dir(self):
+        dirname = self._config.project_dir
+        if not os.path.exists(dirname):
+            if self.rank == 0:
+                pathlib.Path(dirname).mkdir(parents=True, exist_ok=True)
+        return dirname
+
+    @property
+    def rank(self):
+        return self._rank
+
+    @property
+    def device_id(self):
+        return paddle.distributed.ParallelEnv().device_id
+
+    # TODO Generate complete program with all parts like forward, backward, update
+    # as well as parallelism transformation.
+    def _build_programs_without_optimization(self):
+        serial_main_program = self._baseline_dist_context.serial_main_program
+        serial_startup_program = (
+            self._baseline_dist_context.serial_startup_program
+        )
+        serial_loss = self._baseline_dist_context.serial_loss
+
+        with program_guard(serial_main_program, serial_startup_program):
+            params_grads = append_backward(
+                serial_loss,
+                distop_context=self._baseline_dist_context.dist_op_context,
+            )
+
+        self._baseline_completer.complete_backward_annotation(
+            serial_main_program
+        )
+        self._baseline_dist_context.block_state.parse_backward_blocks(
+            serial_main_program
+        )
+        self._baseline_dist_context._params_grads = params_grads
+
+        if self._config.debug:
+            baseline_dir = os.path.join(self.project_dir, "baseline")
+            if not os.path.exists(baseline_dir):
+                pathlib.Path(baseline_dir).mkdir(parents=True, exist_ok=True)
+            debug_program(
+                self._baseline_dist_context._serial_main_program,
+                baseline_dir,
+                "main",
+            )
+            debug_program(
+                self._baseline_dist_context._serial_startup_program,
+                baseline_dir,
+                "startup",
+            )
+
+    def _select_tuning_algorithm(self):
+        selected_passes_set = self._config.tuning_passes_name
+        algorithm_name = "_".join(sorted(selected_passes_set))
+        self._algorithm = new_algorithm(algorithm_name, self._config)
+
+    def _apply_optimization(self, trial):
+        new_strategy = trial.space
+        dist_context = _copy_context(self._baseline_dist_context)
+        pass_context = PassContext()
+        completer = Completer(dist_context)
+
+        main_program = dist_context.serial_main_program
+        startup_program = dist_context.serial_startup_program
+
+        # applying optimization pass
+        if new_strategy.amp.enable:
+            config = copy.deepcopy(new_strategy.amp.to_dict())
+            config["dist_context"] = dist_context
+            config["params_grads"] = dist_context._params_grads
+            # TODO AMP Pass should not use loss var
+            config["loss"] = dist_context.serial_loss
+            config["input_data"] = (
+                self._baseline_dist_context.serial_feed_vars["inputs"]
+                + self._baseline_dist_context.serial_feed_vars["labels"]
+            )
+            if config["dtype"] == "float16" and config["level"] == "o2":
+                config["base_opt"] = dist_context.serial_optimizer
+                auto_parallel_fp16_pass = new_pass("auto_parallel_fp16", config)
+                auto_parallel_fp16_pass.apply(
+                    [main_program], [startup_program], pass_context
+                )
+                dist_context._serial_loss = auto_parallel_fp16_pass.get_loss()
+            else:
+                auto_parallel_amp_pass = new_pass("auto_parallel_amp", config)
+                auto_parallel_amp_pass.apply(
+                    [main_program], [startup_program], pass_context
+                )
+                dist_context._serial_loss = auto_parallel_amp_pass.get_loss()
+
+        if new_strategy.recompute.enable:
+            config = copy.deepcopy(new_strategy.recompute.to_dict())
+            config["dist_context"] = dist_context
+            config["no_grad_set"] = None
+            config["loss"] = dist_context.serial_loss
+            auto_parallel_recompute_pass = new_pass(
+                "auto_parallel_recompute", config
+            )
+            auto_parallel_recompute_pass.apply(
+                [main_program], [startup_program], pass_context
+            )
+
+        # Do logical partition
+        partitioner = Partitioner(dist_context, self.rank)
+        (
+            dist_main_prog,
+            dist_startup_prog,
+            dist_params_grads,
+        ) = partitioner.partition(
+            main_program, startup_program, dist_context._params_grads
+        )
+
+        # Generate optimizer
+        # FIXME should be remove from apply pass after pass support optimizers
+        with (
+            program_guard(dist_main_prog, dist_startup_prog),
+            dist_main_prog.switch_name_generator_guard("opt_"),
+        ):
+            optimizer_ops = dist_context.serial_optimizer.apply_gradients(
+                dist_params_grads
+            )
+        completer.complete_update_annotation(dist_main_prog)
+
+        resharder = Resharder(
+            dist_main_prog,
+            dist_startup_prog,
+            self.rank,
+            dist_context,
+            dist_params_grads,
+        )
+        resharder.reshard()
+
+        config = {}
+        config["dist_context"] = dist_context
+        config["global_rank"] = self.rank
+        config["use_sharding"] = new_strategy.sharding.enable
+        dp_pass = new_pass("auto_parallel_data_parallel_optimization", config)
+        dp_pass.apply([dist_main_prog], [dist_startup_prog], pass_context)
+
+        if new_strategy.sharding.enable:
+            config = copy.deepcopy(new_strategy.sharding.to_dict())
+            config["dist_context"] = dist_context
+            config["params_grads"] = dist_params_grads
+            config["global_rank"] = self.rank
+            auto_parallel_sharding_pass = new_pass(
+                "auto_parallel_sharding", config
+            )
+            auto_parallel_sharding_pass.apply(
+                [dist_main_prog], [dist_startup_prog], pass_context
+            )
+            dist_params_grads = pass_context.get_attr("params_grads")
+
+        # gradient clip
+        config = copy.deepcopy(new_strategy.sharding.to_dict())
+        config["dist_context"] = dist_context
+        config["params_grads"] = dist_params_grads
+        config["rank_id"] = self.rank
+        auto_parallel_clip_pass = new_pass("auto_parallel_grad_clip", config)
+        auto_parallel_clip_pass.apply(
+            [dist_main_prog], [dist_startup_prog], pass_context
+        )
+
+        if new_strategy.gradient_merge.enable:
+            config = copy.deepcopy(new_strategy.gradient_merge.to_dict())
+            config["dist_context"] = dist_context
+            config["params_grads"] = dist_params_grads
+            auto_parallel_gradient_merge_pass = new_pass(
+                "auto_parallel_gradient_merge_pass", config
+            )
+            auto_parallel_gradient_merge_pass.apply(
+                [dist_main_prog], [dist_startup_prog], pass_context
+            )
+        trial.main_program, trial.startup_program = (
+            dist_main_prog,
+            dist_startup_prog,
+        )
+        return trial
+
+    def _get_profile_context(self, trial, result_path):
+        profile_ctx = {}
+
+        profile_ctx['distributed_env'] = copy.deepcopy(
+            paddle.distributed.ParallelEnv()
+        )
+        profile_ctx['group_map'] = parse_process_groups()
+        profile_ctx["loss_var_name"] = (
+            self._baseline_dist_context.serial_loss.name
+        )
+        profile_ctx["main_program_decs"] = (
+            trial.main_program.desc.serialize_to_string()
+        )
+        profile_ctx["startup_program_decs"] = (
+            trial.startup_program.desc.serialize_to_string()
+        )
+        self._dataset.batch_size = self._batch_size
+        self._dataset.input_names = self._get_input_names()
+
+        profile_ctx["dataset"] = self._dataset
+        profile_ctx["result_filename"] = result_path
+
+        return profile_ctx
+
+    def _get_input_names(self):
+        input_names = []
+        for input_spec in self._inputs_spec[:] + self._labels_spec[:]:
+            input_names.append(input_spec.name)
+        return input_names
+
+    def _launch_profile(self, ctx_path, trial_dir):
+        if os.environ.get("WITH_COVERAGE", "OFF") == "ON":
+            coverage_args = ["-m", "coverage", "run", "--branch", "-p"]
+        else:
+            coverage_args = []
+
+        profile_args = " ".join(
+            [
+                "--rank",
+                str(self.rank),
+                "--device_id",
+                str(self.device_id),
+                "--ctx_filename",
+                ctx_path,
+                "--profile_start_step",
+                str(self._config.profile_start_step),
+                "--profile_end_step",
+                str(self._config.profile_end_step),
+            ]
+        )
+        cmd_args = (
+            "-m paddle.distributed.auto_parallel.static.tuner.profiler"
+            + " "
+            + profile_args
+        )
+        cmd = [sys.executable, "-u", *coverage_args, *shlex.split(cmd_args)]
+
+        parent_env = copy.copy(os.environ.copy())
+        # env flags need for profile
+        new_env = {}
+        new_env.update(parent_env)
+
+        # TODO if any rank hang or fail, kill all processes
+        self._logger.debug("Executing cmd:\n{} .".format(" ".join(cmd)))
+        # new_process = subprocess.Popen(cmd, env=new_env)
+        with (
+            open(
+                os.path.join(trial_dir, "stdout.log" + str(self.rank)), "wb"
+            ) as out,
+            open(
+                os.path.join(trial_dir, "stderr.log" + str(self.rank)), "wb"
+            ) as err,
+        ):
+            result = subprocess.Popen(cmd, stdout=out, stderr=err, env=new_env)
+            result.wait()
+            out.flush()
+            err.flush()
+            os.fsync(out)
+            os.fsync(err)
+
+    def _profile_trial(self, trial):
+        # Making working directory
+        trial_dir = self._get_trial_dir(trial)
+        if not os.path.exists(trial_dir):
+            if self.rank == 0:
+                pathlib.Path(trial_dir).mkdir(parents=True, exist_ok=True)
+            else:
+                while not os.path.exists(trial_dir):
+                    pass
+        ctx_filename = "profile_ctx." + str(self.rank)
+        ctx_path = os.path.join(trial_dir, ctx_filename)
+        result_path = os.path.join(trial_dir, "result.json")
+
+        # Prepare Profile Context
+        profile_ctx = self._get_profile_context(trial, result_path)
+        with open(ctx_path, 'wb') as f:
+            pickle.dump(profile_ctx, f, protocol=4)
+
+        if self._config.debug:
+            debug_program(trial.main_program, trial_dir, "main_program")
+            debug_program(trial.startup_program, trial_dir, "startup_program")
+
+        # Run
+        self._launch_profile(ctx_path, trial_dir)
+
+        # Load results
+        try:
+            with open(result_path, 'r') as fp:
+                results = json.load(fp)
+            return results
+        except FileNotFoundError:
+            Error_results = {"Throughput": -1, "ErrorType": 'FatalError'}
+            return Error_results
+
+    def _evaluate_trial(self, trial):
+        self._logger.info(f"Trial {trial.name} evaluation start.")
+        self._apply_optimization(trial)
+
+        if self._config.mode == "PROFILE":
+            results = self._profile_trial(trial)
+
+        elif self._config.mode == "COSTMODEL":
+            raise NotImplementedError(
+                "COSTMODEL mode for optimization tuning is not supported yet!"
+            )
+        else:
+            raise NotImplementedError(
+                f"invalid evaluation mode: {self._config.mode}"
+            )
+
+        self._logger.info(
+            f"Trial {trial.name} evaluation finish with {parse_results(results)}."
+        )
+        return results
+
+    def _update(self, i, trial, results):
+        self._finished_trials.append(trial)
+
+        cur_metric = get_metric(results)
+        if self._best_metric is None or cur_metric > self._best_metric:
+            self._best_metric = cur_metric
+            self._best_iter = i
+
+    def _get_trial_dir(self, trial):
+        return os.path.join(self.project_dir, trial.name)
+
+    def get_best_config(self):
+        """
+        Return the best optimization configuration found in the tuning.
+
+        Returns:
+            A object of fleet.DistributedStrategy with best configuration.
+        """
+        assert self._best_iter >= 0, "The best configuration is not found yet !"
+        best_trial = self._finished_trials[self._best_iter]
+        return self._algorithm.get_config_from_trial(best_trial)
+
+    def summary(self):
+        """
+        Display tuning result summary.
+        """
+        # TODO summary with the trial_name with metric_of_trial
+        best_trial = self._finished_trials[self._best_iter]
+        summary_ = f"""
+Tuning Result Summary
+Run total {len(self._finished_trials)} trials with {(time.time() - self._tuning_start_time) / 60} min.
+The best trial is: [{best_trial.name}], whose configuration is following:
+        """
+        summary_ += "\n" + best_trial.summary() + "\n"
+        self._logger.info(summary_)
+        with open(os.path.join(self.project_dir, "summary.txt"), "w+") as fw:
+            fw.writelines(line + "\n" for line in summary_.split("\n"))
+
+        # full_strategy = self.get_best_config()
+        # path = os.path.join(self.project_dir, "tuned_dist_strategy.yaml")
+        # with open(path, 'w') as outfile:
+        #     yaml.dump(full_strategy, outfile, default_flow_style=False)
+
+    def clear(self):
+        """
+        Clear the temporary file generated in tuning procedure.
+        """
+        # TODO clear up zombie process created by tuning
+        if not self._config.debug:
+            for trial in self._finished_trials:
+                trial_dir = self._get_trial_dir(trial)
+                shutil.rmtree(trial_dir, ignore_errors=True)
+
+    def tune(self):
+        """
+        Performs the search for best hyperparameter configurations
+        for the selected optimization pass(es).
+        """
+
+        # step1: collect model info which might be used for
+        # pruning the search space of the algorithm
+        self._tuning_start_time = time.time()
+        self._algorithm.collect_model_info(
+            self._baseline_dist_context.serial_main_program,
+            self._baseline_dist_context.serial_startup_program,
+        )
+
+        # main search loop
+        i = 0
+        while i < self._config.max_num_trial:
+            # step2: create a new trial
+            trial = self._algorithm.next_trial()
+
+            if trial.status == TrialStatus.STOPPED:
+                break
+
+            # step3: evaluate the trial
+            results = self._evaluate_trial(trial)
+
+            # step4: update the algorithm with last result,
+            # which could be used by algorithm to pruning the
+            # remaining search space.
+            self._algorithm.update(results)
+            self._update(i, trial, results)
+
+            # early stop
+            i += 1
+            if (
+                self._config.early_stop
+                and self._config.early_stop <= i - self._best_iter
+            ):
+                self._logger.info(
+                    f"Early stop the Tuning since there is no better trial found within [{self._config.early_stop}] trials"
+                )
+                break
+
+        # step5: summary the best config and return
+        self.summary()
+
+        self.clear()

python/paddle/distributed/auto_parallel/static/tuner/profiler.py

@@ -0,0 +1,296 @@
+# Copyright (c) 2022 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 argparse
+import json
+import os
+import sys
+import time
+import traceback
+
+import paddle
+from paddle.distributed.auto_parallel.static.dist_loader import (
+    DistributedDataLoaderFromGenerator,
+)
+from paddle.distributed.auto_parallel.static.process_group import (
+    get_all_process_groups,
+    new_process_group,
+)
+from paddle.distributed.collective import _get_global_env
+from paddle.framework import Program, _current_expected_place
+from paddle.static import Operator
+
+paddle.enable_static()
+
+
+def str2bool(v):
+    if v.lower() in ('yes', 'true', 't', 'y', '1'):
+        return True
+    elif v.lower() in ('no', 'false', 'f', 'n', '0'):
+        return False
+    else:
+        raise argparse.ArgumentTypeError('Unsupported value encountered.')
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--profile_start_step",
+        default=10,
+        type=int,
+        help="integer indicates the warmup step before starting profile.",
+    )
+    parser.add_argument(
+        "--profile_end_step",
+        default=30,
+        type=int,
+        help="integer indicates at the end step of profile.",
+    )
+    parser.add_argument(
+        "--rank",
+        type=int,
+        required=True,
+        help="the rank id of the this process.",
+    )
+    parser.add_argument(
+        "--device_id",
+        type=int,
+        required=True,
+        help="the device id of the this process.",
+    )
+    parser.add_argument(
+        "--ctx_filename",
+        type=str,
+        required=True,
+        help="the filename to the profile context file saved by optimization tuner",
+    )
+
+    args = parser.parse_args()
+
+    return args
+
+
+def init_process_groups(group_map, rank):
+    for group_id, ranks in group_map.items():
+        if group_id == 0:
+            continue
+        new_process_group(ranks=ranks, group_id=group_id)
+
+    # TODO should instantiate global group first
+    all_process_groups = get_all_process_groups()
+    for process_group in all_process_groups:
+        print(process_group)
+        process_group.instantiate()
+
+
+def get_cpp_error_type(error):
+    msg = str(error).splitlines()
+    cpp_error_types = [
+        'InvalidArgumentError',
+        'NotFoundError',
+        'OutOfRangeError',
+        'AlreadyExistsError',
+        'ResourceExhaustedError',
+        'PreconditionNotMetError',
+        'PermissionDeniedError',
+        'ExecutionTimeoutError',
+        'UnimplementedError',
+        'UnavailableError',
+        'FatalError',
+        'ExternalError',
+    ]
+    error_type = 'FatalError'
+    for et in cpp_error_types:
+        for line in msg:
+            if et in line:
+                return et
+    return error_type
+
+
+def create_dataloader(
+    main_program, startup_program, profile_ctx, epochs=1, steps_per_epoch=None
+):
+    dataset = profile_ctx["dataset"]
+    main_block = main_program.global_block()
+    feed_list = []
+    for name in dataset.input_names:
+        if name in main_block.vars:
+            feed_list.append(main_block.vars[name])
+
+    # remove the first three ops if multi run fit/evaluate/predict
+    op_size = len(main_block.ops)
+    if main_block.ops[0].type == 'create_py_reader':
+        op_size -= 3
+        for _ in range(3):
+            main_block._remove_op(0, sync=False)
+
+    # insert read op at the end of program
+    places = paddle.static.cuda_places()
+    with paddle.static.program_guard(main_program, startup_program):
+        dataloader = DistributedDataLoaderFromGenerator(
+            dataset=dataset,
+            feed_list=feed_list,
+            capacity=70,
+            places=places,
+            batch_size=dataset.batch_size,
+            epochs=epochs,
+            steps_per_epoch=steps_per_epoch,
+            data_parallel_world_size=dataset.dp_world_size,
+            data_parallel_rank=dataset.dp_rank,
+        )
+
+    # move read op from the end of program to the start of program
+    new_op_size = len(main_block.ops)
+    for _ in range(new_op_size - 1, op_size - 1, -1):
+        op = main_block.ops[new_op_size - 1]
+        new_op_desc = main_block.desc._prepend_op()
+        new_op_desc.copy_from(op.desc)
+        new_op = Operator(main_block, new_op_desc, type=new_op_desc.type())
+        main_block.ops.insert(0, new_op)
+    for _ in range(new_op_size - op_size):
+        main_block._remove_op(new_op_size, sync=False)
+    main_block._sync_with_cpp()
+    return dataloader
+
+
+def init_comm(profile_ctx):
+    # override the env for current process
+    dist_env = profile_ctx['distributed_env']
+    genv = _get_global_env()
+    genv = dist_env
+    print(
+        f"current process rank: {genv.rank}, device_id: {genv.device_id}, ip: {genv.current_endpoint}."
+    )
+
+    # init nccl comm
+    group_map = profile_ctx['group_map']
+    init_process_groups(group_map, args.rank)
+
+
+def load_programs(profile_ctx):
+    main_program_desc_str = profile_ctx['main_program_decs']
+    main_program = Program.parse_from_string(main_program_desc_str)
+
+    startup_program_decs_str = profile_ctx['startup_program_decs']
+    startup_program = Program.parse_from_string(startup_program_decs_str)
+
+    loss_var_name = profile_ctx["loss_var_name"]
+    assert main_program.global_block().has_var(loss_var_name)
+    loss_var = main_program.global_block().var(loss_var_name)
+
+    return main_program, startup_program, loss_var
+
+
+def get_executor():
+    place_type = _current_expected_place()
+    if not isinstance(place_type, paddle.CUDAPlace):
+        raise RuntimeError("OptimizationTuner only support CUDA GPU right now.")
+
+    genv = _get_global_env()
+    place = paddle.CUDAPlace(genv.device_id)
+    exe = paddle.static.Executor(place)
+    return exe
+
+
+def profiler(args):
+    """
+    main function to profile experiment for each pass hyper-parameter.
+    """
+    # load ctx
+    if not os.path.isfile(args.ctx_filename):
+        raise ValueError(
+            f"There is no profile context named {args.ctx_filename}."
+        )
+    with open(args.ctx_filename, 'rb') as f:
+        from paddle.framework.restricted_unpickler import safe_load_pickle
+
+        profile_ctx = safe_load_pickle(f, encoding='latin1')
+
+    init_comm(profile_ctx)
+
+    main_program, startup_program, loss_var = load_programs(profile_ctx)
+
+    data_loader = create_dataloader(main_program, startup_program, profile_ctx)
+
+    result_path = profile_ctx["result_filename"]
+
+    exe = get_executor()
+
+    try:
+        exe.run(startup_program)
+        # profile main
+        duration = 0
+        eval_step = 0
+        data_loader._inner_dataloader.start()
+        while eval_step < args.profile_end_step:
+            start_time = time.time()
+
+            loss = exe.run(
+                main_program,
+                fetch_list=[loss_var],
+                use_program_cache=True,
+            )
+
+            end_time = time.time()
+
+            if eval_step >= args.profile_start_step:
+                duration += end_time - start_time
+
+            print(f"step: {eval_step}, loss_print: {loss[0]:f}")
+            eval_step += 1
+
+        avg_tput = (
+            1.0 * (args.profile_end_step - args.profile_start_step) / duration
+        )
+
+        result_dict = {
+            "Throughput": avg_tput,
+            "ErrorType": None,
+        }
+
+        if paddle.distributed.get_rank() == 0:
+            with open(result_path, 'w') as fp:
+                json.dump(result_dict, fp)
+
+        print(f"profile done! avg speed : {avg_tput} step / s.")
+
+    except paddle.framework.core.EOFException:
+        data_loader._inner_dataloader.reset()
+
+    except Exception as e:
+        error_type = get_cpp_error_type(e)
+        result_dict = {
+            "Throughput": -1,
+            "ErrorType": error_type,
+        }
+        if not os.path.isfile(result_path):
+            with open(result_path, 'w') as fp:
+                json.dump(result_dict, fp)
+
+        print(f"profile failed with error: [{error_type}]")
+        print(e)
+        print(traceback.format_exc())
+
+        data_loader._inner_dataloader.reset()
+        del data_loader._inner_dataloader
+        sys.exit(1)
+
+    data_loader._inner_dataloader.reset()
+    del data_loader._inner_dataloader
+
+
+if __name__ == "__main__":
+    paddle.framework.set_flags({'FLAGS_new_executor_sequential_run': 1})
+    args = parse_args()
+    profiler(args)

python/paddle/distributed/auto_parallel/static/tuner/recorder.py

@@ -0,0 +1,216 @@
+#   Copyright (c) 2022 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.
+
+# Notice that the following codes are modified from KerasTuner for a different purpose.
+# Please refer to https://github.com/keras-team/keras-tuner/blob/master/keras_tuner/engine/metrics_tracking.py.
+
+import numpy as np
+
+
+class MetricRecord:
+    """
+    One record for a single metric at a given execution step.
+    """
+
+    def __init__(self, value, step):
+        self._value = value
+        self._step = step
+
+    @property
+    def value(self):
+        return self._value
+
+    @value.setter
+    def value(self, value):
+        self._value = value
+
+    @property
+    def step(self):
+        return self._step
+
+    @step.setter
+    def step(self, step):
+        self._step = step
+
+    def mean(self):
+        return np.mean(self.value)
+
+    def get_state(self):
+        return {"value": self.value, "step": self.step}
+
+    @classmethod
+    def from_state(cls, state):
+        return cls(**state)
+
+    def __eq__(self, other):
+        if not isinstance(other, MetricRecord):
+            return False
+        return other.value == self.value and other.step == self.step
+
+    def __repr__(self):
+        return f"MetricRecord(value={self.value}, step={self.step})"
+
+
+class MetricRecords:
+    """
+    Records of a single metric across different executions.
+    """
+
+    def __init__(self, direction="min"):
+        if direction not in {"min", "max"}:
+            raise ValueError(
+                f"direction should be one of {{min, max}}, but got: {direction}."
+            )
+        self._direction = direction
+        self._records = {}
+
+    @property
+    def records(self):
+        return sorted(self._records.values(), key=lambda r: r.step)
+
+    @records.setter
+    def records(self, records):
+        for r in records:
+            self.update(r.value, step=r.step)
+
+    @property
+    def direction(self):
+        return self._direction
+
+    @direction.setter
+    def direction(self, direction):
+        self._direction = direction
+
+    def update(self, value, step=0):
+        if step in self._records:
+            self._records[step].set_value(value)
+        else:
+            self._records[step] = MetricRecord(value, step=step)
+
+    def get_best_value(self):
+        values = [r.mean() for r in self._records.values()]
+        if not values:
+            return None
+        if self._direction == "min":
+            return np.nanmin(values)
+        return np.nanmax(values)
+
+    def get_best_step(self):
+        best_value = self.get_best_value()
+        if best_value is None:
+            return None
+        for r in self._records.values():
+            if r.mean() == best_value:
+                return r.step
+
+    def get_statistics(self):
+        records = self.records
+        records_values = [r.mean() for r in records]
+        if not len(records_values):
+            return {}
+        return {
+            "min": float(np.nanmin(records_values)),
+            "max": float(np.nanmax(records_values)),
+            "mean": float(np.nanmean(records_values)),
+            "median": float(np.nanmedian(records_values)),
+            "var": float(np.nanvar(records_values)),
+            "std": float(np.nanstd(records_values)),
+        }
+
+    def get_state(self):
+        state = {}
+        state["direction"] = self._direction
+        state["records"] = [r.get_state() for r in self.records]
+        return state
+
+    @classmethod
+    def from_state(cls, state):
+        records = cls(state["direction"])
+        records.records = [MetricRecord.from_state(r) for r in state["records"]]
+        return records
+
+
+class MetricsRecorder:
+    """
+    Record the values for all metrics.
+    """
+
+    def __init__(self, metrics=None):
+        self._records = {}
+        self.register_metrics(metrics)
+
+    @property
+    def records(self):
+        return self._records
+
+    def exists(self, name):
+        return name in self._records
+
+    def register_metrics(self, metrics=None):
+        metrics = metrics or []
+        for metric in metrics:
+            self.register(metric.name)
+
+    def register(self, name, direction=None):
+        if self.exists(name):
+            raise ValueError(f"Metric {name} have been registered.")
+        if direction is None:
+            direction = "min"
+        self._records[name] = MetricRecords(direction)
+
+    def update(self, name, value, step=0):
+        value = float(value)
+        if not self.exists(name):
+            self.register(name)
+
+        prev_best = self._records[name].get_best_value()
+        self._records[name].update(value, step=step)
+        new_best = self._records[name].get_best_value()
+
+        improved = new_best != prev_best
+        return improved
+
+    def get_records(self, name):
+        return self._records[name].records
+
+    def set_records(self, name, records):
+        if not self.exists(name):
+            self.register(name)
+        self._records[name].records = records
+
+    def get_best_value(self, name):
+        return self._records[name].get_best_value()
+
+    def get_best_step(self, name):
+        return self._records[name].get_best_step()
+
+    def get_statistics(self, name):
+        return self._records[name].get_statistics()
+
+    def get_state(self):
+        return {
+            "metrics": {
+                name: metric_records.get_state()
+                for name, metric_records in self._records.items()
+            }
+        }
+
+    @classmethod
+    def from_state(cls, state):
+        recorder = cls()
+        recorder._records = {
+            name: MetricRecords.from_state(metric_records)
+            for name, metric_records in state["metrics"].items()
+        }
+        return recorder

python/paddle/distributed/auto_parallel/static/tuner/storable.py

@@ -0,0 +1,42 @@
+#   Copyright (c) 2022 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.
+
+# Notice that the following codes are modified from KerasTuner for a different purpose.
+# Please refer to https://github.com/keras-team/keras-tuner/blob/master/keras_tuner/engine/metrics_tracking.py.
+
+import json
+
+
+class Storable:
+    def get_state(self):
+        raise NotImplementedError
+
+    def set_state(self, state):
+        raise NotImplementedError
+
+    def save(self, path):
+        state = self.get_state()
+        state_json = json.dumps(state)
+        try:
+            with open(path, "w") as f:
+                f.write(state_json)
+            return str(path)
+        except OSError as e:
+            raise OSError(f"Failed to save file at {path}: {e}") from e
+
+    def load(self, path):
+        with open(path, "r") as f:
+            state_data = f.read()
+        state = json.loads(state_data)
+        self.set_state(state)

python/paddle/distributed/auto_parallel/static/tuner/trial.py

@@ -0,0 +1,169 @@
+#   Copyright (c) 2022 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.
+
+# Notice that the following codes are modified from KerasTuner to implement our own tuner.
+# Please refer to https://github.com/keras-team/keras-tuner/blob/master/keras_tuner/engine/trial.py.
+
+import hashlib
+import random
+import time
+
+from .recorder import MetricsRecorder
+from .storable import Storable
+from .tunable_space import TunableSpace
+
+
+class TrialStatus:
+    RUNNING = "RUNNING"
+    COMPLETED = "COMPLETED"
+    STOPPED = "STOPPED"
+    INVALID = "INVALID"
+
+
+class Trial(Storable):
+    def __init__(
+        self, tunable_space, trial_id=None, status=TrialStatus.RUNNING
+    ):
+        self._id = _generate_trial_id() if trial_id is None else trial_id
+        self._space = tunable_space
+        self._recorder = MetricsRecorder()
+        self._score = None
+        self._best_step = None
+        self._status = status
+
+    @property
+    def id(self):
+        return self._id
+
+    @property
+    def space(self):
+        return self._space
+
+    @property
+    def recorder(self):
+        return self._recorder
+
+    @property
+    def score(self):
+        return self._score
+
+    @score.setter
+    def score(self, score):
+        self._score = score
+
+    @property
+    def best_step(self):
+        return self._best_step
+
+    @best_step.setter
+    def best_step(self, best_step):
+        self._best_step = best_step
+
+    @property
+    def status(self):
+        return self._status
+
+    @status.setter
+    def status(self, status):
+        self._status = status
+
+    def summary(self):
+        print("Tunable space:")
+        if self.space.values:
+            for tv, value in self.space.values.items():
+                print(tv + ":", value)
+
+        if self.score is not None:
+            print(f"Score: {self.score}")
+
+    def get_state(self):
+        return {
+            "id": self.id,
+            "space": self.space.get_state(),
+            "recorder": self.recorder.get_state(),
+            "score": self.score,
+            "best_step": self.best_step,
+            "status": self.status,
+        }
+
+    def set_state(self, state):
+        self._id = state["id"]
+        self._space = TunableSpace.from_state(state["space"])
+        self._recorder = MetricsRecorder.from_state(state["recorder"])
+        self._score = state["score"]
+        self._best_step = state["best_step"]
+        self._status = state["status"]
+
+    @classmethod
+    def from_state(cls, state):
+        trial = cls(tunable_space=None)
+        trial.set_state(state)
+        return trial
+
+
+class OptimizationTunerTrial(Trial):
+    def __init__(
+        self,
+        config,
+        name,
+        changed_configs,
+        trial_id=None,
+        status=TrialStatus.RUNNING,
+    ):
+        super().__init__(config, trial_id, status)
+        self._name = name
+        self._changed_configs = changed_configs
+
+    @property
+    def name(self):
+        return self._name
+
+    def summary(self):
+        spacing = 2
+        max_k = 38
+        max_v = 38
+
+        length = max_k + max_v + spacing
+
+        h1_format = "    " + f"|{{:^{length}s}}|\n"
+        h2_format = "    " + "|{{:>{}s}}{}{{:^{}s}}|\n".format(
+            max_k, " " * spacing, max_v
+        )
+
+        border = "    +" + "".join(["="] * length) + "+"
+        line = "    +" + "".join(["-"] * length) + "+"
+
+        draws = border + "\n"
+        draws += h1_format.format("")
+        draws += h1_format.format("Tuned Configurations Overview")
+        draws += h1_format.format("")
+
+        for name in self._changed_configs:
+            draws += border + "\n"
+            draws += h1_format.format(f"{name} auto=True <-> {name}")
+            draws += line + "\n"
+            my_configs = getattr(self.space, name)
+            keys = my_configs.to_dict().keys()
+            for key in keys:
+                draws += h2_format.format(
+                    key, str(my_configs.to_dict().get(key, None))
+                )
+
+        result_res = draws + border
+        return result_res
+
+
+def _generate_trial_id():
+    s = str(time.time()) + str(random.randint(1, int(1e7)))
+    return hashlib.sha256(s.encode("utf-8")).hexdigest()[:32]

python/paddle/distributed/auto_parallel/static/tuner/tunable_space.py

@@ -0,0 +1,156 @@
+#   Copyright (c) 2022 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.
+
+# Notice that the following codes are modified from KerasTuner to implement our own tuner.
+# Please refer to https://github.com/keras-team/keras-tuner/blob/master/keras_tuner/engine/hyperparameters.py.
+
+from .tunable_variable import Boolean, Choice, Fixed, FloatRange, IntRange
+
+
+class TunableSpace:
+    """
+    A TunableSpace is constructed by the tunable variables.
+    """
+
+    def __init__(self):
+        # Tunable variables for this tunable variables
+        self._variables = {}
+        # Specific values corresponding to each tunable variable
+        self._values = {}
+
+    @property
+    def variables(self):
+        return self._variables
+
+    @variables.setter
+    def variables(self, variables):
+        self._variables = variables
+
+    @property
+    def values(self):
+        return self._values
+
+    @values.setter
+    def values(self, values):
+        self._values = values
+
+    def get_value(self, name):
+        if name in self.values:
+            return self.values[name]
+        else:
+            raise KeyError(f"{name} does not exist.")
+
+    def set_value(self, name, value):
+        if name in self.values:
+            self.values[name] = value
+        else:
+            raise KeyError(f"{name} does not exist.")
+
+    def _exists(self, name):
+        if name in self._variables:
+            return True
+        return False
+
+    def _retrieve(self, tv):
+        tv = tv.__class__.from_state(tv.get_state())
+        if self._exists(tv.name):
+            return self.get_value(tv.name)
+        return self._register(tv)
+
+    def _register(self, tv):
+        self._variables[tv.name] = tv
+        if tv.name not in self.values:
+            self.values[tv.name] = tv.default
+        return self.values[tv.name]
+
+    def __getitem__(self, name):
+        return self.get_value(name)
+
+    def __setitem__(self, name, value):
+        self.set_value(name, value)
+
+    def __contains__(self, name):
+        try:
+            self.get_value(name)
+            return True
+        except (KeyError, ValueError):
+            return False
+
+    def fixed(self, name, default):
+        tv = Fixed(name=name, default=default)
+        return self._retrieve(tv)
+
+    def boolean(self, name, default=False):
+        tv = Boolean(name=name, default=default)
+        return self._retrieve(tv)
+
+    def choice(self, name, values, default=None):
+        tv = Choice(name=name, values=values, default=default)
+        return self._retrieve(tv)
+
+    def int_range(self, name, start, stop, step=1, default=None):
+        tv = IntRange(
+            name=name, start=start, stop=stop, step=step, default=default
+        )
+        return self._retrieve(tv)
+
+    def float_range(self, name, start, stop, step=None, default=None):
+        tv = FloatRange(
+            name=name, start=start, stop=stop, step=step, default=default
+        )
+        return self._retrieve(tv)
+
+    def get_state(self):
+        return {
+            "variables": [
+                {"class_name": v.__class__.__name__, "state": v.get_state()}
+                for v in self._variables.values()
+            ],
+            "values": dict(self.values.items()),
+        }
+
+    @classmethod
+    def from_state(cls, state):
+        ts = cls()
+        for v in state["variables"]:
+            v = _deserialize_tunable_variable(v)
+            ts._variables[v.name] = v
+        ts._values = dict(state["values"].items())
+        return ts
+
+
+def _deserialize_tunable_variable(state):
+    classes = (Boolean, Fixed, Choice, IntRange, FloatRange)
+    cls_name_to_cls = {cls.__name__: cls for cls in classes}
+
+    if isinstance(state, classes):
+        return state
+
+    if (
+        not isinstance(state, dict)
+        or "class_name" not in state
+        or "state" not in state
+    ):
+        raise ValueError(
+            f"Expect state to be a python dict containing class_name and state as keys, but found {state}"
+        )
+
+    cls_name = state["class_name"]
+    cls = cls_name_to_cls[cls_name]
+    if cls is None:
+        raise ValueError(f"Unknown class name {cls_name}")
+
+    cls_state = state["state"]
+    deserialized_object = cls.from_state(cls_state)
+    return deserialized_object

python/paddle/distributed/auto_parallel/static/tuner/tunable_variable.py

@@ -0,0 +1,240 @@
+#   Copyright (c) 2022 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.
+
+# Notice that the following codes are modified from KerasTuner to implement our own tuner.
+# Please refer to https://github.com/keras-team/keras-tuner/blob/master/keras_tuner/engine/hyperparameters.py.
+
+import numpy as np
+
+
+class TunableVariable:
+    """
+    TunableVariable base class.
+    """
+
+    def __init__(self, name, default=None):
+        self.name = name
+        self._default = default
+
+    @property
+    def default(self):
+        return self._default
+
+    def get_state(self):
+        return {"name": self.name, "default": self.default}
+
+    @classmethod
+    def from_state(cls, state):
+        return cls(**state)
+
+
+class Fixed(TunableVariable):
+    """
+    Fixed variable which cannot be changed.
+    """
+
+    def __init__(self, name, default):
+        super().__init__(name=name, default=default)
+        self.name = name
+        if not isinstance(default, (str, int, float, bool)):
+            raise ValueError(
+                f"Fixed must be an str, int, float or bool, but found {default}"
+            )
+        self._default = default
+
+    def random(self, seed=None):
+        return self._default
+
+    def __repr__(self):
+        return f"Fixed(name: {self.name}, value: {self.default})"
+
+
+class Boolean(TunableVariable):
+    """
+    Choice between True and False.
+    """
+
+    def __init__(self, name, default=False):
+        super().__init__(name=name, default=default)
+        if default not in {True, False}:
+            raise ValueError(
+                f"default must be a Python boolean, but got {default}"
+            )
+
+    def random(self, seed=None):
+        rng = np.random.default_rng(seed)
+        return rng.choice((True, False))
+
+    def __repr__(self):
+        return f'Boolean(name: "{self.name}", default: {self.default})'
+
+
+class Choice(TunableVariable):
+    def __init__(self, name, values, default=None):
+        super().__init__(name=name, default=default)
+
+        types = {type(v) for v in values}
+        if len(types) > 1:
+            raise TypeError(
+                f"Choice can contain only one type of value, but found values: {values} with types: {types}."
+            )
+        self._is_unknown_type = False
+
+        if isinstance(values[0], str):
+            values = [str(v) for v in values]
+            if default is not None:
+                default = str(default)
+        elif isinstance(values[0], int):
+            values = [int(v) for v in values]
+            if default is not None:
+                default = int(default)
+        elif isinstance(values[0], float):
+            values = [float(v) for v in values]
+            if default is not None:
+                default = float(default)
+        elif isinstance(values[0], bool):
+            values = [bool(v) for v in values]
+            if default is not None:
+                default = bool(default)
+        else:
+            self._is_unknown_type = True
+            self._indices = list(range(len(values)))
+        self.values = values
+
+        if default is not None and default not in values:
+            raise ValueError(
+                f"The default value should be one of the choices {values}, but found {default}"
+            )
+        self._default = default
+
+    @property
+    def default(self):
+        if self._default is None:
+            if None in self.values:
+                return None
+            return self.values[0]
+        return self._default
+
+    def random(self, seed=None):
+        rng = np.random.default_rng(seed)
+        if self._is_unknown_type:
+            indice = rng.choice(self._indices)
+            return self.values[indice]
+        else:
+            return rng.choice(self.values)
+
+    def get_state(self):
+        state = super().get_state()
+        state["values"] = self.values
+        return state
+
+    def __repr__(self):
+        return f'Choice(name: "{self.name}", values: {self.values}, default: {self.default})'
+
+
+class IntRange(TunableVariable):
+    """
+    Integer range.
+    """
+
+    def __init__(self, name, start, stop, step=1, default=None, endpoint=False):
+        super().__init__(name=name, default=default)
+        self.start = self._check_int(start)
+        self.stop = self._check_int(stop)
+        self.step = self._check_int(step)
+        self._default = default
+        self.endpoint = endpoint
+
+    @property
+    def default(self):
+        if self._default is not None:
+            return self._default
+        return self.start
+
+    def random(self, seed=None):
+        rng = np.random.default_rng(seed)
+        value = (self.stop - self.start) * rng.random() + self.start
+        if self.step is not None:
+            if self.endpoint:
+                values = np.arange(self.start, self.stop + 1e-7, step=self.step)
+            else:
+                values = np.arange(self.start, self.stop, step=self.step)
+            closest_index = np.abs(values - value).argmin()
+            value = values[closest_index]
+        return int(value)
+
+    def get_state(self):
+        state = super().get_state()
+        state["start"] = self.start
+        state["stop"] = self.stop
+        state["step"] = self.step
+        state["default"] = self._default
+        return state
+
+    def _check_int(self, val):
+        int_val = int(val)
+        if int_val != val:
+            raise ValueError(f"Expects val is an int, but found: {val}.")
+        return int_val
+
+    def __repr__(self):
+        return f"IntRange(name: {self.name}, start: {self.start}, stop: {self.stop}, step: {self.step}, default: {self.default})"
+
+
+class FloatRange(TunableVariable):
+    """
+    Float range.
+    """
+
+    def __init__(
+        self, name, start, stop, step=None, default=None, endpoint=False
+    ):
+        super().__init__(name=name, default=default)
+        self.stop = float(stop)
+        self.start = float(start)
+        if step is not None:
+            self.step = float(step)
+        else:
+            self.step = None
+        self._default = default
+        self.endpoint = endpoint
+
+    @property
+    def default(self):
+        if self._default is not None:
+            return self._default
+        return self.start
+
+    def random(self, seed=None):
+        rng = np.random.default_rng(seed)
+        value = (self.stop - self.start) * rng.random() + self.start
+        if self.step is not None:
+            if self.endpoint:
+                values = np.arange(self.start, self.stop + 1e-7, step=self.step)
+            else:
+                values = np.arange(self.start, self.stop, step=self.step)
+            closest_index = np.abs(values - value).argmin()
+            value = values[closest_index]
+        return value
+
+    def get_state(self):
+        state = super().get_state()
+        state["start"] = self.start
+        state["stop"] = self.stop
+        state["step"] = self.step
+        state["endpoint"] = self.endpoint
+        return state
+
+    def __repr__(self):
+        return f"FloatRange(name: {self.name}, start: {self.start}, stop: {self.stop}, step: {self.step}, default: {self.default}, endpoint: {self.endpoint})"