dist_embedding.py 26.5 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 16 17
from paddle.distributed.auto_parallel.cost.comm_op_cost import (
    AllreduceSumOpCost,
    IdentityOpCost,
18
)
19
from paddle.distributed.fleet.meta_optimizers.common import OP_ROLE_KEY, OpRole
20 21 22 23 24 25 26 27 28 29 30 31
from paddle.fluid import core, unique_name
from paddle.fluid.data_feeder import check_dtype, check_variable_and_dtype

from ..cost import (
    EmbeddingGradOpCost,
    EmbeddingOpCost,
    build_comm_costs_from_descs,
    build_comm_desc_from_dist_op,
    build_comp_costs_from_descs,
    build_comp_desc_from_dist_op,
    build_dp_costs,
)
32
from ..dist_attribute import OperatorDistAttr
33
from ..process_group import new_process_group
34 35 36
from ..utils import (
    _get_comm_group,
    _get_corresponding_rank,
37 38 39 40
    _get_idx_in_axis,
    compute_compatible_and_update_dim_mapping,
    is_dim_replicate,
    is_dim_shard,
41 42
    set_var_dist_attr,
)
43 44 45 46 47 48 49 50 51
from .common import (
    DistributedOperatorImpl,
    DistributedOperatorImplContainer,
    gradient_synchronization,
    infer_shape,
    naive_copy_op_dist_attr_for_program,
    register_distributed_operator_impl,
    register_distributed_operator_impl_container,
    set_comm_op_dist_attr_for_program,
52
)
53 54


55
class DistributedEmbedding(DistributedOperatorImplContainer):
56
    def __init__(self, op_type):
57
        super().__init__(op_type)
58 59


60
register_distributed_operator_impl_container(
61 62
    DistributedEmbedding("lookup_table_v2")
)
63
register_distributed_operator_impl_container(
64 65
    DistributedEmbedding("c_embedding")
)
66
register_distributed_operator_impl_container(
67 68
    DistributedEmbedding("lookup_table")
)
69 70 71 72


def adopt_lookup_table_v1(ctx, main_block, src_op, Ids_var):

73 74 75 76 77
    assert (
        len(Ids_var.shape) == 3
    ), "input Ids to lookup_table should have 3 dimensions but got [{}] with shape [{}]".format(
        Ids_var.name, Ids_var.shape
    )
78 79 80 81 82 83 84
    if not Ids_var.stop_gradient:
        raise NotImplementedError(
            'Requiring the gradient of Ids of lookup_table(v1)dist op is not currently supported. Please open an issue with details on your use case so that we can prioritize adding this (for instance, adversarial training for language model).'
        )

    target_shape = list(Ids_var.shape[:-1])
    intermediate_var_0 = main_block.create_var(
85 86 87
        name=unique_name.generate_with_ignorable_key(
            ".".join(["dist_reshape", 'tmp'])
        ),
88 89 90 91
        dtype=Ids_var.dtype,
        shape=target_shape,
        type=core.VarDesc.VarType.LOD_TENSOR,
        persistable=False,
92 93
        stop_gradient=True,
    )
94 95 96

    target_shape = [0] + list(Ids_var.shape[:-1])
    xshape_var = main_block.create_var(
97 98 99
        name=unique_name.generate_with_ignorable_key(
            ".".join(["dist_Xshape", 'tmp'])
        ),
100 101 102 103
        dtype=Ids_var.dtype,
        shape=target_shape,
        type=core.VarDesc.VarType.LOD_TENSOR,
        persistable=False,
104 105
        stop_gradient=True,
    )
106 107

    # TODO use inplace reshape for memory saving
108 109 110 111 112 113 114 115
    reshape_op = main_block.append_op(
        type='reshape2',
        inputs={'X': [Ids_var]},
        outputs={'Out': [intermediate_var_0], 'XShape': [xshape_var]},
        attrs={
            "shape": [0, -1],
        },
    )
116 117 118 119 120 121

    # set dist attr
    op_dist_attr = ctx.get_op_dist_attr_for_program(src_op)
    Ids_var_dist_attr = op_dist_attr.get_input_dist_attr(Ids_var.name)
    assert Ids_var_dist_attr is not None
    intermediate_var_0_dist_attr = set_var_dist_attr(
122 123 124 125 126 127 128 129 130 131 132
        ctx,
        intermediate_var_0,
        Ids_var_dist_attr.dims_mapping,
        Ids_var_dist_attr.process_mesh,
    )
    set_var_dist_attr(
        ctx,
        xshape_var,
        [-1] + list(Ids_var_dist_attr.dims_mapping),
        Ids_var_dist_attr.process_mesh,
    )
133
    op_dist_attr.del_input_dist_attr(Ids_var.name)
134 135 136
    op_dist_attr.set_input_dist_attr(
        intermediate_var_0.name, intermediate_var_0_dist_attr
    )
137

138
    new_op_dist_attr = OperatorDistAttr()
139 140 141
    new_op_dist_attr.process_mesh = Ids_var_dist_attr.process_mesh
    new_op_dist_attr.impl_type = "default"
    new_op_dist_attr.impl_idx = 0
142 143 144 145 146 147
    new_op_dist_attr.set_input_dims_mapping(
        Ids_var.name, Ids_var_dist_attr.dims_mapping
    )
    new_op_dist_attr.set_output_dims_mapping(
        intermediate_var_0.name, Ids_var_dist_attr.dims_mapping
    )
148
    new_op_dist_attr.set_output_dims_mapping(
149 150
        xshape_var.name, [-1] + list(Ids_var_dist_attr.dims_mapping)
    )
151 152 153
    ctx.set_op_dist_attr_for_program(reshape_op, new_op_dist_attr)

    return intermediate_var_0
154 155 156 157 158


# RowParallel
class DistributedEmbeddingImpl(DistributedOperatorImpl):
    def __init__(self, name):
159
        super().__init__(name)
160
        self._forward_implemented = True
161
        self._backward_implemented = True
162

C
caozhou 已提交
163 164 165 166 167 168 169 170 171 172 173 174
    def calc_cost(self, op_role, dist_op, ctx, cluster):
        """Calculate the cost by the op role."""
        cost = None
        if int(op_role) == int(OpRole.Forward):
            cost = self.calc_fwd_cost(dist_op, ctx, cluster)
        elif int(op_role) == int(OpRole.Backward):
            cost = self.calc_bwd_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
175 176 177
        desc_mapping = build_comp_desc_from_dist_op(
            dist_op=dist_op, dist_context=ctx
        )
178
        processes = dist_op.dist_attr.process_mesh.process_ids
C
caozhou 已提交
179
        # embedding need start_index
180 181 182
        cost_mapping = build_comp_costs_from_descs(
            EmbeddingOpCost, ctx, processes, desc_mapping, cluster
        )
C
caozhou 已提交
183 184 185

        serial_op = dist_op.serial_op
        parallel_axis = dist_op.dist_attr.get_input_dims_mapping(
186 187
            serial_op.input("W")[0]
        )[0]
C
caozhou 已提交
188 189 190 191 192 193 194 195
        attrs = {"use_calc_stream": True, "use_model_parallel": True}
        var_names = serial_op.output("Out")
        c_allreduce_sum_desc_mapping = build_comm_desc_from_dist_op(
            "c_allreduce_sum",
            dist_op,
            ctx,
            var_names,
            attrs=attrs,
196 197
            parallel_axis=parallel_axis,
        )
C
caozhou 已提交
198 199

        comm_op_cost_list = build_comm_costs_from_descs(
200 201 202 203 204 205
            AllreduceSumOpCost,
            ctx,
            processes,
            c_allreduce_sum_desc_mapping,
            cluster,
        )
C
caozhou 已提交
206 207 208 209 210 211 212 213 214 215 216 217 218

        res_cost = [cost_mapping, comm_op_cost_list]

        return res_cost

    def calc_bwd_cost(self, dist_op, ctx, cluster):
        # by now the backward function only insert the gradient allreduce for dist op itself
        res = []
        backward_op = dist_op.serial_op
        main_block = backward_op.block
        dist_attr = dist_op.dist_attr

        embedding_row_dim_mapping = dist_attr.get_input_dims_mapping(
219 220
            backward_op.input("W")[0]
        )[0]
C
caozhou 已提交
221 222 223 224 225 226 227 228 229
        parallel_axis = embedding_row_dim_mapping
        attrs = {"use_calc_stream": True, "use_model_parallel": True}
        var_names = [backward_op.input("Out@GRAD")[0]]
        c_identity_desc_mapping = build_comm_desc_from_dist_op(
            "c_identity",
            dist_op,
            ctx,
            var_names,
            attrs=attrs,
230 231
            parallel_axis=parallel_axis,
        )
C
caozhou 已提交
232 233

        process_mesh = dist_attr.process_mesh
234
        processes = process_mesh.process_ids
C
caozhou 已提交
235
        comm_op_cost_list = build_comm_costs_from_descs(
236 237
            IdentityOpCost, ctx, processes, c_identity_desc_mapping, cluster
        )
C
caozhou 已提交
238 239 240
        res.append(comm_op_cost_list)

        # calc comp op cost
241 242 243 244 245 246
        desc_mapping = build_comp_desc_from_dist_op(
            dist_op=dist_op, dist_context=ctx
        )
        cost_mapping = build_comp_costs_from_descs(
            EmbeddingGradOpCost, ctx, processes, desc_mapping, cluster
        )
C
caozhou 已提交
247 248 249 250
        res.append(cost_mapping)

        # need gradient allreduce
        var_dim_mapping = dist_attr.get_input_dims_mapping(
251 252
            backward_op.input("Ids")[0]
        )
253
        mesh_shape = process_mesh.shape
C
caozhou 已提交
254 255 256 257 258
        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 = [backward_op.output('W@GRAD')[0]]
259 260 261
            build_dp_costs(
                res, dist_op, ctx, var_names, attrs, parallel_axis, cluster
            )
C
caozhou 已提交
262 263 264

        return res

265 266 267
    def is_input_compatible(self, dist_op):
        op_desc = dist_op.serial_op.desc
        op_dist_attr = dist_op.dist_attr
268 269 270 271
        ids_name = op_desc.input('Ids')[0]
        w_name = op_desc.input('W')[0]
        ids_dims_mapping = op_dist_attr.get_input_dims_mapping(ids_name)
        w_dims_mapping = op_dist_attr.get_input_dims_mapping(w_name)
272
        if is_dim_replicate(w_dims_mapping[-2]) or is_dim_shard(
273 274
            w_dims_mapping[-1]
        ):
275 276 277 278 279 280 281
            return False
        # Other dimensions must be replicate except the batch dimension
        for mapping in ids_dims_mapping[1:]:
            if is_dim_shard(mapping):
                return False
        return True

282 283 284
    def is_output_compatible(self, dist_op):
        op_desc = dist_op.serial_op.desc
        op_dist_attr = dist_op.dist_attr
285 286 287 288 289 290 291 292
        out_name = op_desc.output('Out')[0]
        out_dims_mapping = op_dist_attr.get_output_dims_mapping(out_name)
        # Other dimensions must be replicate except the batch dimension
        for mapping in out_dims_mapping[1:]:
            if is_dim_shard(mapping):
                return False
        return True

沉潜的鱼儿's avatar
沉潜的鱼儿 已提交
293
    def is_auto_compatible(self, dist_op):
294 295 296
        if (not self.is_input_compatible(dist_op)) or (
            not self.is_output_compatible(dist_op)
        ):
297 298
            return False

沉潜的鱼儿's avatar
沉潜的鱼儿 已提交
299 300 301 302 303 304 305 306
        op_desc = dist_op.serial_op.desc
        op_dist_attr = dist_op.dist_attr
        ids_name = op_desc.input('Ids')[0]
        w_name = op_desc.input('W')[0]
        out_name = op_desc.output('Out')[0]
        out_dims_mapping = op_dist_attr.get_output_dims_mapping(out_name)
        ids_dims_mapping = op_dist_attr.get_input_dims_mapping(ids_name)
        w_dims_mapping = op_dist_attr.get_input_dims_mapping(w_name)
307

308
        if ids_dims_mapping != out_dims_mapping[: len(ids_dims_mapping)]:
沉潜的鱼儿's avatar
沉潜的鱼儿 已提交
309 310 311 312
            return False

        return True

313
    def update_dims_mapping(self, dist_op):
314
        changed = False
315 316
        op_desc = dist_op.serial_op.desc
        op_dist_attr = dist_op.dist_attr
317 318 319 320 321 322 323 324 325
        ids_name = op_desc.input('Ids')[0]
        w_name = op_desc.input('W')[0]
        out_name = op_desc.output('Out')[0]
        ids_dims_mapping = op_dist_attr.get_input_dims_mapping(ids_name)
        w_dims_mapping = op_dist_attr.get_input_dims_mapping(w_name)
        out_dims_mapping = op_dist_attr.get_output_dims_mapping(out_name)

        for i in range(len(ids_dims_mapping)):
            dim_changed = compute_compatible_and_update_dim_mapping(
326 327
                [ids_dims_mapping, out_dims_mapping], [i, i]
            )
328 329 330 331
            if dim_changed:
                changed = True

        dim_changed = compute_compatible_and_update_dim_mapping(
332 333
            [w_dims_mapping, out_dims_mapping], [-1, -1]
        )
334 335 336
        if dim_changed:
            changed = True

337 338 339 340 341
        if changed:
            op_dist_attr.set_input_dims_mapping(ids_name, ids_dims_mapping)
            op_dist_attr.set_input_dims_mapping(w_name, w_dims_mapping)
            op_dist_attr.set_output_dims_mapping(out_name, out_dims_mapping)

342 343
        return changed

344 345 346 347 348 349
    @staticmethod
    def forward(ctx, *args, **kwargs):
        """
        kwargs: inputname_mapping & outputname_mapping
        """

350
        dist_op_context = ctx.dist_op_context
351 352 353 354
        main_block = dist_op_context.work_block
        startup_block = dist_op_context.startup_block
        src_op = dist_op_context.cur_src_op
        rank_id = dist_op_context.rank_id
355
        op_dist_attr = ctx.get_op_dist_attr_for_program(src_op)
356 357 358
        assert (
            op_dist_attr is not None
        ), "backward op [{}] don't have dist attribute !".format(str(src_op))
359

360
        # check validation of inputs / outputs
361 362 363 364
        assert 'Ids' in kwargs, "input [{}] is not given".format('Ids')
        assert 'W' in kwargs, "input [{}] is not given".format('W')
        assert 'Out' in kwargs, "output [{}] is not given".format('Out')

365 366 367
        assert (
            len(kwargs['Ids']) == 1
        ), "row_parallel_embedding input Ids take 1 variable but got {}".format(
368
            kwargs['Ids']
369 370 371 372
        )
        assert (
            len(kwargs['W']) == 1
        ), "row_parallel_embedding input W take 1 variable but got {}".format(
373
            kwargs['W']
374 375 376 377
        )
        assert (
            len(kwargs['Out']) == 1
        ), "row_parallel_embedding output Out take 1 variable but got {}".format(
378
            kwargs['Out']
379
        )
380

Z
zhaoyingli 已提交
381
        Ids_var = main_block._var_recursive(kwargs['Ids'][0])
382
        Weight_var = main_block._var_recursive(kwargs['W'][0])
Z
zhaoyingli 已提交
383
        Out_var = main_block._var_recursive(kwargs['Out'][0])
384

385 386 387 388
        # support lookup_table_v1
        if src_op.type == 'lookup_table':
            Ids_var = adopt_lookup_table_v1(ctx, main_block, src_op, Ids_var)

389 390
        # got dist attribute info
        embedding_row_dim_mapping = op_dist_attr.get_input_dims_mapping(
391 392 393 394 395 396 397
            Weight_var.name
        )[0]
        assert (
            embedding_row_dim_mapping >= 0
        ), "row_parallel_embedding's row should be divided by a specific mesh axis, but got [{}]".format(
            embedding_row_dim_mapping
        )
398 399
        process_mesh_shape = op_dist_attr.process_mesh.shape
        process_mesh_group = op_dist_attr.process_mesh.process_ids
400 401 402

        # FIXME (JZ-LIANG) Remove this hack to support any op mesh group for Pipeline Parallelism
        if rank_id not in process_mesh_group:
403 404 405
            rank_id = _get_corresponding_rank(
                ctx, op_dist_attr.process_mesh, rank_id
            )
406 407

        # A generalized method to caculate embedding offset using cartisian product
408 409 410 411 412 413
        relative_idx = _get_idx_in_axis(
            process_mesh_group,
            process_mesh_shape,
            embedding_row_dim_mapping,
            rank_id,
        )
414 415 416 417

        per_part_size = Weight_var.shape[0]
        relative_idx = relative_idx * per_part_size

418
        # TODO caculate ring id
419
        parallel_axis = embedding_row_dim_mapping
420 421 422
        group_ranks = _get_comm_group(
            process_mesh_group, process_mesh_shape, parallel_axis, rank_id
        )
423 424 425
        group = new_process_group(group_ranks)

        # append op
426 427 428
        check_variable_and_dtype(
            Ids_var, 'input', ['int32', 'int64'], 'c_embedding'
        )
429

Z
zhaoyingli 已提交
430 431 432 433 434
        # infer new var shape with op dist attr
        out_tensor_dist_attr = ctx.get_tensor_dist_attr_for_program(Out_var)
        assert out_tensor_dist_attr is not None
        out_var_dist_attr = op_dist_attr.get_output_dist_attr(Out_var.name)
        assert out_var_dist_attr is not None
435 436 437
        ref_shape = infer_shape(
            main_block, Out_var, out_tensor_dist_attr, out_var_dist_attr
        )
Z
zhaoyingli 已提交
438

439
        intermediate_var_0 = main_block.create_var(
440 441 442
            name=unique_name.generate_with_ignorable_key(
                ".".join(["c_embedding", 'tmp'])
            ),
443 444 445 446
            dtype=Weight_var.dtype,
            shape=Out_var.shape,
            type=core.VarDesc.VarType.LOD_TENSOR,
            persistable=False,
447 448
            stop_gradient=Out_var.stop_gradient,
        )
Z
zhaoyingli 已提交
449
        # set intermediate_var_0's dist_attr with Out_var's dist_attr
450 451 452
        ctx.set_tensor_dist_attr_for_program(
            intermediate_var_0, out_var_dist_attr
        )
453 454

        check_variable_and_dtype(
455 456
            Out_var,
            'tensor',
457
            ['float16', 'float32', 'float64', 'int32', 'int64'],
458 459
            'c_allreduce_sum',
        )
460 461 462

        c_embedding_op = main_block.append_op(
            type='c_embedding',
463
            inputs={'Ids': [Ids_var], 'W': [Weight_var]},
464
            outputs={'Out': [intermediate_var_0]},
465 466
            attrs={
                "start_index": relative_idx,
467 468 469
                OP_ROLE_KEY: src_op.attr('op_role'),
            },
        )
Z
zhaoyingli 已提交
470 471
        if intermediate_var_0.shape != ref_shape:
            intermediate_var_0.desc.set_shape(ref_shape)
472 473 474 475 476 477 478 479 480 481

        # use_model_parallel
        c_allreduce_sum_op = main_block.append_op(
            type='c_allreduce_sum',
            inputs={'X': [intermediate_var_0]},
            outputs={'Out': [Out_var]},
            attrs={
                'ring_id': group.id,
                'use_calc_stream': True,
                'use_model_parallel': True,
482 483 484
                OP_ROLE_KEY: src_op.attr('op_role'),
            },
        )
Z
zhaoyingli 已提交
485 486 487 488 489
        if Out_var.shape != ref_shape:
            Out_var.desc.set_shape(ref_shape)

        # set dist op's dist_attr with serial op's dist_attr
        # matmulv2
490
        embedding_op_dist_attr = OperatorDistAttr()
Z
zhaoyingli 已提交
491
        embedding_op_dist_attr.process_mesh = op_dist_attr.process_mesh
492
        embedding_op_dist_attr.impl_type = op_dist_attr.impl_type
Z
zhaoyingli 已提交
493 494 495 496
        embedding_op_dist_attr.impl_idx = op_dist_attr.impl_idx
        for input_varname in c_embedding_op.desc.input_arg_names():
            input_dist_attr = op_dist_attr.get_input_dist_attr(input_varname)
            assert input_dist_attr is not None, "dist_attr is {}".format(
497 498 499 500 501
                op_dist_attr
            )
            embedding_op_dist_attr.set_input_dist_attr(
                input_varname, input_dist_attr
            )
Z
zhaoyingli 已提交
502 503 504
        output_varname = c_embedding_op.desc.output_arg_names()[0]
        output_dist_attr = op_dist_attr.get_output_dist_attr(Out_var.name)
        assert output_dist_attr is not None, "dist_attr is {}".format(
505 506 507 508 509
            op_dist_attr
        )
        embedding_op_dist_attr.set_output_dist_attr(
            output_varname, output_dist_attr
        )
Z
zhaoyingli 已提交
510
        ctx.set_op_dist_attr_for_program(c_embedding_op, embedding_op_dist_attr)
511

Z
zhaoyingli 已提交
512
        # allreduce
513
        allreduce_op_dist_attr = OperatorDistAttr()
Z
zhaoyingli 已提交
514
        allreduce_op_dist_attr.process_mesh = op_dist_attr.process_mesh
515
        allreduce_op_dist_attr.impl_type = op_dist_attr.impl_type
Z
zhaoyingli 已提交
516 517
        allreduce_op_dist_attr.impl_idx = op_dist_attr.impl_idx
        for input_varname in c_allreduce_sum_op.desc.input_arg_names():
Z
zhaoyingli 已提交
518
            input_var = main_block._var_recursive(input_varname)
Z
zhaoyingli 已提交
519 520
            tensor_dist_attr = ctx.get_tensor_dist_attr_for_program(input_var)
            assert tensor_dist_attr is not None
521 522 523
            allreduce_op_dist_attr.set_input_dist_attr(
                input_varname, tensor_dist_attr
            )
Z
zhaoyingli 已提交
524 525 526
        for output_varname in c_allreduce_sum_op.desc.output_arg_names():
            output_dist_attr = op_dist_attr.get_output_dist_attr(output_varname)
            assert output_dist_attr is not None, "dist_attr is {}".format(
527 528 529 530 531 532 533 534
                op_dist_attr
            )
            allreduce_op_dist_attr.set_output_dist_attr(
                output_varname, output_dist_attr
            )
        ctx.set_op_dist_attr_for_program(
            c_allreduce_sum_op, allreduce_op_dist_attr
        )
535 536

        # param initialization sync
537
        if Weight_var.is_parameter and not op_dist_attr.is_recompute:
538 539
            if Weight_var.name in dist_op_context.already_init_sync_vars:
                return
J
JZ-LIANG 已提交
540 541 542 543 544 545 546
            dist_op_context.already_init_sync_vars.add(Weight_var.name)
            param = startup_block.var(Weight_var.name)
            param_dist_attr = ctx.get_tensor_dist_attr_for_program(param)
            process_mesh = param_dist_attr.process_mesh
            dim_mapping = param_dist_attr.dims_mapping

            # NOTE all not splited axis should be presented in mesh
547
            for axis, size in enumerate(process_mesh.shape):
J
JZ-LIANG 已提交
548 549 550
                if size <= 1 or axis in dim_mapping:
                    pass
                else:
551
                    group_ranks = _get_comm_group(
552 553
                        process_mesh.process_ids,
                        process_mesh.shape,
554 555 556
                        axis,
                        rank_id,
                    )
J
JZ-LIANG 已提交
557 558
                    sync_group = new_process_group(group_ranks)

559 560 561 562 563 564 565 566 567 568 569
                    startup_block.append_op(
                        type='c_broadcast',
                        inputs={'X': param},
                        outputs={'Out': param},
                        attrs={
                            'ring_id': sync_group.id,
                            'root': 0,
                            'use_calc_stream': True,
                            OP_ROLE_KEY: OpRole.Forward,
                        },
                    )
570 571 572 573 574

    @staticmethod
    def backward(ctx, *args, **kwargs):

        # by now the backward function only insert the gradient allreduce for dist op itself
575
        dist_op_context = ctx.dist_op_context
576 577 578
        main_block = dist_op_context.work_block
        backward_op = dist_op_context.cur_src_op
        rank_id = dist_op_context.rank_id
579
        dist_attr = ctx.get_op_dist_attr_for_program(backward_op)
580 581 582 583 584
        assert (
            dist_attr is not None
        ), "backward op [{}] don't have dist attribute !".format(
            str(backward_op)
        )
585

586
        # FIXME (JZ-LIANG) Remove this hack to support any op mesh group for Pipeline Parallelism
587
        if rank_id not in dist_attr.process_mesh.process_ids:
588 589 590
            rank_id = _get_corresponding_rank(
                ctx, dist_attr.process_mesh, rank_id
            )
591 592 593 594 595 596

        assert 'Ids' in kwargs, "input [{}] is not given".format('Ids')
        assert 'W' in kwargs, "input [{}] is not given".format('W')
        assert 'Out@GRAD' in kwargs, "input [{}] is not given".format('Out')
        assert 'W@GRAD' in kwargs, "output [{}] is not given".format('W@GRAD')

597 598 599
        assert (
            len(kwargs['Ids']) == 1
        ), "row_parallel_embedding input Ids take 1 variable but got {}".format(
600
            kwargs['Ids']
601 602 603 604
        )
        assert (
            len(kwargs['W']) == 1
        ), "row_parallel_embedding input Ids take 1 variable but got {}".format(
605
            kwargs['W']
606 607 608 609 610 611 612 613 614
        )
        assert (
            len(kwargs['Out@GRAD']) == 1
        ), "row_parallel_embedding input Ids take 1 variable but got {}".format(
            kwargs['Out']
        )
        assert (
            len(kwargs['W@GRAD']) == 1
        ), "row_parallel_embedding output Ids take 1 variable but got {}".format(
615
            kwargs['W@GRAD']
616
        )
617

Z
zhaoyingli 已提交
618 619 620 621
        Ids_var = main_block._var_recursive(kwargs['Ids'][0])
        Weight_var = main_block._var_recursive(kwargs['W'][0])
        Out_grad = main_block._var_recursive(kwargs['Out@GRAD'][0])
        Weight_grad = main_block._var_recursive(kwargs['W@GRAD'][0])
622 623

        embedding_row_dim_mapping = dist_attr.get_input_dims_mapping(
624 625 626 627 628 629 630
            Weight_var.name
        )[0]
        assert (
            embedding_row_dim_mapping >= 0
        ), "row_parallel_embedding's row should be divided by a specific mesh axis, but got [{}]".format(
            embedding_row_dim_mapping
        )
631 632
        process_mesh_shape = dist_attr.process_mesh.shape
        process_mesh_group = dist_attr.process_mesh.process_ids
633 634

        # A generalized method to caculate embedding offset using cartisian product
635 636 637 638 639 640
        relative_idx = _get_idx_in_axis(
            process_mesh_group,
            process_mesh_shape,
            embedding_row_dim_mapping,
            rank_id,
        )
641 642 643 644
        per_part_size = Weight_var.shape[0]
        relative_idx = relative_idx * per_part_size

        check_variable_and_dtype(
645 646 647 648 649
            Out_grad,
            'tensor',
            ['float16', 'float32', 'float64', 'int32', 'int64'],
            '_c_identity',
        )
650 651

        intermediate_var_0 = main_block.create_var(
652 653 654
            name=unique_name.generate_with_ignorable_key(
                ".".join(["c_embedding", '@tmp_0@GRAD'])
            ),
655 656 657 658
            dtype=Out_grad.dtype,
            shape=Out_grad.shape,
            type=core.VarDesc.VarType.LOD_TENSOR,
            persistable=False,
659 660
            stop_gradient=Out_grad.stop_gradient,
        )
661 662 663 664

        # copy X_var's dist_attr to intermediate_var_0's dist_attr
        out_grad_dist_attr = dist_attr.get_input_dist_attr(Out_grad.name)
        assert out_grad_dist_attr is not None
665 666 667
        ctx.set_tensor_dist_attr_for_program(
            intermediate_var_0, out_grad_dist_attr
        )
668

669 670 671 672 673 674
        group_ranks = _get_comm_group(
            process_mesh_group,
            process_mesh_shape,
            embedding_row_dim_mapping,
            rank_id,
        )
675 676 677 678 679 680 681 682 683 684 685
        group = new_process_group(group_ranks)

        c_identity_op = main_block.append_op(
            type='c_identity',
            inputs={'X': [Out_grad]},
            outputs={'Out': intermediate_var_0},
            attrs={
                'ring_id': group.id,
                'use_calc_stream': True,
                'use_model_parallel': True,
                OP_ROLE_KEY: OpRole.Backward,
686 687 688 689 690 691 692 693 694 695 696
            },
        )
        check_variable_and_dtype(
            intermediate_var_0, 'x', ['float16', 'float32', 'float64'], 'linear'
        )
        check_dtype(
            intermediate_var_0.dtype,
            'dtype',
            ['float16', 'float32', 'float64'],
            'linear',
        )
697

698 699 700
        set_comm_op_dist_attr_for_program(
            c_identity_op, dist_attr.process_mesh, out_grad_dist_attr, ctx
        )
701

702
        c_embedding_grad_op_desc = main_block.append_op(type='nop').desc
703 704 705
        c_embedding_grad_op_desc.set_type("c_embedding_grad")
        c_embedding_grad_op_desc.set_input('Ids', [Ids_var.name])
        c_embedding_grad_op_desc.set_input('W', [Weight_var.name])
706 707 708
        c_embedding_grad_op_desc.set_input(
            'Out@GRAD', [intermediate_var_0.name]
        )
709 710 711 712 713 714
        c_embedding_grad_op_desc.set_output('W@GRAD', [Weight_grad.name])
        c_embedding_grad_op_desc._set_attr('start_index', relative_idx)
        c_embedding_grad_op_desc._set_attr(OP_ROLE_KEY, OpRole.Backward)

        c_embedding_grad_op = main_block.ops[-1]
        assert c_embedding_grad_op.type == "c_embedding_grad"
715 716 717
        naive_copy_op_dist_attr_for_program(
            c_embedding_grad_op, backward_op, ctx
        )
718

719 720 721
        # data parallel gradient synchronization
        act_grad_names = [Ids_var.name]
        out_grad_names = [kwargs['W@GRAD'][0]]
722

723 724 725
        gradient_synchronization(
            ctx, backward_op, act_grad_names, out_grad_names, rank_id
        )
726

727

728 729 730 731 732 733 734 735 736
register_distributed_operator_impl(
    "lookup_table_v2", DistributedEmbeddingImpl("row_parallel")
)
register_distributed_operator_impl(
    "c_embedding", DistributedEmbeddingImpl("row_parallel")
)
register_distributed_operator_impl(
    "lookup_table", DistributedEmbeddingImpl("row_parallel")
)