dist_embedding.py 26.4 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

Z
zhaoyingli 已提交
15
from .common import infer_shape
16
from .common import DistributedOperatorImplContainer
17
from .common import DistributedOperatorImpl
18
from .common import register_distributed_operator_impl_container
19
from .common import gradient_synchronization
20 21 22 23 24
from .common import (
    naive_copy_op_dist_attr_for_program,
    register_distributed_operator_impl,
    set_comm_op_dist_attr_for_program,
)
25 26 27
from ..utils import is_dim_shard
from ..utils import is_dim_replicate
from ..utils import compute_compatible_and_update_dim_mapping
28
from ..dist_attribute import OperatorDistributedAttribute
29 30
from paddle.fluid import core, unique_name
from paddle.fluid.data_feeder import check_variable_and_dtype, check_dtype
31
from paddle.distributed.fleet.meta_optimizers.common import OP_ROLE_KEY, OpRole
32
from ..process_group import new_process_group
33 34 35 36 37 38
from ..utils import (
    _get_comm_group,
    _get_idx_in_axis,
    _get_corresponding_rank,
    set_var_dist_attr,
)
C
caozhou 已提交
39
from ..cost import build_comp_desc_from_dist_op, build_comm_desc_from_dist_op
40 41 42 43 44
from ..cost import (
    build_comm_costs_from_descs,
    build_comp_costs_from_descs,
    build_dp_costs,
)
45
from ..cost import EmbeddingOpCost, EmbeddingGradOpCost
46 47 48 49
from paddle.distributed.auto_parallel.cost.comm_op_cost import (
    AllreduceSumOpCost,
    IdentityOpCost,
)
50 51


52
class DistributedEmbedding(DistributedOperatorImplContainer):
53
    def __init__(self, op_type):
54
        super().__init__(op_type)
55 56


57
register_distributed_operator_impl_container(
58 59
    DistributedEmbedding("lookup_table_v2")
)
60
register_distributed_operator_impl_container(
61 62
    DistributedEmbedding("c_embedding")
)
63
register_distributed_operator_impl_container(
64 65
    DistributedEmbedding("lookup_table")
)
66 67 68 69


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

70 71 72 73 74
    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
    )
75 76 77 78 79 80 81
    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(
82 83 84
        name=unique_name.generate_with_ignorable_key(
            ".".join(["dist_reshape", 'tmp'])
        ),
85 86 87 88
        dtype=Ids_var.dtype,
        shape=target_shape,
        type=core.VarDesc.VarType.LOD_TENSOR,
        persistable=False,
89 90
        stop_gradient=True,
    )
91 92 93

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

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

    # 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(
119 120 121 122 123 124 125 126 127 128 129
        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,
    )
130
    op_dist_attr.del_input_dist_attr(Ids_var.name)
131 132 133
    op_dist_attr.set_input_dist_attr(
        intermediate_var_0.name, intermediate_var_0_dist_attr
    )
134 135 136 137 138

    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
139 140 141 142 143 144
    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
    )
145
    new_op_dist_attr.set_output_dims_mapping(
146 147
        xshape_var.name, [-1] + list(Ids_var_dist_attr.dims_mapping)
    )
148 149 150
    ctx.set_op_dist_attr_for_program(reshape_op, new_op_dist_attr)

    return intermediate_var_0
151 152 153 154 155


# RowParallel
class DistributedEmbeddingImpl(DistributedOperatorImpl):
    def __init__(self, name):
156
        super().__init__(name)
157
        self._forward_implemented = True
158
        self._backward_implemented = True
159

C
caozhou 已提交
160 161 162 163 164 165 166 167 168 169 170 171
    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
172 173 174
        desc_mapping = build_comp_desc_from_dist_op(
            dist_op=dist_op, dist_context=ctx
        )
C
caozhou 已提交
175 176
        processes = dist_op.dist_attr.process_mesh.processes
        # embedding need start_index
177 178 179
        cost_mapping = build_comp_costs_from_descs(
            EmbeddingOpCost, ctx, processes, desc_mapping, cluster
        )
C
caozhou 已提交
180 181 182

        serial_op = dist_op.serial_op
        parallel_axis = dist_op.dist_attr.get_input_dims_mapping(
183 184
            serial_op.input("W")[0]
        )[0]
C
caozhou 已提交
185 186 187 188 189 190 191 192
        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,
193 194
            parallel_axis=parallel_axis,
        )
C
caozhou 已提交
195 196

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

        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(
216 217
            backward_op.input("W")[0]
        )[0]
C
caozhou 已提交
218 219 220 221 222 223 224 225 226
        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,
227 228
            parallel_axis=parallel_axis,
        )
C
caozhou 已提交
229 230 231 232

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

        # calc comp op cost
238 239 240 241 242 243
        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 已提交
244 245 246 247
        res.append(cost_mapping)

        # need gradient allreduce
        var_dim_mapping = dist_attr.get_input_dims_mapping(
248 249
            backward_op.input("Ids")[0]
        )
C
caozhou 已提交
250 251 252 253 254 255
        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]]
256 257 258
            build_dp_costs(
                res, dist_op, ctx, var_names, attrs, parallel_axis, cluster
            )
C
caozhou 已提交
259 260 261

        return res

262 263 264
    def is_input_compatible(self, dist_op):
        op_desc = dist_op.serial_op.desc
        op_dist_attr = dist_op.dist_attr
265 266 267 268
        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)
269
        if is_dim_replicate(w_dims_mapping[-2]) or is_dim_shard(
270 271
            w_dims_mapping[-1]
        ):
272 273 274 275 276 277 278
            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

279 280 281
    def is_output_compatible(self, dist_op):
        op_desc = dist_op.serial_op.desc
        op_dist_attr = dist_op.dist_attr
282 283 284 285 286 287 288 289
        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
沉潜的鱼儿 已提交
290
    def is_auto_compatible(self, dist_op):
291 292 293
        if (not self.is_input_compatible(dist_op)) or (
            not self.is_output_compatible(dist_op)
        ):
294 295
            return False

沉潜的鱼儿's avatar
沉潜的鱼儿 已提交
296 297 298 299 300 301 302 303
        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)
304

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

        return True

310
    def update_dims_mapping(self, dist_op):
311
        changed = False
312 313
        op_desc = dist_op.serial_op.desc
        op_dist_attr = dist_op.dist_attr
314 315 316 317 318 319 320 321 322
        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(
323 324
                [ids_dims_mapping, out_dims_mapping], [i, i]
            )
325 326 327 328
            if dim_changed:
                changed = True

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

        return changed

336 337 338 339 340 341
    @staticmethod
    def forward(ctx, *args, **kwargs):
        """
        kwargs: inputname_mapping & outputname_mapping
        """

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

352
        # check validation of inputs / outputs
353 354 355 356
        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')

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

Z
zhaoyingli 已提交
373
        Ids_var = main_block._var_recursive(kwargs['Ids'][0])
374
        Weight_var = main_block._var_recursive(kwargs['W'][0])
Z
zhaoyingli 已提交
375
        Out_var = main_block._var_recursive(kwargs['Out'][0])
376

377 378 379 380
        # support lookup_table_v1
        if src_op.type == 'lookup_table':
            Ids_var = adopt_lookup_table_v1(ctx, main_block, src_op, Ids_var)

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

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

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

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

410
        # TODO caculate ring id
411
        parallel_axis = embedding_row_dim_mapping
412 413 414
        group_ranks = _get_comm_group(
            process_mesh_group, process_mesh_shape, parallel_axis, rank_id
        )
415 416 417
        group = new_process_group(group_ranks)

        # append op
418 419 420
        check_variable_and_dtype(
            Ids_var, 'input', ['int32', 'int64'], 'c_embedding'
        )
421

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

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

        check_variable_and_dtype(
447 448
            Out_var,
            'tensor',
449
            ['float16', 'float32', 'float64', 'int32', 'int64'],
450 451
            'c_allreduce_sum',
        )
452 453 454

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

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

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

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

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

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

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

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

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

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

Z
zhaoyingli 已提交
610 611 612 613
        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])
614 615

        embedding_row_dim_mapping = dist_attr.get_input_dims_mapping(
616 617 618 619 620 621 622
            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
        )
623 624 625 626
        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
627 628 629 630 631 632
        relative_idx = _get_idx_in_axis(
            process_mesh_group,
            process_mesh_shape,
            embedding_row_dim_mapping,
            rank_id,
        )
633 634 635 636
        per_part_size = Weight_var.shape[0]
        relative_idx = relative_idx * per_part_size

        check_variable_and_dtype(
637 638 639 640 641
            Out_grad,
            'tensor',
            ['float16', 'float32', 'float64', 'int32', 'int64'],
            '_c_identity',
        )
642 643

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

        # 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
657 658 659
        ctx.set_tensor_dist_attr_for_program(
            intermediate_var_0, out_grad_dist_attr
        )
660

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

690 691 692
        set_comm_op_dist_attr_for_program(
            c_identity_op, dist_attr.process_mesh, out_grad_dist_attr, ctx
        )
693

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

711 712 713
        # data parallel gradient synchronization
        act_grad_names = [Ids_var.name]
        out_grad_names = [kwargs['W@GRAD'][0]]
714

715 716 717
        gradient_synchronization(
            ctx, backward_op, act_grad_names, out_grad_names, rank_id
        )
718

719

720 721 722 723 724 725 726 727 728
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")
)