Method forward

ernie/longcontext_ops.py:145–303

Balance the input tensor between groups based on token_type_ids

(
        ctx,
        tensor: paddle.Tensor,
        token_type_ids: paddle.Tensor,
        group=None,
        is_tensor_sharded=True,
        axis=0,
        is_token_type_ids_sharded=False,
        unique_tokens_type=None,
    )

    @staticmethod
    def forward(
        ctx,
        tensor: paddle.Tensor,
        token_type_ids: paddle.Tensor,
        group=None,
        is_tensor_sharded=True,
        axis=0,
        is_token_type_ids_sharded=False,
        unique_tokens_type=None,
    ):
        """Balance the input tensor between groups based on token_type_ids"""
        ctx.is_tensor_sharded = is_tensor_sharded
        ctx.axis = axis
        ctx.tensor_shape = tensor.shape
        ctx.tensor_dtype = tensor.dtype
        ctx.token_type_ids_shape = token_type_ids.shape
        ctx.token_type_ids_dtype = token_type_ids.dtype
        ctx.group = fleet.get_hybrid_communicate_group().get_model_parallel_group() if group is None else group
        ctx.rank = ctx.group.rank
        ctx.world_size = ctx.group.nranks

        if len(ctx.tensor_shape) == 1:
            if not ctx.is_tensor_sharded:
                tensor = tensor.split(num_or_sections=ctx.world_size, axis=0)[ctx.rank]
            tensor = tensor.reshape([-1, 1])
        else:
            if (
                len(ctx.tensor_shape) == 2
                and not ctx.is_tensor_sharded
                and (axis == -1 or axis == len(ctx.tensor_shape) - 1)
            ):
                raise ValueError(
                    "Do not support len(ctx.tensor_shape) == 2 and not ctx.is_tensor_sharded"
                    + " and (axis == -1 or axis == len(ctx.tensor_shape) -1)"
                )
            assert len(ctx.tensor_shape) <= 3, f"len(tensor.shape) must <= 3, but got {len(tensor.shape)}"
            if len(ctx.tensor_shape) == 3:
                assert ctx.tensor_shape[0] == 1, "only support tensor.shape[0] == 1"

            if not ctx.is_tensor_sharded:
                tensor = tensor.split(num_or_sections=ctx.world_size, axis=ctx.axis)[ctx.rank]
            tensor = tensor.reshape([-1, tensor.shape[-1]])

        ctx.tensor_grad_shape = tensor.shape

        if is_token_type_ids_sharded:
            assert (
                unique_tokens_type is not None and len(unique_tokens_type) > 0
            ), "require len(unique_tokens_type) > 0 when is_token_type_ids_sharded=True"
            ctx.unique_tokens_type = unique_tokens_type
            token_type_ids_per_rank = token_type_ids.flatten()
        else:
            ctx.unique_tokens_type = token_type_ids.unique().tolist()
            token_type_ids_per_rank = token_type_ids.flatten().split(ctx.world_size)
        assert tensor.shape[0] == token_type_ids_per_rank[ctx.rank].shape[0], (
            f"tensor.shape[0]:{tensor.shape[0]} != "
            + f"token_type_ids_per_rank[ctx.rank].shape[0]:{token_type_ids_per_rank[ctx.rank].shape[0]}"
        )