strategy_compiler.py 5.5 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32
#   Copyright (c) 2020 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 maximum_path_len_algo(optimizer_list):
    max_idx = 0
    max_len = 0
    candidates = []
    for idx, opt in enumerate(optimizer_list):
        local_buffer = [opt]
        for opt_inner in optimizer_list:
            if opt._can_update(opt_inner):
                local_buffer.append(opt_inner)
        if len(local_buffer) > max_len:
            max_idx = idx
            max_len = len(local_buffer)
        candidates.append(local_buffer)
    if len(candidates) == 0:
        return None
    for idx, opt in enumerate(candidates[max_idx][:-1]):
        opt._update_inner_optimizer(candidates[max_idx][idx + 1])
D
Dong Daxiang 已提交
33
    return candidates[max_idx]
34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53


class StrategyCompilerBase(object):
    def __init__(self):
        pass


class StrategyCompiler(StrategyCompilerBase):
    """
    StrategyCompiler is responsible for meta optimizers combination
    Generally, a user can define serveral distributed strategies that
    can generate serveral meta optimizer. The combination of these 
    meta optimizers should have the right order to apply the optimizers'
    minimize function.
    This class is responsible for the executable distributed optimizer
    generation.
    """

    def __init__(self):
        super(StrategyCompiler, self).__init__()
D
Dong Daxiang 已提交
54 55 56 57 58 59 60 61 62
        self._meta_optimizer = None
        self._graph_optimizer = None
        self._valid_optimizer_list = None
        self._user_defined_strategy = None
        self._meta_optimizer_candidates = []
        self._graph_optimizer_candidates = []

    def _get_valid_strategy(self, dist_strategy, can_not_apply_optimizer_list):
        import copy
M
mapingshuo 已提交
63
        valid_strategy = copy.deepcopy(dist_strategy)
D
Dong Daxiang 已提交
64 65 66 67 68 69 70 71 72 73 74 75 76 77
        invalid_optimizers = []
        for candidate in self._meta_optimizer_candidates:
            is_valid = False
            for valid in self._meta_optimizers:
                if candidate.__class__.__name__ == valid.__class__.__name__:
                    is_valid = True
                    break
            if not is_valid:
                invalid_optimizers.append(candidate)
        for opt in invalid_optimizers:
            opt._disable_strategy(valid_strategy)
        for opt in can_not_apply_optimizer_list:
            opt._disable_strategy(valid_strategy)
        return valid_strategy
78

79 80 81 82 83 84 85 86 87 88 89 90
    """
    Meta Optimizer Type A: rewrite forward, backward. e.g. recompute, async, sync, pipeline.
                           results will be splitted in async, sync, pipeline
    Meta Optimizer Type B: rewrite forward, 
                           e.g. AMP and the corresponding backward is generated by rewritten forward
    Meta Opitmizer Type B: rewrite backward. e.g. gradient fusion
    Meta Optimizer Type D: rewrite optimize. e.g. lars, lamb, localsgd, gradient merge, dgc
    Meta Optimizer Type E: only transpile to Graph structure for runtime,
                           currently, grad fusion and kernel fusion, sync batch-norm included.
                           we will remove grad fusion and sync batch-norm
    """

91
    def generate_optimizer(self, loss, role_maker, optimizer,
D
Dong Daxiang 已提交
92
                           user_defined_strategy, meta_optimizer_list,
93
                           graph_optimizer_list):
D
Dong Daxiang 已提交
94 95 96 97
        self._user_defined_strategy = user_defined_strategy
        self._meta_optimizer_candidates = meta_optimizer_list
        self._graph_optimizer_candidates = graph_optimizer_list

98 99 100 101 102
        if len(meta_optimizer_list) == 0 and len(graph_optimizer_list) == 0:
            return optimizer, None
        else:
            # currently, we use heuristic algorithm to select
            # meta optimizers combinations
D
Dong Daxiang 已提交
103 104
            meta_optimizers = maximum_path_len_algo(meta_optimizer_list)
            graph_optimizers = maximum_path_len_algo(graph_optimizer_list)
105 106 107 108
            # should design a distributed strategy update interface
            # when we have finally decided the combination of meta_optimizer
            # and graph_optimizer, the corresponding distributed strategy
            # should be updated.
D
Dong Daxiang 已提交
109 110 111 112 113 114 115 116

            self._meta_optimizers = meta_optimizers
            self._graph_optimizers = graph_optimizers

            return_meta = None if meta_optimizers == None else meta_optimizers[
                0]
            return_graph = None if graph_optimizers == None else graph_optimizers[
                0]
117 118 119 120 121 122 123 124 125 126 127 128 129 130

            if meta_optimizers == None or graph_optimizers == None:
                return return_meta, return_graph

            # do heuristic filter here, if any meta optimizer in graph optimizers is in 
            # any meta optimizers' black list, set return_graph to None
            need_graph_opt = True
            for graph_opt in graph_optimizers:
                for program_opt in meta_optimizers:
                    if graph_opt.__class__.__name__ in program_opt.meta_optimizers_black_list:
                        need_graph_opt = False
            if not need_graph_opt:
                return_graph = None

D
Dong Daxiang 已提交
131
            return return_meta, return_graph