作者:Mantavers,AGI独角兽
声明:本文只做分享,版权归原作者,侵权私信删除!
原文:https://zhuanlan.zhihu.com/p/706097271编辑:青稞AI
引言
最近,Andrew大神发布了一个全新的视频教程,讲解了从零开始预训练GPT-2的全过程。这个四小时的视频详细介绍了模型的构建、训练数据的加载、评估方法以及在分布式框架下的DDP训练。受到此视频的启发,我决定使用LLaMA3架构,从零开始预训练一个大型语言模型,并对比不同模型参数下模型能力的提升。本文将开源所有相关代码在:
https://github.com/hengjiUSTC/learn-llm/tree/main/pretrain
接下来让我们进入正题。
模型构建和评估
为了能够有一个对照效果,同时保证我们之后自己从零实现的LLaMA模型的正确性,我们首先通过加载Huggingface的官方LLaMA3模型,对模型进行HellaSwag评估。
代码开源在:
https://github.com/hengjiUSTC/learn-llm/blob/main/pretrain/play_with_llama.ipynb
模型加载
![null b957d5df8607673db864e9e3a82c97b1.jpeg](https://img-blog.csdnimg.cn/img_convert/b957d5df8607673db864e9e3a82c97b1.jpeg)
hellaswag 评估
![null 2f477748ac998fa291a0163f86814b2c.jpeg](https://img-blog.csdnimg.cn/img_convert/2f477748ac998fa291a0163f86814b2c.jpeg)
我们可以看到,原始LLaMA3-8B模型的得分是70分。这一结果与官方公开的模型结果基本一致:
https://paperswithcode.com/sota/sentence-completion-on-hellaswag
![null a4df3f126514c9a57cf368de48b8aaf2.jpeg](https://img-blog.csdnimg.cn/img_convert/a4df3f126514c9a57cf368de48b8aaf2.jpeg)
有了这一基准,我们就可以开始从零实现自己的LLaMA模型了。
![null b51dfba6e27a8fe0c2dface1782dc7d7.jpeg](https://img-blog.csdnimg.cn/img_convert/b51dfba6e27a8fe0c2dface1782dc7d7.jpeg)
RMSNorm层
RMSNorm(Root Mean Square Normalization)是一种归一化方法,用于稳定和加速训练过程。
• 输入:
一个形状为(batch_size, sequence_length, hidden_dim)的张量,表示输入的隐藏状态。
• 输出:
一个形状与输入相同的张量,但每个隐藏状态经过归一化处理。
• 作用:
RMSNorm通过计算输入张量每个元素的均方根值,并使用它来归一化输入。这可以有效地控制每层的输出范围,防止梯度爆炸或消失。
![null dee1113c8fbe424798c41601b54c2dd1.jpeg](https://img-blog.csdnimg.cn/img_convert/dee1113c8fbe424798c41601b54c2dd1.jpeg)
RotaryEmbedding
RotaryEmbedding是一种位置编码技术,用于引入位置信息。
• 输入:
一个形状为(batch_size, sequence_length, hidden_dim)的张量,表示输入的查询(q)或键(k)向量。
• 输出:
一个形状与输入相同的张量,但每个向量经过旋转编码处理。
• 作用:
通过生成与位置ID对应的旋转向量,RotaryEmbedding可以在不显式使用位置编码的情况下,引入位置信息,使模型能够更好地理解序列中的位置信息。
RotaryEmbedding本质上是生成一组与position ID对应的旋转向量,代表了不同位置的数据应该怎么旋转。
![null 7ee2aca304b6d32ad2e19714e5468721.jpeg](https://img-blog.csdnimg.cn/img_convert/7ee2aca304b6d32ad2e19714e5468721.jpeg)
Attention Layer
Attention Layer用于计算输入序列中各个元素之间的相关性,并根据相关性对输入进行加权。
• 输入:
一个形状为(batch_size, sequence_length, hidden_dim)的张量,表示输入的隐藏状态。
• 输出:
一个形状为(batch_size, sequence_length, hidden_dim)的张量,表示经过注意力机制处理后的输出。
![null 4d35e2221e3b39e50832f31868f545f5.jpeg](https://img-blog.csdnimg.cn/img_convert/4d35e2221e3b39e50832f31868f545f5.jpeg)
计算过程:
• 投影:将输入数据分别使用q_proj、k_proj、v_proj进行线性投影,生成查询、键、值向量。
• 旋转投影:对查询和键进行旋转投影,引入位置信息。
• 计算Attention值:通过查询向量和键向量的点积(q @ k.transpose())计算Attention值。
• 加权求和:使用Attention值对值向量进行加权求和。
• 线性投影:最后通过o_proj进行线性投影输出结果。
![null 911b5379109109e2628f1fd5f3b17d92.jpeg](https://img-blog.csdnimg.cn/img_convert/911b5379109109e2628f1fd5f3b17d92.jpeg)
MLP
MLP层(多层感知机)用于对输入进行非线性变换。
• 输入:
一个形状为(batch_size, sequence_length, hidden_dim)的张量,表示经过Attention处理后的隐藏状态。
• 输出:
一个形状为(batch_size, sequence_length, hidden_dim)的张量,表示经过非线性变换后的输出。
• 作用:
MLP层通常包含两个线性投影和一个非线性激活函数(如SiLU),用于增加模型的表达能力。具体步骤是:
• 输入通过第一个线性投影,生成一个中间表示。
• 中间表示经过SiLU激活函数。
• 激活后的表示通过第二个线性投影,生成最终输出。
![null 8cb224cde1bda892bb55a4aa2dd44038.jpeg](https://img-blog.csdnimg.cn/img_convert/8cb224cde1bda892bb55a4aa2dd44038.jpeg)
Block组装
完成一次组装,这里的计算过程包含了一次residual connection目的是为了提高梯度传播的质量和效率。
![null b0ca502c2075dd1fb6567912247b8263.jpeg](https://img-blog.csdnimg.cn/img_convert/b0ca502c2075dd1fb6567912247b8263.jpeg)
整体模型结构
最终的LLaMA3模型结构包含:
• 输入embedding层,将输入词汇转换为隐藏状态。
• 多个重复的Block,每个Block包含RMSNorm层、RotaryEmbedding层、Attention Layer和MLP层。
• 输出线性层,将隐藏状态转换为输出。
通过这种结构,模型能够逐层处理输入数据,并逐步生成最终的预测结果。
![null 971358989b3fe5d71665071719e7db17.jpeg](https://img-blog.csdnimg.cn/img_convert/971358989b3fe5d71665071719e7db17.jpeg)
整个过程之后,我们需要验证模型构建的正确性。我们将Huggingface的LLaMA3模型的所有参数复制到我们的模型中,然后通过HellaSwag评估我们自己的LLaMA3模型是否能够达到相同的效果。
在我们的llama3模型初始化时,copy huggingface原始llama3模型的权重。
![null e3a24c6572476948129031d9965ac32a.jpeg](https://img-blog.csdnimg.cn/img_convert/e3a24c6572476948129031d9965ac32a.jpeg)
通过评估,我们发现得分与Huggingface模型完全一致!这证明我们自己的实现是正确的。
![null 25625a31ec2c27cebd17d26089095d2f.jpeg](https://img-blog.csdnimg.cn/img_convert/25625a31ec2c27cebd17d26089095d2f.jpeg)
从零开始预训练
为了能够预训练模型,我们使用了FineWeb 10TB的数据进行训练,训练时长为一个epoch。为了适应预训练场景,我们创建了自定义的数据Dataloader,并通过DDP(Distributed Data Parallel)加速训练。否则,训练过程将会耗费大量时间。
训练代码开源在:
https://github.com/hengjiUSTC/learn-llm/blob/main/pretrain/train_llama.py
使用多GPU进行并行训练
为了同时使用多张GPU进行并行训练,我们使用了PyTorch提供的DDP模块,将模型加载到多个GPU上并行训练。
DDP(Distributed Data Parallel):DDP是PyTorch中常用的分布式训练模块,用于在多个GPU之间并行训练模型。DDP通过在每个GPU上运行一个训练进程,并在每次梯度计算后进行参数同步,从而实现并行加速训练。DDP解决的问题包括:
• 数据并行:将数据切分成多个batch,分发到不同的GPU上并行处理,提升训练速度。
• 梯度同步:在每次梯度更新后同步各个GPU的参数,保证模型的一致性。
首先,需要初始化DDP环境。针对我们运行代码时
torchrun --standalone --nproc_per_node=4 train_llama.py
指定的机器数量,PyTorch会自动分配对应的环境变量给每一个训练子进程。
![null 6e7f718b97843ce53bc5ed6788f031fb.jpeg](https://img-blog.csdnimg.cn/img_convert/6e7f718b97843ce53bc5ed6788f031fb.jpeg)
在加载模型后,使用DDP封装模型
![null e42a9e78f04b53eaaad1a7fb95c16088.jpeg](https://img-blog.csdnimg.cn/img_convert/e42a9e78f04b53eaaad1a7fb95c16088.jpeg)
在训练阶段,我们可以按照单一模型的训练方式构建代码。在获得loss之后,使用loss.backward()积累梯度。这里的require_backward_grad_sync本质上是在梯度积累步长不为1时,通过在前几步中减少跨GPU的梯度同步,来优化训练速度。
![null 014e9d5df1d86bb57062e6d8020865cf.jpeg](https://img-blog.csdnimg.cn/img_convert/014e9d5df1d86bb57062e6d8020865cf.jpeg)
之后,依旧正常利用optimizer迭代模型权重。使用torch.cuda.synchronize()的目的是为了同步并行的GPU,使每次iteration时所有GPU都完成这一轮的训练。
![null 5b48ba72d75ea8cb111fec99c406626a.jpeg](https://img-blog.csdnimg.cn/img_convert/5b48ba72d75ea8cb111fec99c406626a.jpeg)
由于使用了多GPU训练,需要调整DataLoader让每个GPU获得不一样的训练数据。我们的DataLoader会根据GPU的编号,使用B * T * self.process_rank的方式加载数据。例如,如果我们有4个GPU,第0个DataLoader加载的数据就是第0,4,8,12个chunk,第1个GPU加载的是1,5,9,13个chunk,以此类推。
![null 18c27cd54ce514b5fb9ae671a1dc00d8.jpeg](https://img-blog.csdnimg.cn/img_convert/18c27cd54ce514b5fb9ae671a1dc00d8.jpeg)
在这样的结构下,我们就可以开始训练了。当然,不要忘记提前通过fineweb.py脚本生成我们需要的10TB训练数据。
训练结果
在对比中,我们首先与Andrew的GPT-2 demo进行对比。在相同的模型配置下(12层,12个head),经过一个epoch的FineWeb数据训练,andrew的GPT-2模型达到了2.9的training loss,超过了OpenAI原始版本的GPT-2。同时,在HellaSwag评测任务中,模型得分为0.305,也超过了GPT-2的0.294,但尚未达到GPT-3的0.337。我们认为FineWeb数据集起到了重要作用,它的质量已经超过了GPT-2的训练数据,但与GPT-3仍有差距。
![null 91afacd45fc178b542b8afd4c65c90f3.jpeg](https://img-blog.csdnimg.cn/img_convert/91afacd45fc178b542b8afd4c65c90f3.jpeg)
对于我们自己的LLaMA架构模型,训练结果如下:在相同的12个layer、12个head配置下,经过一个epoch的训练,LLaMA3架构在得分上超越了andrew的GPT-2,训练loss为0.269(Andrew的gpt2为0.29),hellaswag得分为0.33(Andrew的gpt2为0.305),效果更加接近GPT-3。我们认为这主要是因为FineWeb数据集的质量虽不如GPT-3,但LLaMA架构的优势在于其使用了Rotary Position Embedding,而GPT-2使用的是Learnable Position Embedding。
![null 547f066947b2e80055726186d09ad84b.jpeg](https://img-blog.csdnimg.cn/img_convert/547f066947b2e80055726186d09ad84b.jpeg)
进一步测试不同大小模型的效果可以让我们浅窥Scaling Law。下图显示了不同模型参数量在相同训练配置下的HellaSwag效果。可以看到,当模型参数量较小时,架构细节并不太重要;参数量大的模型训练效果更好。在参数量接近的情况下,更合理的架构可能会带来细微的差别。然而,当前行业趋势还是倾向于通过更大规模的模型提升LLM的能力。我也认为,在小规模模型上过度优化架构不如寻找更基础的研究方向。
![null 0e84c3a5ee91f0be30c7dbbac58faf57.jpeg](https://img-blog.csdnimg.cn/img_convert/0e84c3a5ee91f0be30c7dbbac58faf57.jpeg)
总结
在这次探索中,我们从零开始预训练了一个语言模型。尽管本文中涉及的知识点还有很多没有详细讲解的地方,但我会在之后推出更多相关的实战文章,以补充这些内容。
• DDP训练与FSDP和TP:我们讨论了如何使用DDP进行训练。我简单的调整了DDP的llama训练代码,有一个FSDP实现[1]。FSDP会通过模型分片极大程度地减少每个GPU内存占用。此外,还有TP(Tensor Parallelism)等大规模并发训练结构,这些方法都可以显著提升训练效率和模型规模。
• MoE模型和多模态模型:目前在MoE(Mixture of Experts)模型和多模态模型方面有许多新进展,这也是一个非常有趣的讨论方向。我们可以通过实践进一步探索这些模型的应用和优化。
引用链接
[1]
FSDP实现: ttps://github.com/hengjiUSTC/learn-llm/blob/main/pretrain/train_llama_fsdp.py
备注:昵称-学校/公司-方向/会议(eg.ACL),进入技术/投稿群
id:DLNLPer,记得备注呦