大模型在AI芯片如何训练？AWS AI等KDD2023最新《在新兴AI芯片上训练大规模基础模型》教程，147页pdf

基础模型如ChatGPT和GPT-4由于其新兴的能力，如少量提示、多步推理、指令跟随和模型校准，在学术界和工业界都受到了极大的关注。这样的能力过去只能通过特别设计的模型获得，例如使用知识图谱的模型，但现在可以通过基础模型在更大的规模上实现。

随着基础模型的能力的增加，它们的大小也以远远超过摩尔定律的速度增长。例如，2018年发布的BERT大型模型是一个334M参数模型。而2022年发布的Pathways Language Model (PaLM)是一个540B参数模型，这在短短4年内增加了超过三个数量级。基础模型的训练需要大量的计算能力。例如，使用多个A100芯片的单一最先进的GPU机器训练BERT模型可能需要几天，而在大型多实例GPU集群上训练GPT-3模型可能需要几个月的时间来完成估计的3*10^23 flops。

本教程提供了对支持新AI芯片的基础模型训练和推理的最新进展的概述。它回顾了建模方面的进展，重点是Transformer架构，并介绍了支持训练和服务基础模型的系统架构。这包括编程语言框架如PyTorch和TensorFlow、图编译器、3D并行性，以及像GPU H100、TPU和Trainium这样的加速器。最后，该教程介绍了我们使用不同系统训练基础模型的经验。

1.1 从模型到比特

1.1.1 算法概览。我们涵盖了广泛使用的深度学习（DL）架构，包括递归和卷积神经网络，以及稳定的扩散模型和Transformer架构（例如，仅编码器、仅解码器、编码器-解码器和多模态模型）。我们概述了常见的训练和推理任务，如预训练、微调、推理和评估。我们介绍了高级建模技术，如不同的注意机制、训练目标及其实现细节。 1.1.2 系统概览。有效地训练大规模的基础模型需要在所有层面上都有强大的系统支持。在前端，我们需要在DL框架中适当地表示模型和分布式处理策略。在中端，模型执行需要广泛的性能优化，如图重写、内核融合、内核优化、计算和通信重叠等。在后端，模型将在某种硬件上运行，该硬件按照某种网络拓扑连接。在这个环节中，我们将概述每一层中出现的挑战。

1.2 使用模型并行化训练基础模型

在这一节中，我们将讨论用于跨多台机器训练拥有数十亿参数的模型的现代分布式训练技术，包括完全分片数据并行(FSDP)[16, 15]、Deepspeed Zero [11, 10]、张量并行化[4, 7, 14]和流水线并行化[3, 6]。这些技术在使大型神经网络的训练成为可能方面起到了关键的作用。我们将研究和比较与特定模型架构紧密集成的方法的局限性，以及依赖于底层ML框架快速演变的内部接口的技术。最后，我们将介绍可以用来使用这些技术训练模型的开源库。

1.3 案例研究

在四个案例研究中，我们将介绍不同基础模型的特点以及在新兴AI加速器上运行它们时出现的挑战。我们将介绍包括BERT[2]、RoBERTa[5]和DistilBERT[13]在内的“BERT系列”模型、GPT2[8]和GPT-3[1]、文本到文本转移变换器(T5)[9]和稳定扩散[12]。在对每种架构简要介绍之后，我们将分享我们在不同加速器上训练这些模型的经验。我们将以对有兴趣在新兴AI芯片上训练或评估模型的实践者的实际建议和注意事项来结束本教程。

讲者:

Jun (Luke) Huan 是AWS AI实验室的首席科学家。焕博士的研究领域是AI和数据科学。他已经发表了超过160篇的同行评审的文章并指导了11名博士学生完成学业。在加入AWS之前，他在百度研究院工作，并曾创建了一个AI初创公司StylingAI Inc。在进入工业界之前，他在堪萨斯大学的EECS部门担任教授。焕博士还在美国的NSF工作过，负责其大数据项目。 Yida Wang 是亚马逊的AWS AI团队的首席科学家。他的研究领域包括系统、高性能计算和大数据分析。他目前专注于深度学习系统的研究，特别是编译和优化深度学习模型。 Youngsuk Park 是AWS AI实验室的高级应用科学家。他的研究兴趣集中在机器学习、基础模型、优化和决策制定等方面。在加入AWS之前，他在Stanford University完成了电气工程的硕士和博士学位。 Aashiq Muhamed 在AWS AI实验室工作，专注于优化深度学习系统。在加入AWS AI实验室之前，他在Amazon Search工作，并在Stanford University进行研究生学习。 Rahul Solanki 在亚马逊的AWS Neuron团队工作。他目前的工作重点是构建可以高效地在AI加速器上进行深度学习模型训练和推理的框架和工具。他从Georgia Institute of Technology获得了硕士学位。 Christian Bock 是AWS AI实验室的应用科学家。他的研究兴趣主要集中在语言模型的可解释性以及其在各个领域的应用上。他在ETH Zurich完成了博士学位。