HLAT: High-quality Large Language Model Pre-trained on AWS Trainium

📄 arXiv: 2404.10630v2 📥 PDF

作者: Haozheng Fan, Hao Zhou, Guangtai Huang, Parameswaran Raman, Xinwei Fu, Gaurav Gupta, Dhananjay Ram, Yida Wang, Jun Huan

分类: cs.CL, cs.LG

发布日期: 2024-04-16 (更新: 2024-11-23)

🔗 代码/项目: GITHUB

💡 一句话要点

提出HLAT以解决AWS Trainium上大语言模型训练的挑战

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 大语言模型 AWS Trainium 分布式训练 NeuronX 深度学习

📋 核心要点

  1. 现有方法在AWS Trainium上训练大语言模型面临软件生态系统不成熟的挑战，影响了训练效率和效果。
  2. 论文提出HLAT，通过4096个AWS Trainium加速器在1.8万亿个标记上进行预训练，展示了其在大语言模型训练中的有效性。
  3. 实验结果表明，HLAT在多个评估任务上达到了与基线模型相当的性能，证明了其高效性和成本效益。

📝 摘要（中文）

为了使大型语言模型（LLMs）在下游任务中表现良好，通常需要在数万亿个标记上进行预训练。这通常需要大量强大的计算设备以及稳定的分布式训练框架。随着AI/ML应用的增加，传统加速器（如GPU）的稀缺性凸显了对可扩展且具有成本效益的替代专用加速器的需求。AWS Trainium是专为训练大型深度学习模型而设计的第二代机器学习加速器。本文展示了HLAT：一个使用4096个AWS Trainium加速器在1.8万亿个标记上预训练的7B和70B解码器模型系列。我们将HLAT的性能与在NVIDIA GPU和Google TPU上训练的开源模型进行基准测试，结果表明HLAT在相似模型规模的基线模型上达到了相当的模型质量。我们还开源了所有训练脚本和配置，并分享了使用NeuronX分布式训练库的最佳实践。

🔬 方法详解

问题定义：本文旨在解决在AWS Trainium上训练大型语言模型时面临的软件生态系统不成熟和计算资源稀缺的问题。现有方法依赖于传统GPU，难以满足日益增长的计算需求。

核心思路：HLAT通过利用AWS Trainium加速器的强大计算能力，结合NeuronX分布式训练库，优化了大语言模型的预训练过程，以实现高效且经济的训练。

技术框架：整体架构包括数据预处理、模型设计、分布式训练和评估四个主要模块。数据预处理负责处理和准备训练数据，模型设计则基于解码器架构，分布式训练模块利用NeuronX库进行高效训练，最后通过评估模块对模型性能进行验证。

关键创新：HLAT的主要创新在于其在AWS Trainium上实现的高效预训练，克服了传统GPU训练的局限性，展示了在新硬件上进行大规模模型训练的可行性。

关键设计：在训练过程中，HLAT采用了特定的超参数设置和损失函数，优化了模型的收敛速度和性能，同时确保了在大规模数据集上的训练稳定性。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

HLAT在多个评估任务中表现出色，其性能与在NVIDIA GPU和Google TPU上训练的开源模型相当。具体而言，HLAT在相似模型规模下达到了与基线模型相同的质量，证明了其在高性能和成本效益方面的优势。

🎯 应用场景

该研究的潜在应用领域包括自然语言处理、对话系统和文本生成等。HLAT的高效训练能力使其能够在资源受限的环境中实现大规模模型的应用，推动AI技术的普及和发展。未来，HLAT可能会在更多AI应用场景中发挥重要作用，促进智能化服务的提升。

📄 摘要（原文）

Getting large language models (LLMs) to perform well on the downstream tasks requires pre-training over trillions of tokens. This typically demands a large number of powerful computational devices in addition to a stable distributed training framework to accelerate the training. The growing number of applications leveraging AI/ML led to a scarcity of the expensive conventional accelerators (such as GPUs), which emphasizes the need for the alternative specialized-accelerators that are scalable and cost-efficient. AWS Trainium is the second-generation machine learning accelerator purposely built for training large deep learning models. However, training LLMs with billions of parameters on AWS Trainium is challenging due to its relatively nascent software ecosystem. In this paper, we showcase HLAT: a family of 7B and 70B decoder-only LLMs pre-trained using 4096 AWS Trainium accelerators over 1.8 trillion tokens. The performance of HLAT is benchmarked against popular open source models including LLaMA and OpenLLaMA, which have been trained on NVIDIA GPUs and Google TPUs, respectively. On various evaluation tasks, we show that HLAT achieves model quality on par with the baselines of similar model size. We also open-source all the training scripts and configurations of HLAT (https://github.com/awslabs/HLAT) and share the best practice of using the NeuronX Distributed Training (NxDT), a customized distributed training library for AWS Trainium. Our work demonstrates that AWS Trainium powered by NxDT is able to successfully pre-train state-of-the-art LLM models with high performance and cost-effectiveness.