chore: import upstream snapshot with attribution

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
+                )