最近，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。
进一步测试不同大小模型的效果可以让我们浅窥Scaling Law。下图显示了不同模型参数量在相同训练配置下的HellaSwag效果。可以看到，当模型参数量较小时，架构细节并不太重要；参数量大的模型训练效果更好。在参数量接近的情况下，更合理的架构可能会带来细微的差别。然而，当前行业趋势还是倾向于通过更大规模的模型提升LLM的能力。我也认为，在小规模模型上过度优化架构不如寻找更基础的研究方向。
在这次探索中，我们从零开始预训练了一个语言模型。尽管本文中涉及的知识点还有很多没有详细讲解的地方，但我会在之后推出更多相关的实战文章，以补充这些内容。
• DDP训练与FSDP和TP：我们讨论了如何使用DDP进行训练。我简单的调整了DDP的llama训练代码，有一个FSDP实现[1]。FSDP会通过模型分片极大程度地减少每个GPU内存占用。此外，还有TP（Tensor Parallelism）等大规模并发训练结构，这些方法都可以显著提升训练效率和模型规模。
• MoE模型和多模态模型：目前在MoE（Mixture of Experts）模型和多模态模型方面有许多新进展，这也是一个非常有趣的讨论方向。我们可以通过实践进一步探索这些模型的应用和优化。
