deepspeedai--deepspeed/tests/unit/sequence_parallelism/test_ulysses.py

# Copyright (c) Microsoft Corporation.
# SPDX-License-Identifier: Apache-2.0

# DeepSpeed Team

import pytest
import torch
import torch.nn.functional as F
import deepspeed.comm as dist
from deepspeed import initialize
from transformers import AutoModel
from unit.common import DistributedTest
from deepspeed.sequence.layer import _SeqAllToAll
from deepspeed.sequence.fpdt_layer import _FPDTGPUOffloadingAttentionImpl_, FPDT_InputConstruct
from unit.util import skip_on_arch
from unit.simple_model import *
from deepspeed.utils import groups
from deepspeed.module_inject.tp_shard import get_shard_size_list
#Use mesh device to create data and sequence parallel group


class TestUlyssesUtils(DistributedTest):
    world_size = 4

    def test_mesh_device_creation(self) -> None:
        skip_on_arch(min_arch=8)
        model = AutoModel.from_pretrained('bert-base-uncased')
        sp_size = 2
        dp_size = 2
        ds_engine, _, _, _ = initialize(
            model=model,
            config_params={
                "train_batch_size": 8,
                "data_parallel_size": dp_size,
                "sequence_parallel_size": sp_size
            },
        )
        assert ds_engine.seq_parallel_group is not None
        assert ds_engine.data_parallel_group is not None
        assert dist.get_world_size(group=ds_engine.seq_parallel_group) == sp_size
        assert dist.get_world_size(group=ds_engine.data_parallel_group) == dp_size
        assert dist.get_world_size() == sp_size * dp_size


#Sweep b,s,h,d to test all2all consistency
@pytest.mark.parametrize("d0", [2, 4])  #batch or sequence dimension
@pytest.mark.parametrize("d1", [4, 8])  #batch or sequence dimension
@pytest.mark.parametrize("num_heads", [4, 8])
@pytest.mark.parametrize("head_dim", [16, 32])
class TestUlyssesAll2All(DistributedTest):
    world_size = 4

    def test_alltoall_output_consistency(self, d0: int, d1: int, head_dim: int, num_heads: int) -> None:
        skip_on_arch(min_arch=8)
        model = AutoModel.from_pretrained('bert-base-uncased')
        ds_engine, _, _, _ = initialize(model=model, config_params={"train_batch_size": 8}, mesh_param=(2, 2))
        #4D tensor : b,s,h,d or s,b,h,d
        input_tensor = torch.randn(d0, d1, num_heads, head_dim, device=ds_engine.device)
        scatter_idx = 2
        batch_dim_idx = 0
        outputs = []
        seq_dims = [0]  #seq first API
        #TODO: Add support for batch first (that seq_dims=[0,1]) after PR for bs>1 issue with batch first is fixed
        ## See discussion in : https://github.com/deepspeedai/DeepSpeed/issues/5808
        for seq_dim in seq_dims:
            gather_idx = seq_dim
            #first all2all: sequence parallel to head parallel
            s2h_tensor = _SeqAllToAll.apply(ds_engine.seq_parallel_group, input_tensor, scatter_idx, gather_idx,
                                            batch_dim_idx)

            #No op
            # second all2all: head parallel to sequence parallel
            h2s_tensor = _SeqAllToAll.apply(ds_engine.seq_parallel_group, s2h_tensor, gather_idx, scatter_idx,
                                            batch_dim_idx)
            print(
                f'[{dist.get_rank()}] s={seq_dim} input: {input_tensor.shape} s2h: {s2h_tensor.shape} h2s_tensor: {h2s_tensor.shape}'
            )
            outputs.append(h2s_tensor)

        # Check outputs are the same as input
        for i in range(1, len(outputs)):
            assert torch.allclose(input_tensor, outputs[i]), f"Outputs differ for sequence dim {seq_dims[i]}"


@pytest.mark.parametrize("d0", [2, 4])  #batch or sequence dimension
@pytest.mark.parametrize("d1", [4, 8])  #batch or sequence dimension
@pytest.mark.parametrize("num_heads", [3, 7])
@pytest.mark.parametrize("head_dim", [16])
class TestUlyssesAll2All_odd(DistributedTest):
    world_size = 4

    def test_alltoall_output_consistency(self, d0: int, d1: int, head_dim: int, num_heads: int) -> None:

        data_parallel_size = 2
        seq_parallel_size = self.world_size // data_parallel_size
        skip_on_arch(min_arch=8)

        def seq_batch_heads_hash(d0, d1, h, offset_d0=0, offset_d1=0, offset_h=0):
            d0 += offset_d0
            d1 += offset_d1
            h += offset_h
            return d0 * 10 + h + d1 * 0.1

        hidden_dim = 10
        model = SimpleModel(hidden_dim)
        ds_engine, _, _, _ = initialize(model=model,
                                        config_params={"train_batch_size": 8},
                                        mesh_param=(data_parallel_size, seq_parallel_size))

        scatter_idx = 2
        outputs = []
        inputs = []
        batch_dims = [0, 1]
        seq_dims = [1, 0]

        for idx, seq_dim in enumerate(seq_dims):
            gather_idx = seq_dim
            batch_dim_idx = batch_dims[idx]

            #4D tensor : b,s,h,d or s,b,h,d
            #create a hash tensor from pos_id, head_id, and batch_id
            d0_indices = torch.arange(d0).reshape(-1, 1, 1, 1)
            d1_indices = torch.arange(d1).reshape(1, -1, 1, 1)
            h_indices = torch.arange(num_heads).reshape(1, 1, -1, 1)
            input_tensor = torch.randn(d0, d1, num_heads, head_dim, device=ds_engine.device)
            if batch_dim_idx == 1:  #seq_len_dim : 0(d0)
                input_tensor[:] = seq_batch_heads_hash(d0_indices, d1_indices, h_indices,
                                                       d0 * groups._get_sequence_parallel_rank(), 0)
            elif batch_dim_idx == 0:  #seq_len_dim : 1(d1)
                input_tensor[:] = seq_batch_heads_hash(d0_indices, d1_indices, h_indices, 0,
                                                       d1 * groups._get_sequence_parallel_rank())
            inputs.append(input_tensor)

            ### first all2all: sequence parallel to head parallel
            s2h_tensor = _SeqAllToAll.apply(ds_engine.seq_parallel_group, input_tensor, scatter_idx, gather_idx,
                                            batch_dim_idx)

            # s2h_tensor check for the first all2all: compare with the expected ground truth
            d0_indices = torch.arange(s2h_tensor.shape[0]).reshape(-1, 1, 1, 1)
            d1_indices = torch.arange(s2h_tensor.shape[1]).reshape(1, -1, 1, 1)
            h_indices = torch.arange(s2h_tensor.shape[2]).reshape(1, 1, -1, 1)
            shard_list = get_shard_size_list(num_heads, groups._get_sequence_parallel_world_size())
            head_offset = sum(shard_list[:groups._get_sequence_parallel_rank()])
            s2h_truth = torch.zeros_like(s2h_tensor)
            s2h_truth[:] = seq_batch_heads_hash(d0_indices, d1_indices, h_indices, 0, 0, head_offset)

            assert torch.allclose(s2h_truth,
                                  s2h_tensor), f"s2h_tensor differs from the expected for sequence dim: {seq_dim}"
            #No op
            ### second all2all: head parallel to sequence parallel
            h2s_tensor = _SeqAllToAll.apply(ds_engine.seq_parallel_group, s2h_tensor, gather_idx, scatter_idx,
                                            batch_dim_idx)
            print(
                f'[{dist.get_rank()}] s={seq_dim} input: {input_tensor.shape} s2h: {s2h_tensor.shape} h2s_tensor: {h2s_tensor.shape}'
            )
            outputs.append(h2s_tensor)

        # Check outputs for the second all2all
        for i in range(0, len(outputs)):
            assert torch.allclose(inputs[i],
                                  outputs[i]), f"[{dist.get_rank()}]Outputs differ for sequence dim {seq_dims[i]}"


@pytest.mark.parametrize("d0", [4, 1])  #batch dimension
@pytest.mark.parametrize("d1", [2048, 8192])  #sequence dimension
@pytest.mark.parametrize("chunk_size", [128, 256])  #size of chunk
@pytest.mark.parametrize("num_heads", [8, 4])
@pytest.mark.parametrize("head_dim", [32])
class TestFPDTAttention(DistributedTest):

    def test_FPDT_attention_offloading_output_consistency(self, d0: int, d1: int, chunk_size: int, head_dim: int,
                                                          num_heads: int) -> None:
        skip_on_arch(min_arch=8)
        world_size = 2

        try:
            from flash_attn.flash_attn_interface import _flash_attn_forward, _flash_attn_backward
        except ImportError:
            _flash_attn_forward = None
            _flash_attn_backward = None

        if _flash_attn_forward is None or _flash_attn_backward is None:
            pytest.skip("Flash Attention is not available.")

        model = AutoModel.from_pretrained('bert-base-uncased')
        ds_engine, _, _, _ = initialize(
            model=model,
            config_params={
                "train_batch_size": 8,
                "data_parallel_size": 1,
                "sequence_parallel_size": world_size
            },
        )
        #3D tensor : l, b, d
        dim = head_dim * num_heads

        seed = 42
        torch.manual_seed(seed)
        get_accelerator().manual_seed_all(seed)

        input_tensor = torch.randn(d1, d0, dim, device=ds_engine.device, dtype=torch.half)  # l, b, d
        spg = ds_engine.seq_parallel_group

        dist.broadcast(input_tensor, src=0, group=spg)

        class args:

            def __init__(self):
                self.ds_sequence_parallel_fpdt_chunk_size = chunk_size

        fpdt_input_tensor = FPDT_InputConstruct(input_tensor.permute(1, 0, 2), None, None, None, None, args(),
                                                world_size, dist.get_rank()).generate()[0].permute(1, 0, 2)

        if dist.get_rank() == 0:
            qkv_linear_weight = torch.nn.Parameter(
                torch.empty(dim + 2 * dim, dim, device=dist.get_rank(), dtype=torch.half))
            torch.nn.init.normal_(qkv_linear_weight, mean=0.0, std=0.02)

            qkv_linear_bias = torch.nn.Parameter(torch.empty(dim + 2 * dim, device=dist.get_rank(), dtype=torch.half))
            torch.nn.init.normal_(qkv_linear_bias, mean=0.0, std=0.02)
        else:
            qkv_linear_weight = torch.nn.Parameter(
                torch.empty(dim + 2 * dim, dim, device=dist.get_rank(), dtype=torch.half))
            qkv_linear_bias = torch.nn.Parameter(torch.empty(dim + 2 * dim, device=dist.get_rank(), dtype=torch.half))

        dist.broadcast(qkv_linear_weight, src=0, group=spg)
        dist.broadcast(qkv_linear_bias, src=0, group=spg)

        num_chunks_attn = fpdt_input_tensor.shape[0] * dist.get_world_size(spg) // chunk_size
        fpdt_output = _FPDTGPUOffloadingAttentionImpl_.apply(fpdt_input_tensor, None, None, None, spg, 2, 0, dim, dim,
                                                             head_dim, dim, qkv_linear_weight, qkv_linear_bias, 0,
                                                             num_chunks_attn, True)

        # baseline
        qkv = torch.matmul(input_tensor, qkv_linear_weight.t()) + qkv_linear_bias
        q = qkv[:, :, :dim].contiguous().reshape(qkv.shape[0], qkv.shape[1], -1, head_dim).permute(1, 2, 0,
                                                                                                   3).contiguous()
        k = qkv[:, :, dim:dim * 2].contiguous().reshape(qkv.shape[0], qkv.shape[1], -1,
                                                        head_dim).permute(1, 2, 0, 3).contiguous()
        v = qkv[:, :, dim * 2:dim * 3].contiguous().reshape(qkv.shape[0], qkv.shape[1], -1,
                                                            head_dim).permute(1, 2, 0,
                                                                              3).contiguous()  # b, nhead, l, d

        scores = torch.matmul(q, k.transpose(-2, -1)) / torch.sqrt(torch.tensor(dim, dtype=torch.half))

        causal_mask = torch.triu(torch.ones(d1, d1, device=ds_engine.device), diagonal=1).bool()
        causal_mask = causal_mask.unsqueeze(0).unsqueeze(0)
        scores = scores.masked_fill(causal_mask, float('-inf'))
        attn_weights = F.softmax(scores, dim=-1)
        output = torch.matmul(attn_weights, v).permute(0, 2, 1, 3)

        baseline_output_shuffled = FPDT_InputConstruct(output, None, None, None, None, args(), world_size,
                                                       dist.get_rank()).generate()[0]  # b, l, n, d

        assert torch.allclose(
            fpdt_output, baseline_output_shuffled, rtol=0.01, atol=0.1
        ), f"rank {dist.get_rank()}, sp size: {dist.get_world_size(spg)}, input_tensor: {input_tensor.shape}, fpdt_input_tensor: {fpdt_input_tensor.shape}, fpdt_output: {fpdt_output.shape},            baseline_output_shuffled: {baseline_output_shuffled.shape},{torch.max(torch.abs(fpdt_output - baseline_output_shuffled))}"


@pytest.mark.parametrize("sp_size", [2])
class TestUlyssesLossBackward(DistributedTest):
    world_size = 4

    def test_sp_loss_backward_stability(self, sp_size: int) -> None:
        """
        Regression test for Issue #7672.
        Verifies that using all_reduce for loss aggregation is stable
        when sequence_parallel_size < world_size, preventing IndexError.
        """
        skip_on_arch(min_arch=8)

        # Setup
        dp_size = self.world_size // sp_size
        model = SimpleModel(4)
        ds_engine, _, _, _ = initialize(
            model=model,
            config_params={
                "train_batch_size": 8,
                "data_parallel_size": dp_size,
                "sequence_parallel_size": sp_size
            },
        )

        sp_group = ds_engine.seq_parallel_group

        # Simulate Loss on each rank
        rank = dist.get_rank()
        local_loss = torch.tensor(float(rank + 1), device=ds_engine.device, requires_grad=True)
        local_weight = torch.tensor(1.0, device=ds_engine.device)

        # Numerator: Weighted Loss summation
        weighted_loss = local_loss * local_weight
        dist.all_reduce(weighted_loss, op=dist.ReduceOp.SUM, group=sp_group)

        # B. Denominator: Sum of total weights
        total_weight = local_weight.clone()
        dist.all_reduce(total_weight, op=dist.ReduceOp.SUM, group=sp_group)

        # C. Calculate the final loss
        dist_loss = weighted_loss / total_weight

        # Backward Pass verification
        try:
            dist_loss.backward()
        except IndexError as e:
            pytest.fail(f"Backward crashed with IndexError: {e}")

        # Verify Gradients
        # Loss = (L1*1 + L2*1) / 2 = 0.5*L1 + 0.5*L2
        expected_grad = 0.5
        assert torch.allclose(local_loss.grad, torch.tensor(expected_grad, device=ds_engine.device)), \
            f"Gradient mismatch! Expected {expected_grad}, got {local_loss.grad}"