从零预训练LLAMA3的完整指南：一个文件，探索Scaling Law

最新推荐文章于 2024-09-27 21:16:51 发布

AI学习不迷路

最新推荐文章于 2024-09-27 21:16:51 发布

阅读量1.1k

点赞数 8

文章标签：深度学习人工智能自然语言处理大模型 Llama3 AI大模型大语言模型

本文链接：https://blog.csdn.net/qkh1234567/article/details/140989199

版权

引言

最近，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

模型加载

hellaswag 评估

我们可以看到，原始LLaMA3-8B模型的得分是70分。这一结果与官方公开的模型结果基本一致：

https://paperswithcode.com/sota/sentence-completion-on-hellaswag

有了这一基准，我们就可以开始从零实现自己的LLaMA模型了。

RMSNorm层

RMSNorm（Root Mean Square Normalization）是一种归一化方法，用于稳定和加速训练过程。

• 输入：

一个形状为(batch_size, sequence_length, hidden_dim)的张量，表示输入的隐藏状态。

• 输出：

一个形状与输入相同的张量，但每个隐藏状态经过归一化处理。

• 作用：

RMSNorm通过计算输入张量每个元素的均方根值，并使用它来归一化输入。这可以有效地控制每层的输出范围，防止梯度爆炸或消失。

RotaryEmbedding

RotaryEmbedding是一种位置编码技术，用于引入位置信息。

• 输入：

一个形状为(batch_size, sequence_length, hidden_dim)的张量，表示输入的查询（q）或键（k）向量。

• 输出：

一个形状与输入相同的张量，但每个向量经过旋转编码处理。

• 作用：

通过生成与位置ID对应的旋转向量，RotaryEmbedding可以在不显式使用位置编码的情况下，引入位置信息，使模型能够更好地理解序列中的位置信息。

RotaryEmbedding本质上是生成一组与position ID对应的旋转向量，代表了不同位置的数据应该怎么旋转。

Attention Layer

Attention Layer用于计算输入序列中各个元素之间的相关性，并根据相关性对输入进行加权。

• 输入：

一个形状为(batch_size, sequence_length, hidden_dim)的张量，表示输入的隐藏状态。

• 输出：

一个形状为(batch_size, sequence_length, hidden_dim)的张量，表示经过注意力机制处理后的输出。

计算过程：

• 投影：将输入数据分别使用q_proj、k_proj、v_proj进行线性投影，生成查询、键、值向量。
• 旋转投影：对查询和键进行旋转投影，引入位置信息。
• 计算Attention值：通过查询向量和键向量的点积（q @ k.transpose()）计算Attention值。
• 加权求和：使用Attention值对值向量进行加权求和。
• 线性投影：最后通过o_proj进行线性投影输出结果。

MLP

MLP层（多层感知机）用于对输入进行非线性变换。

• 输入：

一个形状为(batch_size, sequence_length, hidden_dim)的张量，表示经过Attention处理后的隐藏状态。

• 输出：

一个形状为(batch_size, sequence_length, hidden_dim)的张量，表示经过非线性变换后的输出。

• 作用：

MLP层通常包含两个线性投影和一个非线性激活函数（如SiLU），用于增加模型的表达能力。具体步骤是：

• 输入通过第一个线性投影，生成一个中间表示。
• 中间表示经过SiLU激活函数。
• 激活后的表示通过第二个线性投影，生成最终输出。

Block组装

完成一次组装，这里的计算过程包含了一次residual connection目的是为了提高梯度传播的质量和效率。

整体模型结构

最终的LLaMA3模型结构包含：

• 输入embedding层，将输入词汇转换为隐藏状态。
• 多个重复的Block，每个Block包含RMSNorm层、RotaryEmbedding层、Attention Layer和MLP层。
• 输出线性层，将隐藏状态转换为输出。

通过这种结构，模型能够逐层处理输入数据，并逐步生成最终的预测结果。

整个过程之后，我们需要验证模型构建的正确性。我们将Huggingface的LLaMA3模型的所有参数复制到我们的模型中，然后通过HellaSwag评估我们自己的LLaMA3模型是否能够达到相同的效果。

在我们的llama3模型初始化时，copy huggingface原始llama3模型的权重。

通过评估，我们发现得分与Huggingface模型完全一致！这证明我们自己的实现是正确的。

从零开始预训练

为了能够预训练模型，我们使用了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会自动分配对应的环境变量给每一个训练子进程。

在加载模型后，使用DDP封装模型

在训练阶段，我们可以按照单一模型的训练方式构建代码。在获得loss之后，使用loss.backward()积累梯度。这里的require_backward_grad_sync本质上是在梯度积累步长不为1时，通过在前几步中减少跨GPU的梯度同步，来优化训练速度。

之后，依旧正常利用optimizer迭代模型权重。使用torch.cuda.synchronize()的目的是为了同步并行的GPU，使每次iteration时所有GPU都完成这一轮的训练。

由于使用了多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，以此类推。

在这样的结构下，我们就可以开始训练了。当然，不要忘记提前通过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仍有差距。

对于我们自己的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。

同等规模LLAMA3架构训练结果

进一步测试不同大小模型的效果可以让我们浅窥Scaling Law。下图显示了不同模型参数量在相同训练配置下的HellaSwag效果。可以看到，当模型参数量较小时，架构细节并不太重要；参数量大的模型训练效果更好。在参数量接近的情况下，更合理的架构可能会带来细微的差别。然而，当前行业趋势还是倾向于通过更大规模的模型提升LLM的能力。我也认为，在小规模模型上过度优化架构不如寻找更基础的研究方向。

总结

在这次探索中，我们从零开始预训练了一个语言模型。尽管本文中涉及的知识点还有很多没有详细讲解的地方，但我会在之后推出更多相关的实战文章，以补充这些内容。

• DDP训练与FSDP和TP：我们讨论了如何使用DDP进行训练。我简单的调整了DDP的llama训练代码，有一个FSDP实现[1]。FSDP会通过模型分片极大程度地减少每个GPU内存占用。此外，还有TP（Tensor Parallelism）等大规模并发训练结构，这些方法都可以显著提升训练效率和模型规模。
• MoE模型和多模态模型：目前在MoE（Mixture of Experts）模型和多模态模型方面有许多新进展，这也是一个非常有趣的讨论方向。我们可以通过实践进一步探索这些模型的应用和优化。

如何学习AI大模型？

我在一线互联网企业工作十余年里，指导过不少同行后辈。帮助很多人得到了学习和成长。

我意识到有很多经验和知识值得分享给大家，也可以通过我们的能力和经验解答大家在人工智能学习中的很多困惑，所以在工作繁忙的情况下还是坚持各种整理和分享。但苦于知识传播途径有限，很多互联网行业朋友无法获得正确的资料得到学习提升，故此将并将重要的AI大模型资料包括AI大模型入门学习思维导图、精品AI大模型学习书籍手册、视频教程、实战学习等录播视频免费分享出来。

在这里插入图片描述