#

GPT

これは PyTorch における 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 ロジットを生成する最後の完全接続層です。

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でも使用されるため、2番目の値は状態用です）

84        return x, None

#

コンフィギュレーション

これは以下から継承されます NLPAutoRegressionConfigs

87class Configs(NLPAutoRegressionConfigs):

#

GPT モデル

96    model: GPT

#

変圧器

98    transformer: TransformerConfigs

#

体重減少

100    weight_decay: float = 0.1

#

ウォームアップ用のトークン数

102    warmup_steps: int = 128 * 128 * 20

#

カスタムオプティマイザー

105    optimizer = 'transformer_optimizer'

#

変圧器構成

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

#

設定可能なトランス実装を使用しています

116    conf = TransformerConfigs()

#

埋め込みやロジットの生成に使用するボキャブラリーサイズを設定

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

#

ウェイトを初期化

線形レイヤーと埋め込みレイヤーの重みは、デフォルトの Xavier $N (0, 0.02)$ 初期化の代わりに初期化されます。

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()