#

GPT

这是 P yTorch 中 OpenAI GPT 架构的教程/实现。@karpathy 从 MinGpt 那里得到了很多实现细节。此实现还使用角色小莎士比亚数据集。

GPT 模型本质上是一个标准的变压器，但有一些调整。GPT-2，尤其是 GPT-3 模型非常大，不适合单个 GPU，需要模型并行处理。此实现甚至不使用数据并行性，旨在更像是一个教程。

与@@

简单的自回归转换器相比，其主要区别在于参数初始化、权重衰减和学习速率调度。对于变压器，我们重用了现有的 labml/nn 变换器实现。

这是一本用于在 Tiny Shakespeare 数据集上训练 GPT 模型的笔记本。

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

#

GPT 型号

这包括令牌嵌入层、变压器编码器和给出令牌日志的最终线性层。

46class GPT(Module):

#

encoder 是变压器编码器
src_embed 是令牌嵌入模块（带有位置编码）
generator 是给出 logit 的最后一个完全连接的层。

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

#

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

#

掩码将在第一次调用时初始化

67        self.mask = None

#

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

#

如果掩码未初始化或掩码大小不同，则创建后续掩码

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

#

后续的掩码，将掩盖令牌以免看到未来的代币

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

#

使用位置编码获取令牌嵌入

76        x = self.src_embed(x)

#

变压器编码

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

#

获取日志

80        x = self.generator(x)

#

返回结果（第二个值用于状态，因为我们的训练器也与 RNN 一起使用）

84        return x, None

#

配置

这继承自 NLPAutoRegressionConfigs

87class Configs(NLPAutoRegressionConfigs):

#

GPT 型号

96    model: GPT

#

变压器

98    transformer: TransformerConfigs

#

体重衰减

100    weight_decay: float = 0.1

#

wamup 的代币数量

102    warmup_steps: int = 128 * 128 * 20

#

自定义优化器

105    optimizer = 'transformer_optimizer'

#

变压器配置

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

#

我们使用我们的可配置变压器实现

116    conf = TransformerConfigs()

#

设置嵌入和生成 logit 的词汇量大小

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

#

GPT 使用 GELU 激活进行位置明智前馈

121    conf.ffn.activation = 'GELU'

#

124    return conf

#

初始化权重

线性层和嵌入层的权重初始化为， $N (0, 0.02)$ 而不是默认的 Xavier 初始化。

127def _init_weights(module):

#

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

#

将偏差初始化为 $0$

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

#

创建 GPT 模型并初始化权重

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

#

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

#

应用自定义权重初始化

156    m.apply(_init_weights)
157
158    return m

#

创建具有权重衰减的自定义优化器

此代码取自 MingPT。这仅将权重衰减应用于线性图层的权重。

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

#

收集参数名称以应用权重衰减

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

#

获取所有参数

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

#

未衰减的参数

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

#

创建 pytorch 优化器对象

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

#

创建一个可配置的优化器，这样我们就可以通过传递配置字典来更改它们。

192    optimizer = OptimizerConfigs()

#

设置参数组以进行优化。

195    optimizer.parameters = opt_groups

#

使用余弦衰减优化器。这就是 GPT 使用的。

198    optimizer.optimizer = 'AdamWarmupCosineDecay'

#

设置模型嵌入大小，如果我们使用具有指数衰减的 Noam 优化器，则需要设置模型嵌入大小。

201    optimizer.d_model = c.d_model

#

设置默认权重衰减。这不是必需的，因为我们在参数组中设置了权重衰减。

204    optimizer.weight_decay = c.weight_decay

#

GPT 使用的最大学习速率为 $6 \times 10^{- 4}$ 。

206    optimizer.learning_rate = 6e-4

#

$β_{1} = 0 .9, β_{2} = 0 .95$

208    optimizer.betas = (0.9, 0.95)

#

$ϵ = 10^{- 8}$

210    optimizer.eps = 1e-8

#

权重衰减与梯度分离

212    optimizer.weight_decouple = True

#

学习速率余弦衰减的优化步骤总数

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

#

预热优化步骤数

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

#

221def main():

#

创建实验

223    experiment.create(name="gpt")

#

创建配置

225    conf = Configs()

#

覆盖配置

227    experiment.configs(conf, {

#

使用角色等级分词器

229        'tokenizer': 'character',

#

提示分隔符为空

231        'prompt_separator': '',

#

开始采样提示

233        'prompt': 'It is ',

#

使用小莎士比亚数据集

235        'text': 'tiny_shakespeare',

#

使用上下文大小为 $128$

238        'seq_len': 128,

#

为 $32$ 时代而训练

240        'epochs': 32,

#

批量大小 $128$

242        'batch_size': 128,

#

在训练和验证之间切换每个纪元的 $10$ 次数

245        'inner_iterations': 10,

#

变压器配置

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

#

设置用于保存和加载的模型

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

#

开始实验

258    with experiment.start():

#

跑步训练

260        conf.run()

#

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