wehub-resource-sync
242 lines
8.5 KiB
Python
242 lines
8.5 KiB
Python
# 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
|
|
)
|