dist_softmax.py 7.0 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14
# Copyright (c) 2021 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

15
from .common import DistributedOperatorImplContainer
16
from .common import DistributedOperatorImpl
17
from .common import register_distributed_operator_impl_container
18
from .common import register_distributed_operator_impl
C
caozhou 已提交
19
from .common import is_parameter_related
20 21 22 23 24 25
from ..utils import is_dim_shard
from ..utils import is_dim_replicate
from ..utils import is_valid_list_index
from ..utils import compute_compatible_dim_mapping
from ..utils import compute_compatible_dims_mapping
from ..utils import compute_compatible_and_update_dim_mapping
26
from .dist_default import DistributedDefaultImpl0
27
from ..cost import _g_op_cost_factory
C
caozhou 已提交
28 29 30 31
from ..cost import build_comp_desc_from_dist_op, build_dp_costs
from ..cost import build_comp_costs_from_descs
from ..cost import SoftmaxOpCost, SoftmaxGradOpCost
from paddle.distributed.fleet.meta_optimizers.common import OpRole
32
from paddle.distributed.auto_parallel.cost.comm_op_cost import AllreduceSumOpCost
33 34


35
class DistributedSoftmax(DistributedOperatorImplContainer):
36

37 38
    def __init__(self, op_type):
        super(DistributedSoftmax, self).__init__(op_type)
39 40


41
register_distributed_operator_impl_container(DistributedSoftmax("softmax"))
42 43 44


class DistributedSoftmaxImpl(DistributedOperatorImpl):
45

46
    def __init__(self, name):
47
        super(DistributedSoftmaxImpl, self).__init__(name)
48
        self._forward_implemented = False
49
        self._backward_implemented = False
50

C
caozhou 已提交
51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106
    def calc_cost(self, op_role, dist_op, ctx, cluster):
        cost = None
        if int(op_role) == int(OpRole.Backward):
            cost = self.calc_bwd_cost(dist_op, ctx, cluster)
        else:
            cost = self.calc_fwd_cost(dist_op, ctx, cluster)
        assert cost is not None
        return cost

    def calc_fwd_cost(self, dist_op, ctx, cluster):
        # calc comp op cost
        desc_mapping = build_comp_desc_from_dist_op(dist_op=dist_op,
                                                    dist_context=ctx)
        processes = dist_op.dist_attr.process_mesh.processes
        cost_mapping = build_comp_costs_from_descs(SoftmaxOpCost, ctx,
                                                   processes, desc_mapping,
                                                   cluster)

        res_cost = [cost_mapping]
        return res_cost

    def calc_bwd_cost(self, dist_op, ctx, cluster):
        # calc comp op cost
        res = []
        desc_mapping = build_comp_desc_from_dist_op(dist_op=dist_op,
                                                    dist_context=ctx)
        dist_attr = dist_op.dist_attr
        process_mesh = dist_attr.process_mesh
        processes = process_mesh.processes
        cost_mapping = build_comp_costs_from_descs(SoftmaxGradOpCost, ctx,
                                                   processes, desc_mapping,
                                                   cluster)
        res.append(cost_mapping)

        backward_op = dist_op.serial_op
        main_block = backward_op.block
        need_gradient_allreduce = False
        vars = main_block.vars
        for input_name in backward_op.desc.input_names():
            for varname in backward_op.desc.input(input_name):
                if "@GRAD" not in varname and is_parameter_related(
                        varname, main_block):
                    # NOTE input var's dim_mapping of backward op should be the same with input var instead of corresponding varname of forward op
                    var_dim_mapping = dist_attr.get_input_dims_mapping(varname)

                    mesh_shape = process_mesh.topology
                    batch_size_axis = var_dim_mapping[0]
                    if batch_size_axis > -1 and mesh_shape[batch_size_axis] > 1:
                        parallel_axis = batch_size_axis
                        attrs = {"use_calc_stream": True}
                        var_names = [varname + "@GRAD"]
                        build_dp_costs(res, dist_op, ctx, var_names, attrs,
                                       parallel_axis, cluster)

        return res

107 108 109
    def is_input_compatible(self, dist_op):
        op_desc = dist_op.serial_op.desc
        op_dist_attr = dist_op.dist_attr
110 111 112 113
        x_name = op_desc.input('X')[0]
        axis = op_desc.attr('axis')
        x_dims_mapping = op_dist_attr.get_input_dims_mapping(x_name)

114 115
        # if axis != -1 and axis != len(x_dims_mapping) - 1:
        #     return False
116 117 118 119 120 121

        if is_dim_shard(x_dims_mapping[axis]):
            return False

        return True

122 123 124
    def is_output_compatible(self, dist_op):
        op_desc = dist_op.serial_op.desc
        op_dist_attr = dist_op.dist_attr
125 126 127 128
        out_name = op_desc.output('Out')[0]
        axis = op_desc.attr('axis')
        out_dims_mapping = op_dist_attr.get_output_dims_mapping(out_name)

129 130
        # if axis != -1 and axis != len(out_dims_mapping) - 1:
        #     return False
131 132 133 134 135 136

        if is_dim_shard(out_dims_mapping[axis]):
            return False

        return True

沉潜的鱼儿's avatar
沉潜的鱼儿 已提交
137
    def is_auto_compatible(self, dist_op):
138 139 140 141
        if (not self.is_input_compatible(dist_op)) or \
            (not self.is_output_compatible(dist_op)):
            return False

沉潜的鱼儿's avatar
沉潜的鱼儿 已提交
142 143 144 145 146 147 148
        op_desc = dist_op.serial_op.desc
        op_dist_attr = dist_op.dist_attr
        x_name = op_desc.input('X')[0]
        axis = op_desc.attr('axis')
        out_name = op_desc.output('Out')[0]
        x_dims_mapping = op_dist_attr.get_input_dims_mapping(x_name)
        out_dims_mapping = op_dist_attr.get_output_dims_mapping(out_name)
149 150
        # if axis != -1 and axis != len(x_dims_mapping) - 1:
        #     return False
沉潜的鱼儿's avatar
沉潜的鱼儿 已提交
151 152 153 154 155 156

        if x_dims_mapping != out_dims_mapping:
            return False

        return True

157
    def update_dims_mapping(self, dist_op):
158
        changed = False
159 160
        op_desc = dist_op.serial_op.desc
        op_dist_attr = dist_op.dist_attr
161 162 163 164 165 166 167 168 169 170 171 172 173
        x_name = op_desc.input('X')[0]
        out_name = op_desc.output('Out')[0]
        x_dims_mapping = op_dist_attr.get_input_dims_mapping(x_name)
        out_dims_mapping = op_dist_attr.get_output_dims_mapping(out_name)

        for i in range(len(x_dims_mapping)):
            dim_changed = compute_compatible_and_update_dim_mapping(
                [x_dims_mapping, out_dims_mapping], [i, i])
            if dim_changed:
                changed = True

        return changed

174 175 176 177
    @staticmethod
    def forward(ctx, *args, **kwargs):
        DistributedDefaultImpl0.forward(ctx, *args, **kwargs)

178 179
    @staticmethod
    def backward(ctx, *args, **kwargs):
180
        DistributedDefaultImpl0.backward(ctx, *args, **kwargs)
181

182 183 184

register_distributed_operator_impl(
    "softmax", DistributedSoftmaxImpl("replicate_last_axis"))