test_faster_schedule.py 2.72 KB
Newer Older
Rick Ho's avatar
Rick Ho committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
import pytest

import os
import sys
import json
import math

import torch
import torch.distributed as dist
import torch.nn.functional as F
from fmoe.functions import ensure_comm
from test_ddp import _ensure_initialized, _run_distributed
from test_numerical import _assert_numerical
from fmoe.fastermoe.schedule import _fmoe_general_global_forward as smart_fwd
from fmoe.layers import _fmoe_general_global_forward as naive_fwd


@pytest.mark.parametrize("n_process", [8])
@pytest.mark.parametrize("d_model", [1024])
@pytest.mark.parametrize("batch_size", [16])
Rick Ho's avatar
Rick Ho committed
21
@pytest.mark.parametrize("n_expert", [1, 4])
Rick Ho's avatar
Rick Ho committed
22
23
@pytest.mark.parametrize("group_sz", [1, 2, 4])
def test_faster_schedule(n_process, d_model, batch_size, n_expert, group_sz):
Rick Ho's avatar
Rick Ho committed
24
25
26
27
28
29
30
31
    _run_distributed('_test_faster_schedule',
            n_process,
            {
                'd_model': d_model,
                'batch_size': batch_size,
                'n_expert': n_expert
            },
            script=__file__,
Rick Ho's avatar
Rick Ho committed
32
33
34
            env=dict(
                FMOE_FASTER_GROUP_SIZE=str(group_sz)
            )
Rick Ho's avatar
Rick Ho committed
35
36
37
38
39
40
41
42
    )


def _test_faster_schedule(d_model, batch_size, n_expert):
    _ensure_initialized()
    rank = dist.get_rank()
    world_size = dist.get_world_size()

Rick Ho's avatar
Rick Ho committed
43
44
45
46
    x1 = torch.rand(batch_size, d_model).cuda()
    x1.requires_grad = True
    x2 = x1.data.clone()
    x2.requires_grad = True
Rick Ho's avatar
Rick Ho committed
47
    topk_idx = torch.randint(0, world_size * n_expert, (batch_size, 2)).cuda()
Rick Ho's avatar
Rick Ho committed
48
49
    m1s = [torch.nn.Linear(d_model, d_model).cuda() for _ in range(n_expert)]
    m2s = [torch.nn.Linear(d_model, d_model).cuda() for _ in range(n_expert)]
Rick Ho's avatar
Rick Ho committed
50
    with torch.no_grad():
Rick Ho's avatar
Rick Ho committed
51
52
53
54
55
56
        for m1, m2 in zip(m1s, m2s):
            m2.weight.copy_(m1.weight)
            m2.bias.copy_(m1.bias)

    def ef1(x, fec, eidx):
        return m1s[eidx](x)
Rick Ho's avatar
Rick Ho committed
57

Rick Ho's avatar
Rick Ho committed
58
    def ef2(x, fec):
Rick Ho's avatar
Rick Ho committed
59
60
61
62
63
64
65
        o = 0
        ys = []
        for m, i in zip(m2s, fec):
            if i > 0:
                ys.append(m(x[o:o + i]))
                o += i
        y = torch.cat(ys)
Rick Ho's avatar
Rick Ho committed
66
67
68
        return y

    ensure_comm(x1, None)
Rick Ho's avatar
Rick Ho committed
69
    y1 = smart_fwd(x1, topk_idx, ef1, n_expert, world_size, experts=m1s)
Rick Ho's avatar
Rick Ho committed
70
    y1.sum().backward()
Rick Ho's avatar
Rick Ho committed
71

Rick Ho's avatar
Rick Ho committed
72
    y2 = naive_fwd(x2, topk_idx, ef2, n_expert, world_size, experts=m2s)
Rick Ho's avatar
Rick Ho committed
73
    y2.sum().backward()
Rick Ho's avatar
Rick Ho committed
74
75
76
77
78
79
80
    _assert_numerical(['out', 'grad_in'],
            [y1, x1.grad],
            [y2, x2.grad], rank)
    for i in range(n_expert):
        _assert_numerical([f'grad_bias_{i}', f'grad_weight_{i}'],
                          [m1s[i].bias.grad, m1s[i].weight.grad],
                          [m2s[i].bias.grad, m2s[i].weight.grad], rank)
Rick Ho's avatar
Rick Ho committed
81
82
83
84
85
86
87


if __name__ == '__main__':
    if len(sys.argv) >= 3:
        args = json.loads(sys.argv[2])
        locals()[sys.argv[1]](**args)
    else:
Rick Ho's avatar
Rick Ho committed
88
        # test_faster_schedule(8, 16, 16, 1, 2)
Rick Ho's avatar
Rick Ho committed
89
        _test_faster_schedule(4, 2, 4)