#

DeepNorm Experiment

13import copy
14
15import torch
16import torch.nn as nn
17
18from labml import experiment
19from labml.configs import option
20from labml_helpers.module import Module
21from labml_nn.experiments.nlp_autoregression import NLPAutoRegressionConfigs
22from labml_nn.normalization.deep_norm import DeepNormTransformerLayer
23from labml_nn.transformers import MultiHeadAttention
24from labml_nn.transformers.feed_forward import FeedForward

#

Auto-Regressive model

This is a autoregressive transformer model that uses DeepNorm.

27class AutoregressiveTransformer(Module):

#

n_tokens is the number of tokens in the vocabulary
d_model is the embedding size
n_layers is the number of transformer layers
layer is the layer. We use n_layers copies of this for the tranformer.

34    def __init__(self, n_tokens: int, d_model: int, n_layers: int, layer: DeepNormTransformerLayer):

#

41        super().__init__()

#

Transformer with n_layers layers

43        self.transformer = nn.Sequential(*[copy.deepcopy(layer) for _ in range(n_layers)])

#

Token embedding layer

46        self.emb = nn.Embedding(n_tokens, d_model)

#

Readout layer

48        self.readout = nn.Linear(d_model, n_tokens)

#

x are the input tokens of shape [seq_len, batch_size]

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

#

Get the token embeddings

55        x = self.emb(x)

#

Transformer encoder

57        x = self.transformer(x)

#

Get logits

59        x = self.readout(x)

#

Return results

62        return x, None

#

Configurations

This inherits from NLPAutoRegressionConfigs

65class Configs(NLPAutoRegressionConfigs):

#

Model

74    model: AutoregressiveTransformer

#

Number of layers

77    n_layers: int = 32

#

$α$ and $β$ for DeepNorm

80    deep_norm_alpha: float
81    deep_norm_beta: float

#

Number of heads in the attention

84    n_heads: int = 4

#

Embedding size

86    d_model: int = 64

#

Size of each attention head

88    d_k: int = 16

#

Calculate $α$

$α = (2 M)^{\frac{1}{4}}$

91@option(Configs.deep_norm_alpha)
92def _deep_norm_alpha(c: Configs):

#

98    return (2. * c.n_layers) ** (1. / 4.)

#

Calculate $β$

$β = (8 M)^{- \frac{1}{4}}$

101@option(Configs.deep_norm_beta)
102def _deep_norm_beta(c: Configs):

#

108    return (8. * c.n_layers) ** -(1. / 4.)

#

Initialize the model

111@option(Configs.model)
112def _model(c: Configs):

#

116    m = AutoregressiveTransformer(c.n_tokens, c.d_model, c.n_layers,
117                                  DeepNormTransformerLayer(d_model=c.d_model,
118                                                           deep_norm_alpha=c.deep_norm_alpha,
119                                                           deep_norm_beta=c.deep_norm_beta,
120                                                           feed_forward=FeedForward(d_model=c.d_model,
121                                                                                    d_ff=c.d_model * 4),
122                                                           self_attn=MultiHeadAttention(c.n_heads, c.d_model,
123                                                                                        dropout_prob=0.0)))
124
125    return m.to(c.device)

#

Create and run the experiment

128def main():

#

Create experiment

133    experiment.create(name="deep_norm", writers={'screen', 'web_api'})

#

Create configs

135    conf = Configs()

#

Override configurations

137    experiment.configs(conf, {

#

Use character level tokenizer

139        'tokenizer': 'character',

#

Prompt separator is blank

141        'prompt_separator': '',

#

Starting prompt for sampling

143        'prompt': 'It is ',

#

Use Tiny Shakespeare dataset

145        'text': 'tiny_shakespeare',

#

Use a context size of $256$

148        'seq_len': 256,

#

Train for 32 epochs

150        'epochs': 32,

#

Batch size $16$

152        'batch_size': 16,

#

Switch between training and validation for $10$ times per epoch

154        'inner_iterations': 10,

#

Number of layers

157        'n_layers': 50,

#

Adam optimizer with no warmup

161        'optimizer.optimizer': 'Adam',
162        'optimizer.learning_rate': 1.25e-4,
163    })

#

Set model(s) for saving and loading

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

#

Start the experiment

169    with experiment.start():

#

Run training

171        conf.run()

#

175if __name__ == '__main__':
176    main()