#

Switch Transformer Experiment

This is an annotated PyTorch experiment to train a switch transformer.

12import torch
13import torch.nn as nn
14
15from labml import experiment, tracker
16from labml.configs import option
17from labml_helpers.module import Module
18from labml_helpers.train_valid import BatchIndex
19from labml_nn.experiments.nlp_autoregression import NLPAutoRegressionConfigs

#

Auto regressive model

22class AutoregressiveModel(Module):

#

27    def __init__(self, n_vocab: int, d_model: int, transformer: Module):
28        super().__init__()

#

Token embedding module

30        self.src_embed = nn.Embedding(n_vocab, d_model)

#

Transformer

32        self.transformer = transformer

#

Final layer

34        self.generator = nn.Linear(d_model, n_vocab)
35        self.mask = None

#

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

#

Initialize the subsequent mask

39        if self.mask is None or self.mask.size(0) != len(x):
40            from labml_nn.transformers.utils import subsequent_mask
41            self.mask = subsequent_mask(len(x)).to(x.device)

#

Token embeddings

43        x = self.src_embed(x)

#

Run it through the transformer

45        res, counts, route_prob, n_dropped = self.transformer(x, self.mask)

#

Generate logits of the next token

47        res = self.generator(res)

#

49        return res, counts, route_prob, n_dropped

#

Configurations

This extends NLPAutoRegressionConfigs.

The default configs can and will be over-ridden when we start the experiment

52class Configs(NLPAutoRegressionConfigs):

#

61    model: AutoregressiveModel
62    transformer: Module

#

Token embedding size

65    d_model: int = 128

#

Number of attention heads

67    heads: int = 4

#

Dropout probability

69    dropout: float = 0.0

#

Number of features in FFN hidden layer

71    d_ff: int = 256

#

Number of transformer layers

73    n_layers: int = 6

#

Number of experts

75    n_experts: int = 4

#

Load balancing coefficient

77    load_balancing_loss_ceof = 0.01

#

Whether to scale the chosen expert outputs by the routing probability

79    is_scale_prob: bool = True

#

Whether to drop tokens

81    drop_tokens: bool = False

#

Capacity factor to determine capacity of each model

83    capacity_factor: float = 1.0

#

85    def init(self):
86        super().init()

#

Initialize tracking indicators

88        tracker.set_scalar("lb_loss.*", False)
89        tracker.set_scalar("route.*", False)
90        tracker.set_scalar("dropped.*", False)

#

Training or validation step

92    def step(self, batch: any, batch_idx: BatchIndex):

#

Move data to the device

98        data, target = batch[0].to(self.device), batch[1].to(self.device)

#

Update global step (number of tokens processed) when in training mode

101        if self.mode.is_train:
102            tracker.add_global_step(data.shape[0] * data.shape[1])

#

Whether to capture model outputs

105        with self.mode.update(is_log_activations=batch_idx.is_last):

#

Get model outputs.

107            output, counts, route_prob, n_dropped = self.model(data)

#

Calculate and cross entropy loss

110        cross_entropy_loss = self.loss_func(output, target)

#

Total number of tokens processed, $T$, in the current batch $\mathscr{B}$

112        total = counts.sum(dim=-1, keepdims=True)

#

Fraction of tokens routed to each expert $f_i = \frac{1}{T} \sum_{x \in \mathscr{B}} \unicode{x1D7D9} \{ \mathop{argmax} p(x), i \}$ $f_i$ is the count of tokens where the argmax of $p(x)$ is equal to $i$.

116        route_frac = counts / total

#

Mean routing probability $P_i = \frac{1}{T} \sum_{x \in \mathscr{B}} p_i (x)$

119        route_prob = route_prob / total

#

Load balancing loss $\mathscr{L} = N \sum_{i=1}^N f_i \cdot P_i$

122        load_balancing_loss = self.n_experts * (route_frac * route_prob).sum()

#

Track stats

125        tracker.add('dropped.', total.new_tensor(n_dropped) / total)
126        tracker.add('route.min.', route_frac.min())
127        tracker.add('route.max.', route_frac.max())
128        tracker.add('route.std.', route_frac.std())
129        tracker.add("loss.", cross_entropy_loss)
130        tracker.add("lb_loss.", load_balancing_loss)

#

Combined loss. The load balancing loss is multiplied by a coefficient $\alpha$ which is set to something small like $\alpha = 0.01$.

135        loss = cross_entropy_loss + self.load_balancing_loss_ceof * load_balancing_loss

#

Calculate and log accuracy

138        self.accuracy(output, target)
139        self.accuracy.track()

#

Train the model

142        if self.mode.is_train:

#

Calculate gradients

144            loss.backward()

#

Clip gradients

146            torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=self.grad_norm_clip)

#

Take optimizer step

148            self.optimizer.step()

#

Log the model parameters and gradients on last batch of every epoch

150            if batch_idx.is_last:
151                tracker.add('model', self.model)

#

Clear the gradients

153            self.optimizer.zero_grad()

#

Save the tracked metrics

156        tracker.save()

#

Initialize the auto-regressive model

159@option(Configs.model)
160def autoregressive_model(c: Configs):

#

164    m = AutoregressiveModel(c.n_tokens, c.d_model, c.transformer)
165    return m.to(c.device)

#

Initialize the switch transformer

168@option(Configs.transformer)
169def switch_transformer(c: Configs):

#

173    from labml_nn.transformers.switch import SwitchTransformer, SwitchTransformerLayer, SwitchFeedForward
174    from labml_nn.transformers import MultiHeadAttention
175    from labml_nn.transformers.feed_forward import FeedForward
176
177    return SwitchTransformer(
178        SwitchTransformerLayer(d_model=c.d_model,
179                               attn=MultiHeadAttention(c.heads, c.d_model, c.dropout),
180                               feed_forward=SwitchFeedForward(capacity_factor=c.capacity_factor,
181                                                              drop_tokens=c.drop_tokens,
182                                                              is_scale_prob=c.is_scale_prob,
183                                                              n_experts=c.n_experts,
184                                                              expert=FeedForward(c.d_model, c.d_ff, c.dropout),
185                                                              d_model=c.d_model),
186                               dropout_prob=c.dropout),
187        c.n_layers)

#

Run the experiment

190def main():

#

Create experiment

195    experiment.create(name="switch_transformer", comment='')

#

Create configs

197    conf = Configs()

#

Load configurations

199    experiment.configs(conf,

#

A dictionary of configurations to override

201                       {'tokenizer': 'character',
202                        'text': 'tiny_shakespeare',
203                        'optimizer.learning_rate': 1.,
204                        'optimizer.optimizer': 'Noam',
205                        'prompt': 'It is',
206                        'prompt_separator': '',
207
208                        'transformer': 'switch_transformer',
209                        'is_scale_prob': False,
210                        'n_experts': 4,
211
212                        'drop_tokens': True,
213                        'capacity_factor': 1.2,
214
215                        'train_loader': 'shuffled_train_loader',
216                        'valid_loader': 'shuffled_valid_loader',
217
218                        'seq_len': 64,
219                        'epochs': 128,
220                        'batch_size': 32,
221                        'inner_iterations': 25,
222                        })

#

Set models for saving and loading

225    experiment.add_pytorch_models({'model': conf.model})

#

Start the experiment

228    with experiment.start():

#

TrainValidConfigs.run

230        conf.run()

#

234if __name__ == '__main__':
235    main()