dist_embedding.py 26.3 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 32
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,
)
from ..dist_attribute import OperatorDistributedAttribute
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 139 140 141

    new_op_dist_attr = OperatorDistributedAttribute()
    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
        )
C
caozhou 已提交
178 179
        processes = dist_op.dist_attr.process_mesh.processes
        # 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 234 235

        process_mesh = dist_attr.process_mesh
        processes = process_mesh.processes
        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]
        )
C
caozhou 已提交
253 254 255 256 257 258
        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 = [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 337 338
        if dim_changed:
            changed = True

        return changed

339 340 341 342 343 344
    @staticmethod
    def forward(ctx, *args, **kwargs):
        """
        kwargs: inputname_mapping & outputname_mapping
        """

345
        dist_op_context = ctx.dist_op_context
346 347 348 349
        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
350
        op_dist_attr = ctx.get_op_dist_attr_for_program(src_op)
351 352 353
        assert (
            op_dist_attr is not None
        ), "backward op [{}] don't have dist attribute !".format(str(src_op))
354

355
        # check validation of inputs / outputs
356 357 358 359
        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')

360 361 362
        assert (
            len(kwargs['Ids']) == 1
        ), "row_parallel_embedding input Ids take 1 variable but got {}".format(
363
            kwargs['Ids']
364 365 366 367
        )
        assert (
            len(kwargs['W']) == 1
        ), "row_parallel_embedding input W take 1 variable but got {}".format(
368
            kwargs['W']
369 370 371 372
        )
        assert (
            len(kwargs['Out']) == 1
        ), "row_parallel_embedding output Out take 1 variable but got {}".format(
373
            kwargs['Out']
374
        )
375

Z
zhaoyingli 已提交
376
        Ids_var = main_block._var_recursive(kwargs['Ids'][0])
377
        Weight_var = main_block._var_recursive(kwargs['W'][0])
Z
zhaoyingli 已提交
378
        Out_var = main_block._var_recursive(kwargs['Out'][0])
379

380 381 382 383
        # support lookup_table_v1
        if src_op.type == 'lookup_table':
            Ids_var = adopt_lookup_table_v1(ctx, main_block, src_op, Ids_var)

384 385
        # got dist attribute info
        embedding_row_dim_mapping = op_dist_attr.get_input_dims_mapping(
386 387 388 389 390 391 392
            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
        )
393 394
        process_mesh_shape = op_dist_attr.process_mesh.topology
        process_mesh_group = op_dist_attr.process_mesh.processes
395 396 397

        # FIXME (JZ-LIANG) Remove this hack to support any op mesh group for Pipeline Parallelism
        if rank_id not in process_mesh_group:
398 399 400
            rank_id = _get_corresponding_rank(
                ctx, op_dist_attr.process_mesh, rank_id
            )
401 402

        # A generalized method to caculate embedding offset using cartisian product
403 404 405 406 407 408
        relative_idx = _get_idx_in_axis(
            process_mesh_group,
            process_mesh_shape,
            embedding_row_dim_mapping,
            rank_id,
        )
409 410 411 412

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

413
        # TODO caculate ring id
414
        parallel_axis = embedding_row_dim_mapping
415 416 417
        group_ranks = _get_comm_group(
            process_mesh_group, process_mesh_shape, parallel_axis, rank_id
        )
418 419 420
        group = new_process_group(group_ranks)

        # append op
421 422 423
        check_variable_and_dtype(
            Ids_var, 'input', ['int32', 'int64'], 'c_embedding'
        )
424

Z
zhaoyingli 已提交
425 426 427 428 429
        # 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
430 431 432
        ref_shape = infer_shape(
            main_block, Out_var, out_tensor_dist_attr, out_var_dist_attr
        )
Z
zhaoyingli 已提交
433

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

        check_variable_and_dtype(
450 451
            Out_var,
            'tensor',
452
            ['float16', 'float32', 'float64', 'int32', 'int64'],
453 454
            'c_allreduce_sum',
        )
455 456 457

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

        # 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,
477 478 479
                OP_ROLE_KEY: src_op.attr('op_role'),
            },
        )
Z
zhaoyingli 已提交
480 481 482 483 484 485 486
        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
        embedding_op_dist_attr = OperatorDistributedAttribute()
        embedding_op_dist_attr.process_mesh = op_dist_attr.process_mesh
487
        embedding_op_dist_attr.impl_type = op_dist_attr.impl_type
Z
zhaoyingli 已提交
488 489 490 491
        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(
492 493 494 495 496
                op_dist_attr
            )
            embedding_op_dist_attr.set_input_dist_attr(
                input_varname, input_dist_attr
            )
Z
zhaoyingli 已提交
497 498 499
        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(
500 501 502 503 504
            op_dist_attr
        )
        embedding_op_dist_attr.set_output_dist_attr(
            output_varname, output_dist_attr
        )
Z
zhaoyingli 已提交
505
        ctx.set_op_dist_attr_for_program(c_embedding_op, embedding_op_dist_attr)
506

Z
zhaoyingli 已提交
507 508 509
        # allreduce
        allreduce_op_dist_attr = OperatorDistributedAttribute()
        allreduce_op_dist_attr.process_mesh = op_dist_attr.process_mesh
510
        allreduce_op_dist_attr.impl_type = op_dist_attr.impl_type
Z
zhaoyingli 已提交
511 512
        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 已提交
513
            input_var = main_block._var_recursive(input_varname)
Z
zhaoyingli 已提交
514 515
            tensor_dist_attr = ctx.get_tensor_dist_attr_for_program(input_var)
            assert tensor_dist_attr is not None
516 517 518
            allreduce_op_dist_attr.set_input_dist_attr(
                input_varname, tensor_dist_attr
            )
Z
zhaoyingli 已提交
519 520 521
        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(
522 523 524 525 526 527 528 529
                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
        )
530 531

        # param initialization sync
532
        if Weight_var.is_parameter and not op_dist_attr.is_recompute:
533 534
            if Weight_var.name in dist_op_context.already_init_sync_vars:
                return
J
JZ-LIANG 已提交
535 536 537 538 539 540 541 542 543 544 545
            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
            for axis, size in enumerate(process_mesh.topology):
                if size <= 1 or axis in dim_mapping:
                    pass
                else:
546 547 548 549 550 551
                    group_ranks = _get_comm_group(
                        process_mesh.processes,
                        process_mesh.topology,
                        axis,
                        rank_id,
                    )
J
JZ-LIANG 已提交
552 553
                    sync_group = new_process_group(group_ranks)

554 555 556 557 558 559 560 561 562 563 564
                    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,
                        },
                    )
565 566 567 568 569

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

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

581
        # FIXME (JZ-LIANG) Remove this hack to support any op mesh group for Pipeline Parallelism
582
        if rank_id not in dist_attr.process_mesh.processes:
583 584 585
            rank_id = _get_corresponding_rank(
                ctx, dist_attr.process_mesh, rank_id
            )
586 587 588 589 590 591

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

592 593 594
        assert (
            len(kwargs['Ids']) == 1
        ), "row_parallel_embedding input Ids take 1 variable but got {}".format(
595
            kwargs['Ids']
596 597 598 599
        )
        assert (
            len(kwargs['W']) == 1
        ), "row_parallel_embedding input Ids take 1 variable but got {}".format(
600
            kwargs['W']
601 602 603 604 605 606 607 608 609
        )
        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(
610
            kwargs['W@GRAD']
611
        )
612

Z
zhaoyingli 已提交
613 614 615 616
        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])
617 618

        embedding_row_dim_mapping = dist_attr.get_input_dims_mapping(
619 620 621 622 623 624 625
            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
        )
626 627 628 629
        process_mesh_shape = dist_attr.process_mesh.topology
        process_mesh_group = dist_attr.process_mesh.processes

        # A generalized method to caculate embedding offset using cartisian product
630 631 632 633 634 635
        relative_idx = _get_idx_in_axis(
            process_mesh_group,
            process_mesh_shape,
            embedding_row_dim_mapping,
            rank_id,
        )
636 637 638 639
        per_part_size = Weight_var.shape[0]
        relative_idx = relative_idx * per_part_size

        check_variable_and_dtype(
640 641 642 643 644
            Out_grad,
            'tensor',
            ['float16', 'float32', 'float64', 'int32', 'int64'],
            '_c_identity',
        )
645 646

        intermediate_var_0 = main_block.create_var(
647 648 649
            name=unique_name.generate_with_ignorable_key(
                ".".join(["c_embedding", '@tmp_0@GRAD'])
            ),
650 651 652 653
            dtype=Out_grad.dtype,
            shape=Out_grad.shape,
            type=core.VarDesc.VarType.LOD_TENSOR,
            persistable=False,
654 655
            stop_gradient=Out_grad.stop_gradient,
        )
656 657 658 659

        # 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
660 661 662
        ctx.set_tensor_dist_attr_for_program(
            intermediate_var_0, out_grad_dist_attr
        )
663

664 665 666 667 668 669
        group_ranks = _get_comm_group(
            process_mesh_group,
            process_mesh_shape,
            embedding_row_dim_mapping,
            rank_id,
        )
670 671 672 673 674 675 676 677 678 679 680
        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,
681 682 683 684 685 686 687 688 689 690 691
            },
        )
        check_variable_and_dtype(
            intermediate_var_0, 'x', ['float16', 'float32', 'float64'], 'linear'
        )
        check_dtype(
            intermediate_var_0.dtype,
            'dtype',
            ['float16', 'float32', 'float64'],
            'linear',
        )
692

693 694 695
        set_comm_op_dist_attr_for_program(
            c_identity_op, dist_attr.process_mesh, out_grad_dist_attr, ctx
        )
696

697
        c_embedding_grad_op_desc = main_block.append_op(type='nop').desc
698 699 700
        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])
701 702 703
        c_embedding_grad_op_desc.set_input(
            'Out@GRAD', [intermediate_var_0.name]
        )
704 705 706 707 708 709
        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"
710 711 712
        naive_copy_op_dist_attr_for_program(
            c_embedding_grad_op, backward_op, ctx
        )
713

714 715 716
        # data parallel gradient synchronization
        act_grad_names = [Ids_var.name]
        out_grad_names = [kwargs['W@GRAD'][0]]
717

718 719 720
        gradient_synchronization(
            ctx, backward_op, act_grad_names, out_grad_names, rank_id
        )
721

722

723 724 725 726 727 728 729 730 731
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")
)