LLama 3 405B 技术报告解读
Meta于2024年7月24日发布了Llama 3的405B版本，这一模型在技术报告中引起了广泛关注。报告开头特别强调了一个关键概念：“Managing complexity”，即管控复杂度。Meta没有使用Mixture of Experts (MoE)架构，而是选择了405B参数的密集模型（dense model），并且在优化过程中采用了Direct Preference Optimization (DPO)而非Proximal Policy Optimization (PPO)。下面将详细解读这份技术报告，涵盖预训练、模型结构、规模定律、硬件和网络、并行策略等方面的亮点。
一、预训练
数据部分
- 数据清洗：对个人隐私数据（PII）进行清洗、去重和去黄。
- 数据配比：通过模型进行细粒度打标签，最终确定了50%通用数据，25%数理数据，17%代码数据，8%多语言数据。
- 数据退火：在大模型训练的最后阶段，使用高质量数据学习并逐渐将学习率衰减到0。
数据配比和Scaling Law
- 数据配比：作者在不同的小模型上进行了不同的配比实验，用于预测大模型的最优配比。
- Scaling Law：通过在小模型上探索数据配比的Scaling Law，预测了大模型的最优配比。
二、模型结构
GQA和8-kv head
Llama 3使用了分组查询注意力（GQA）和8个键值对头，提高了推理速度并减少了解码过程中的键值对缓存大小。
长文拼接
使用注意力掩码防止不同来源数据串味，这在扩长序列时尤为重要。
词表和RoPE
- 词表大小：128K，增强了非英语能力。
- RoPE theta：调整到500000，支持更长的上下文窗口。
三、Scaling Law
两步方法
- 计算最优模型：在下游任务上的负对数似然与训练FLOPs之间的相关性。
- Benchmark表现：利用Scaling Law模型和使用更高计算FLOPs训练的旧模型，将负对数似然与Benchmark的准确率关联。
结果
在ARC Challenge基准上，拟合结果表现良好。
四、硬件和网络
GPU资源
16K H100 80GB with NVLink，采用Meta的生产集群，提供240 PB SSD和7500台机器，支持2TB-7TB/s的存储吞吐。
网络部分
- RoCE和400Gb/s网络接口：采用3层网络架构。
- 负载均衡和拥塞控制：使用16个流减少单流流量，在网络包头部增加特殊区域，通过hash使得流的选路更均衡。
五、并行策略
4D并行
结合张量并行（TP）、流水线并行（PP）、上下文并行（CP）和数据并行（DP），多了一个上下文并行CP或称为序列并行。
FSDP
使用全分片数据并行（FSDP），优化模型权重的拉取次数，减少反向梯度计算时的通信。
PP并行策略改进
- Batch Size限制：当前的流水线并行策略限制了micro batch个数为流水线stage的整数倍。
- 显存和计算不均衡：首尾stage做了padding，以平衡显存和计算。
CP并行策略的改进
基于allgather的通信方法，支持不同类型的注意力掩码。
六、预训练策略
三阶段训练法
- 初始训练：余弦调度，逐步增加上下文长度和batch size。
- 长上下文训练：从8K逐步增加到128K。
- 退火训练：最后40M token，用128K长度，逐渐线性缩减学习率到0。
七、后训练（对齐）
对话数据格式
设置多种角色的输入输出格式和特殊token。
奖励建模
增加了人工编辑样本，进行更精细的训练和权重平均。
SFT和DPO
拒绝采样，使用DPO+NLL的loss进行优化，迭代式训练，增强数据选择和训练效果。
八、基础设施、扩展性和效率
训练基础设施
Meta的生产集群提供了大规模的计算资源和存储网络，支持高效的模型训练和调度。
模型扩展的并行性
通过4D并行性将计算分布到多个GPU上，实现高效的模型训练。
总结
Llama 3 405B版本在数据处理、模型结构、并行策略和训练基础设施等方面进行了全面优化。尽管其性能未达到压倒性的领先地位，但在各个方面都展示了许多值得学习的技术亮点。特别是对数据的细粒度处理和大规模并行训练策略，对于提升模型性能和效率具有重要意义。
