colossalai

1
import pytest
2
import torch
3
import torch.distributed as dist
4

5
import colossalai
6
from colossalai.accelerator import get_accelerator
7
from colossalai.moe import SparseMLP
8
from colossalai.moe.manager import MOE_MANAGER
9
from colossalai.testing import rerun_if_address_is_in_use, spawn
10

11
BATCH_SIZE = 4
12
NUM_EXPERTS = 4
13

14

15
def check_equal(tensor_a, tensor_b, atol=1e-06):
16
    assert torch.allclose(tensor_a, tensor_b, rtol=0, atol=atol) is True
17

18

19
def run_routing(rank, world_size, port, rs=2, hidden_size=128, data_type=torch.float32, topk=1):
20
    # Here we do not need TF32, since it brings absolute error on results
21
    torch.backends.cuda.matmul.allow_tf32 = False
22

23
    colossalai.launch(config=dict(), rank=rank, world_size=world_size, host="localhost", port=port, backend="nccl")
24
    local_rank = dist.get_rank()
25

26
    MOE_MANAGER.setup(parallel="EP")  # MOE environment initialization
27
    MOE_MANAGER.reset_loss()
28
    torch.manual_seed(rs + local_rank)  # set each process has different random seed
29

30
    # get randomized data
31
    tokens = torch.randn(
32
        BATCH_SIZE, hidden_size, dtype=data_type, device=get_accelerator().get_current_device(), requires_grad=True
33
    )
34

35
    layer = SparseMLP(
36
        hidden_size=hidden_size,
37
        intermediate_size=hidden_size * 2,
38
        num_experts=NUM_EXPERTS,
39
        router_top_k=topk,
40
        router_capacity_factor_train=1.0,
41
    )
42
    layer = layer.to(get_accelerator().get_current_device())
43
    if data_type == torch.float16:
44
        layer = layer.half()
45

46
    # use matrix multiplication instead of COL_MOE_KERNEL in MOE dispatch and combine
47
    layer.enable_kernel = False
48
    old_out = layer(tokens)
49
    ech = old_out.shape
50
    grad = torch.randn(ech, device=get_accelerator().get_current_device())
51
    old_out.backward(grad)  # get gradient
52

53
    # save all results
54
    o_tk_grad = tokens.grad.data.clone()
55
    o_gt_grad = layer.gate_weight.grad.data.clone()
56

57
    # reset all gradients
58
    tokens.grad.zero_()
59
    layer.gate_weight.grad.zero_()
60

61
    layer.enable_kernel = True
62
    new_out = layer(tokens)  # get outputs through colossal kernel
63

64
    if data_type == torch.float32:
65
        check_equal(old_out, new_out)
66
    else:
67
        check_equal(old_out, new_out, 1e-2)
68
    # forward function passed
69

70
    new_out.backward(grad)  # get new type gradient
71
    n_tk_grad = tokens.grad.data.clone()
72
    n_gt_grad = layer.gate_weight.grad.data.clone()
73

74
    if data_type == torch.float32:
75
        check_equal(o_tk_grad, n_tk_grad)
76
    else:
77
        check_equal(o_tk_grad, o_tk_grad, 1e-2)
78
    # tokens gradient is correct
79

80
    if data_type == torch.float32:
81
        check_equal(o_gt_grad, n_gt_grad, 5e-05)
82
    else:
83
        check_equal(o_gt_grad, n_gt_grad, 2e-01)
84
    # bias gradient is correct
85

86

87
@pytest.mark.dist
88
@pytest.mark.parametrize("rs", [131])
89
@pytest.mark.parametrize("hidden_size", [32, 144])
90
@pytest.mark.parametrize("data_type", [torch.float32, torch.float16])
91
@pytest.mark.parametrize("topk", [1, 2])
92
@rerun_if_address_is_in_use()
93
def test_moe_kernel(rs, hidden_size, data_type, topk):
94
    spawn(run_routing, 4, rs=rs, hidden_size=hidden_size, data_type=data_type, topk=topk)
95

96

97
if __name__ == "__main__":
98
    test_moe_kernel(2, 256, torch.float16, 2)
99