これは 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
46class GPT(Module):
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
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
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)
バイアスを初期化
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
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 の最大学習率はです。
206 optimizer.learning_rate = 6e-4
208 optimizer.betas = (0.9, 0.95)
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',
コンテキストサイズを次の値にしてください
238 'seq_len': 128,
時代に合わせた列車
240 'epochs': 32,
バッチサイズ
242 'batch_size': 128,
エポックごとにトレーニングと検証を切り替える
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()