行程式化地重现 GPT-2：第三部分

原文：towardsdatascience.com/line-by-line-lets-reproduce-gpt-2-section-3-training-f2fef87880fc

在第一篇博客文章中，我们编写了 Transformer 架构的代码。在第二篇博客文章中，我们优化了 NVIDIA 交互。为了完成我们的系列，我们现在将预训练 GPT-2 模型。如果你还没有看过，我强烈推荐查看 Andrej Karpathy 的 “Let’s reproduce GPT-2 (124M)” 视频（https://www.youtube.com/watch?v=l8pRSuU81PU），以及我们系列中的前两篇文章。

行程式化地重现 GPT-2：第一部分
行程式化地重现 GPT-2：第二部分 – 硬件优化

不再拖延，让我们深入探讨！

在我们开始预训练之前，我们希望对代码进行一些修改，以提高性能，并在训练过程中提高可见性。请注意，我们将多次修改代码的许多部分，所以如果你在跟随，请注意你将无法运行代码，直到最后。

AdamW 超参数

我们在这里做的第一件事是调整我们的 AdamW 超参数，使其更接近 GPT-3 的训练。虽然我们主要遵循 GPT-2 的设置，但 GPT-2 论文没有指定其 AdamW beta 或 epsilon，所以我们将遵循 GPT-3 的设置。

GPT-3 将 beta 设置为 0.9 和 0.95（请注意，截至写作时 PyTorch 的默认值是 0.9 和 0.999），以及一个 epsilon 为 1e-8（请注意，这等于截至写作时的 PyTorch 默认值）。

AdamW 优化器很复杂，所以在这里我不会深入细节。就我们的目的而言，beta 是两个我们用来告诉 AdamW 看多远的历史的值。Beta-1 通知梯度移动平均的衰减率，Beta-2 通知平方梯度的移动平均的衰减率。优化器使用这些值（以及学习率）来确定给定其在损失函数中的当前位置的最佳下一个值。

通过将第二个 beta 从 0.99 改为 0.95，我们告诉优化器不要像以前那样看那么远的历史。同时，epsilon 告诉模型至少要移动多少权重。1e-8 似乎是基于经验结果设定的值。

梯度裁剪

接下来，从 GPT-3 论文中，我们将剪辑梯度。这里的想法是规避不良批次的风险。如果你有一个不良批次，你最终会通过调整比实际应该调整的更多参数来过度补偿。这可能导致覆盖模型中的良好参数值，实际上逆转了之前的步骤。为了避免这种过度补偿，我们剪辑梯度，使其具有最大范数。这在数学上意味着我们的梯度不能超过一定的大小（在这个例子中是 1.0）。

norm=torch.nn.utils.clip_grad_norm_(model.parameters(),1.0)

虽然这防止了我们的模型在开始时改变太多，但也防止了它从一个特定的批次改变太多。因为在开始时，它将连续跨多个批次改变，这不应该是一个问题。

然而，这更像是对我们训练数据集中不良数据问题的一种“临时修补”。当数据集中出现不良批次时，它可能存在这个问题只能减轻。完全移除它并拥有高质量的数据集才是最好的。

学习率调度器

让我们更深入地探讨我们希望模型何时改变其参数。我们使用学习率值来告诉优化器它应该改变多少值。如果有一个超参数会极大地改变行为，那就是这个。这个值决定了优化器在每次训练步骤后损失函数上采取的步长有多大。太大，你会过度纠正；太小，你纠正不够。

在过去，我们将其设置为恒定值，但因为我们预计开始时的变化比结束时更多，我们将编写一个调度器来调整训练过程中的学习率。

我们下面的代码从预热阶段开始，那时它将使用最大学习率。这是有道理的，因为大多数存储的权重是随机的，因此对这些权重的重大变化可能是必需的。

在预热完成之后，我们进入一个阶段，希望降低学习率，以停止改变许多权重，从而在回合之间保持更多的权重相同或相似。这里的理论是，随着模型训练，权重越来越接近良好值，因此由于高学习率导致的任何重大变化都是适得其反的。我们使用余弦衰减来确定衰减到最小学习率的速度。

# inside of our training loop# ...max_lr=6e-4min_lr=max_lr*0.1warmup_steps=715max_steps=19073ifit<warmup_steps:returnmax_lr*(it+1)/warmup_stepsifit>max_steps:returnmin_lr decay_ratio=(it-warmup_steps)/(max_steps-warmup_steps)assert0<=decay_ratio<=1coeff=0.5*(1.0+math.cos(math.pi*decay_ratio))returnmin_lr+coeff*(max_lr-min_lr)# ...

注意，在 GPT-3 论文中，他们有几个步骤，其中学习率简单地是最低学习率，而在这上面，我们只是在训练的末尾达到那个值。

https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/c6499b0eb786728c147a703723c3b4de.png

表 2.1 来自 "Language Models are Few-Shot Learners

在这里还有一个需要注意的事项，那就是随着模型规模的增大，学习率往往会变小。这里的想法是，随着模型规模的增大，你需要记住更多的权重，因此在开始时不必频繁更新它们是可以接受的。在我看来，这是一个值得进一步研究的有趣地方，因为模型蒸馏已经表明，大型模型（如 Llama 3.1 405B）的大部分性能可以在显著更小的模型（如 Llama3.1 70B）中保留。

权重衰减和融合 AdamW

接下来，我们将设置优化器以具有权重衰减，或者在每个训练步骤后，我们将一个衰减计算应用到我们的权重上。这里的想法是避免过拟合。当我们衰减这些权重时，我们迫使优化器不断增加它们的值以保持它们较高。这种动态使得一些权重变得非常高并开始主导的可能性非常低。AdamW 的默认值是 0.01，这比我们打算用于训练的值小 10 倍。

在下面，我们将参数分成我们想要衰减的和不会衰减的类别。我们在这里的重点是对涉及矩阵乘法和嵌入的权重进行权重衰减。我们可以通过选择维度大于 1 的权重来过滤这些权重。

classGPT(nn.Module):# ...defconfigure_optimizers(self,weight_decay,learning_rate,device):param_dict={pn:pforpn,pinself.named_parameters()}param_dict={pn:pforpn,pinparam_dict.items()ifp.requires_grad}decay_params=[pforn,pinparam_dict.items()ifp.dim()>=2]nodecay_params=[pforn,pinparam_dict.items()ifp.dim()<2]optim_groups=[{'params':decay_params,'weight_decay':weight_decay},{'params':nodecay_params,'weight_decay':0.0}]num_decay_params=sum(p.numel()forpindecay_params)num_nodecay_params=sum(p.numel()forpinnodecay_params)print(f"num decayed parameter tensors:{len(decay_params)}, with{num_decay_params:,}parameters")print(f"num non-decayed parameter tensors:{len(nodecay_params)}, with{num_nodecay_params:,}parameters")fused_available='fused'ininspect.signature(torch.optim.AdamW).parameters use_fused=fused_availableand'cuda'indeviceprint(f"using fused AdamW:{use_fused}")optimizer=torch.optim.AdamW(optim_groups,lr=learning_rate,betas=(0.9,0.95),eps=1e-8,fused=use_fused)returnoptimizer

另一个 Karpathy 指出的项目是用于 AdamW 的核的使用。就像我们上次讨论的那样，融合核是一种优化 GPU 内存移动和计算效率的方法，以便在相关数据位于 SRAM 时进行更多计算。通过使用 AdamW 的已使用版本，我们在代码中获得了重大的加速。

梯度累积

返回到我们的 GPT-3 表格，我们应该注意，在这个表格中，我们还没有硬编码的唯一变量是批量大小。在这篇论文中，他们使用了 500,000 个标记的批量大小。如果我们天真地处理这么多，考虑到我们的最大序列长度（T）等于 1024，我们需要有 488 个 B 值（500,000/1024 ≈ 488）。即使使用今天的优质 GPU，我们也没有足够的内存来处理这个问题。

https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/c380af8cb10299004716e52dd89776fe.png

表 2.1 来自"Language Models are Few-Shot Learners

尽管如此，我们确实希望使用确切的批量大小，因此我们转向梯度累积来解决此问题。这里的想法是进行许多“微批次”，然后将它们的梯度相加，以获得运行一个常规批次的等效结果。

注意，为了避免对我们梯度计算的数学干扰，我们必须将我们的损失按我们将要采取的总梯度步数进行缩放（损失 = 损失 / grad_accum_steps）。

loss_accum=0.0formicro_stepinrange(grad_accum_steps):x,y=train_loader.next_batch()x,y=x.to(device),y.to(device)withtorch.autocast(device_type=device,dtype=torch.bfloat16):logits,loss=model(x,y)loss=loss/grad_accum_steps loss_accum+=loss.detach()loss.backward()

分布式数据并行化

这是一个变革性的改变。在我们做出这个改变之前，我们的代码只在单个 GPU 上运行。通过我们即将实施的分布式数据并行化（ddp），它现在将开始使用我们配置的所有 8 个 A100 GPU。自然地，这也是我们代码复杂性的大部分来源。

Rank vs Local Rank

我们第一次重大的改变是将必要的变量设置为使用 ddp。在我们的操作系统环境中，有一个名为 RANK 的环境变量。你不会在我们的代码中看到它被设置。相反，它是在我们通过命令torchrun - standalone - nproc_per_node=8 train_gpt2.py运行我们的程序时设置的。这个命令告诉我们，我们将在集群中的每个节点上运行 8 个进程。

ddp=int(os.environ.get('RANK',-1))!=-1# is this a ddp run?ifddp:# use of DDP atm demands CUDA, we set the device appropriately according to rankasserttorch.cuda.is_available(),"for now i think we need CUDA for DDP"init_process_group(backend='nccl')ddp_rank=int(os.environ['RANK'])ddp_local_rank=int(os.environ['LOCAL_RANK'])ddp_world_size=int(os.environ['WORLD_SIZE'])device=f'cuda:{ddp_local_rank}'torch.cuda.set_device(device)master_process=ddp_rank==0# this process will do logging, checkpointing etc.else:# vanilla, non-DDP runddp_rank=0ddp_local_rank=0ddp_world_size=1master_process=True# attempt to autodetect devicedevice="cpu"iftorch.cuda.is_available():device="cuda"elifhasattr(torch.backends,"mps")andtorch.backends.mps.is_available():device="mps"print(f"using device:{device}")

要开始遍历我们的代码，我们确定我们是否在运行‘ddp’，如果我们的 RANK 不等于-1（未设置时的默认值）。一旦我们知道我们正在使用 ddp，我们就告诉 PyTorch 使用 nccl（Nvidia 集体通信库）来处理 GPU 之间的连接。这有助于我们在之后同步梯度累加。然后，我们捕获与我们相关的环境变量。WORLD_SIZE告诉我们有多少 GPU 参与（在我们的情况下是 8）。RANK帮助我们识别在整个作业中运行的每个进程，而LOCAL_RANK帮助我们识别在特定节点上运行的进程。

下面的例子有助于看到RANK和LOCAL_RANK之间的区别：

https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/d70a6a7b726e57dbccced41f283f3d58.png

作者图片 – 两个节点之间的 Rank 与 Local Rank 对比

我们任意地将RANK设置为 0 分配给主进程，该进程将处理我们训练中的所有累加和检查点。

将数据加载到进程

现在我们正在运行并行进程，我们需要确保每个进程都有其独特的数据进行训练。为此，我们正在调整我们的next_batch()函数，以便我们用来为每个批次预留标记的current_position变量对于每个运行的进程都是独特的。

classDataLoaderLite:def__init__(self,B,T,process_rank,num_processes):self.B=B self.T=T self.process_rank=process_rank self.num_processes=num_processeswithopen('shakespeare.txt',"r")asf:text=f.read()enc=tiktoken.get_encoding('gpt2')tokens=enc.encode(text)self.tokens=torch.tensor(tokens)ifmaster_process:print(f"loaded{len(self.tokens)}tokens")self.current_position=self.B*self.T*self.process_rank#change heredefnext_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)y=(buf[1:]).view(B,T)self.current_position+=B*T*self.num_processesifself.current_position+(B*T*self.num_processes+1)>len(self.tokens):self.current_position=self.B*self.T*self.process_rank# change herereturnx,y

在我们的训练循环中，我们需要小心避免我们的累加器和梯度发生扭曲。我们首先确保只有在批次的最后一个微步时才运行反向梯度。

ifddp:model.require_backward_grad_sync=(micro_step==grad_accum_steps-1)loss.backward()

批次完成后，我们运行以下减少操作以平均化所有找到的损失。这有助于确保进程不会失去同步，并且它们的损失值与其他进程适当地计算。

ifddp:dist.all_reduce(loss_accum,op=dist.ReduceOp.AVG)

这样，我们已经准备好了我们的代码以并行运行！大部分的主要复杂性都是由nccl处理的，但在整个训练循环中，我们需要小心不要意外地给我们的计算带来副作用。

训练语料库改进

数据集

到目前为止，我们一直在使用莎士比亚文本进行训练。如果机器学习有一个最重要的真理，那就是：数据质量至关重要。因此，我们将升级我们的数据集以升级我们的模型。

卡帕西指出了几个高质量的集合。红睡衣是通过尝试重建 Meta 用于训练 Llama1 的数据集而创建的。它从维基百科、GitHub、ArXiv、StackExchange、清洗后的 Common Crawl (C4)和 Common Crawl 本身中抽取数据。还进行了大量的去重工作。

瘦睡衣对红睡衣进行了进一步的清洗和去重，从而降低了总标记数。虽然更多的数据通常是好的，但你能使数据质量越高越好。

考虑到这一点，我们还有FineWeb和FineWeb Edu。FineWeb 有 1500 万亿个标记长，并且有详细的文档说明其作者是如何组装它的。FineWeb Edu 是将 FineWeb 精炼成只有高教育价值的示例——虽然这是一个主观的判断，但作者仍然指出 FineWeb Edu 优于其同类。我们将使用 FineWeb EDU，它使用 Open Data Commons License。

添加分块

尽管 FineWeb Edu 比 FineWeb 小，但我们使用的版本仍然非常庞大——10 亿个标记。自然地，我们无法像加载 Shakespeare.txt 那样简单地一次性将这么多标记加载到内存中。相反，我们将通过分块加载标记。

classDataLoaderLite:def__init__(self,B,T,process_rank,num_processes):self.B=B self.T=T self.process_rank=process_rank self.num_processes=num_processesassertsplitin{'train','val'}# get the shard filenamesdata_root="edu_fineweb10B"shards=os.listdir(data_root)shards=[sforsinshardsifsplitins]shards=sorted(shards)shards=[os.path.join(data_root,s)forsinshards]self.shards=shardsassertlen(shards)>0,f"no shards found for split{split}"ifmaster_process:print(f"found{len(shards)}shards for split{split}")self.current_shard=0self.tokens=load_tokens(self.shards[self.current_shard])self.current_position=self.B*self.T*self.process_rankdefnext_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)y=(buf[1:]).view(B,T)self.current_position+=B*T*self.num_processesifself.current_position+(B*T*self.num_processes+1)>len(self.tokens):self.current_shard=(self.current_shard+1)%len(self.shards)self.tokens=load_tokens(self.shards[self.current_shard])self.current_position=B*T*self.process_rankreturnx,y

分块是数据语料库的一部分。我们经常在数据库术语中听到它们的讨论，但这个术语在这里也适用，因为我们正在加载标记。对于上述内容，我们正在修改next_batch，以便我们加载的标记来自正确的分块。除此之外，我们在分块内的位置逻辑保持不变。

加载标记

defload_tokens(filename):npt=np.load(filename)npt=npt.astype(np.int32)ptt=torch.tensor(npt,dtype=torch.long)returnptt

最后，我们创建了一个load_tokens函数，它实际上会遵循我们为 Shakespeare.txt 所做的方法，只是我们告诉numpy以整数 32 位加载标记，以便torch.compile()可以优化。

这里的大致情况是我们正在限制加载到分块中的内存量。如果我们一次性加载所有内容，就会耗尽内存，所以我们一次只加载一部分。

在开始训练运行之前，我们必须确保我们有我们的数据集，因此我们将运行python fineweb.py来拉取所有数据。这花了我相当多的时间，但你的速度可能会根据你的连接速度而有所不同。

验证循环

就像所有机器学习一样，我们想要一种方法来了解模型性能接近最优的程度。虽然损失是一个重要的度量，但如果没有在之前未见过的数据上测试模型，我们就无法知道我们是否过拟合。

为了避免这个问题，下面的代码将在之前未见过的数据上测试模型，并计算其损失。

ifstep%100==0:model.eval()val_loader.reset()withtorch.no_grad():val_loss_accum=0.0val_loss_steps=20for_inrange(val_loss_steps):x,y=val_loader.next_batch()x,y=x.to(device),y.to(device)withtorch.autocast(device_type=device,dtype=torch.bfloat16):logits,loss=model(x,y)loss=loss/val_loss_steps val_loss_accum+=loss.detach()ifddp:dist.all_reduce(val_loss_accum,op=dist.ReduceOp.AVG)ifmaster_process:print(f"validation loss:{val_loss_accum.item():.4f}")# training loopmodel.train()

在这里，我们运行了 20 步的验证，并取其平均值。记住，我们使用的数据集与训练数据集不同，所以我们进行 20 次是为了避免一个糟糕的批次给我们带来糟糕的验证值。

在这里，最终的结果是我们应该通过训练来获得模型泛化程度的感觉。

HellaSwag 评估

我们确定性能的另一种方式是在特定的测试数据集上运行。HellaSwag 就是这样一种数据集，它曾经被用于基准测试。最近，最先进模型在数据集上的准确率达到了~95%，这使得仅根据数据集本身很难真正区分质量。尽管如此，在 GPT-2 时代，这是一个难以攻克的难题，所以我们将在这里用它来衡量性能。

尽管评估 HellaSwag 并不像计算熵损失那样简单。HellaSwag 要求 LLM 从 4 种不同的可能性中选择一种来完成一个句子（想想多项选择）。我们的模型实际上无法做到这一点，因此我们不得不在模型“选择”选项的方式上发挥创意。

if(step%250==0orlast_step)and(notuse_compile):num_correct_norm=0num_total=0fori,exampleinenumerate(iterate_examples("val")):# only process examples for select processesifi%ddp_world_size!=ddp_rank:continue# render the example into tokens and labels_,tokens,mask,label=render_example(example)tokens=tokens.to(device)mask=mask.to(device)# get the logitswithtorch.no_grad():withtorch.autocast(device_type=device_type,dtype=torch.bfloat16):logits,loss=model(tokens)pred_norm=get_most_likely_row(tokens,mask,logits)num_total+=1num_correct_norm+=int(pred_norm==label)# reduce the stats across all processesifddp:num_total=torch.tensor(num_total,dtype=torch.long,device=device)num_correct_norm=torch.tensor(num_correct_norm,dtype=torch.long,device=device)dist.all_reduce(num_total,op=dist.ReduceOp.SUM)dist.all_reduce(num_correct_norm,op=dist.ReduceOp.SUM)num_total=num_total.item()num_correct_norm=num_correct_norm.item()acc_norm=num_correct_norm/num_totalifmaster_process:print(f"HellaSwag accuracy:{num_correct_norm}/{num_total}={acc_norm:.4f}")withopen(log_file,"a")asf:f.write(f"{step}hella{acc_norm:.4f}n")

在上述代码中，模型会拉取每个可能的答案，然后从每个答案中获取 logits。哪个答案的 logits 概率最高，那就是模型“选择”的答案，然后我们可以评估其性能。

这段代码通过利用模型训练循环的批处理特性来工作。我们从每个示例中获取标记、一个掩码，它让我们知道应该从哪里生成新的标记，以及一个标签，这样我们就可以知道每个示例的正确答案。我们从每个选项中获取 logits，然后使用这些 logits 为每个选项获取一个概率。我们的函数get_most_likely_row将输出模型“选择”的行，然后我们可以将其与标签进行比较，以评估模型是否选择了正确答案。

我们只是任意地每 250 步做一次。显然，每一步都进行评估将会非常昂贵，而只在最后进行评估则不会让你看到模型在训练过程中的进展。因此，选择一个好的节奏来运行这个操作是很困难的。

Torchrun

就这样！我们现在准备好训练了！为了利用我们的数据并行化工作，我们将使用torchrun命令行工具：

torchrun --standalone --nproc_per_node=8 train_gpt2.py

为了解释这个命令，standalone表示我们将在一个节点上运行训练过程。这很合理，因为我们只使用 1 个 CPU 而不是集群。如果这是在集群上运行，我们需要提供一个外部进程来协调训练。nproc_per_node告诉系统每个节点上有 8 个进程要运行，对应于我们连接的 8 个 GPU，这意味着我们的 LOCAL_RANK 将从 0 运行到 7，同时我们的 RANK 也在相同的范围内运行。如果我们有一个集群，那么这些数字的范围就不会相同。

训练时间估计

我使用每个具有 40GB 内存的 8 个 A100s 训练了模型。由于可用的内存较少，我不得不将批大小从视频中使用的 64 减少到适合我 GPU 的 16。

https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/cad75c00d7e4af5f9382c6158e3c587d.png

图片由作者提供 – NVIDIA 训练期间内存使用情况截图

如上图所示，批大小为 16 的批次并没有完全使用 GPU 的内存，然而，32 的批大小刚刚好导致 CUDAOutOfMemoryError，因此我别无选择，只能在这里低效使用我的 GPU。对于未来的优化，这是一个关键改进领域。

以每步大约半秒的速度进行，我的模型大约在两个半小时后完成训练。

结果

运行后，我们可以将损失和 HellaSwag 评估与 1.24 亿参数的 OpenAI 模型进行比较。

观察我们的损失，我们看到一个漂亮的曲线，显示了损失随时间逐渐降低 – 我们的曲线在大约 6500 步时低于 OpenAI GPT-2 的损失。显然，OpenAI 在当时拥有高质量的硬件和良好的优化，这里的关键区别可能是我们手头的数据更好。这是展示高质量数据如何产生差异的绝佳方式。

https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/64804252ff7e60bf89548ac11ed21930.png

图片由作者提供 – 与 OpenAI 的训练相比，我们的训练损失

接下来，对于 HellaSwag 评估，我们看到我们的模型比 GPT-2 表现得好，但并不比 GPT-3 好得多。记住，HellaSwag 是一个多项选择题，预期结果是 0.25。我们大约 0.31 的结果是好的，但当然远不及现代 SoTA 模型所展现的 0.95 的性能。为了提高我们模型在 HellaSwag 上的性能，我们可以从向我们的语料库添加更多标记开始 – GPT-3 在 3000 亿个标记上进行了训练，这比我们的数据多出大约 30 倍！

https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/a52062da819dbc58434df3da16e365ed.png

图片由作者提供 – HellaSwag 评估与 OpenAI 训练步骤对比