ජීපීටී

35import torch
36from torch import nn
37
38from labml import experiment
39from labml.configs import option
40from labml_helpers.module import Module
41from labml_nn.experiments.nlp_autoregression import NLPAutoRegressionConfigs
42from labml_nn.optimizers.configs import OptimizerConfigs
43from labml_nn.transformers import TransformerConfigs, Encoder
44from labml_nn.transformers.utils import subsequent_mask

47class GPT(Module):

55    def __init__(self, encoder: Encoder, src_embed: Module, generator: Module):

62        super().__init__()
63        self.src_embed = src_embed
64        self.encoder = encoder
65        self.generator = generator

68        self.mask = None

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

73        if self.mask is None or self.mask.size(0) != len(x):

75            self.mask = subsequent_mask(len(x)).to(x.device)

77        x = self.src_embed(x)

79        x = self.encoder(x, self.mask)

81        x = self.generator(x)

85        return x, None

88class Configs(NLPAutoRegressionConfigs):

97    model: GPT

99    transformer: TransformerConfigs

101    weight_decay: float = 0.1

103    warmup_steps: int = 128 * 128 * 20

106    optimizer = 'transformer_optimizer'

109@option(Configs.transformer, 'GPT')
110def _transformer_configs(c: Configs):

117    conf = TransformerConfigs()

119    conf.n_src_vocab = c.n_tokens
120    conf.n_tgt_vocab = c.n_tokens

122    conf.ffn.activation = 'GELU'

125    return conf

128def _init_weights(module):

137    if not isinstance(module, (nn.Linear, nn.Embedding)):
138        return
139
140    module.weight.data.normal_(mean=0.0, std=0.02)

143    if isinstance(module, nn.Linear) and module.bias is not None:
144        module.bias.data.zero_()

147@option(Configs.model)
148def _model(c: Configs):

152    m = GPT(c.transformer.encoder,
153            c.transformer.src_embed,
154            c.transformer.generator).to(c.device)

157    m.apply(_init_weights)
158
159    return m

162@option(NLPAutoRegressionConfigs.optimizer)
163def transformer_optimizer(c: NLPAutoRegressionConfigs):

171    decay = set()
172    for mn, m in c.model.named_modules():
173        for pn, p in m.named_parameters():
174            fpn = f'{mn}.{pn}' if mn else pn  # full param name
175
176            if fpn.endswith('weight') and isinstance(m, nn.Linear):
177                decay.add(fpn)

180    param_dict = {pn: p for pn, p in c.model.named_parameters()}

182    no_decay = set(param_dict.keys()) - decay

185    opt_groups = [
186        {"params": [param_dict[pn] for pn in sorted(list(decay))], "weight_decay": c.weight_decay},
187        {"params": [param_dict[pn] for pn in sorted(list(no_decay))], "weight_decay": 0.0},
188    ]

193    optimizer = OptimizerConfigs()

196    optimizer.parameters = opt_groups

199    optimizer.optimizer = 'AdamWarmupCosineDecay'

202    optimizer.d_model = c.d_model

205    optimizer.weight_decay = c.weight_decay

207    optimizer.learning_rate = 6e-4

209    optimizer.betas = (0.9, 0.95)

211    optimizer.eps = 1e-8

213    optimizer.weight_decouple = True

215    optimizer.total_steps = c.epochs * len(c.text.train) // (c.batch_size * c.seq_len)

217    optimizer.warmup = c.warmup_steps // (c.batch_size * c.seq_len)
218
219    return optimizer

222def main():

224    experiment.create(name="gpt")

226    conf = Configs()

228    experiment.configs(conf, {

230        'tokenizer': 'character',

232        'prompt_separator': '',

234        'prompt': 'It is ',

236        'text': 'tiny_shakespeare',

239        'seq_len': 128,

241        'epochs': 32,

243        'batch_size': 128,

246        'inner_iterations': 10,

249        'transformer.d_model': 512,
250        'transformer.ffn.d_ff': 2048,
251        'transformer.n_heads': 8,
252        'transformer.n_layers': 6
253    })

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

259    with experiment.start():

261        conf.run()

265if __name__ == '__main__':
266    main()