40import torch
41from torch import nn
42
43from labml_helpers.module import Module
44from labml_nn.transformers.feed_forward import FeedForward
45from labml_nn.transformers.mha import MultiHeadAttention
46from labml_nn.utils import clone_module_list

49class SwitchFeedForward(Module):

54    def __init__(self, *,
55                 capacity_factor: float,
56                 drop_tokens: bool,
57                 is_scale_prob: bool,
58                 n_experts: int,
59                 expert: FeedForward,
60                 d_model: int):

71        super().__init__()
72
73        self.capacity_factor = capacity_factor
74        self.is_scale_prob = is_scale_prob
75        self.n_experts = n_experts
76        self.drop_tokens = drop_tokens

79        self.experts = clone_module_list(expert, n_experts)

81        self.switch = nn.Linear(d_model, n_experts)
82        self.softmax = nn.Softmax(dim=-1)

84    def forward(self, x: torch.Tensor):

90        seq_len, batch_size, d_model = x.shape

92        x = x.view(-1, d_model)

98        route_prob = self.softmax(self.switch(x))

102        route_prob_max, routes = torch.max(route_prob, dim=-1)

105        indexes_list = [torch.eq(routes, i).nonzero(as_tuple=True)[0] for i in range(self.n_experts)]

108        final_output = x.new_zeros(x.shape)

114        capacity = int(self.capacity_factor * len(x) / self.n_experts)

116        counts = x.new_tensor([len(indexes_list[i]) for i in range(self.n_experts)])

119        dropped = []

121        if self.drop_tokens:

123            for i in range(self.n_experts):

125                if len(indexes_list[i]) <= capacity:
126                    continue

128                indexes_list[i] = indexes_list[i][torch.randperm(len(indexes_list[i]))]

130                dropped.append(indexes_list[i][capacity:])

132                indexes_list[i] = indexes_list[i][:capacity]

135        expert_output = [self.experts[i](x[indexes_list[i], :]) for i in range(self.n_experts)]

138        for i in range(self.n_experts):
139            final_output[indexes_list[i], :] = expert_output[i]

142        if dropped:
143            dropped = torch.cat(dropped)
144            final_output[dropped, :] = x[dropped, :]
145
146        if self.is_scale_prob:

148            final_output = final_output * route_prob_max.view(-1, 1)
149        else:

152            final_output = final_output * (route_prob_max / route_prob_max.detach()).view(-1, 1)

155        final_output = final_output.view(seq_len, batch_size, d_model)

166        return final_output, counts, route_prob.sum(0), len(dropped), route_prob_max

169class SwitchTransformerLayer(Module):

177    def __init__(self, *,
178                 d_model: int,
179                 attn: MultiHeadAttention,
180                 feed_forward: SwitchFeedForward,
181                 dropout_prob: float):

188        super().__init__()
189        self.size = d_model
190        self.attn = attn
191        self.feed_forward = feed_forward
192        self.dropout = nn.Dropout(dropout_prob)
193        self.norm_self_attn = nn.LayerNorm([d_model])
194        self.norm_ff = nn.LayerNorm([d_model])

196    def forward(self, *,
197                x: torch.Tensor,
198                mask: torch.Tensor):

200        z = self.norm_self_attn(x)

202        self_attn = self.attn(query=z, key=z, value=z, mask=mask)

204        x = x + self.dropout(self_attn)

207        z = self.norm_ff(x)

209        ff, counts, route_prob, n_dropped, route_prob_max = self.feed_forward(z)

211        x = x + self.dropout(ff)
212
213        return x, counts, route_prob, n_dropped, route_prob_max

216class SwitchTransformer(Module):

221    def __init__(self, layer: SwitchTransformerLayer, n_layers: int):
222        super().__init__()

224        self.layers = clone_module_list(layer, n_layers)

226        self.norm = nn.LayerNorm([layer.size])

228    def forward(self, x: torch.Tensor, mask: torch.Tensor):

230        counts, route_prob, n_dropped, route_prob_max = [], [], [], []
231        for layer in self.layers:
232            x, f, p, n_d, p_max = layer(x=x, mask=mask)
233            counts.append(f)
234            route_prob.append(p)
235            n_dropped.append(n_d)
236            route_prob_max.append(p_max)

238        x = self.norm(x)

240        return x, torch.stack(counts), torch.stack(route_prob), n_dropped, torch.stack(route_prob_max)