layers.py 9.18 KB
Newer Older
Sengxian's avatar
Sengxian committed
1
r"""
Rick Ho's avatar
Rick Ho committed
2
FMoE core layer
Sengxian's avatar
Sengxian committed
3
"""
4
import tree
Rick Ho's avatar
Rick Ho committed
5
import torch
Rick Ho's avatar
Rick Ho committed
6
7
import torch.nn as nn

8
from .functions import prepare_forward, ensure_comm
Rick Ho's avatar
Rick Ho committed
9
from .functions import MOEScatter, MOEGather
Sengxian's avatar
Sengxian committed
10
from .functions import AllGather, Slice
Rick Ho's avatar
Rick Ho committed
11
from .gates import NaiveGate
Rick Ho's avatar
Rick Ho committed
12

Rick Ho's avatar
Rick Ho committed
13

Rick Ho's avatar
Rick Ho committed
14
def mark_module_parallel_comm(module, comm):
Sengxian's avatar
Sengxian committed
15
    r"""
Rick Ho's avatar
Rick Ho committed
16
17
    Mark all parameters in `module` as doing data parallel in `comm`, where
    `comm` may be one of `'world', 'dp', 'none'`.
Sengxian's avatar
Sengxian committed
18
    """
Rick Ho's avatar
Rick Ho committed
19
    for p in module.parameters():
Sengxian's avatar
Sengxian committed
20
        setattr(p, "dp_comm", comm)
Rick Ho's avatar
Rick Ho committed
21
22


23
def _fmoe_general_global_forward(inp, gate, expert_fn, num_expert, world_size):
Sengxian's avatar
Sengxian committed
24
    r"""
Rick Ho's avatar
Rick Ho committed
25
26
27
28
29
    A private function that performs the following steps to complete the MoE
    computation.
    * Count the number of tokens from each worker to each expert.
    * Send the features to their target position so that input features to each
    expert are contiguous in memory.
Rick Ho's avatar
Rick Ho committed
30
    * Perform the forward computation of the experts using `expert_fn`
Rick Ho's avatar
Rick Ho committed
31
32
33
    * Gather the output features of experts back, and reorder them as sentences.
    Intermediate results like expert counts are hidden from users by this
    function.
Sengxian's avatar
Sengxian committed
34
    """
35
    (
Sengxian's avatar
Sengxian committed
36
37
38
39
40
        pos,
        local_expert_count,
        global_expert_count,
        fwd_expert_count,
        fwd_batch_size,
41
    ) = prepare_forward(gate, num_expert, world_size)
Rick Ho's avatar
Rick Ho committed
42
43
44
    topk = 1
    if len(gate.shape) == 2:
        topk = gate.shape[1]
45

46
47
48
    def scatter_func(tensor):
        return MOEScatter.apply(
            tensor,
49
50
51
52
53
54
55
56
57
            pos // topk,
            local_expert_count,
            global_expert_count,
            fwd_batch_size,
            world_size,
        )

    x = tree.map_structure(scatter_func, inp)

Rick Ho's avatar
Rick Ho committed
58
    x = expert_fn(x, fwd_expert_count)
59

60
    out_batch_size = tree.flatten(inp)[0].shape[0]
61
62
63
    if len(gate.shape) == 2:
        out_batch_size *= gate.shape[1]

64
65
66
    def gather_func(tensor):
        return MOEGather.apply(
            tensor,
67
68
69
70
71
72
73
            pos,
            local_expert_count,
            global_expert_count,
            out_batch_size,
            world_size,
        )

74
    outp = tree.map_structure(gather_func, x)
75
    return outp
Rick Ho's avatar
Rick Ho committed
76
77


Rick Ho's avatar
Rick Ho committed
78
class FMoE(nn.Module):
Sengxian's avatar
Sengxian committed
79
    r"""
Rick Ho's avatar
Rick Ho committed
80
81
    A general moe implementation that supports an arbitrary module as the
    expert.
Rick Ho's avatar
Rick Ho committed
82
83
84
    * `num_expert` stands for the number of experts on **each** worker.
    * `world_size` stands for the total number of workers that contains
    different experts.
Rick Ho's avatar
Rick Ho committed
85
86
87
88
89
    * `slice_group` can be a torch's communication group, indicating that
    specific model parallel is applied across the group, and workers in the
    group hold the same copy of input feature, and requires the same copy of
    the output. For each worker, FMoE only computes the output of a certain
    slice of the input batch, and will all-gather the outputs after
Rick Ho's avatar
Rick Ho committed
90
    computation.
Rick Ho's avatar
Rick Ho committed
91
    * `top_k` stands for the number of experts each token is going to.
Rick Ho's avatar
Rick Ho committed
92
93
94
    * `gate` is a gate class which can found in `fmoe.gates`.
    * `expert` can be specified as a module class, it is used to generate
    `num_expert` expert modules.
Sengxian's avatar
Sengxian committed
95
96
97
98
99
100
101
    """

    def __init__(
        self,
        num_expert=32,
        d_model=1024,
        world_size=1,
102
        mp_group=None,  # being deprecated
Rick Ho's avatar
Rick Ho committed
103
        slice_group=None,
Rick Ho's avatar
Rick Ho committed
104
        moe_group=None,
Sengxian's avatar
Sengxian committed
105
106
107
        top_k=2,
        gate=NaiveGate,
        expert=None,
108
        gate_hook=None,
Colin's avatar
Colin committed
109
110
        mask=None,
        mask_dict=None,
Sengxian's avatar
Sengxian committed
111
    ):
Rick Ho's avatar
Rick Ho committed
112
        super().__init__()
Rick Ho's avatar
Rick Ho committed
113
114
115
        self.num_expert = num_expert
        self.d_model = d_model
        self.world_size = world_size
Rick Ho's avatar
Rick Ho committed
116
117
118

        self.slice_group = slice_group
        if mp_group is not None:
119
            print("[Warning] mp_group is being deprecated")
Rick Ho's avatar
Rick Ho committed
120
121
122
123
            self.slice_group = mp_group
        if self.slice_group is None:
            self.slice_size = 1
            self.slice_rank = 0
Rick Ho's avatar
Rick Ho committed
124
        else:
Rick Ho's avatar
Rick Ho committed
125
126
            self.slice_size = self.slice_group.size()
            self.slice_rank = self.slice_group.rank()
Rick Ho's avatar
Rick Ho committed
127

Rick Ho's avatar
Rick Ho committed
128
        self.top_k = top_k
Colin's avatar
Colin committed
129
130
131
132
133
        if type(expert) is list:
            self.experts = nn.ModuleList([e(d_model) for e in expert])
            self.experts_fused = False
            self.num_expert = num_expert = len(expert)
        elif expert is not None:
134
            self.experts = nn.ModuleList([expert(d_model) for _ in range(num_expert)])
135
136
137
            self.experts_fused = False
        else:
            self.experts_fused = True
Rick Ho's avatar
Rick Ho committed
138

Colin's avatar
Colin committed
139
        self.gate = gate(d_model, num_expert, world_size, top_k)
140
        self.gate_hook = gate_hook
Colin's avatar
Colin committed
141
142
        self.mask = mask
        self.mask_dict = mask_dict
Rick Ho's avatar
Rick Ho committed
143
        self.moe_group = moe_group
Rick Ho's avatar
Rick Ho committed
144
145

    def expert_fn(self, inp, fwd_expert_count):
Sengxian's avatar
Sengxian committed
146
        r"""
Rick Ho's avatar
Rick Ho committed
147
148
        The default expert function which either calls the experts as a whole
        or as separate experts.
Sengxian's avatar
Sengxian committed
149
        """
150
        if self.experts_fused:
Rick Ho's avatar
Rick Ho committed
151
152
153
154
155
            return self.experts(inp, fwd_expert_count)
        outputs = []
        base_idx = 0
        for i in range(self.num_expert):
            batch_size = fwd_expert_count[i].item()
Sengxian's avatar
Sengxian committed
156
            inp_slice = inp[base_idx : base_idx + batch_size]
Rick Ho's avatar
Rick Ho committed
157
158
159
            outputs.append(self.experts[i](inp_slice))
            base_idx += batch_size
        return torch.cat(outputs, dim=0)
Rick Ho's avatar
Rick Ho committed
160

Sengxian's avatar
Sengxian committed
161
162
    def mark_parallel_comm(self, expert_dp_comm="none"):
        r"""
Rick Ho's avatar
Rick Ho committed
163
164
165
        Automatically mark the data parallel comms of the parameters within the
        module. This can be typically called at the end of the __init__ function
        in child classes.
Sengxian's avatar
Sengxian committed
166
        """
Rick Ho's avatar
Rick Ho committed
167
        if self.experts is not None:
168
            comm = expert_dp_comm
Rick Ho's avatar
Rick Ho committed
169
170
171
172
173
            if isinstance(self.experts, list):
                for e in self.experts:
                    mark_module_parallel_comm(e, comm)
            else:
                mark_module_parallel_comm(self.experts, comm)
Rick Ho's avatar
Rick Ho committed
174
        mark_module_parallel_comm(self.gate, "gate")
Rick Ho's avatar
Rick Ho committed
175

Jiezhong Qiu's avatar
Jiezhong Qiu committed
176
    def forward(self, moe_inp):
Sengxian's avatar
Sengxian committed
177
        r"""
Rick Ho's avatar
Rick Ho committed
178
179
180
        The FMoE module first computes gate output, and then conduct MoE forward
        according to the gate.  The score of the selected gate given by the
        expert is multiplied to the experts' output tensors as a weight.
Sengxian's avatar
Sengxian committed
181
        """
182
183
184
185
186
187
188
189

        moe_inp_batch_size = tree.flatten(
            tree.map_structure(lambda tensor: tensor.shape[0], moe_inp)
        )
        assert all(
            [batch_size == moe_inp_batch_size[0] for batch_size in moe_inp_batch_size]
        ), "MoE inputs must have the same batch size"

190
        if self.world_size > 1:
191
192
193
194
195

            def ensure_comm_func(tensor):
                ensure_comm(tensor, self.moe_group)

            tree.map_structure(ensure_comm_func, moe_inp)
Rick Ho's avatar
Rick Ho committed
196
        if self.slice_size > 1:
Sengxian's avatar
Sengxian committed
197

198
199
200
201
202
203
204
205
            def slice_func(tensor):
                return Slice.apply(
                    tensor, self.slice_rank, self.slice_size, self.slice_group
                )

            moe_inp = tree.map_structure(slice_func, moe_inp)

        gate_top_k_idx, gate_score = self.gate(moe_inp)
Rick Ho's avatar
Rick Ho committed
206

207
208
209
        if self.gate_hook is not None:
            self.gate_hook(gate_top_k_idx, gate_score, None)

Colin's avatar
Colin committed
210
        # delete masked tensors
Colin's avatar
Colin committed
211
        if self.mask is not None and self.mask_dict is not None:
212
213
214
215
216
217
            # TODO: to fix
            def delete_mask_func(tensor):
                # to: (BxL') x d_model
                tensor = tensor[mask == 0, :]
                return tensor

Colin's avatar
Colin committed
218
            mask = self.mask.view(-1)
219
            moe_inp = tree.map_structure(delete_mask_func, moe_inp)
Colin's avatar
Colin committed
220
221
222
            gate_top_k_idx = gate_top_k_idx[mask == 0, :]

        fwd = _fmoe_general_global_forward(
223
            moe_inp, gate_top_k_idx, self.expert_fn, self.num_expert, self.world_size
Sengxian's avatar
Sengxian committed
224
        )
225

Colin's avatar
Colin committed
226
        # recover deleted tensors
Colin's avatar
Colin committed
227
        if self.mask is not None and self.mask_dict is not None:
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247

            def recover_func(tensor):
                # to: (BxL') x top_k x dim
                dim = tensor.shape[-1]
                tensor = tensor.view(-1, self.top_k, dim)
                # to: (BxL) x top_k x d_model
                x = torch.zeros(
                    mask.shape[0],
                    self.top_k,
                    dim,
                    device=tensor.device,
                    dtype=tensor.dtype,
                )
                # recover
                x[mask == 0] = tensor
                for k, v in self.mask_dict.items():
                    x[mask == k] = v
                return x

            moe_outp = tree.map_structure(recover_func, fwd)
Colin's avatar
Colin committed
248
249
        else:

250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
            def view_func(tensor):
                dim = tensor.shape[-1]
                tensor = tensor.view(-1, self.top_k, dim)
                return tensor

            moe_outp = tree.map_structure(view_func, fwd)

        gate_score = gate_score.view(-1, 1, self.top_k)

        def bmm_func(tensor):
            dim = tensor.shape[-1]
            tensor = torch.bmm(gate_score, tensor).reshape(-1, dim)
            return tensor

        moe_outp = tree.map_structure(bmm_func, moe_outp)
Sengxian's avatar
Sengxian committed
265

Rick Ho's avatar
Rick Ho committed
266
        if self.slice_size > 1:
267
268
269
270
271
272
273

            def all_gather_func(tensor):
                return AllGather.apply(
                    tensor, self.slice_rank, self.slice_size, self.slice_group
                )

            moe_outp = tree.map_structure(all_gather_func, moe_outp)
274
275
276
277
278
279
280

        moe_outp_batch_size = tree.flatten(
            tree.map_structure(lambda tensor: tensor.shape[0], moe_outp)
        )
        assert all(
            [batch_size == moe_outp_batch_size[0] for batch_size in moe_outp_batch_size]
        ), "MoE outputs must have the same batch size"
281
        return moe_outp