test_moe_layer.py 6.9 KB
Newer Older
1
2
3
4
5
# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
#
# This source code is licensed under the BSD license found in the
# LICENSE file in the root directory of this source tree.

6
import functools
7
import tempfile
8

9
10
import pytest
import torch
11
import torch.distributed as dist
12
import torch.multiprocessing as mp
13
14

from fairscale.nn import MOELayer, Top2Gate
15
from fairscale.utils.testing import torch_version
16

17
18
19
pytestmark = pytest.mark.skipif(
    not (torch.cuda.is_available() and torch_version() >= (1, 8, 0)), reason="cuda and torch>=1.8.0 required"
)
20

21
devices = ["cuda"]
22

23

24
25
26
27
28
29
30
def pg_worker(rank, world_size, init_file, func, *args):
    init_url = "file://" + init_file
    dist.init_process_group(backend=dist.Backend.NCCL, rank=rank, world_size=world_size, init_method=init_url)
    torch.cuda.set_device(rank)
    dist.all_reduce(torch.zeros(1).cuda())
    func(*args)
    dist.destroy_process_group()
31

32

33
34
35
36
37
38
def pg_test(world_size=torch.cuda.device_count()):
    def decorator(func):
        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            tempfile_name = tempfile.mkstemp()[1]
            mp.spawn(pg_worker, args=(world_size, tempfile_name, func, *kwargs.values()), nprocs=world_size)
39

40
41
        globals()["test_" + func.__name__] = wrapper
        return func
42

43
    return decorator
44
45


46
@pg_test(world_size=1)
47
@pytest.mark.parametrize("device", devices)
48
def create(device):
49
50
51
52
    model_dim = 8
    num_experts = 4
    gate = Top2Gate(model_dim, num_experts)
    expert = torch.nn.Linear(model_dim, model_dim)
53
    moe = MOELayer(gate, expert).to(device)
54
55


56
@pg_test(world_size=1)
57
@pytest.mark.parametrize("device", devices)
58
def expert_params(device):
59
60
61
62
63
64
    model_dim = 8
    num_experts = 4
    gate = Top2Gate(model_dim, num_experts)
    expert = torch.nn.Linear(model_dim, model_dim)
    moe = MOELayer(gate, expert).to(device)
    for p in expert.parameters():
65
        assert p.expert is True, str(p.expert)
66
67


68
69
70
@pg_test()
@pytest.mark.parametrize("device", devices)
def forward(device):
71
    model_dim = 8
72
    num_experts = dist.get_world_size(dist.group.WORLD)
73
    input = torch.randn(4, 16, model_dim).to(device)
74
75
76
77
78
79
    gate = Top2Gate(model_dim, num_experts)
    expert = torch.nn.Linear(model_dim, model_dim, bias=False)
    # Use identity matrix
    expert.weight = torch.nn.Parameter(torch.eye(model_dim))
    moe = MOELayer(gate, expert).to(device)
    output = moe(input)
80
    assert output.shape == input.shape, f"{output.shape} != {input.shape}"
81
    # Re-assembled output should match input due to identity expert.
82
    torch.testing.assert_allclose(input, output)
83
84


85
86
87
@pg_test()
@pytest.mark.parametrize("device", devices)
def forward_multi(device):
88
89
90
91
    torch.set_printoptions(threshold=5000)
    num_local_experts = 4
    model_dim = 4
    num_experts = dist.get_world_size(dist.group.WORLD) * num_local_experts
92
    input = torch.randn(4 * num_local_experts, 16, model_dim).to(device)
93
94
95
96
97
98
99
100
101
    gate = Top2Gate(model_dim, num_experts)
    experts = []
    for i in range(num_local_experts):
        expert = torch.nn.Linear(model_dim, model_dim, bias=False)
        # Use identity matrix
        expert.weight = torch.nn.Parameter(torch.eye(model_dim))
        experts += [expert]
    moe = MOELayer(gate, torch.nn.ModuleList(experts)).to(device)
    output = moe(input)
102
    assert output.shape == input.shape, f"{output.shape} != {input.shape}"
103
    # 90% of the input should have gone to an expert
104
105
106
    assert (
        len(output.nonzero(as_tuple=False)) / output.numel() > 0.90
    ), f"{len(output.nonzero(as_tuple=False))} / {output.numel()}"
107
    # Except for zeros, re-assembled output should match input due to identity expert.
108
    torch.testing.assert_allclose(input, torch.where(output > 0, output, input))
109
110


111
112
113
114
115
116
117
118
# Test Gate which round-robin routes tokens to experts
class RoundRobinGate(torch.nn.Module):
    def __init__(self, model_dim, num_experts):
        super().__init__()
        self.model_dim = model_dim
        self.num_experts = num_experts

    def forward(self, input):
119
        s = input.shape[0]
120
        assert s % self.num_experts == 0, f"{s} % {self.num_experts} != 0"
121
        capacity = 2 * s // self.num_experts
122
        output = torch.zeros(s, self.num_experts, capacity, dtype=input.dtype, device=input.device)
123
        for i in range(s):
124
            output[i, i % self.num_experts, i // self.num_experts] = 1.0
125
126
127
        return 0.0, output, output.bool()


128
129
130
@pg_test()
@pytest.mark.parametrize("device", devices)
def forward_routing(device):
131
132
    model_dim = 8
    num_experts = dist.get_world_size()
133
    input = torch.randn(4, 16, model_dim).to(device)
134
135
136
137
138
139
140
    gate = RoundRobinGate(model_dim, num_experts)
    expert = torch.nn.Linear(model_dim, model_dim, bias=False)
    # Use scaling matrix (each rank has a different scale)
    scale = dist.get_rank() + 1
    expert.weight = torch.nn.Parameter(torch.eye(model_dim) * scale)
    moe = MOELayer(gate, expert).to(device)
    output = moe(input)
141
    assert output.shape == input.shape, f"{output.shape} != {input.shape}"
142
    # Verify that each token was sent to the correct expert by checking its scale.
143
    t = input.shape[1]
144
145
    for i in range(t):
        expert = i % num_experts
146
        torch.testing.assert_allclose(input[:, i] * (expert + 1), output[:, i])
147
148


149
150
151
@pg_test()
@pytest.mark.parametrize("device", devices)
def forward_routing_multi(device):
152
153
154
155
156
157
158
159
160
161
162
163
164
165
    model_dim = 8
    num_local_experts = 4
    num_experts = dist.get_world_size(dist.group.WORLD) * num_local_experts
    input = torch.randn(4 * num_local_experts, 16, model_dim).to(device)
    gate = RoundRobinGate(model_dim, num_experts)
    experts = []
    for i in range(num_local_experts):
        expert = torch.nn.Linear(model_dim, model_dim, bias=False)
        # Use scaling matrix (each rank has a different scale)
        scale = dist.get_rank() * num_local_experts + i + 1
        expert.weight = torch.nn.Parameter(torch.eye(model_dim) * scale)
        experts += [expert]
    moe = MOELayer(gate, torch.nn.ModuleList(experts)).to(device)
    output = moe(input)
166
    assert output.shape == input.shape, f"{output.shape} != {input.shape}"
167
168
169
170
    # Verify that each token was sent to the correct expert by checking its scale.
    t = input.shape[1]
    for i in range(t):
        expert = i % num_experts
171
        torch.testing.assert_allclose(input[:, i] * (expert + 1), output[:, i])
172
173


174
175
176
@pg_test()
@pytest.mark.parametrize("device", devices)
def backward(device):
177
178
179
    loss = torch.nn.MSELoss()
    model_dim = 8
    num_experts = dist.get_world_size(dist.group.WORLD)
180
    input = torch.randn(4, 16, model_dim).to(device)
181
182
183
184
185
186
    gate = Top2Gate(model_dim, num_experts)
    expert = torch.nn.Linear(model_dim, model_dim, bias=False)
    # Use identity matrix
    expert.weight = torch.nn.Parameter(torch.eye(model_dim))
    moe = MOELayer(gate, expert).to(device)
    output = moe(input)
187
    assert output.shape == input.shape, f"{output.shape} != {input.shape}"
188
189
    output = loss(output, input)
    output.backward()
190
    torch.testing.assert_allclose(expert.weight.grad, torch.zeros_like(expert.weight))