使用AWS Trainium进行快速且经济高效的LLaMA 2微调
利用AWS Trainium，实现快速高效的LLaMA 2微调，省时省力
大型语言模型（LLM）已经引起了开发者、科学家、技术专家、企业家以及各个行业的高管们的想象和关注。这些模型可以用于问答、摘要、翻译等多个应用场景，比如客户支持的对话代理、市场营销的内容创作以及编码助手等。
最近，Meta发布了Llama 2，面向研究人员和商业实体，增加了其他LLM模型的列表，包括MosaicML的MPT和Falcon。在本文中，我们将介绍如何在AWS Trainium上对Llama 2进行微调，这是一个专为LLM训练而建的加速器，可以减少训练时间和成本。我们将回顾AWS Neuron SDK提供的微调脚本（使用NeMo Megatron-LM），我们使用的各种配置以及我们观察到的吞吐量结果。
关于Llama 2模型
与之前的Llama 1模型和其他像GPT这样的模型一样，Llama 2采用了Transformer的仅解码器架构。它有三种规模：70亿、130亿和700亿个参数。与Llama 1相比，Llama 2将上下文长度从2000增加到4000，并使用了分组查询注意力（仅适用于700B）。Llama 2的预训练模型是在2万亿个标记上训练的，微调模型已经经过100多万个人类注释的训练。
Llama 2的分布式训练
为了适应具有2000和4000序列长度的Llama 2，我们使用支持数据并行性（DP）、张量并行性（TP）和管道并行性（PP）的Trainium，并使用新的实现功能（如解除单词嵌入、旋转嵌入、RMSNorm和Swiglu激活），我们使用GPT Neuron Megatron-LM的通用脚本来支持Llama 2训练脚本。
我们的高级训练流程如下：对于我们的训练环境，我们使用由SLURM系统管理的多实例集群进行分布式训练和调度，使用NeMo框架。
首先，下载Llama 2模型和训练数据集，并使用Llama 2的分词器进行预处理。例如，要使用RedPajama数据集，请使用以下命令：
wget https://data.together.xyz/redpajama-data-1T/v1.0.0/book/book.jsonlpython nemo/scripts/nlp_language_modeling/preprocess_data_for_megatron.py有关下载模型和预处理脚本参数的详细指南，请参阅下载LlamaV2数据集和分词器。
接下来，编译模型：
sbatch --nodes 4 compile.slurm ./llama_7b.sh模型编译完成后，使用以下已经针对Llama 2进行了最佳配置和超参数优化的脚本启动训练作业（包含在示例代码中）：
sbatch --nodes 4 run.slurm ./llama_7b.sh最后，我们使用TensorBoard监控训练进度：
tensorboard --logdir ./有关完整的示例代码和我们提到的脚本，请参阅Llama 7B教程和NeMo代码中的Neuron SDK，详细了解更多步骤。
微调实验
我们在 Neuron 2.12 环境（PyTorch）中，对 OSCAR（Open Super-large Crawled ALMAnaCH coRpus）和 QNLI（Question-answering NLI）数据集上的 7B 模型进行了微调。对于每个 2,000 和 4,000 序列长度，我们优化了一些配置，如 batchsize 和 gradient_accumulation，以提高训练效率。作为微调策略，我们采用了对所有参数进行完全微调（约 500 步），这可以扩展到更长步数和更大数据集的预训练（例如，1T RedPajama）。序列并行性也可启用，以允许 NeMo Megatron 成功微调 4,000 的更长序列长度的模型。下表显示了 Llama 7B 微调实验的配置和吞吐量结果。吞吐量几乎按照实例数量的增加呈线性扩展，直到 4 个为止。
| 分布式库 | 数据集 | 序列长度 | 实例数量 | 张量并行 | 数据并行 | 管道并行 | 全局批量大小 | 吞吐量（seq/s） |
| Neuron NeMo Megatron | OSCAR | 4096 | 1 | 8 | 4 | 1 | 256 | 3.7 |
| . | . | 4096 | 2 | 8 | 4 | 1 | 256 | 7.4 |
| . | . | 4096 | 4 | 8 | 4 | 1 | 256 | 14.6 |
| . | QNLI | 4096 | 4 | 8 | 4 | 1 | 256 | 14.1 |
最后一步是使用基准模型验证准确性。我们为 GPU 实验实施了一个参考脚本，并确认 GPU 和 Trainium 的训练曲线匹配，如下图所示。该图显示了在 QNLI 数据集上的训练步数和损失曲线。GPU 采用了混合精度（蓝色），Trainium 采用了 bf16 和默认随机舍入（橙色）。
结论
在本篇文章中，我们展示了Trainium在Llama 2的高性能和经济高效的微调能力。对于使用Trainium进行分布式预训练和微调生成型AI模型的更多资源，请参考AWS Neuron参考资料 – NeMo Megatron。
