[LLM]从零开始搭建GPT-2(1)

本篇将介绍一下GPT-2(124M)的搭建和初始化，主要是一些关键点的讲解，完整代码点击这里 。

因为OpneAI发布的gpt2源代码使用的tensorflow框架，我们会使用pytorch，使用Huggingface提供的转化过的权重参数。

搬出下面这张图镇楼，作为一切的一切的起点，我们将围绕这张图来搭建gpt2（下面统称gpt2）的模型框架。

首先gpt2没有图的左边这一部分，包括接入的交叉注意力头，并且根据gpt2的官方文档有两个变动点：1. 交换了归一化层的位置 2. 在最后的自注意力块后面又添加了一个新的归一化层。

gpt2的一些配置信息：

block_size: int = 1024 # 最大序列长度 vocab_size: int = 50257 # 词汇表大小 50000 + 256 + 1
n_layer: int = 12
n_head: int = 12
n_embd: int = 768 # 嵌入维度

开始搭建框架：

self.transformer = nn.ModuleDict(dict(
wte = nn.Embedding(config.vocab_size, config.n_embd), # 对应图中的Output Embedding
wpe = nn.Embedding(config.block_size, config.n_embd), # 对应图中的Positional Embedding
h = nn.ModuleList([Block(config) for _ in range(config.n_layer)]), # 对应主体的模块
ln_f = nn.LayerNorm(config.n_embd), # gpt添加的额外的归一化层
)) self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False) # 对应图中的linear

在Block块里面的attn()里tokens相互交流信息，而在mlp()里进行信息的消化理解，完成mlp和多头自注意力块后gpt2的雏形就算完成了，接着来加载gpt2的配置信息和参数。

在模型的初始化以及后面的训练中一般是要把模型移到gpu训练的，但就算没有gpu也可以跟随着步骤走一段路，下面的代码可以用来自动检测设备：

device = 'cpu'
if torch.cuda.is_available():
device = 'cuda'
elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
device = 'mps'
print(f'using device: {device}')

要注意的是在模型移到gpu的同时不要忘记张量也是需要移到gpu的，防止张量出现在多个计算设备上面

buf.to(device) # 这种方法是错误的，有漏洞而且不稳定
buf = buf.to(device) # 使用这种方法

在debug阶段使用的数据集推荐tiny shakespeare(link)

gpt2的Tokenizer压缩率在300%左右，比如1000个字符会生成300个token左右

在处理tokens可以这样做：

buf = torch.tensor(tokens[:6 + 1])
x = buf[:-1].view(2, 3) # input
y = buf[1:].view(2, 3) # target
x = ([56, 22, 12],
[82, 14, 77])
y = ([22, 12, 82],
[14, 77, 11])
这样x中的每一个token的下一个token就可以在y里面的对应位置找见，并且也解决了x中最后一个token的问题

有了y就可以计算loss了

loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))

在初始化阶段的loss我们是希望每一个token有着相似的可能性，loss = -ln(1/50257) = 10.82 (约)

optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4) # 默认是Adam，建议使用AdamW，修复了一些小bug；
另外这个初始的学习率在现阶段就可以
for i in range(50):
x, y =train_loader.next_batch()
x, y = x.to(device), y.to(device) # 把x,y张量移动到device计算
optimizer.zero_grad() # 这里容易忘记
logits, loss = model(x, y)
loss.backward()
optimizer.step()
print(f'step{i}, loss: {loss.item()}')

接下来就可以把数据集分批次送入模型中了，这需要一个数据装入模块，见源码的DataLoaderLite，建议输出的信息包括：

print(f'loaded {len(self.tokens)} tokens') # 总计数据集的token数
print(f'1 epoch = {len(self.tokens) // (B * T)} batches') # 一次完整的循环要的批次数
def next_batch(self):
B, T = self.B, self.T
buf = self.tokens[self.current_position : self.current_position+B*T+1]
x = (buf[:-1]).view(B, T) # inputs
y = (buf[1:]).view(B, T) # targets
# advance the position in the tensor
self.current_position += B * T # 更新当前的批次位置
# if loading the next batch would be out of bounds, advance to next shard
if self.current_position + (B * T + 1) > len(self.tokens):
self.current_position = 0
return x, y

现在我们在不断送入新的批次数据，可这也涵盖不了所有的词汇表token，一些生僻的多语言的token是不会出现在数据集里面的，所以loss应该会下降一些。根据gpt2的官方文档，Token Emdebdding 的权重参数其实和Transformer的最上层的Linear层的权重参数是完全相同一致的，添加：

self.transformer.wte.weight = self.lm_head.weight

别小看仅仅添加了这一行代码，这行代码帮我们节省了近4百万的训练参数，相当于gpt2模型30%的参数量

初始化参数：

self.apply(self._init_weights)
def _init_weights(self, module):
if isinstance(module, nn.Linear):
std = 0.02
if hasattr(module, 'NANOGPT_SCALE_INIT'):
std *= (2 * self.config.n_layer) ** -0.5 # 控制残差层内的激活增长
torch.nn.init.normal_(module.weight, mean=0.0, std=std)
if module.bias is not None:
torch.nn.init.zeros_(module.bias)
elif isinstance(module, nn.Embedding):
torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)

最后可以设一下种子

torch.manual_seed(1337)
if torch.cuda.is_available():
torch.cuda.manual_seed(1337)

运行，看到下面的输出就说明你已经成功完成gpt2的初始化了

using device: mps
loaded 338025 tokens
1 epoch = 2640 batches
step0, loss: 10.960028648376465
step1, loss: 9.687705993652344
step2, loss: 9.08289909362793
step3, loss: 9.145987510681152
step4, loss: 8.626202583312988
.
.
.
step48, loss: 6.953254699707031
step49, loss: 6.799217224121094