大语言模型之OpenAI GPT架构代码详解

这是Pytorch中OpenAI GPT架构的教程/实现。
此实现使用了小莎士比亚数据集。
GPT 模型本质上是一个标准的Transformer，但有一些调整。
GPT-2，尤其是 GPT-3 模型非常大，不适合单个 GPU，需要模型并行处理。
此实现甚至不使用数据并行性，旨在更像是一个教程。
与简单的自回归转换器相比，其主要区别在于参数初始化、权重衰减和学习速率调度。
对于Transformer，我们重用了现有的 labml/nn 变换器实现。

```python
import torch
from torch import nn
from labml import experiment
from labml.configs import option
from labml_helpers.module import Module
from labml_nn.experiments.nlp_autoregression import NLPAutoRegressionConfigs
from labml_nn.optimizers.configs import OptimizerConfigs
from labml_nn.transformers import TransformerConfigs, Encoder
from labml_nn.transformers.utils import subsequent_mask


class GPT(Module):


	def __init__(self, encoder: Encoder, src_embed: Module, generator: Module):
		super().__init__()
		
		self.src_embed = src_embed
		
		self.encoder = encoder
		
		self.generator = generator
		
		self.mask = None  #掩码将在第一次调用时初始化
	
	
	def forward(self, x: torch.Tensor):
	
		#如果掩码未初始化或掩码大小不同，则创建后续掩码
		if self.mask is None or self.mask.size(0) != len(x):
		
			#后续的掩码，将掩盖令牌以免看到未来的代替
			self.mask = subsequent_mask(len(x)).to(x.device)
			
		x = self.src_embed(x) #使用位置编码获取令牌嵌入
		
		x = self.encoder(x, self.mask) #变换器编码
		
		x = self.generator(x) #获取日志
		
		return x, None #返回结果（第二个值用于状态，因为我们的训练器也与RNN一起使用）


class Configs(NLPAutoRegressionConfigs):

	model: GPT
	
	transformer: TransformerConfigs
	
	weight_decay: float = 0.1
	
	warmup_steps: int = 128 * 128 * 20
	
	optimizer = 'transformer_optimizer'
	
	@option(Configs.transformer, 'GPT')
	def _transformer_configs(c: Configs):
	
		conf = TransformerConfigs()
		
		conf.n_src_vocab = c.n_tokens
		
		conf.n_tgt_vocab = c.n_tokens
		
		conf.ffn.activation = 'GELU'
		
		return conf


def _init_weights(module):

	if not isinstance(module, (nn.Linear, nn.Embedding)):
		return
		
	module.weight.data.normal_(mean=0.0, std=0.02)
	
	if isinstance(module, nn.Linear) and module.bias is not None:
		module.bias.data.zero_() #偏差初始化为0


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

	m = GPT(c.transformer.encoder,
			c.transformer.src_embed,
			c.transformer.generator).to(c.device)
			
	m.apply(_init_weights) #自定义权重初始化
	
	return m

@option(NLPAutoRegressionConfigs.optimizer)
def transformer_optimizer(c: NLPAutoRegressionConfigs):
	decay = set()
	
	for mn, m in c.model.named_modules():
		for pn, p in m.named_parameters():
		
			fpn = f'{mn}.{pn}' if mn else pn  # full param name
			
			if fpn.endswith('weight') and isinstance(m, nn.Linear):
				decay.add(fpn)
				
	# 获取所有参数
	param_dict = {pn: p for pn, p in c.model.named_parameters()}
	
	no_decay = set(param_dict.keys()) - decay
	
	# 创建pytorch优化器对象
	opt_groups = [{"params": [param_dict[pn] for pn in sorted(list(decay))], "weight_decay": c.weight_decay},
     			{"params": [param_dict[pn] for pn in sorted(list(no_decay))], "weight_decay": 0.0},]
     			
    # 创建一个可配置的优化器，这样我们就可以通过传递配置字典来更改它们
    optimizer = OptimizerConfigs()
    
    # 设置参数组以进行优化
    optimizer.parameters = opt_groups
    
    # 使用余弦衰减优化器，这就是GPT使用的
    optimizer.optimizer = 'AdamWarmupCosineDecay'
    
    # 如果我们使用具有指数衰减的Noam优化器，则需要设置模型嵌入大小
    optimizer.d_model = c.d_model
    
    # 设置默认权重衰减。这不是必需的，因为我们在参数组中设置了权重衰减
    optimizer.weight_decay = c.weight_decay
    
    # GPT使用的最大学习率为6×10^-4
    optimizer.learning_rate = 6e-4
    
    # β1=0.9,β2=0.95
    optimizer.betas = (0.9, 0.95)
    
    optimizer.eps = 1e-8
    
    # 权重衰减与梯度分离
    optimizer.weight_decouple = True
    
    # 学习速率余弦衰减的优化步骤总数
    optimizer.total_steps = c.epochs * len(c.text.train) // (c.batch_size * c.seq_len)
    
    # 预热优化步骤数
    optimizer.warmup = c.warmup_steps // (c.batch_size * c.seq_len)
    
    return optimizer

def main():

    experiment.create(name='GPT', comment='Train GPT model')
    
    # 初始化配置
    conf = Configs()
    
    # 覆盖配置
    experiment.configs(conf, {
    'tokenizer': 'character',  # 使用角色等级分词器
    'prompt_separator': '',    # 提示分隔符为空
    'prompt': 'It is ',        # 开始采样提示
    'text': 'tiny_shakespeare',# 使用小莎士比亚数据集
    'seq_len': 128,			   # 使用上下文大小为128
    'epochs': 32,
    'batch_size': 128,
    'inner_iterations': 10,    # 在训练和验证之间切换每个epoch的次数
    'transformer.d_model': 512,
    'transformer.ffn.d_ff': 2048,
    'transformer.n_heads': 8,
    'transformer.n_layers': 6
    })
    
    # 设置用于保存和加载的模型
    experiment.add_pytorch_models({'model': conf.model})
    
    #with experiment.start():
    #	conf.run()
    
    # 初始化模型
    conf.model.to(conf.device)
    
    # 设置优化器
    optimizer = conf.optimizer
    with experiment.start():
        for epoch in range(conf.epochs):
            for batch in conf.text.train:
            
                # Forward
                data = batch['tokens'].to(conf.device)
                
                output, _ = conf.model(data)
                
                # 计算损失
                loss = nn.CrossEntropyLoss()(output.view(-1, output.size(-1)), data.view(-1))
                
                # Backward
                optimizer.zero_grad()
                
                loss.backward()
                
                optimizer.step()
                
                experiment.update({
                    'loss': loss.item()
                })


if __name__ == '__main__':
    main()

``