Efficient Training of Large Language Models on Distributed Infrastructures: A Survey
高效训练分布式基础设施上的大型语言模型:一篇综述
作者/机构: Jiangfei Duan, Shuo Zhang, Zerui Wang*, Lijuan Jiang, Wenwen Qu, Qinghao Hu, Guoteng Wang, Qizhen Weng, Hang Yan, Xingcheng Zhang, Xipeng Qiu, Dahua Lin, Yonggang Wen, Xin Jin, Tianwei Zhang and Peng SunB
A1 主要贡献
本文旨在全面综述大型语言模型(LLM)训练系统和基础设施的最新进展,以应对其在“可扩展性、效率和可靠性”(SER)方面带来的重大挑战。
* 核心问题: 训练如GPT和LLaMA等先进的大型语言模型需要庞大的GPU集群和大量的计算时间,这在可扩展性、效率和可靠性方面构成了巨大挑战。
* 研究目标: 本综述旨在提供关于LLM训练系统和基础设施进展的全面概述。这包括对分布式训练基础设施(AI加速器、网络、存储、调度)的创新方法,以及分布式LLM训练中的并行策略、计算、通信和内存优化技术,同时还涵盖了在长时间训练中维持系统可靠性的方法。
* 创新点: 本文通过系统性地整合最新的研究进展,为研究人员和从业者提供了改进LLM训练系统的宝贵见解,并指明了未来的研究方向。文章的组织结构清晰,从背景知识、基础设施、并行方案,到计算、内存、通信优化,再到容错机制,全面覆盖了LLM高效训练的各个层面。特别地,文章还强调了传统计算系统在满足LLM计算需求方面的局限性,并指出了光计算和光网络等创新解决方案的必要性。
Fig. 1: 本综述的整体结构。
A3 背景知识
2.1 基于Transformer的大型语言模型
当前最先进的LLM主要基于Transformer架构。这些模型的核心是注意力机制【索引10,Attention is all you need,2017,Advances in neural information processing systems】,它允许模型动态地权衡句子中不同单词的重要性。典型的Transformer层架构如图2所示,可以多次堆叠以构建一个LLM。输入文本首先被分词为独立的token,然后通过嵌入层转换为token向量X。为了保留文本的序列性,token向量会嵌入位置信息。最终得到的token向量被送入Transformer层,该层由一个注意力(Attention)块和一个前馈神经网络(FFN)块组成。
Fig. 2: 一个典型的Transformer层包含一个注意力(Attention)块和一个前馈网络(FFN)块。
注意力机制的计算过程。假设输入token向量为$X = [x1, x2, · · · , xn]$。这些token首先通过线性变换转换为查询(Query)Q、键(Key)K和值(Value)V张量。注意力机制计算注意力输出如下:
其中d是键张量的维度。这个公式确保LLM能够通过计算值的加权和来关注输入序列的相关部分,权重则源自查询和键之间的相似度。在注意力层之后,输出会经过一个FFN进行进一步处理。
现有LLM对原始Transformer架构的改进。如今的LLM通常遵循原始的仅解码器(decoder-only)Transformer架构,但对注意力机制和FFN进行了修改以提升效率和性能。原始的注意力机制被称为多头注意力(Multi-Head Attention, MHA)【索引10,Attention is all you need,2017,Advances in neural information processing systems】,由于键值缓存(key-value cache)的存在,其计算复杂度为二次方且内存消耗高。为解决这些问题,提出了多种变体,如多查询注意力(Multi-Query Attention, MQA)【索引11,Fast transformer decoding: One write-head is all you need,2019,arXiv】、分组查询注意力(GroupQuery Attention, GQA)【索引12,GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints,2023,arXiv.org】和多潜注意力(Multi-Latent Attention, MLA)【索引13,Deepseek-v2: A strong, economical, and efficient mixture-of-experts language model,2024,arXiv】。FFN组件的一个显著进步是专家混合(Mixture-of-Experts, MoE)【索引14,Adaptive mixtures of local experts,1991,Neural computation;索引15,Gshard: Scaling giant models with conditional computation and automatic sharding,2020,arXiv】架构,它采用稀疏激活的FFN。在MoE中,每个输入只激活FFN层(或称专家)的一个子集,从而在保持高模型容量的同时显著减少计算负载。
2.2 LLM训练工作负载的特性
LLM训练工作负载与传统深度学习工作负载的显著差异。LLM训练工作负载的特性因其复杂性和规模而与传统深度学习工作负载有很大不同。这些独特的特性影响了训练系统的设计、性能、可扩展性和资源利用率。以下是LLM的关键差异和要求。
(1) 同构模型架构。与先前使用多种模型架构(如LSTM【索引16,Neural machine translation by jointly learning to align and translate,2015,ICLR ’15】、CNN【索引17,Gradient-based learning applied to document recognition,1998,Proceedings of the IEEE】)处理不同任务的深度学习工作负载不同,LLM主要使用Transformer架构【索引10,Attention is all you need,2017,Advances in neural information processing systems】。像GPT【索引5,Language models are unsupervised multitask learners,2019,OpenAI blog】、LLaMA【索引6,Llama 2: Open foundation and fine-tuned chat models,2023,arXiv】、InternLM【索引18,Internlm2 technical report,2024,arXiv】和MOSS【索引19,MOSS: An open conversational large language model,2024,Machine Intelligence Research】等模型都共享这一共同基础。这种架构的统一性为针对特定模型架构优化系统性能提供了巨大潜力。
(2) 前所未有的规模和训练时长。LLM训练的规模空前,通常使用TB级的训练数据集更新拥有数千亿参数的模型。如此巨大的规模需要在庞大的GPU集群上进行分布式训练,并给保持高效率带来了挑战。此外,LLM的训练可能持续数周或数月,要求系统具备强大的容错机制和高效的检查点策略,以防止数据丢失并方便中断的训练会话恢复。
(3) 专门的软件优化。为了适应LLM庞大的模型尺寸,专门的系统实现了先进的技术来优化执行。例如,Megatron【索引20,Megatron-lm: Training multi-billion parameter language models using model parallelism,2019,arXiv】和Alpa【索引21,Alpa: Automating interand {Intra-Operator} parallelism for distributed deep learning,2022,OSDI 22】通过混合并行加速训练。DeepSpeed【索引22,Deepspeed: System optimizations enable training deep learning models with over 100 billion parameters,2020,KDD】通过集成状态分片优化器来减少内存消耗。
(4) 训练范式的转变。传统的深度学习工作负载遵循任务特定的范式,在特定领域的数据上为翻译等特定任务训练模型。相比之下,LLM采用在海量数据集上进行自监督训练的方法来创建基础模型,然后将这些模型适配到各种下游任务。这种范式转变为模型开发流程带来了重大变化,包括预训练和对齐阶段,并导致其工作负载特性与之前的深度学习工作负载相比有显著不同。从数据中心的角度来看,LLM开发涉及大量与预训练相关的小规模工作负载,包括对齐(即微调)和定期的评估工作负载【索引23,Characterization of large language model development in the datacenter,2024,NSDI 24】。
2.3 LLM训练的挑战
LLM训练工作负载特性带来的挑战。LLM训练工作负载的独特性给开发高效的训练系统和基础设施带来了重大挑战。这些挑战主要体现在三个关键领域:可扩展性、效率和可靠性。每个挑战都直接源于LLM的巨大规模及其训练过程的复杂性,需要创新的解决方案来推动分布式计算和机器学习系统的边界。
(1) 可扩展性。LLM的成功很大程度上归功于其规模,性能通常随着LLM的增大而提高【索引8,Scaling laws for neural language models,2020,arXiv】。然而,模型规模的扩大带来了巨大的可扩展性挑战,因为训练LLM需要越来越大的GPU或专用AI加速器集群。首先,构建能够提供大规模计算和内存容量的可扩展基础设施至关重要。这涉及设计和部署大型GPU或专用AI加速器集群、连接这些设备的高性能网络,以及能够处理海量数据集和模型检查点的分布式存储系统。挑战在于确保这些组件在大规模下高效协同工作,同时管理大规模部署中的散热、功耗和硬件故障。其次,设计能够有效并行利用大规模加速器的可扩展训练系统至关重要。这包括设计并行化策略和通信算法,使其能够在数千个加速器上实现近线性的可扩展性,同时保持一致的LLM准确性。
(2) 效率。LLM训练巨大的计算需求转化为高昂的训练成本,因此最大化硬件和软件系统的效率势在必行。效率可以通过模型FLOPs利用率(MFU)来衡量,它量化了系统利用可用计算资源的有效性。然而,在大规模下实现高效率仍然是一个重大挑战。例如,LLaMA3在16K个GPU上仅实现了38%到41%的MFU【索引9,The llama 3 herd of models,2024,ai.meta.com】,这凸显了随着系统规模扩大,维持高利用率的难度。最大化效率需要在并行、计算、通信和内存方面进行优化。首先,分布式LLM训练的并行性需要精心设计以最小化通信需求。其次,优化的计算算子和较低精度的算术对于实现高GPU FLOPS利用率至关重要。第三,需要最小化通信开销以减少GPU空闲时间。最后,需要高效的内存优化来在现有硬件中容纳LLM,并减少重计算造成的FLOPs浪费。
(3) 可靠性。确保LLM训练在长时间内的可靠性至关重要。由于训练任务可能在数万个GPU的大型集群上持续数周或数月,训练失败的概率增加,因此需要快速的故障检测和恢复机制以实现有弹性的LLM训练。首先,LLM训练任务可能因各种错误而崩溃,这使得在数万个GPU中快速确定确切的故障原因变得困难。其次,由于训练的同步性,LLM训练任务的挂起会导致所有GPU空闲,造成巨大浪费。此外,一些复杂的异常,如冗余链路故障或掉队者(stragglers),可能不会立即导致崩溃,但会导训练减速。这种不稳定性可能导致训练效率降低。为应对这些挑战,能够检测灾难性故障和性能退化的强大异常检测系统至关重要。此外,实现能够无缝处理节点故障和网络问题的容错训练框架也至关重要。
2.4 相关综述
本综述的重点和范围。这项工作专注于基于Transformer的LLM的高效训练系统和基础设施,包括底层分布式基础设施的设计、并行范式、计算和通信的优化、高效的内存管理以及训练系统的弹性。我们还研究了新兴工作负载(如MoE,一种有前途的高效LLM变体)和微调(对齐LLM能力的必要阶段)的高效训练系统。然而,这项工作不包括有前途的LLM架构的演变【索引24,A comprehensive survey on pretrained foundation models: A history from bert to chatgpt,2023,arXiv;索引25,Advancing transformer architecture in longcontext large language models: A comprehensive survey,2023,arXiv】以及用于训练【索引26,A survey of large language models,2023,arXiv】、指令微调【索引27,Instruction Tuning for Large Language Models: A Survey,2023,arXiv】和对齐【索引28,Aligning large language models with human: A survey,2023,arXiv】的算法,以实现强大且安全的LLM。虽然以前的工作【索引29,Efficient Large Language Models: A Survey,2023,arXiv;索引30,Understanding llms: A comprehensive overview from training to inference,2024,arXiv;索引31,A Survey of Resource-efficient LLM and Multimodal Foundation Models,2024,arXiv】已经讨论了LLM训练系统的某些方面,但它们的主要重点不是高效训练系统和基础设施的设计。Wan等人【索引29,Efficient Large Language Models: A Survey,2023,arXiv】旨在提供一个关于以模型和数据为中心的高效LLM进展的整体视图。Liu等人【索引30,Understanding llms: A comprehensive overview from training to inference,2024,arXiv】涵盖了LLM的训练和推理部署技术。Xu等人【索引31,A Survey of Resource-efficient LLM and Multimodal Foundation Models,2024,arXiv】旨在讨论LLM开发中算法和系统方面的资源高效策略。这项工作也讨论了量化LLM训练和高效LLM微调的方法,但我们专注于系统方法。压缩和微调LLM的算法方法由Zhu等人【索引32,A survey on model compression for large language models,2023,arXiv】和Han等人【索引33,Parameter-efficient finetuning for large models: A comprehensive survey,2024,arXiv】讨论。这项工作的讨论范围不包括高级优化算法【索引34,Large-scale deep learning optimizations: A comprehensive survey,2021,arXiv】和分布式DNN训练系统【索引35,Scalable deep learning on distributed infrastructures: Challenges, techniques, and tools,2020,ACM Computing Surveys (CSUR)】。虽然Liang等人【索引36,A survey on auto-parallelism of large-scale deep learning training,2023,IEEE Transactions on Parallel and Distributed Systems】广泛回顾了自动并行方法,但他们的重点是通用DNN,而不是专门针对LLM。
A2 方法细节
3 LLM训练基础设施
本节我们探讨用于训练LLM的基础设施设计,包括加速器、网络、存储和调度系统(图3)。
Fig. 3: 分布式LLM训练的基础设施概览。
Fig. 4: 分布式LLM训练的基础设施优化研究。
3.1 AI加速器
AI加速器对LLM发展的推动作用。LLM的快速发展显著受益于GPU和AI加速器的演进,它们对于提高模型训练性能至关重要。
3.1.1 NVIDIA图形处理单元(GPU)
NVIDIA GPU在LLM训练中的核心地位。由于其卓越的并行计算能力,NVIDIA GPU已成为分布式LLM训练的重要组成部分。这些处理器内置了大量紧凑、高效的核心,可以同时执行众多任务。GPU的设计非常适合LLM训练中的矩阵和向量运算。它们支持多种数值精度格式,如FP32、TF32、FP16、BF16、FP8、INT8,甚至FP4。这使得研究人员能够在训练速度和准确性之间取得良好平衡,从而使LLM训练更高效【索引110,Towards efficient generative large language model serving: A survey from algorithms to systems,2023,arXiv】。NVIDIA的GPU编程语言(即CUDA)使得在GPU上管理任务的分割和并行处理变得更加容易,帮助研究人员充分利用GPU的强大能力来训练先进的LLM。
GPU架构的演进。典型的GPU包含一个流式多处理器(SMs)阵列,每个SM内置多个核心,这些核心共享一个指令单元但能并行执行不同的线程。每个SM内的共享内存允许线程间有效的数据交换和同步,这对于优化LLM计算所需的内存访问模式至关重要。此外,GPU配备了高带宽内存(HBM),可加速数据传输并减轻计算密集型任务中的内存访问瓶颈。最新的GPU架构,如NVIDIA的Ampere【索引37,NVIDIA Ampere Architecture,,NVIDIA】、Hopper【索引38,NVIDIA Hopper Architecture,,NVIDIA】和Blackwell【索引39,NVIDIA Blackwell Architecture,,NVIDIA】,不断推动LLM计算的边界。它们提供增强的内存带宽和容量、更高的每秒浮点运算次数(FLOPS),以及像Tensor Cores这样的专用混合精度计算单元。值得注意的是,NVIDIA的Hopper架构引入了一项重大进步,即Transformer引擎【索引111,Nvidia hopper gpu: Scaling performance,2022,IEEE Hot Chips 34 Symposium (HCS)】,该功能利用混合FP8和FP16精度来加速基于Transformer的LLM的训练。
3.1.2 其他AI加速器
AMD GPU及其他AI加速器的应用。在AMD GPU上进行分布式LLM训练已成为现实,尤其是在全球首台百亿亿次超级计算机Frontier【索引112,The exscale era is upon us: The frontier supercomputer may be the first to reach 1,000,000,000,000,000,000 operations per second,2022,IEEE spectrum】上。每个Frontier节点配备8个MI250X【索引40,Amd instinct tm mi250x accelerator enabled by elevated fanout bridge advanced packaging architecture,2023,IEEE Symposium on VLSI Technology and Circuits (VLSI Technology and Circuits)】 AMD GPU,每个GPU拥有64GB的HBM,理论FP16峰值性能为191.5 TFLOPS。这种配置为高效训练万亿参数模型提供了前所未有的机会。释放这一潜力的关键在于将现有的基于CUDA的工具和框架适配到ROCm平台【索引113,Optimizing distributed training on frontier for large language models,2024,ISC High Performance 2024 Research Paper Proceedings (39th International Conference);索引114,Evaluation of pre-training large language models on leadership-class supercomputers,2023,The Journal of Supercomputing】。值得注意的是,支持ROCm的FlashAttention【索引115,Flashattention: Fast and memory-efficient exact attention with io-awareness,2022,Advances in Neural Information Processing Systems】和FlashAttention2【索引116,FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning,2023,arxiv】版本已经被开发出来,从而实现了注意力机制的高效执行。
专用AI加速器的发展。各种具有强大计算能力和软件优化的AI加速器已被开发用于训练LLM。GAUDI【索引41,Gaudi training platform white paper,2023,intel.com】提供了一个异构计算架构,包括两个矩阵乘法引擎和一个全可编程张量处理器核心集群,能够高效处理LLM训练操作。该处理器可以使用384个GAUDI2卡训练一个拥有1750亿参数的GPT-3模型【索引117,Benchmarking and in-depth performance study of large language models on habana gaudi processors,2023,SC’23 Workshops】。Google TPUv4【索引42,Tpu v4: An optically reconfigurable supercomputer for machine learning with hardware support for embeddings,2023,ISCA】超级计算机拥有4096个芯片,支持LLM训练时平均达到约60%的峰值FLOPS。Graphcore Bow Pod64【索引43,A comprehensive performance study of large language models on novel ai accelerators,2023,arXiv】,一个包含64个Bow级IPU的单机架配置,达到22 petaFLOPS,支持使用256个IPU进行GPT-3模型训练。Cerebras CS-2【索引44,The cerebras cs-2: Designing an ai accelerator around the world’s largest 2.6 trillion transistor chip,2022,ISPD】是一个晶圆级深度学习加速器,包含850,000个处理核心,每个核心提供48KB的专用SRAM内存。它被用来训练Cerebras-GPT,一个开放的计算最优语言模型家族【索引118,Cerebras-gpt: Open compute-optimal language,2023,arXiv:2304.03208】。
3.2 网络基础设施
通信开销是扩展LLM训练的主要障碍。通信开销是扩展LLM训练的主要障碍【索引119,Is network the bottleneck of distributed training?,2020,Workshop on Network Meets AI & ML;索引120,High-speed data communication with advanced networks in large language model training,2024,IEEE Micro】。例如,在训练过程中减少模型梯度可能导致超过90%的训练时间用于通信【索引121,Multi-resource interleaving for deep learning training,2022,SIGCOMM】。为解决此问题,研究界致力于改进LLM训练的通信基础设施。
3.2.1 芯片间通信
芯片间通信对LLM训练效率至关重要。芯片间通信对于节点内AI加速器之间的数据传输至关重要,显著影响LLM训练的效率。传统上,这种通信依赖于PCI Express(PCIe)【索引122,Using pci express® as the primary system interconnect in multiroot compute, storage, communications and embedded systems,2008,Integrated Device Technology】,它采用树形拓扑——一种多个设备连接到单个根复合体的分层结构。多年来,PCIe的带宽得到了提升:PCIe 3.0每通道提供约1 GB/s,16通道配置总共约16 GB/s;PCIe 4.0将带宽翻倍至每通道2 GB/s,而PCIe 5.0进一步增至每通道4 GB/s。尽管有这些增强,PCIe在带宽、延迟和可扩展性方面的固有局限性使其不适合LLM训练【索引123,Evaluating modern gpu interconnect: Pcie, nvlink, nvsli, nvswitch and gpudirect,2019,IEEE Transactions on Parallel and Distributed Systems】。为解决这些限制,像NVLink【索引45,Nvidia dgx-1 system architecture white paper,2017,NVIDIA】这样的专用芯片间互连技术越来越受LLM训练青睐。与PCIe相比,这些先进的互连技术通过使用各种拓扑结构(立方体网格、全连接和3D环面)提供了显著更高的带宽和更低的延迟。此外,共享内存模型、专用通信协议和同步机制也发挥着关键作用。
Fig. 5: 五种芯片间拓扑:树形拓扑、立方体网格拓扑、基于交换机的全连接拓扑、基于P2P的全连接拓扑和2D环面拓扑。
立方体网格拓扑(Cube-Mesh Topology)。NVLink-1.0【索引45,Nvidia dgx-1 system architecture white paper,2017,NVIDIA】提供了GPU之间的直接高速连接,每条链路提供160 GB/s的双向带宽。该架构支持为四个GPU形成平面网格结构,为八个GPU形成立方体网格拓扑,可配置为DGX-1服务器。这种立方体网格配置虽然不是全连接,但显著提高了GPU上的数据通信效率和训练性能。
全连接拓扑(Fully-Connected Topology)。许多互连技术利用基于交换机或基于P2P的全连接拓扑来提高芯片间通信性能。NVIDIA使用NVSwitch【索引46,3.2 the a100 datacenter gpu and ampere architecture,2021,ISSCC】实现GPU间的基于交换机的全连接。在DGX-2【索引124,Nvidia dgx-2: The world’s most powerful system for the most complex ai challenges,2018,nvidia.com】系统中,六个NVSwitch将16个GPU中的每一个都完全连接到所有其他GPU,任意两个GPU之间的双向带宽为300 GB/s。该带宽在NVSwitch 2.0中增加到600 GB/s,在NVSwitch 3.0中进一步增加到900 GB/s。Intel、AMD和华为Ascend为其加速器采用基于P2P的全连接拓扑,其中每个芯片通过以太网或Infinity Fabric【索引47,2.2 amd chiplet architecture for high-performance server and desktop products,2020,ISSCC】直接连接到同一节点内的其他每个芯片。与基于交换机的拓扑相比,基于P2P拓扑中两个GPU之间的带宽受到直接连接链路带宽的限制。
2D/3D环面拓扑(2D/3D-Torus Topology)。Google的TPU系统利用环面网络拓扑【索引59,Interconnection networks,2003,Morgan Kaufmann】进行芯片间通信。它通过将每个TPU芯片连接到其在网格中的四个相邻邻居(带有环绕边缘)来建立连接,从而形成一个环面结构。这种架构设计由于芯片之间存在多条直接路径,确保了低延迟和高带宽。具体来说,TPUv2【索引48,A domain-specific supercomputer for training deep neural networks,2020,Communications of the ACM】超级计算机采用16x16的2D环面配置,包含256个芯片,通过高速芯片间互连(ICI)链路连接。TPUv3【索引49,Google’s training chips revealed: Tpuv2 and tpuv3,2020,Hot Chips Symposium】超级计算机使用32x32的2D环面,包含1024个芯片。在2D环面设计的基础上,TPUv4【索引42,Tpu v4: An optically reconfigurable supercomputer for machine learning with hardware support for embeddings,2023,ISCA】超级计算机将计算资源组织成具有3D环面拓扑的多机立方体。每个TPU机器包含四个呈2x2x1网格排列的芯片,通过ICI链路互连。16个这样的TPU机器组合成一个数据中心机架,机架内的ICI链路互连形成一个4x4x4的网格,从而构成一个3D环面结构。这种先进配置显著提高了通信效率和可扩展性,特别有益于LLM训练。
3.2.2 节点间通信
RDMA技术在节点间通信中的作用。远程直接内存访问(RDMA)【索引54,Architectural Specifications for RDMA over TCP/IP ,,rdmaconsortium.org】实现了节点间高速、低延迟的数据传输。RDMA允许一台计算机的内存直接访问另一台计算机的内存,而无需涉及任何一个节点的操作系统。GPUDirect-RDMA【索引50,The development of mellanox/nvidia gpudirect over infiniband—a new model for gpu to gpu communications,2011,Computer Science-Research and Development】通过实现不同节点上GPU之间的直接通信,完全绕过CPU,从而增强了这一过程。这项技术对LLM训练特别有益,因为它加速了模型参数和梯度的同步。两种最普遍的RDMA技术是InfiniBand【索引51,An introduction to the infiniband architecture,2001,High performance mass storage and parallel I/O】和RoCE(RDMA over Converged Ethernet)【索引52,Supplement to infiniband architecture specification volume 1 release 1.2.2 annex a16,2010,Infiniband Trade Association】。
InfiniBand与RoCE的对比。InfiniBand是一种高速、低延迟的网络技术,广泛用于HPC(高性能计算)环境,例如Eagle超级计算机【索引125,Nvidia hopper gpu and grace cpu highlights,2022,Computing in Science & Engineering】。该技术需要专用的网络基础设施,这反映了其专注于提供卓越性能的设计。多年来,InfiniBand在带宽能力方面有了显著发展,从100 Gbps的EDR(增强数据速率)发展到200 Gbps的HDR(高动态范围),最近又发展到每链路400 Gbps的NDR(下一代数据速率)【索引126,Infiniband product guide,2020,network.nvidia.com】。RoCE利用现有的以太网基础设施提供RDMA能力。这种方法提供了一种更具成本效益且易于部署的解决方案,尤其是在已经使用以太网的数据中心。RoCE有两个版本:RoCE-v1【索引52,Supplement to infiniband architecture specification volume 1 release 1.2.2 annex a16,2010,Infiniband Trade Association】,作为以太网链路层协议运行;RoCE-v2【索引53,Supplement to infiniband architecture specification volume 1 release 1.2.2 annex a17,2010,Infiniband Trade Association】,在UDP上运行。字节跳动和Meta等行业领导者已采用这些技术来扩展LLM训练。另一种RDMA协议是iWARP(Internet Wide Area RDMA Protocol)【索引54,Architectural Specifications for RDMA over TCP/IP ,,rdmaconsortium.org】,它支持在TCP/IP网络上实现RDMA。然而,由于其性能相对有限,iWARP不常用于分布式LLM训练【索引127,Roce vs. iwarp competitive analysis,2017,network.nvidia.com】。
3.2.3 网络拓扑
LLM训练集群的网络架构。在LLM训练集群中,网络架构分为前端和后端组件(图3)。前端网络处理各种流量,如作业管理、模型推理和存储活动,而后端网络则专用于训练过程中产生的大量流量。我们优化LLM训练的主要重点是提高这个后端网络的性能和效率,以便将AI加速器扩展到数万个。
Fig. 6: 大规模GPU集群中四种典型的网络拓扑:Clos拓扑、Dragonfly+拓扑、rail-optimization拓扑和rail-only拓扑。
HPC网络拓扑。用于HPC环境的传统拓扑也可用于AI集群进行分布式训练,例如Clos【索引55,A study of non-blocking switching networks,1953,Bell System Technical Journal】、BCube【索引56,Bcube: a high performance, server-centric network architecture for modular data centers,2009,SIGCOMM】、DCell【索引57,Dcell: a scalable and fault-tolerant network structure for data centers,2008,SIGCOMM】、Jellyfish【索引58,Jellyfish: Networking data centers randomly,2012,NSDI 12】、Torus【索引59,Interconnection networks,2003,Morgan Kaufmann】、Dragonfly【索引60,Technology-driven, highly-scalable dragonfly topology,2008,SIGARCH Computer Architecture News】和Dragonfly+【索引61,Dragonfly+: Low cost topology for scaling datacenters,2017,HiPINEB】。Clos网络架构,通常称为胖树(Fat-Tree)拓扑,在LLM训练集群中被广泛使用。在基于Clos的集群中,每个配备一个或多个NIC的服务器被组织成连接到叶交换机的机架。这些叶交换机连接到脊交换机,提供机架间连接并形成一个pod。这些pod再与核心交换机互连,从而实现集群内服务器之间的任意通信。例如,Meta的上一代GPU集群架构支持多达24,000个GPU,由八个pod组成,它们之间具有全胖树带宽,并在核心层使用7:1的超订购比【索引9,The llama 3 herd of models,2024,ai.meta.com】。Meta使用24,000个GPU在该集群上训练Llama 3 405B。
训练优化拓扑。许多网络拓扑是与分布式训练算法协同设计的。rail-optimized拓扑【索引62,Doubling all2all performance with nvidia collective communication library 2.12,2018,developer.nvidia.com】增强了从GPU到叶交换机的连接。在每个rail内,跨不同服务器具有相同索引的GPU通过同一个叶交换机互连。这种配置通过减少数据流之间的网络干扰来提高集体通信性能。SuperPod架构利用rail-optimized网络,能够连接超过16,000个GPU【索引128,Dgx superpod architecture,2024,docs.nvidia.com】。字节跳动在其MegaScale系统设计中采用三层rail-optimized网络连接超过10,000个GPU【索引71,{MegaScale}: Scaling large language model training to more than 10,000 {GPUs},2024,NSDI 24】。然而,rail-optimized网络设计可能效率较低,因为它们需要将GPU连接到远处的交换机,这需要昂贵且耗电的光收发器。这些光学组件增加了功耗和热量,导致网络故障率更高,这对于分布式LLM训练至关重要。阿里巴巴通过一个名为HPN【索引63,Alibaba hpn: A data center network for large language model training,2024,SIGCOMM】的两层双平面架构进一步优化了rail-optimized拓扑。该架构采用最新的51.2Tbps单芯片交换机,支持在一个tier-1网络中连接1,000个GPU,在一个pod内连接多达15,000个GPU。
基于流量分析的拓扑优化。GPT/OPT-175B模型训练的网络流量分析显示,99%的GPU对不承载任何流量,而处理流水线/张量并行和数据并行流量的GPU对不到0.25%【索引64,Rail-only: a cost-effective and high-performance network topology for large-scale llm training with billions of parameters,2023,arXiv】。基于这些发现,rail-only拓扑【索引64,Rail-only: a cost-effective and high-performance network topology for large-scale llm training with billions of parameters,2023,arXiv】取消了rail-optimized网络上不同rail之间的连接。每个rail由一个专用但独立的Clos网络连接。跨不同rail上GPU的通信通过内部芯片间互连转发数据来管理。这种方法可以在保持性能的同时有效降低成本。HammingMesh【索引65,Hammingmesh: a network topology for large-scale deep learning,2022,SC22】将GPU组织成具有2D环面拓扑的组,并通过稀疏连接的交换机连接这些2D环面组。该设计旨在在不牺牲训练性能的情况下节省成本。对于仅通过PCIe连接的GPU,BiGraph【索引129,Eflops: Algorithm and system co-design for a high performance distributed training platform,2020,HPCA】提出了一种新的网络架构,将节点内GPU通信导出到节点外,绕过PCIe带宽瓶颈。它采用一个通过Clos架构互连的两层网络,具有支持应用控制流量路由的独特最短通信路径。
可重构拓扑。可重构网络可以动态调整以优化通信模式,从而提高训练性能。它们通常利用光交换和定制配置来提高网络基础设施的带宽利用率、灵活性和可扩展性。在硅光子(SiP)接口的推动下,SiP-ML【索引66,Sip-ml: highbandwidth optical network interconnects for machine learning training,2021,SIGCOMM】提出了两种主要架构:SiP-OCS和SiP-Ring。SiP-OCS采用全连接配置,通过使用商用光路交换机(OCS)将GPU连接到所有交换机,从而最大化带宽。相反,SiP-Ring采用无交换机的环形配置,通过在SiP接口内集成微环谐振器来减少重构延迟。Wang等人提出了TopoOpt【索引67,{TopoOpt}: Co-optimizing network topology and parallelization strategy for distributed training jobs,2023,NSDI 23】,用于协同优化分布式训练中的网络拓扑和并行策略。这种方法不仅优化了计算和通信需求,还解决了网络拓扑的物理层问题。TPUv4【索引42,Tpu v4: An optically reconfigurable supercomputer for machine learning with hardware support for embeddings,223,ISCA】具有光路交换机(OCS),允许动态重构基于3D环面的互连拓扑,为LLM训练中多样化且密集的通信模式优化数据流。例如,对于512个芯片,TPUv4在3D环面拓扑中提供了灵活性,如4x4x32或8x8x8。
3.2.4 负载均衡与拥塞控制
负载均衡。LLM训练的网络流量特点是少数大象流(elephant flows)。具体来说,由于梯度同步,LLM训练表现出周期性的网络流量爆发【索引63,Alibaba hpn: A data center network for large language model training,2024,SIGCOMM】。每次爆发都需要巨大的网络带宽。此外,参与LLM训练的每个计算节点产生的连接非常少【索引63,Alibaba hpn: A data center network for large language model training,2024,SIGCOMM】。传统的负载均衡技术ECMP(等价多路径路由)【索引68,Analysis of an equal-cost multi-path algorithm,2000,rfc-editor.org】使用哈希算法将流量均匀分布在等价路径上,例如在Clos拓扑中从叶交换机到脊交换机的路径。然而,这种基于哈希的方案在处理由少数大象流组成的LLM训练流量时效率低下。当多个大象流被路由到同一条链路时,可能导致拥塞和高延迟。
针对LLM训练的负载均衡策略。为解决大规模GPU集群中的负载均衡挑战,已开发了多种策略。在Llama 3 405B训练期间,集合通信库在两个GPU之间建立16个网络流,而不是单个流,从而减少了每个流的流量,并增强了负载均衡的机会【索引9,The llama 3 herd of models,2024,ai.meta.com】。此外,增强型ECMP(E-ECMP)协议通过对RoCE报头中的附加字段进行哈希,有效地将这16个流分布在不同的网络路径上。包喷射(Packet spraying)【索引69,On the impact of packet spraying in data center networks,2013,ieee infocom】将一个流的包分布在所有可用的并行链路上,这可能导致包乱序。NIC需要处理乱序的RDMA包。基于LLM训练的流量模式,Ethereal【索引70,Challenging the need for packet spraying in large-scale distributed training,2024,arXiv】证明了贪婪地为每个流分配路径可以均匀地将负载分布在所有网络路径上,并解决ECMP哈希冲突问题。在一个大规模GPU集群中,HPN【索引63,Alibaba hpn: A data center network for large language model training,2024,SIGCOMM】通过识别精确的不相交等价路径并在集合通信库内平衡负载来实现高效的负载均衡。MegaScale【索引71,{MegaScale}: Scaling large language model training to more than 10,000 {GPUs},2024,NSDI 24】表明,rail-optimized拓扑也可以缓解ECMP哈希冲突。
拥塞控制。在RDMA集群中,无损传输至关重要。基于优先级的流控制(PFC)【索引72,802.1qbb – priority-based flow control,,ieee802.org】是一种防止数据包丢失的流控制机制。当 下游设备上启用PFC的队列发生拥塞时,该设备会指示上游设备暂停该队列中的流量,从而确保零丢包。由于PFC是一种粗粒度机制,它可能导致队头阻塞(head-of-line blocking)【索引130,Rdma over commodity ethernet at scale,2016,SIGCOMM】,从而显著降低网络吞吐量。为应对这些挑战,已开发了多种通用拥塞控制方案。这些技术包括TIMELY【索引73,Timely: Rtt-based congestion control for the datacenter,2015,ACM SIGCOMM Computer Communication Review】、数据中心量化拥塞通知(DCQCN)【索引75,Congestion control for large-scale rdma deployments,2015,ACM SIGCOMM Computer Communication Review;索引76,Ecn or delay: Lessons learnt from analysis of dcqcn and timely,2016,CoNEXT】、Swift【索引74,Swift: Delay is simple and effective for congestion control in the datacenter,2020,SIGCOMM】、高精度拥塞控制(HPCC)【索引77,Hpcc: High precision congestion control,2019,SIGCOMM】、边缘排队数据报服务(EQDS)【索引78,An edge-queued datagram service for all datacenter traffic,2022,NSDI 22】和鲁棒拥塞控制(RoCC)【索引79,Rocc: robust congestion control for rdma,2020,CoNEXT】。这些方案监控网络拥塞,调整数据速率以缓解拥塞,并恢复速率以最小化吞吐量下降。
针对并发训练作业的拥塞控制。当存在并发训练作业时,许多拥塞控制方案利用突发性和周期性的流量模式来有效地交错网络流量。MLTCP【索引80,Mltcp: Congestion control for dnn training,2024,arXiv】根据一个关键洞察,即训练流应根据每次训练迭代中发送的字节数来调整其拥塞窗口大小,从而交错竞争带宽的作业的通信阶段。CASSINI【索引81,{CASSINI}:{Network-Aware} job scheduling in machine learning clusters,2024,NSDI 24】通过考虑不同作业的通信模式来优化网络链路上的作业放置。MLT【索引82,Towards {Domain-Specific} network transport for distributed {DNN} training,2024,NSDI 24】利用LLM训练的特性,即早期层的梯度不如后期层的重要,且较大的梯度比小梯度更重要。因此,在发生通信拥塞时,MLT在交换机层面根据梯度在数据包中的重要性来优先排队或丢弃数据包,以缓解通信拥塞问题。
3.3 存储
存储系统在LLM训练中的关键作用。存储系统在分布式LLM训练中扮演着关键角色,需要满足几个关键要求。首先,它应与GPU的计算能力相匹配,以最大化其利用率并避免由存储瓶颈造成的资源浪费。其次,它应支持大规模结构化和非结构化训练数据集的存储,并在分布式处理环境中具有可扩展性。此外,模型检查点的存储和检索在LLM训练中提出了挑战,要求系统满足由模型大小和训练时长决定的读写带宽。最后,存储系统应满足传统企业级要求,如数据保护、高可用性和安全性。
3.3.1 用于检查点的存储系统
LLM训练中检查点存储的挑战和解决方案。在LLM训练中,模型检查点的大小是巨大的。随着参数数量的增加,需要写入的数据量也随之增加,要求存储系统提供更大的写带宽。例如,一个70B参数的LLM的检查点大小为980GB。许多存储系统已部署在大型GPU数据中心以管理模型检查点。Meta的分布式文件系统Tectonic【索引83,Facebook’s tectonic filesystem: Efficiency from exascale,2021,FAST 21】使数千个GPU能够同时保存和加载模型检查点,为大规模训练操作提供高效且可扩展的存储解决方案【索引131,Building Meta’s GenAI Infrastructure,,engineering.fb.com】。在字节跳动,HDFS【索引84,The hadoop distributed file system,2010,MSST】用于集中维护模型检查点,确保大规模下的一致性和可靠性【索引71,{MegaScale}: Scaling large language model training to more than 10,000 {GPUs},2024,NSDI 24】。为缓解检查点恢复期间的带宽瓶颈,一种常见方法是指定一个工作节点从HDFS读取检查点分区,然后将其广播给共享相同数据的其他工作节点。分布式对象存储,如Ceph对象存储【索引85,Ceph: A scalable, high-performance distributed file system,2006,OSDI’06】,更容易扩展。这一优势源于它们没有分层目录树或命名空间,简化了一致性维护。由于这些优点,对象存储已广泛用于模型检查点存储。
3.3.2 用于训练数据的存储系统
LLM训练数据的规模和处理。LLM训练的原始数据集非常庞大。LLaMA 3在超过15万亿个token上进行训练,这比LLaMA 2的数据集大了七倍多【索引6,Llama 2: Open foundation and fine-tuned chat models,2023,arXiv】。每个token约需2字节,总计约30TB数据。为训练准备数据集涉及广泛的预处理步骤,包括数据抓取和清洗,需要大量实验。通常,在这些步骤中处理的数据量是最终训练数据集大小的100倍以上【索引132,Wanjuan-cc: A safe and high-quality opensourced english webtext dataset,2024,arXiv】。例如,WanJuan-CC数据集【索引132,Wanjuan-cc: A safe and high-quality opensourced english webtext dataset,2024,arXiv】选择性地提取了约680亿份文档,在丢弃了99%的原始数据后,生成了约1万亿个高质量token,相当于2TB的数据量。因此,LLM训练的总数据量预计将超过几十PB。
用于数据加载的存储系统和缓存策略。并行文件系统如Lustre【索引86,Lustre: Building a file system for 1000-node clusters,2003,Linux symposium】、GPFS【索引87,{GPFS}: A {Shared-Disk} file system for large computing clusters,2002,FAST 02】和BeeGFS【索引88,I/o characterization and performance evaluation of beegfs for deep learning,2019,ICPP】经常部署在领先的高性能计算系统上,以确保高效的I/O、持久存储和可扩展的性能。这些系统也广泛用于训练集群的数据加载,为处理大规模训练数据提供了必要的基础设施。此外,文件系统能够让工程师在利用数千个GPU的作业上进行交互式调试至关重要,因为代码更改需要立即对所有节点可见【索引131,Building Meta’s GenAI Infrastructure,,engineering.fb.com】。
数据加载中的缓存技术。在大多数LLM的训练中,每个token通常只遇到一次。然而,使用数据缓存仍然至关重要,以缓解数据加载期间的I/O瓶颈。该策略涉及将训练数据从较慢的后端存储预取到较快的缓存存储中。Alluxio【索引89,Tachyon: Reliable, memory speed storage for cluster computing frameworks,2014,ACM Symposium on Cloud Computing】和JuiceFS【索引90,Juicefs: A High-Performance, Cloud-Native, Distributed File System ,,github.com】通过有效地从底层存储系统(如HDFS或对象存储)缓存训练数据来增强LLM训练。Quiver【索引91,Quiver: An informed storage cache for deep learning,2020,FAST 20】支持在操作相同数据集的多个作业和用户之间透明地重用缓存数据。Fluid【索引92,Fluid: Dataset abstraction and elastic acceleration for cloud-native deep learning training jobs,2022,ICDE】利用Alluxio进行数据缓存,并包含一个机制,可以根据I/O条件动态地自动扩展缓存。
3.4 调度
LLM训练工作负载调度。LLM训练工作负载通常在大型多租户基础设施(如GPU集群、公有云)上运行,用户共享集群资源。有效的调度机制对于管理这些工作负载、确保高效的资源利用和任务执行至关重要【索引133,Deep learning workload scheduling in gpu datacenters: A survey,2024,ACM Computing Surveys】。
集群级调度的分类。与专注于单个作业执行的细粒度优化的任务级调度(如流水线调度【索引134,Pipedream: Fast and efficient pipeline parallel dnn training,2018,arXiv;索引135,Pipedream: generalized pipeline parallelism for dnn training,2019,SOSP;索引136,Chimera: efficiently training large-scale neural networks with bidirectional pipelines,2021,SC】)不同,集群级调度旨在优化整个集群的资源分配和任务调度。我们将现有的集群级调度系统根据其主要优化方面分为两类:工作负载调度和资源调度。
3.4.1 工作负载调度
通用深度学习工作负载调度器。近年来,为深度学习训练工作负载量身定制的调度器得到了积极探索【索引93,{Tiresias}: A {GPU} cluster manager for distributed deep learning,2019,NSDI 19;索引94,{Themis}: Fair and efficient {GPU} cluster scheduling,2020,NSDI 20;索引95,Elasticflow: An elastic serverless training platform for distributed deep learning,2023,ASPLOS】。为了提高资源利用率,通常会实现三个高级功能:(1)异构感知调度器(如Gavel【索引96,Heterogeneity-Aware Cluster Scheduling Policies for Deep Learning Workloads,2020,OSDI ’20】、Gandivafair【索引97,Balancing efficiency and fairness in heterogeneous gpu clusters for deep learning,2020,EuroSys ’20】)专注于优化跨不同代GPU的作业分配;(2)作业打包调度器(如FGD【索引98,Beware of fragmentation: Scheduling GPUSharing workloads with fragmentation gradient descent,2023,USENIX ATC ’23】、Lucid【索引99,Lucid: A nonintrusive, scalable and interpretable scheduler for deep learning training jobs,,SC ’21】)支持细粒度的GPU共享以充分利用硬件能力;(3)自适应伸缩调度器(如Pollux【索引100,Pollux: Co-adaptive cluster scheduling for goodput-optimized deep learning,2021,OSDI ’21】、Sia【索引101,Sia: Heterogeneity-aware, goodput-optimized mlcluster scheduling,2023,SOSP】)动态调整GPU数量以及训练超参数以加速训练进程。然而,这些调度器是为通用深度学习工作负载设计的,由于LLM工作负载的独特性质,可能不直接适用于LLM【索引23,Characterization of large language model development in the datacenter,2024,NSDI 24】。
LLM特定工作负载调度器。为了更好地管理LLM工作负载,最近的一些研究提出了针对LLM的系统。Crius【索引102,A codesign of scheduling and parallelization for large model training in heterogeneous clusters,2024,arXiv】联合考虑了混合并行(§4.1)和异构集群内的硬件亲和性。它研究了在集群调度层面集成自适应并行配置的工作流效率,为同时提高多个LLM的训练效率提供了重要机会。为实现LLM的高效超参数调优,Hydro【索引103,Hydro: Surrogate-Based hyperparameter tuning service in datacenters,2023,OSDI ’23】将模型缩小到一个较小的代理模型进行超参数搜索,然后将多个模型融合成一个单一实体以提高硬件利用率。此外,Hydro通过将调优工作负载与启用流水线的LLM预训练任务交错,有效利用流水线气泡,从而扩展了调优工作负载的资源。Acme【索引23,Characterization of large language model development in the datacenter,2024,NSDI 24】进一步描述了LLM开发工作流的工作负载混合,并提出了一个系统来高效调度与LLM训练相关的作业,包括用于及时模型质量反馈的解耦评估调度以及LLM参与的故障诊断和自动恢复。
3.4.2 资源调度
集群管理中的相关资源调度。除了工作负载调度,相关的资源调度(如CPU、内存和网络)是集群级管理的另一个关键方面。对于网络,Cassini【索引81,{CASSINI}:{Network-Aware} job scheduling in machine learning clusters,2024,NSDI 24】通过使用亲和图确定时间移位值来调整通信阶段,从而实现了在不同作业的上下行阶段交错带宽需求。HIRE【索引104,Switches ¨ for hire: Resource scheduling for data center in-network computing,2021,ASPLOS】为数据中心交换机引入了一种创新的网络内计算调度系统,显著减少了网络绕路和尾部放置延迟。对于存储,SiloD【索引105,Silod: A co-design of caching and scheduling for deep learning clusters,2023,EuroSys】将数据缓存和远程I/O视为一等资源进行联合分配,显示出显著的吞吐量提升。对于CPU和内存,Synergy【索引106,Looking beyond {GPUs} for {DNN} scheduling on {MultiTenant} clusters,2022,OSDI 22】通过优化CPU核心分配而不是依赖于GPU比例分配来提高训练效率。此外,一些工作专注于节能。EnvPipe【索引107,{EnvPipe}: Performance-preserving {DNN} training framework for saving energy,2023,USENIX ATC 23】利用流水线并行中的气泡时间,通过降低SM频率来延长流水线单元的执行时间以节省能源。Zeus【索引108,Zeus: Understanding and Optimizing GPU Energy Consumption of DNN Training,2023,USENIX NSDI】自动配置批量大小和GPU功率限制,以提高训练期间的能源效率。Perseus【索引109,Perseus: Removing energy bloat from large model training,2023,arXiv】引入了一种高效的基于图切割的迭代算法,以获得大型模型训练作业的迭代时间-能量帕累托前沿。
Fig. 7: 分布式LLM训练的并行方案研究。
4 LLM训练的并行方案
分布式训练是高效训练LLM的关键。LLM规模的持续增长需要大量的计算资源和内存容量。利用大规模HPC集群进行分布式训练已成为高效训练这些模型的关键方法。本节我们研究了为增强HPC集群在LLM训练中的利用率而提出的各种并行方案。我们将这些方法分为三类:混合并行、自动并行和异构并行。混合并行结合了多种手工设计的并行策略,如数据并行、张量并行、流水线并行、序列并行和专家并行。自动并行根据模型和硬件特性自动确定最佳并行策略。异构并行则利用硬件或模型的异构性进行高效训练,包括利用不同类型的加速器或模型内部的异构性(如RLHF训练)来提高HPC集群上的整体训练效率。
并行训练的编程模型。当今大多数最先进的并行策略采用单程序多数据(SPMD)编程模型,类似于MPI范式【索引239,The mpi message passing interface standard,1994,Programming Environments for Massively Parallel Distributed Systems】,即同一个程序在多个处理器上运行,每个处理器处理不同的数据片段【索引225,Pathways: Asynchronous distributed dataflow for ml,2022,MLS】。例如,数据并行、模型并行和序列并行都利用了SPMD编程模型。这种方法确保了操作的统一性和一致性,非常适合大规模分布式训练环境。一些策略探索打破SPMD的限制,并利用多程序多数据(MPMD)模型进一步提高资源利用率,即不同的程序(或程序的不同部分)在不同的处理器上运行,处理数据或模型的不同部分【索引225,Pathways: Asynchronous distributed dataflow for ml,2022,MLS】。例如,流水线并行在不同设备上运行LLM的不同部分。此外,自动并行和异构并行可以利用SPMD和MPMD模型来增加资源利用率。因此,我们根据并行发生的维度以及所用计算设备是同构还是异构来讨论这些方法,而不是关注底层的编程模型。
4.1 混合并行
混合并行的概念。混合并行通常结合多种手工设计的并行策略来划分LLM的不同可并行维度。这些策略包括数据并行、张量并行、流水线并行和序列并行,如图8所示。数据并行、张量并行和流水线并行的组合也称为3D并行。
4.1.1 数据并行
数据并行的基本原理和实现。数据并行是分布式训练中最常用的并行策略,因其高可扩展性和易于实现。它遵循单程序多数据(SPMD)编程模型。数据并行沿批次维度划分输入训练数据,每个GPU处理其分配的数据段,如图8(a)所示。在整个训练过程中,数据首先逐层通过完整模型权重进行前向计算,然后以相反顺序进行后向计算。每层产生的梯度将通过集体通信操作在所有GPU之间聚合,用于优化器更新。
数据并行的分片策略。数据并行包含多种分片策略,这些策略显著影响内存占用和通信开销。假设全局设备数量为W,引入分片因子F来控制使用的分片策略【索引146,Pytorch fsdp: Experiences on scaling fully sharded data parallel,2023,VLDB Endowment】,定义为参数被划分到的设备数量(1 ≤ F ≤ W)。我们有以下几种情况。完全复制 (F = 1):此分片策略简化为普通的数据并行。Pytorch-DDP【索引240,Pytorch: An imperative style, high-performance deep learning library,2019,Advances in neural information processing systems】和Horovod【索引143,Horovod: fast and easy distributed deep learning in tensorflow,2018,arXiv】在所有设备上完全复制模型,并使用All-Reduce进行梯度聚合。它们还将梯度分成小桶,以将梯度通信与后向计算重叠。
完全分片 (F = W)。此分片策略内存消耗最低,但通信开销最大(比普通数据并行多1.5倍)。完全分片策略将模型完全分片,每个设备只持有模型参数的1/W。完整的权重和梯度在计算前按需通信和恢复,之后立即丢弃。ZeRO-3【索引145,Zero: Memory optimizations toward training trillion parameter models,2020,SC20】采用逐参数分片来划分整个模型,并分别利用All-Gather和Reduce-Scatter进行取消分片和分片通信。分片权重更新(Sharded Weight Update)【索引144,Automatic cross-replica sharding of weight update in data-parallel training,2020,arXiv】也采用逐参数分片,但更侧重于在所有设备间分片冗余的参数更新计算。FSDP(完全分片数据并行)【索引146,Pytorch fsdp: Experiences on scaling fully sharded data parallel,2023,VLDB Endowment】通过以模块单元为粒度分片模型参数来实现相同功能,并提供更用户友好的API。
混合分片 (1 < F < W)。在此策略【索引146,Pytorch fsdp: Experiences on scaling fully sharded data parallel,2023,VLDB Endowment】中,所有设备被划分为N × M的设备网格。模型参数沿网格的N维度分片,并沿M维度复制。MiCS【索引147,Mics: Near-linear scaling for training gigantic model on public,2022,VLDB Endowment】调用All-Gather集合操作来收集分片的参数,并调用All-Reduce来聚合梯度。FSDP【索引146,Pytorch fsdp: Experiences on scaling fully sharded data parallel,2023,VLDB Endowment】用Reduce-Scatter替换All-Reduce,以减少内存和通信开销。与完全复制和完全分片相比,混合分片更灵活,可以通过根据模型架构和硬件约束调整F来在内存消耗和通信开销之间取得平衡。
Fig. 8: 一个包含数据并行、张量并行和流水线并行的3D并行示例。
4.1.2 张量并行
张量并行的概念和应用。张量并行(图8(b)),也称为层内模型并行,是一种为在多个GPU上训练LLM而提出的技术。它沿多个维度划分每层的参数张量,有效地将模型参数分布到可用的GPU上。张量并行通信的是中间激活张量,其大小远小于数据并行中通信的参数和梯度,除非在长上下文LLM训练场景中。然而,在张量并行中,很难将通信与计算重叠,因此需要使用高带宽连接。因此,张量并行更常用于单个GPU节点内。
不同维度的张量并行。张量并行可以根据划分的维度分为1-D【索引20,Megatron-lm: Training multi-billion parameter language models using model parallelism,2019,arXiv】、2-D【索引148,An efficient 2d method for training superlarge deep learning models,2023,IPDPS】、2.5-D【索引149,Tesseract: Parallelize the tensor parallelism efficiently,2022,ICPP】和3-D【索引150,Maximizing parallelism in distributed training for huge neural networks,2021,arXiv】并行。基于Transformer的LLM的MLP和自注意力模块中都有两个参数矩阵。Megatron-LM【索引20,Megatron-lm: Training multi-billion parameter language models using model parallelism,2019,arXiv】首先采用1-D张量并行,沿列划分第一个参数矩阵,沿行划分第二个参数矩阵。它复制每个划分模块的输入和输出张量,并在所有GPU之间引入两次All-Reduce集体通信,以将一个LLM适配到多个GPU中。受用于2-D并行矩阵乘法的可扩展通用矩阵乘法算法(SUMMA)【索引241,Summa: Scalable universal matrix multiplication algorithm,1997,Concurrency: Practice and Experience】和Cannon算法【索引242,A cellular computer to implement the Kalman filter algorithm,1969,Montana State University】的启发,Optimus【索引148,An efficient 2d method for training superlarge deep learning models,2023,IPDPS】进一步在2个维度上划分输入和参数张量,以提高1-D张量并行的通信和内存效率。Tesseract【索引149,Tesseract: Parallelize the tensor parallelism efficiently,2022,ICPP】扩展了为提高Cannon算法效率而提出的2.5-D矩阵乘法方法【索引243,Communication-optimal parallel 2.5 d matrix multiplication and lu factorization algorithms,2011,Euro-Par】,提出了2.5-D张量并行,以克服因模型尺寸增大而导致的大量不必要通信。3-D张量并行【索引150,Maximizing parallelism in distributed training for huge neural networks,2021,arXiv】采用并改进了用于线性操作的3-D并行矩阵乘法算法【索引244,A three-dimensional approach to parallel matrix multiplication,1995,IBM Journal of Research and Development】,并在LLM训练中实现了跨多个设备的完美负载均衡。
4.1.3 流水线并行
流水线并行的基本原理和适用场景。流水线并行(图8(c))【索引151,Gpipe: Efficient training of giant neural networks using pipeline parallelism,2019,Advances in neural information processing systems】,也称为层间模型并行,旨在将大型模型容纳到多个GPU上,特别是跨不同节点。流水线并行将模型的层划分为多个阶段,每个阶段由模型中一组连续的层组成,并映射到一组GPU上。与通常需要像NVLink这样的高带宽连接进行通信的张量并行不同,流水线并行仅需要在指定的切点交换中间张量,导致通信频率较低。因此,流水线并行适用于在通过小带宽连接的多个GPU节点上扩展LLM训练。例如,Strati等人【索引245,Ml training with cloud gpu shortages: Is cross-region the answer?,2024,MLSys】采用流水线并行来充分利用地理分布式资源以克服GPU短缺。由于不同阶段的数据依赖性,流水线并行通常将输入数据分成多个微批次(micro-batches)进行流水线处理,以实现巨型模型的高效训练。然而,它带来了两个重要问题。首先,流水线气泡问题降低了GPU的利用率,因为需要等待前一阶段的输出。其次,不同阶段之间存在内存消耗不平衡,因为前一阶段需要持有比后一阶段更多的活动微批次以获得更好的流水线和更高的利用率。
流水线气泡问题的解决方案。提出了高效的微批次调度算法来减少流水线气泡。GPipe【索引151,Gpipe: Efficient training of giant neural networks using pipeline parallelism,2019,Advances in neural information processing systems】引入了一种“填充-排空”(fill-drain)调度,一次性注入所有微批次进行前向传递执行,然后是后向传递。由于前向和后向传递的预热和冷却,GPipe会产生显著的流水线气泡。PipeDream【索引134,Pipedream: Fast and efficient pipeline parallel dnn training,2018,arXiv;索引135,Pipedream: generalized pipeline parallelism for dnn training,2019,SOSP】引入了1F1B(1个前向,1个后向)调度,一旦某个微批次的相应前向传递完成,就立即执行其后向传递,以减少异步场景中的流水线气泡。DAPPLE【索引152,Dapple: A pipelined data parallel approach for training large models,2021,PPoPP】采用早期后向调度,在每个阶段开始时首先注入固定数量的微批次,然后以轮询方式交错前向和后向传递。Interleaved 1F1B【索引153,Efficient large-scale language model training on gpu clusters using megatron-lm,2021,SC】对1F1B调度进行了调整,但将多个阶段分配给每个GPU(即循环流水线放置)。流水线气泡以更高的通信和峰值内存消耗为代价而减少。Chimera【索引136,Chimera: efficiently training large-scale neural networks with bidirectional pipelines,2021,SC】引入双向流水线,通过权重复制来减少气泡。Hanayo【索引157,Hanayo: Harnessing wave-like pipeline parallelism for enhanced large model training efficiency,2023,SC】进一步提出了一种波浪式流水线,将多个对称阶段分配给一个GPU以提高流水线利用率。Zero bubble【索引156,Zero bubble pipeline parallelism,2023,ICLR】将后向计算分为两部分:激活梯度和参数梯度计算。它使用1F1B调度前向和激活梯度计算,然后用参数梯度计算填充气泡,从而以更高的峰值内存消耗减少气泡。Breadth-First【索引159,Breadth-first pipeline parallelism,2023,MLS】在循环流水线放置中一次性运行所有微批次,以减少与分片数据并行结合时的通信开销。TeraPipe【索引154,Terapipe: Token-level pipeline parallelism for training largescale language models,2021,ICML】沿序列维度拆分微批次,并利用更细粒度的token并行来减少流水线气泡。然而,由于TeraPipe基于GPipe调度,其内存开销很大。Seq1F1B【索引158,Seq1f1b: Efficient sequence-level pipeline parallelism for large language model training,2024,arXiv】将序列分成块并利用1F1B调度来减少峰值内存消耗,同时实现低流水线气泡率。DynaPipe【索引160,Dynapipe: Optimizing multi-task training through dynamic pipelines,2024,EuroSys】对具有可变长度输入的多任务LLM训练使用动态微批次方法。它引入了内存感知的自适应调度算法和预先的通信规划,以进一步降低流水线气泡率。Tessel【索引155,Tessel: Boosting distributed execution of large dnn models via flexible schedule search,2024,HPCA】是一个两阶段方法,包括重复模式构建和调度完成,为指定的分区策略自动搜索高效的流水线调度。DISTMM【索引161,Distmm: Accelerating distributed multimodal model training,2024,NSDI 24】启动加倍的微批次以绕过多模态训练大批量需求引起的依赖障碍,从而减少空闲周期。GraphPipe【索引162,Graphpipe: Improving performance and scalability of dnn training with graph pipeline parallelism,2024,arXiv】保留DNN图拓扑并将其划分为可并行执行的阶段,以提高流水线利用率并减少内存消耗。
内存不平衡问题的解决方案。流水线并行通常在开始阶段注入更多的微批次以提高流水线利用率,导致这些阶段的激活内存消耗更高。为解决此问题,BPipe【索引163,Bpipe: Memorybalanced pipeline parallelism for training large language models,2023,ICML】和MPress【索引164,Mpress: Democratizing billion-scale model training on multi-gpu servers via memory-saving inter-operator parallelism,2023,HPCA】在运行时采用D2D(设备到设备)传输,将中间激活张量从高负载GPU交换到低负载GPU。MPress还结合了激活重计算来减少内存占用。Chimera【索引136,Chimera: efficiently training large-scale neural networks with bidirectional pipelines,2021,SC】引入了双向流水线,将两个不同方向的流水线结合在一起,以实现更均衡的内存消耗。每个GPU持有两个对称阶段,导致权重复制。Hanayo【索引157,Hanayo: Harnessing wave-like pipeline parallelism for enhanced large model training efficiency,2023,SC】将双向流水线转变为两个数据并行流水线,以消除权重复制,并通过对称地将多个阶段分配给一个GPU来实现均衡的内存消耗。V-Shape【索引166,Pipeline parallelism with controllable memory,2024,arXiv】将模型划分为设备数量两倍的阶段,其中两半阶段以相反顺序放置。通过改变阶段之间的偏移量,V-Shape在峰值内存消耗和气泡利用率之间进行权衡。mCAP【索引165,mcap: Memorycentric partitioning for large-scale pipeline-parallel dnn training,2022,Euro-Par】利用增量分析方法,根据峰值内存使用情况在GPU之间均匀划分模型。
结合激活重计算减少峰值内存。峰值内存消耗限制了流水线并行中活动微批次的数量,从而影响其效率。激活重计算可以有效地减少峰值内存消耗。Varuna【索引167,Varuna: scalable, low-cost training of massive deep learning models,2022,EuroSys】结合流水线并行和激活重计算来实现这一目标。它设计了一个基于静态规则的调度,为给定的流水线枚举,并采用机会主义策略来隐藏抖动和减少气泡。静态调度是基于包括激活重计算时机、激活内存管理和后向计算优先级等约束生成的。为解决内存不平衡和低重计算开销,AdaPipe【索引168,Adapipe: Optimizing pipeline parallelism with adaptive recomputation and partitioning,2024,ASPLOS】采用自适应重计算,支持不同阶段采用不同的重计算策略,并基于1F1B调度进行自适应分区以平衡每个阶段的计算。
4.1.4 序列并行
序列并行以支持长上下文训练。当今LLM的上下文窗口迅速增长,最强大的LLM可以支持数百万个token【索引7,Gemini: a family of highly capable multimodal models,2023,arXiv】。如此超长的序列对LLM训练的内存和计算需求提出了巨大挑战:激活的内存占用线性增加,注意力机制的复杂度为二次方。在后向传播中重计算激活可以减少峰值内存消耗,但也会引入显著的开销(完全重计算时为30%)。大的张量并行度会产生显著的通信开销。序列并行(图8(d))【索引169,Sequence parallelism: Long sequence training from system perspective,2023,ACL;索引170,Reducing activation recomputation in large transformer models,2023,MLS】被提出来适应长序列训练,并在内存容量内有效地将计算分布在多个GPU上。它沿序列维度将输入数据分成多个块,每个块被送到一个GPU进行计算。由于序列并行复制了模型参数,它通常与张量和流水线并行结合使用以扩展LLM训练。当与张量并行一起使用时,序列并行将注意力的内存和计算分布在多个GPU上,但在Transformer层的非张量并行区域会产生冗余的内存消耗和计算。Megatron-SP【索引170,Reducing activation recomputation in large transformer models,2023,MLS】沿序列维度拆分这些计算,以减少冗余的激活计算和内存消耗,而不增加通信。
针对序列并行的注意力机制优化。尽管序列并行划分了内存、计算和通信,二次方的因果注意力在训练效率方面仍然存在显著挑战,包括键值张量通信开销、IO感知的注意力计算开销以及由于因果注意力掩码导致的GPU间负载不均衡。大多数用于注意力的序列并行方法都是基于环(ring-based)的【索引169,Sequence parallelism: Long sequence training from system perspective,2023,ACL;索引173,Megatron context parallelism,2023,docs.nvidia.com;索引175,Ring attention with blockwise transformers for near-infinite context,2024,ICLR;索引176,DISTFLASHATTN: Distributed Memory-efficient Attention for Long-context LLMs Training,2023,arxiv;索引178,Striped attention: Faster ring attention for causal transformers,2023,arXiv;索引179,Burstattention: An efficient distributed attention framework for extremely long sequences,2024,arXiv】。Ring Self-Attention【索引169,Sequence parallelism: Long sequence training from system perspective,2023,ACL】利用序列并行,并通过环形通信计算自注意力,以扩展LLM训练的上下文窗口。它首先在GPU之间以循环方式传输键张量以计算注意力分数,然后基于以类似方式传输的注意力分数和值张量计算自注意力输出。DistFlashAttn【索引176,DISTFLASHATTN: Distributed Memory-efficient Attention for Long-context LLMs Training,2023,arxiv】并发传输键值张量块,以利用IO感知的FlashAttention【索引115,Flashattention: Fast and memory-efficient exact attention with io-awareness,2022,Advances in Neural Information Processing Systems】内核,并通过用后来的token填充早期token的空闲周期来平衡不同GPU的计算。Megatron Context Parallel【索引173,Megatron context parallelism,2023,docs.nvidia.com】也利用FlashAttention内核,并移除了由下三角因果掩码导致的不必要计算。它通过与对称GPU交换一半的块来进一步平衡GPU之间的计算。DistFlashAttn和Context Parallel还通过独立的CUDA流将键值张量通信和注意力计算重叠。Striped Attention【索引178,Striped attention: Faster ring attention for causal transformers,2023,arXiv】通过为每个GPU分配在整个序列中均匀分布的token子集,而不是连续的块,来解决不平衡问题。BurstAttention【索引179,Burstattention: An efficient distributed attention framework for extremely long sequences,2024,arXiv】在每个GPU上使用FlashAttention计算注意力,并利用双缓冲来重叠通信和计算。Blockwise Ring Attention【索引175,Ring attention with blockwise transformers for near-infinite context,2024,ICLR】将Ring Self-Attention【索引169,Sequence parallelism: Long sequence training from system perspective,2023,ACL】扩展到块状注意力,该注意力在小块中计算以减少内存占用。受N体模拟的启发,WallFacer【索引180,Wallfacer: Guiding transformer model training out of the long-context dark forest with n-body problem,2024,arXiv】首先将GPU分为子组,并在每个子组内通过异步AllGather复制查询和键值张量。注意力计算利用多个环形P2P通信来提高效率。最后需要一个异步的ReduceScatter来分发注意力输出。
非环形序列并行方法。DeepSpeed-Ulysses【索引172,System optimizations for enabling training of extreme long sequence transformer models,2024,PODC】与之前的基于环的方法不同,它划分头维度而不是序列维度,并利用All-to-All将分区维度从序列转换到头。DeepSpeed-Ulysses可以与现有的注意力实现(如FlashAttention)无缝结合,并且GPU之间的工作负载自然平衡。然而,DeepSpeed-Ulysses的并行度受到头数量的限制,特别是对于使用MQA【索引11,Fast transformer decoding: One write-head is all you need,2019,arXiv】和GQA【索引12,GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints,2023,arXiv.org】的LLM。LoongTrain【索引174,Loongtrain: Efficient training of long-sequence llms with head-context parallelism,2024,arXiv】和USP【索引171,Usp: A unified sequence parallelism approach for long context generative ai,2024,arXiv】是并发工作,它们整合了DeepSpeed-Ulysses和Ring Attention的优点。它们将GPU组织成二维网格,形成混合的ulysses式和环式进程组。在训练期间,它们首先在ulysses组之间执行All-to-All,将分区从序列切换到头维度,然后在环组之间使用Ring-Attention执行注意力计算。LoongTrain进一步提出了Double-Ring-Attention,以充分利用可用的带宽进行节点间通信,并重叠通信与计算。DSP【索引177,Dsp: Dynamic sequence parallelism for multi-dimensional transformers,2024,arXiv】根据多维Transformer(如DiT【索引246,Scalable diffusion models with transformers,2023,ICCV】)中的计算阶段动态切换并行维度。
4.1.5 专家并行
专家混合(MoE)模型及其并行策略。专家混合(Mixture-of-Experts, MoE)是目前LLM中最流行的稀疏模型。MoE在显著增加LLM参数数量的同时,通过条件计算【索引247,Estimating or propa- ´ gating gradients through stochastic neurons for conditional computation,2013,arXiv】并没有大幅增加计算成本。MoE的基本框架如图9所示,包括多个处理训练数据子集的专家网络和一个应用路由算法将输入token分配给不同专家网络的门控网络。MoE使得训练参数超过万亿规模的大型模型成为可能,并据称被用于流行的LLM模型,如Mixtral 8x7B【索引248,Mixtral of experts,2024,arXiv】和DeepSeek2【索引13,Deepseek-v2: A strong, economical, and efficient mixture-of-experts language model,2024,arXiv】。
Fig. 9: 专家并行。虚线突出了Transformer模型中的MoE组件,其中每个设备为一个专家维护一份副本,并基于All-to-All通信进行协作。
稀疏激活与分布式MoE训练。随着模型规模的增加,所有专家无法在单个设备上容纳和训练。因此,GShard【索引15,Gshard: Scaling giant models with conditional computation and automatic sharding,2020,arXiv】将MoE的思想扩展到分布式环境中的Transformer,其中专家分布在不同的工作节点上,并通过All-to-All通信进行协作,如图9所示。后续的专家并行研究通常遵循相同的范式。例如,Switch Transformer【索引181,Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity,2022,JMLR】在T5模型上集成了分布式MoE训练的设计。但与GShard中使用的top-2路由算法不同,Switch Transformer将每个token只路由到top-1专家,以最大化计算效率。此外,DeepSpeed-MoE【索引183,Deepspeed-moe: Advancing mixture-of-experts inference and training to power nextgeneration ai scale,2022,ICML】提出了一种新的分布式MoE架构,该架构在每个工作节点中应用共享专家,并在更深的层中放置更多专家,以平衡通信成本和训练准确性。
MoE与3D并行的结合。专家并行可以有效地与传统的3D并行集成。例如,GShard、Switch Transformer和DeepSpeed-MoE都将专家并行视为混合并行的一个正交维度。为了高效的混合训练,DeepSpeed-TED【索引249,A hybrid tensor-expert-data parallelism approach to optimize mixture-of-experts training,2023,ICS】提出了一种结合数据、张量和专家并行的混合并行算法,以支持MoE模型的训练。作者将MoE参数划分为预定义大小的“瓦片”,以避免优化器内存峰值过高,并提出了通信优化,如重复Token丢弃(DTD)和激活检查点,以消除All-to-All通信中的重复数据。然而,由于MoE的动态性,选择最优的混合并行方案具有挑战性,并且在运行时切换不同的并行策略也会产生巨大的开销。因此,一些研究如Tutel【索引182,Tutel: Adaptive mixture-of-experts at scale,2023,MLS】设计了一种自适应并行切换算法,该算法对所有可能的最优策略应用相同的分布模型布局,并可以在每次迭代中动态切换并行策略而无需任何额外开销。
MoE中的计算优化。由于通用矩阵乘法(GeMMs)要求所有专家输入的大小一致,现有的MoE训练框架通常执行token丢弃和填充以匹配相同的专家容量,这浪费了计算资源。Megablocks【索引185,Megablocks: Efficient sparse training with mixture-of-experts,2023,MLS】通过实现块稀疏矩阵乘法来优化分组GeMMs,并支持在单个内核中为专家计算提供不同的批大小,以避免MoE训练中不必要的token丢弃。另一个支持分组GeMMs的框架是ScatterMoE【索引184,Scattered mixture-ofexperts implementation,2024,arXiv】,它实现了ParallelLinear内核,该内核融合了分组GeMMs和分散的读写操作,以减少top-k(k ≥ 2)门控的内存占用。
MoE中的通信优化。专家并行中的All-to-all通信会严重影响MoE的训练效率,尤其是在网络环境较差的情况下。现有的分布式训练系统试图通过将通信任务与计算任务重叠来优化MoE的性能,从而隐藏部分通信成本。例如,Tutel【索引182,Tutel: Adaptive mixture-of-experts at scale,2023,MLS】沿专家容量维度将输入张量分组,并在不同组之间重叠计算和通信以隐藏All-to-All开销。FasterMoE【索引191,Fastmoe: A fast mixture-of-expert training system,2021,arXiv;索引192,Fastermoe: modeling and optimizing training of large-scale dynamic pre-trained models,2022,PPoPP】使用与Tutel类似的策略,但沿专家维度拆分张量。此外,Tutel【索引182,Tutel: Adaptive mixture-of-experts at scale,2023,MLS】还通过在节点间交换数据之前将小消息聚合成单个大块来优化All-to-All内核实现。这种优化也被FasterMoE和ScheMoe【索引187,Schemoe: An extensible mixture-of-experts distributed training system with tasks scheduling,2024,EuroSys】使用。基于Tutel中的重叠策略,PipeMoE【索引186,Pipemoe: Accelerating mixtureof-experts through adaptive pipelining,2023,IEEE INFOCOM】根据工作负载对通信和计算任务的执行时间进行建模,并设计了一种自适应算法来找到最优的分区数量以最小化训练时间。ScheMoE【索引187,Schemoe: An extensible mixture-of-experts distributed training system with tasks scheduling,2024,EuroSys】在All-to-All通信前考虑数据压缩方法,并将耗时操作模块化,包括数据压缩、集体通信和专家计算。然后,ScheMoE提出了一种自适应最优调度算法来流水线化通信和计算操作以提高训练效率。
通过细粒度任务调度减少通信开销。专家并行通常与MoE训练中的其他并行策略相互作用。通过细粒度的任务调度可以减少通信开销。例如,Lina【索引188,Accelerating distributed {MoE} training and inference with lina,2023,USENIX ATC 23】系统地分析了分布式训练和推理期间MoE的All-to-All开销,并发现当All-to-All与AllReduce操作重叠时,其延迟会延长。Lina提出优先处理All-to-All而不是AllReduce,以提高其带宽并减少其在分布式训练中的阻塞时间。此外,Lina结合了张量分区和流水线技术,以执行类似于Tutel的微操作调度。Lina还在推理期间根据专家热度动态调度资源以最小化开销。Janus【索引189,Janus: A unified distributed training framework for sparse mixture-of-experts models,2023,SIGCOMM】设计了一种以数据为中心的范式,保持数据在原位,并根据参数服务器在GPU之间移动专家。这种以数据为中心的范式使用细粒度的异步通信,并允许专家使用非阻塞通信原语(如pull)在GPU之间移动。Janus实现了一种拓扑感知策略,以有效地在节点之间拉取专家,并支持专家预取,将所有外部专家拉取到本地CPU内存。
模型-系统协同设计优化MoE训练。有一些研究从模型-系统协同设计的角度优化MoE训练。例如,TA-MoE【索引190,Ta-moe: Topologyaware large scale mixture-of-expert training,2022,Advances in Neural Information Processing Systems】为大规模MoE训练提出了一种拓扑感知的路由策略。TA-MoE将分派问题抽象为一个优化目标,以在不同拓扑下获得目标分派模式,并根据分派模式设计了拓扑感知的辅助损失。这种方法自适应地路由数据以适应底层拓扑,而不会牺牲模型准确性。
负载均衡。由于MoE的稀疏和条件计算特性,一个热门的专家可能会比其他专家接收更多的token(通常由较差的路由算法引起),导致严重的负载不均衡,影响MoE的训练效率。FasterMoE【索引192,Fastermoe: modeling and optimizing training of large-scale dynamic pre-trained models,2022,PPoPP】提出了“影子专家”(shadowing experts)方法,根据前几轮迭代的工作负载,动态地将热门专家的参数广播到所有其他GPU。通过将热门专家的工作负载分散到不同设备上,“影子专家”方法减少了专家热度倾斜的影响。SmartMoE【索引193,Smartmoe: Efficiently training {Sparsely-Activated} models through combining offline and online parallelization,2023,USENIX ATC 23】采用两阶段方法来搜索负载均衡的最优并行方案。首先,SmartMoE设计了一个数据敏感的性能模型,将并行方案分为池,池内切换并行模式的成本相对较低。然后,SmartMoE可以在在线阶段切换到适当的并行模式(在SmartMoE中称为专家放置)以保持负载均衡。FlexMoE【索引194,Flexmoe: Scaling large-scale sparse pre-trained model training via dynamic device placement,2023,ACM on Management of Data】发现专家到设备的映射分布在短时间内不会显著变化,因此它引入了细粒度的复制专家并行,将重载的专家复制到多个设备上。FlexMoE监控数据工作负载,并使用三种放置调整原语(即扩展、收缩、迁移)在超过平衡比率时生成最优放置方案。Prophet【索引1195,Prophet: Fine-grained load balancing for parallel training of large-scale moe models,2023,CLUSTER】为大规模MoE模型提出了一种系统化、细粒度且高效的负载均衡训练方法。以MoE模型、设备池和token分布为输入,Prophet的规划器迭代搜索和评估专家放置,最终输出一个均衡的专家放置方案。此外,Prophet使用逐层调度策略隐藏这些资源分配操作的开销。
4.2 自动并行
自动并行的动机。对于任意的DNN模型和GPU集群,存在大量的并行选项,包括单个层的划分及其划分程度。设计能够最大化训练效率的手工混合并行方法是一个耗时且知识密集的过程,需要专家对模型架构、硬件特性以及并行策略中涉及的复杂权衡有深入的理解。此外,高效实现最优并行策略通常需要大量的人力。为解决这些挑战,自动并行作为一种有前途的解决方案应运而生,它旨在为给定的DNN模型在特定的GPU集群上自动确定最有效的并行策略。通过利用复杂的算法和启发式方法,自动并行系统可以分析模型架构、硬件规格和性能特征,以确定数据、张量和流水线并行等并行技术的最佳组合。这种方法简化了跨各种模型和基础设施优化分布式训练的过程,提高了整体效率并减少了手动工作。此外,自动并行可以适应变化的硬件配置和模型架构,自动调整并行策略以保持最佳性能。
4.2.1 通用框架
通用自动并行框架的方法论。通用自动并行框架专注于在特定计算集群上自动并行化通用DNN。这些框架通常遵循一个三步过程:(1)定义并行策略的搜索空间;(2)开发性能模型以衡量不同策略的训练效率;(3)设计算法以高效地识别最优并行策略。
探索混合数据和流水线并行。一些工作探索了用于DNN训练优化的混合数据和流水线并行策略的搜索空间。这些方法专注于自动划分DNN并设计流水线调度以提高流水线利用率。PipeDream【索引134,Pipedream: Fast and efficient pipeline parallel dnn training,2018,arXiv】用最慢阶段的执行时间来衡量流水线分区的效率,并开发了一种动态规划算法,通过最小化最慢阶段来均匀划分DNN。DAPPLE【索引152,Dapple: A pipelined data parallel approach for training large models,2021,PPoPP】建立了一个分析模型来估计一个分区策略的执行时间,并使用动态规划来确定最优的流水线分区。AutoPipe【索引205,Autopipe: A fast pipeline parallelism approach with balanced partitioning and micro-batch slicing,2022,CLUSTER】构建了一个模拟器来模拟流水线执行,并提出了一种启发式算法来获得均衡的分区。AutoPipe还自动拆分微批次以减少预热阶段的延迟。一些设备放置方法【索引206,Device placement optimization with reinforcement learning,2017,ICML;索引207,Spotlight: Optimizing device placement for training deep neural networks,2018,ICML;索引208,Efficient algorithms for device placement of dnn graph operators,2020,Advances in Neural Information Processing Systems】使用强化学习来预测流水线并行的最优算子放置。
探索自动化数据和模型并行。研究人员还通过沿不同维度划分算子来探索自动化数据和模型并行。OptCNN【索引198,Exploring hidden dimensions in parallelizing convolutional neural networks,2018,ICML】沿其输出张量中所有可分维度划分算子,并利用分析性能模型来选择最优的并行策略,包括可并行化的维度和并行度,这定义了如何跨不同设备并行化单个层。FlexFlow【索引199,Beyond data and model parallelism for deep neural networks,2019,MLS】进一步将搜索空间扩展到Sample-Operator-Attribute-Parameter (SOAP),这包括输入和输出张量中几乎所有可分维度,并引入了一种新颖的执行模拟器以进行准确的性能建模。FlexFlow通过MCMC采样高效地找到最优的并行策略。Tofu【索引200,Supporting very large models using automatic dataflow graph partitioning,2019,EuroSys】和HyPar【索引201,Hypar: Towards hybrid parallelism for deep learning accelerator array,2019,HPCA】开发了动态规划算法,最小化总通信成本而不是端到端性能,以确定混合数据和模型并行空间中每个算子的最优分区。TensorOpt【索引204,Tensoropt: Exploring the tradeoffs in distributed dnn training with auto-parallelism,2021,IEEE Transactions on Parallel and Distributed Systems】在给定的内存预算下,使用前沿跟踪算法优化并行策略。AutoMap【索引202,Automap: Towards ergonomic automated parallelism for ml models,2021,arXiv】采用蒙特卡洛树搜索(MCTS)为通过学习的评分器选择的一组重要算子选择一系列由PartIR【索引212,Partir: Composing spmd partitioning strategies for machine learning,2024,arXiv】定义的分区规则。整个并行策略是从通过所选算子的策略中传播出来的。
探索自动化数据、模型和流水线并行。最近的工作还设计了用于自动化数据、模型和流水线并行的方法。Piper【索引209,Piper: Multidimensional planner for dnn parallelization,2021,Advances in Neural Information Processing Systems】设计了一种两级动态规划方法,以找到最优的混合数据、张量和流水线并行,并结合了激活重计算。它首先将模型划分为流水线的小分区,然后在每个分区内拆分算子。Alpa【索引21,Alpa: Automating interand {Intra-Operator} parallelism for distributed deep learning,2022,OSDI 22】通过将并行视为两个层次级别:算子间和算子内并行,构建了一个全面的空间。然后,它在每个并行级别自动推导出高效的并行执行计划。Unity【索引210,Unity: Accelerating {DNN} training through joint optimization of algebraic transformations and parallelization,2022,OSDI 22】通过将并行化和代数变换表示为统一并行计算图上的替换,来联合优化它们。Aceso【索引211,Aceso: Efficient parallel dnn training through iterative bottleneck alleviation,2024,EuroSys】提出了一种迭代瓶颈缓解方法,以显著减少搜索时间。它在每一步识别性能瓶颈,并调整策略以减轻瓶颈,直到收敛。nnScaler【索引213,{nnScaler}:{Constraint-Guided} parallelization plan generation for deep learning training,2024,OSDI 24】引入了三个原语,以支持用任意分区和划分模型的时空调度来组合搜索空间。领域专家可以对原语施加约束,以构建有效且小的搜索空间,这些空间可以以低开销自动探索。AutoDDL【索引215,Autoddl: Automatic distributed deep learning with near-optimal bandwidth cost,2024,IEEE Transactions on Parallel and Distributed Systems】定制了一种坐标下降算法,通过迭代更新每层的SBP【索引214,Oneflow: Redesign the distributed deep learning framework from scratch,2021,arXiv】分布,并快速发现具有近乎最优通信成本的最优策略。
通用自动并行框架的系统支持。通用自动并行框架除了需要快速的优化算法来发现最优并行策略外,还需要对各种并行策略的高效系统支持。这是因为并行通常涉及复杂的计算和通信算子,特别是对于划分算子的模型并行。先前的工作已经开发了支持广泛并行策略的高效系统,要么是基于现代深度学习框架【索引21,Alpa: Automating interand {Intra-Operator} parallelism for distributed deep learning,2022,OSDI 22;索引213,{nnScaler}:{Constraint-Guided} parallelization plan generation for deep learning training,2024,OSDI 24】,要么是从头开始实现【索引199,Beyond data and model parallelism for deep neural networks,2019,MLS】。Mesh-TensorFlow【索引196,Meshtensorflow: Deep learning for supercomputers,2018,Advances in neural information processing systems】观察到实现并行策略的内在复杂性,并首次提出将设备集群抽象为多维网格,并将并行抽象为划分迭代空间(即张量维度)。通过映射张量和网格维度,可以轻松地以高性能实现混合数据和模型并行策略。例如,数据和模型并行分别拆分批次和隐藏维度。GSPMD【索引197,Gspmd: general and scalable parallelization for ml computation graphs,2021,arXiv】进一步提供了一种统一的方式,通过基于JAX【索引250,Jax: composable transformations of python+ numpy programs,2018】和XLA【索引251,Xla: Optimizing compiler for machine learning,2017,github.com】的简单张量分片注释来实现各种通用并行方案。OneFlow【索引214,Oneflow: Redesign the distributed deep learning framework from scratch,2021,arXiv】提出了用于分区的SBP(Split, Broadcast, Partial-value)抽象,并允许用户为张量指定放置和SBP签名以实现不同的并行策略。PartIR【索引212,Partir: Composing spmd partitioning strategies for machine learning,2024,arXiv】将模型与其分区解耦,并设计了一个编译器堆栈,供用户通过调度逐步组合SPMD分片策略。与TVM【索引252,{TVM}: An automated {End-to-End} optimizing compiler for deep learning,2018,OSDI 18】类似,Slapo【索引203,Slapo: A schedule language for progressive optimization of large deep learning model training,2024,ASPLOS】为并行化和子图优化(如算子融合和激活检查点)定义了一套全面的调度原语。这些调度与执行解耦,并保留了原始模型结构以进行渐进式优化。
4.2.2 Transformer特定框架
为Transformer量身定制的自动化系统。由于LLM基于Transformer架构,最近的工作为Transformer量身定制了自动化系统。DeepSpeed-Autotuning【索引216,Deepspeed autotuning,2021,deepspeed.ai】自动调整系统旋钮,以在用户定义的调优空间中找出与性能相关的良好配置,包括并行度。Galvatron【索引217,Galvatron: Efficient transformer training over multiple gpus using automatic parallelism,2022,VLDB Endowment】设计了一种动态规划算法,以生成最有效的混合数据、张量和流水线并行策略。Merak【索引218,Merak: An efficient distributed dnn training framework with automated 3d parallelism for giant foundation models,2023,IEEE】引入了一个用于非侵入式自动并行的自动模型分区器和一个高性能的3D并行运行时引擎,以增强可用资源的利用率。Colossal-AI【索引219,Colossal-Auto: Unified Automation of Parallelization and Activation Checkpoint for Large-scale Models,2023,arXiv;索引253,Colossal-ai: A unified deep learning system for largescale parallel training,2023,ICPP】为混合数据、张量、序列和流水线并行的模块化使用提供了统一的接口。Galvatron-BMW【索引220,Improving automatic parallel training via balanced memory workload optimization,2024,IEEE Transactions on Knowledge and Data Engineering】将Galvatron的空间扩展到包括分片数据并行和激活重计算,并在考虑内存消耗和计算的同时,搜索最大化硬件利用率的最优策略。
4.3 异构并行
异构并行在LLM训练中的应用。LLM训练日益增长的计算需求推动了异构硬件的发展,这些硬件利用了多样化的计算资源和全球分布的设备。这种异构性也反映在模型架构中,特别是在使用人类反馈强化学习(RLHF)时。利用异构硬件和多样化的模型架构已成为高效训练LLM的必要条件。
4.3.1 异构硬件
利用异构和地理分布式硬件。LLM训练的大量计算需求推动了加速器的发展,导致集群中混合了不同类型的设备和不均匀的互连带宽。此外,现代数据和计算集群由于电力短缺等因素,常常是全球分布的。这些现象推动了异构并行的采用,该并行利用多样化的计算资源和地理上分布的设备来加速LLM训练。
利用异构计算资源。一些工作利用CPU、GPU和专用加速器等异构计算资源来提升LLM的性能。这些设备在计算、内存容量和互连带宽上的差异给高效的LLM预训练带来了挑战。HetPipe【索引221,Hetpipe: Enabling large dnn training on (whimpy) heterogeneous gpu clusters through integration of pipelined model parallelism and data parallelism,2020,USENIX ATC 20】将异构集群划分为多个虚拟工作单元。每个虚拟工作单元使用流水线并行处理小批量数据,而不同的虚拟工作单元采用异步数据并行来提高吞吐量。AccPar【索引222,Accpar: Tensor partitioning for heterogeneous deep learning accelerators,2020,HPCA】提出了灵活的张量分区来平衡不同加速器的计算,并使用动态规划为DNN自动决定异构设备间的张量分区。Whale【索引223,Whale: Efficient giant model training over heterogeneous {GPUs},2022,USENIX ATC 22】提出了一个统一的抽象来简化在异构集群上并行训练巨型模型的努力。它通过自动图优化无缝适应异构GPU,并利用硬件信息平衡工作负载。AMP【索引224,Amp: Automatically finding model parallel strategies with heterogeneity awareness,2022,Advances in Neural Information Processing Systems】利用一个异构感知的性能模型来找到最优的混合数据、张量和流水线并行策略。HPH【索引226,Hph: Hybrid parallelism on heterogeneous clusters for accelerating large-scale dnns training,2022,CLUSTER】根据计算通信比按降序排列不同的GPU到各个阶段,并将模型分区公式化为一个整数规划问题,以最小化迭代时间。Pathways【索引225,Pathways: Asynchronous distributed dataflow for ml,2022,MLS】采用分片数据流模型和异步组调度来在异构集群上高效执行ML模型。SDPIPE【索引227,Sdpipe: A semi-decentralized framework for heterogeneity-aware pipelineparallel training,2023,VLDB Endowment】引入了一种半去中心化方案,该方案将通信模型同步去中心化,并将组调度过程中心化,以利用异构设备进行流水线并行。HAP【索引228,Hap: Spmd dnn training on heterogeneous gpu clusters with automated program synthesis,2024,EuroSys】使用基于A*的搜索算法来生成最优的张量分片策略、跨异构设备的分片比例以及分布式训练的通信方法。PipePar【索引229,Pipepar: Enabling fast dnn pipeline parallel training in heterogeneous gpu clusters,223,Neurocomputing】提出了一种动态规划算法,考虑GPU和网络带宽的异构性,将模型划分为流水线阶段。
利用地理分布式设备。其他一些工作探索了具有低网络带宽的地理分布式设备,以提高训练效率。Yuan等人【索引230,Decentralized training of foundation models in heterogeneous environments,2022,Advances in Neural Information Processing Systems】将LLM划分为计算任务单元,并提出了一种新颖的调度算法,以高效利用由慢速异构网络连接的一组异构设备,进行混合数据和流水线并行。SWARM并行【索引231,Swarm parallelism: Training large models can be surprisingly communication-efficient,2023,ICML】将模型划分为大小相等的阶段,并优先将输入路由到具有较低延迟的稳定对等节点以实现工作负载平衡。它还自适应地在阶段间移动设备以最大化训练吞吐量。FusionAI【索引232,Fusionai: Decentralized training and deploying llms with massive consumer-level gpus,2023,arXiv】将训练计算图(DAG)拆分为子图(subDAG),并生成一个负载均衡的任务调度,以利用通过低带宽连接的异构消费级GPU进行流水线训练。通信压缩方法,如CocktailSGD【索引254,Cocktailsgd: Fine-tuning foundation models over 500mbps networks,2023,ICML】,也可以被用来在低带宽集群中高效地训练LLM。
Fig. 10: RLHF示例。推理过程:1 Actor模型根据给定的查询生成响应。2 Critic模型、奖励模型和参考模型使用查询和响应对通过推理生成训练所需的值、分数和KL散度。训练过程:3 Actor模型和Critic模型使用在推理过程中收集的数据通过梯度下降更新它们的权重。
4.3.2 异构模型
模型异构性:以RLHF为例。在LLM训练过程中,异构性不仅体现在硬件上,也体现在模型上。训练可能涉及多个不同模型的交互。一个具体的例子是来自人类反馈的强化学习(RLHF)。RLHF是一种旨在使AI系统更紧密地与人类偏好对齐的训练方法【索引255,Training language models to follow instructions with human feedback,2022,arXiv】,它利用人类在判断适当行为方面的优势,而不是演示。这种方法受到了广泛关注,特别是用于微调大型语言模型。然而,由于近端策略优化(PPO)【索引256,Proximal policy optimization algorithms,2017】算法的特殊性,模型异构性被引入到RLHF训练中,使得RLHF的训练过程与预训练和监督微调有很大不同。
RLHF的PPO阶段的异构性。原则上,RLHF包括三个不同阶段:第一阶段是监督微调,第二阶段是奖励模型的训练,第三阶段是PPO训练。模型异构性呈现在第三阶段,如图10所示。PPO训练阶段包括两个不同的过程,即生成数据的推理过程和更新actor模型与critic模型权重的训练过程。PPO训练是通过这两个过程的协作来执行的。此外,训练阶段引入了更高的内存成本,因为我们需要同时服务多个自回归生成模型和奖励模型的副本,以及更多的时间成本,因为我们必须等待经验生成完成后才能更新权重。
RLHF训练框架。许多框架被提出来用于RLHF训练。例如,DeepSpeed-Chat【索引233,Deepspeedchat: Easy, fast and affordable rlhf training of chatgpt-like models at all scales,2023,arXiv】使用混合引擎(Hybrid Engine)在训练和推理之间无缝切换模型分区,例如在推理期间使用张量并行提高吞吐量,在训练期间使用ZeRO【索引145,Zero: Memory optimizations toward training trillion parameter models,2020,SC20】或LoRA【索引257,Lora: Low-rank adaptation of large language models,2021,arXiv】提高内存利用率,为RLHF训练提供了出色的系统效率。HuggingFace TRL【索引234,Trl - transformer reinforcement learning,,github.com】可以充分利用各种参数高效微调(PEFT)方法,如LoRA或QLoRA【索引258,Qlora: Efficient finetuning of quantized llms,2024,Advances in Neural Information Processing Systems】,以节省内存成本,并使用unsloth【索引259,unsloth - finetune llama 3, mistral, phi-3 & gemma 2-5x faster with 80,,github.com】设计的专用内核来提高RLHF的训练速度。ColossalAI-Chat【索引253,Colossal-ai: A unified deep learning system for largescale parallel training,2023,ICPP】是另一个端到端的RLHF训练框架,它也支持LoRA,并支持使用ZeRO【索引145,Zero: Memory optimizations toward training trillion parameter models,2020,SC20】来减少内存冗余。
RLHF中的模型放置和调度策略。然而,上述工作对模型放置采用扁平化策略,即将RLHF中的四个模型放置在同一设备上,然后使用ZeRO或LoRA等方法来最小化内存成本。但仅使用ZeRO在训练更大模型时会导致内存瓶颈,而使用LoRA等高效参数微调策略会损害模型性能。为解决此问题,OpenRLHF【索引235,Openrlhf: An easy-to-use, scalable and high-performance rlhf framework,2024,arXiv】使用Ray【索引260,Ray: A distributed framework for emerging {AI} applications,2018,OSDI 18】和vLLM【索引261,Efficient memory management for large language model serving with pagedattention,2023,SOSP】将奖励模型分布到不同设备,避免将PPO中的所有四个模型都放置在同一设备上。同样,自适应放置和并行(APP)框架【索引236,An Adaptive Placement and Parallelism Framework for Accelerating RLHF Training,2023,arxiv】提出了另外两种模型放置策略,即交错策略和分离策略。它抓住了PPO训练期间生成部分和训练部分可以独立运行的事实,通过将它们放置在不同设备上可以消除一些串行化,这会引入额外的通信,但可以很好地与计算重叠。
RLHF中的细粒度调度。同时,有一些工作将前两个阶段的并行策略以细粒度调度的方式应用于RLHF的第三阶段。例如,ReaLHF【索引237,Realhf: Optimized rlhf training for large language models through parameter reallocation,2024,arXiv】通过重新分配参数,为第三阶段的不同子阶段切换最合适的并行模式,极大地增加了优化空间。PUZZLE【索引238,{PUZZLE}: Efficiently aligning large language models through {LightWeight} context switch,2024,USENIX ATC 24】根据不同阶段的亲和性重新调度任务执行顺序,使得亲和性更好的阶段可以有效地覆盖执行并提高训练效率。
5 计算优化
有效利用AI加速器的计算能力。当今的AI加速器在FLOPs方面提供了前所未有的计算能力。然而,要有效地将这些FLOPs充分利用,需要复杂的优化技术。本节介绍计算优化的系统和技术,以有效利用GPU FLOPs。我们首先详细阐述算子优化,包括核心的注意力算子优化和通过编译器的自动优化。通过利用大规模并行和针对底层硬件特性的高效多级内存访问,算子和计算图获得了显著的性能提升。其次,详细介绍混合精度训练,其中通过降低精度来加速计算。16位浮点混合训练已成为大多数训练系统中的事实标准。低至1位的低位定点数也已被研究并用于提高训练效率。
Fig. 11: 分布式LLM训练的计算优化研究。
5.1 算子优化
算子优化的分类。算子优化可分为手动优化和自动优化。手动优化主要集中在注意力算子上,而自动优化则应用得更广泛。
5.1.1 手动优化的注意力算子
注意力机制是Transformer的核心。注意力作为Transformer架构的核心,在LLM的训练效率中扮演着至关重要的角色。给定一个查询q,以及键k1, k2, ..., kn和值v1, v2, ..., vn的列表,其中q, ki, vi ∈ Rd,注意力的计算如下:
自注意力相对于序列长度表现出二次方的时间和内存复杂度。自注意力带来的巨大内存消耗和对高带宽内存(HBM)的频繁访问,限制了Transformer模型的性能和上下文长度。已有大量工作致力于优化自注意力。我们专注于精确注意力优化,而像线性注意力这样的有损优化不在我们的讨论范围之内。
内存高效的注意力机制。内存高效的注意力机制主要旨在减轻巨大的内存成本。Rabe等人【索引287,Self-attention does not need o(n2) memory,2021,arXiv】证明了自注意力只需要$O(logn)$的内存复杂度,而不是$O(n^2)$。通过采用懒惰softmax,softmax中的除以$\sum_j e^{s_j}$可以延迟到注意力操作的最后。因此,求和可以增量处理,这只需要一个标量(即$O(1)$)来维持中间结果,但不会改变输出。自注意力需要额外的$O(logn)$内存复杂度来保存查询列表的附加索引,以便顺序计算所有查询的结果。
IO感知的注意力机制(FlashAttention系列)。FlashAttention系列进一步展示了在GPU上具有IO感知、高并行度和均衡工作负载的快速且内存高效的精确注意力。在FlashAttention【索引115,Flashattention: Fast and memory-efficient exact attention with io-awareness,2022,Advances in Neural Information Processing Systems】中,提出了一种IO感知的分块(tiling)算法,以减少慢速HBM和快速片上SRAM之间的内存读写次数,这是基于在线softmax实现的。更具体地说,通过跟踪归一化统计数据,包括最大分数和指数化分数的总和,可以一次计算一个块的softmax。因此,分块算法将自注意力中的所有计算操作链(包括矩阵乘法、softmax、矩阵乘法等)融合到一个CUDA内核中,以减少HBM访问。FlashAttention2【索引116,FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning,2023,arxiv】通过在序列长度维度上增加并行度,并改进线程块内数据共享的warp级调度,进一步改善了FlashAttention中占用率低和不必要的共享内存读写问题。此外,流行的训练系统【索引174,Loongtrain: Efficient training of long-sequence llms with head-context parallelism,2024,arXiv】通常采用FlashAttention-2以获得高性能。FlashAttention-3【索引262,Flashattention-3: Fast and accurate attention with asynchrony and low-precision,2024,arXiv】通过挖掘新硬件的能力,加速了在H100 GPU上的注意力计算,因为之前的FlashAttention实现是基于A100 GPU的。基于FlashAttention-2重新设计了一种交错的块级GEMM和softmax算法,以利用异步WGMMA指令隐藏softmax中非GEMM操作。此外,通过利用Tensor Cores和Tensor Memory Accelerator (TMA)的异步性,通过warp专用的软件流水线方案将整体计算与数据移动重叠。Blockwise Parallel Transformer (BPT)【索引263,Blockwise parallel transformers for large context models,2024,Advances in Neural Information Processing Systems】通过扩展FlashAttention中的分块算法来融合前馈网络,进一步减少了巨大的内存需求。
针对特定硬件架构的注意力优化。注意力操作也通过利用硬件特定功能在各种架构上进行了优化。例如,SWattention【索引264,Swattention: designing fast and memory-efficient attention for a new sunway supercomputer,2024,The Journal of Supercomputing】设计了一种两级分块注意力算法,以利用新神威架构的底层硬件,这是在FlashAttention的基础上构建的。类似地,Bikshand等人【索引265,A case study in cuda kernel fusion: Implementing flashattention-2 on nvidia hopper architecture using the cutlass library,2023,arXiv】使用Cutlass库在H100 GPU上实现了FlashAttention-2。他们利用TMA和WarpGroup矩阵乘累加(WGMMA)指令分别优化数据复制和GEMM操作。此外,还基于Cutlass库精心设计了张量布局转换以及两个GEMM之间数据复制和计算的软件流水线。
针对可变长度序列的注意力优化。注意力机制也针对可变长度序列进行了优化,这在分布式LLM训练中很常见。如果将这些可变长度序列填充到最大长度,会产生显著的内存和计算成本。FlashAttention-2通过不可分割地并行化序列长度维度来高效处理可变长度输入。ByteTransformer【索引2266,Bytetransformer: A high-performance transformer boosted for variable-length inputs,2023,IPDPS】专注于用于可变长度输入的无填充Transformer,在计算过程中维护一个位置数组。该数组记录了原始张量和中间打包张量之间有效token的映射关系。用于长序列的融合多头注意力算法采用了针对未填充张量的优化分组GEMM。这种优化减少了与填充相关的内存和计算开销,从而提高了性能。
5.1.2 通过编译器的自动优化
DNN编译器在LLM训练中的作用。DNN编译器在优化LLM训练中的关键计算方面扮演着重要角色。它们自动生成高效的算子内核,极大地减轻了在不同硬件供应商上基于库进行内核优化的负担。通过在训练过程中自动分析计算图来执行算子融合。
高效的算子内核生成。Halide【索引267,Halide: a language and compiler for optimizing parallelism, locality, and recomputation in image processing pipelines,2013,Acm Sigplan Notices】和TVM【索引252,{TVM}: An automated {End-to-End} optimizing compiler for deep learning,2018,OSDI 18】自动生成高性能的算子实现,依赖于多种有效的调度原语,这些原语在各种后端上利用了并行性和数据局部性。此外,Roller【索引268,{ROLLER}: Fast and efficient tensor compilation for deep learning,2022,OSDI 22】优化了在内核实现的巨大搜索空间中搜索最优替代方案的成本。它主要生成一个由加载、存储和计算接口组成的瓦片(tile)内核,然后通过“先放大后横向扩展”(scale-up-then-scale-out)的方法构建完整的算子内核。Triton【索引269,Triton: an intermediate language and compiler for tiled neural network computations,2019,MLPL】提供了一种基于C的语言和编译器,便于表达和优化瓦片张量程序以获得有竞争力的性能。特别是,通过依赖于机器的编译过程支持了分层分块和共享内存分配等有效优化。ALCOP【索引270,Alcop: Automatic load-compute pipelining in deep learning compiler for ai-gpus,2023,MLS】执行自动的加载-计算流水线,以在GPU上的算子中将高延迟的内存访问与计算重叠。通过流水线缓冲区检测以及在复杂循环结构中复杂的索引分析和替换来利用多级流水线。
用于算子融合的图级优化。随着计算核心和内存带宽速度差距的扩大,现代DNN受限于内存访问。通过编译器使用算子融合来挖掘算子间的数据重用。许多编译器工作【索引288,Taso: optimizing deep learning computation with automatic generation of graph substitutions,2019,SOSP;索引289,Ansor: Generating {High-Performance} tensor programs for deep learning,2020,OSDI 20;索引290,Dnnfusion: accelerating deep neural networks execution with advanced operator fusion,2021,PLDI;索引291,Astitch: enabling a new multidimensional optimization space for memory-intensive ml training and inference on modern simt architectures,2022,ASPLOS】通过设置专家规则来执行算子融合。特别是,Chimera【索引271,Chimera: An analytical optimizing framework for effective compute-intensive operators fusion,2023,HPCA】致力于优化计算密集型算子链。算子链首先被分解为一系列计算块,然后根据分析模型选择最优的块执行顺序以最大化数据重用。此外,还设计了可替换的微内核以利用硬件特定的块内优化。Welder【索引272,Welder: Scheduling deep learning memory access via tile-graph,2023,OSDI 23】将计算图降低为瓦片级数据流图,其节点是算子瓦片,边上标记了连接节点重用的张量数据的内存级别。在瓦片级别搜索最大化跨不同内存层次数据重用的算子融合组合。
现代DL框架中的编译器支持。Pytorch2【索引273,Pytorch 2: Faster machine learning through dynamic python bytecode transformation and graph compilation,2024,ASPLOS】提出了两个扩展,即Python级的JIT编译器TorchDynamo和相应的编译器后端TorchInductor,以在各种后端上实现更稳健的图编译,从而在不牺牲Python灵活性的情况下显著提高性能。Slapo【索引203,Slapo: A schedule language for progressive optimization of large deep learning model training,224,ASPLOS】提出了一种调度语言,以将模型执行与定义解耦。通过声明一组调度原语,用户可以转换为高性能内核的模型。JIT-Q【索引274,Jit-q: Justin-time quantization with processing-in-memory for efficient ml training,2024,MLS】提出了权重的即时量化,这使得在训练期间只存储一份高精度权重副本,并根据商业PIM(内存处理)解决方案的内存中ALU增强功能创建低精度权重副本。
5.2 混合精度训练
低精度训练的优势。低精度训练是减少训练大规模模型中计算、存储和通信成本的有效方法。如今的LLM训练通常利用FP16和BF16数据类型。特别是,BF16可以表示与FP32相同的值范围。BF16训练被用于像BLOOM【索引292,Bloom: A 176b-parameter open-access multilingual language model,2022,arXiv】这样的模型中,因为在FP16中当损失标量变得太低时,损失会缓慢发散【索引293,Reproducible scaling laws for contrastive language-image learning,2023,CVPR】。然而,快速的bfloat16支持仅在TPU或使用NVIDIA Ampere系列及之后开发的GPU上可用。此外,还利用混合精度训练和损失缩放等技术来确保数值稳定性,因为降低精度会限制动态范围。8位甚至更低位的训练也正成为量化研究的焦点。
5.2.1 16位浮点
FP16/BF16混合精度训练。流行的训练系统通常采用FP16/BF16混合精度策略来在训练期间降低精度,如Megatron-LM【索引20,Megatron-lm: Training multi-billion parameter language models using model parallelism,2019,arXiv】和ColossalAI【索引253,Colossal-ai: A unified deep learning system for largescale parallel training,2023,ICPP】等工作所强调的。FP16混合精度训练方案【索引275,Mixed precision training,2017,arXiv】利用IEEE半精度格式来存储权重、激活和梯度,用于前向和后向算术运算。为了在降低精度的情况下保持模型准确性,保留一份单精度权重副本,用于在每个优化器步骤中进行累加。还应用了损失缩放来保留小幅度梯度的值。Campo【索引276,Campo:{Cost-Aware} performance optimization for {Mixed-Precision} neural network training,2022,USENIX ATC 22】通过自动图重写优化了FP32和FP16之间转换产生的转换成本。这一点至关重要,因为转换成本有时会抵消低精度带来的性能优势。Campo还使用离线训练的线性回归模型来预测FP32和FP16操作的转换成本和执行时间。BF16【索引277,A study of bfloat16 for deep learning training,2019,arXiv】也广泛用于各个领域的混合精度训练【索引294,High performance monte carlo simulation of ising model on tpu clusters,2019,SC;索引295,Automatic full compilation of julia programs and ml models to cloud tpus,2018,arXiv】。它具有与FP32相同的表示范围,并且不需要为收敛调整超参数。此外,THC【索引278,{THC}: Accelerating distributed deep learning using tensor homomorphic compression,2024,NSDI 24】通过消除解压缩和压缩的需要,解决了参数服务器架构中的计算开销。THC通过均匀同态压缩属性实现了对压缩梯度值的直接聚合,从而提高了效率。
5.2.2 亚8位浮点
FP8混合精度训练。随着支持如FP8等更低精度数据类型的新芯片的发布,混合精度训练被设计为使用更低的精度进行训练。新设计的数据格式结合确保数值稳定性的技术,主要被用来实现深度学习神经网络的FP8训练。Wang等人【索引279,Training deep neural networks with 8-bit floating point numbers,2018,Advances in neural information processing systems】使用一种新的FP8浮点格式来表示数据和进行计算。在浮点累加和权重更新过程中,分别使用了基于块的计算和随机舍入来保持模型准确性。Sun等人【索引280,Hybrid 8-bit floating point (hfp8) training and inference for deep neural networks,2019,Advances in neural information processing systems】提出了一种混合8位浮点训练方法,适用于整个深度学习模型谱系,且没有准确性下降。新颖的混合FP8格式为前向和后向传播分别使用不同的指数位和尾数位,因为前向和后向传播在范围和精度之间有不同的最佳平衡。此外,还使用了损失缩放等技术来避免准确性下降。随着更多支持FP8数据类型的加速器的成熟,提出了一种基于NVIDIA H100 GPU【索引296,Nvidia hopper h100 gpu: Scaling performance,2023,IEEE Micro】的FP8自动混合精度框架(FP8-LM)【索引281,Fp8-lm: Training fp8 large language models,2023,arXiv】,用于训练LLM。该框架逐步集成了8位梯度、优化器状态和分布式并行训练,并指定了FP8低位并行,包括张量、流水线和序列并行。此外,还设计了精度解耦和自动缩放来解决由于动态范围变窄和精度降低导致的数据下溢或上溢问题。FlashAttention-3也采用了块GEMM量化和非相干处理,利用了H100 GPU上对FP8低精度的硬件支持。此外,Rouhani等人【索引282,Microscaling data formats for deep learning,2023,arXiv】通过使用将缩放因子与张量的细粒度子块关联的微缩放数据格式,以最小的准确性损失,在亚8位权重、激活和梯度下训练LLM。
5.2.3 低位定点
低位定点训练。低位定点训练也被研究用于LLM训练。Jetfire【索引283,Jetfire: Efficient and accurate transformer pretraining with int8 data flow and per-block quantization,2024,arXiv】维持一个INT8数据流,其中输入和输出以INT8数据格式加载和存储,以加速计算密集的线性算子和内存密集的非线性算子。此外,利用分块算法通过每块量化方法来挖掘共享内存数据访问,其中执行更高精度的计算,即线性算子的WMMA张量核心操作使用INT32,非线性操作使用FP32,以保持预训练Transformer的准确性。Xi等人【索引284,Training transformers with 4-bit integers,2023,Advances in Neural Information Processing Systems】为Transformer模型提出了一种新颖的INT4训练算法。在前向传播中,激活矩阵首先被转换为块对角哈达玛矩阵,以减轻激活中异常值引起的准确性下降,然后对转换后的矩阵进行量化。在后向传播中,利用比特分裂和杠杆分数采样来根据激活梯度的结构稀疏性选择信息丰富的梯度进行量化。
1位精度训练。最近,LLM的低精度训练已发展到使用1位精度。BitNet【索引285,Bitnet: Scaling 1-bit transformers for large language models,2023,arXiv】在Transformer块内采用了一种新颖的低位精度矩阵乘法,使用1位权重和8位激活。模型权重围绕零进行中心化,以在有限的数值范围内最大化容量,然后使用signnum函数二值化为+1或-1。为确保训练稳定性和准确性,保留了梯度、优化器状态和一份高精度潜在权重副本用于参数更新。在BitNet的基础上,BitNet b1.58【索引286,The era of 1-bit llms: All large language models are in 1.58 bits,2024,arXiv】通过将模型权重降低到三元值{-1, 0, 1}来进一步增强建模能力。权重矩阵通过其平均绝对值进行缩放,并且每个值都四舍五入到-1、0和+1中最接近的整数。
Fig. 12: 分布式LLM训练的内存优化研究。
6 内存优化
LLM训练中的内存消耗分类。LLM训练期间的内存消耗可分为四个关键部分:模型状态、激活、临时缓冲区和内存碎片。
* 模型状态: 包括优化器状态、梯度和模型参数所消耗的内存。在混合精度训练【索引275,Mixed precision training,2017,arXiv】中,模型参数和激活以16位精度存储。训练一个拥有Φ个参数的模型时,存储参数和梯度需要4Φ字节。32位参数、动量和方差的副本各需要4Φ字节,总计12Φ字节。因此,存储模型状态的总内存需求为16Φ字节。
* 激活: 指在前向传递过程中生成的张量。这些张量对于后向阶段的梯度计算至关重要。
* 临时缓冲区: 用于存储中间结果。例如,梯度AllReduce等操作通常将一个桶中的梯度融合到一个扁平化的缓冲区中,然后应用该操作以提高吞吐量。
* 内存碎片: 可能导致尽管有大量可用内存,内存请求却失败的情况。这是因为可用内存可能变得碎片化,没有足够的连续内存来满足内存请求【索引145,Zero: Memory optimizations toward training trillion parameter models,2020,SC20】。
内存优化技术。为了解决LLM训练的内存限制,已经提出了各种内存高效技术。这些技术包括:激活重计算策略,通过增加计算来减少内存使用;冗余减少方法,最小化训练进程间的数据重复;碎片整理技术,优化内存分配和释放以减少碎片并提高内存利用率;以及交换和卸载方法,利用CPU内存和NVMe SSD来补充GPU内存。图12概述了这些内存高效LLM训练优化的分类。
6.1 激活重计算
激活重计算的原理和目标。在模型训练的后向阶段,激活对于计算梯度至关重要。随着模型规模的增加,训练期间存储这些激活所需的内存可能超过GPU内存容量,从而限制了可以训练的模型规模。激活重计算【索引325,Training deep nets with sublinear memory cost,2016,arXiv】提供了一种解决方案,通过在前向传递期间有策略地丢弃某些激活,并在后向传递期间根据需要重新计算它们。这种方法已成为LLM训练中减少内存消耗的事实标准。有效激活重计算的关键是在内存节省和额外计算开销之间取得平衡。
激活重计算方法的分类。我们将这些方法分为两种主要方法:静态驱逐和动态驱逐。静态驱逐方法通常涉及为特定模型架构或模块量身定制驱逐策略。相比之下,动态驱逐方法在运行时实时做出决策,无需事先了解模型。尽管静态方法需要为新模型进行修改,但大多数LLM的结构共享相似的架构,使得这些策略在LLM训练中可以普遍应用。尽管动态驱逐方法具有固有的灵活性,但尚未在LLM训练中广泛采用。然而,我们仍在本节中探讨一些相关工作以供进一步参考。
6.1.1 静态驱逐
静态驱逐策略。静态驱逐涉及建立一个固定的计划,在前向传递期间丢弃激活,然后在后向传递期间重新计算它们。Checkmate【索引300,Checkmate: Breaking the memory wall with optimal tensor rematerialization,2020,MLS】将此激活重计算问题公式化为一个混合整数线性规划,以确定静态深度学习模型的最佳重物质化计划。然而,由于搜索空间巨大,Checkmate难以扩展到像LLM这样的大型模型。
针对LLM的定制化静态驱逐策略。最近,一些工作提出了为LLM训练量身定制的激活重计算策略。Selective-checkpointing【索引170,Reducing activation recomputation in large transformer models,2023,MLS】选择性地丢弃内存密集型注意力模块的激活。FlashAttention【索引115,Flashattention: Fast and memory-efficient exact attention with io-awareness,2022,Advances in Neural Information Processing Systems】将注意力模块融合到单个内核中,并采用选择性检查点来减少内存消耗。DistFlashAttn【索引176,DISTFLASHATTN: Distributed Memory-efficient Attention for Long-context LLMs Training,2023,arxiv】解决了长序列中由注意力模块重计算引起的高计算开销问题,采用了一种感知重物质化的梯度检查点策略。具体来说,DistFlashAttn将检查点放在FlashAttention内核的输出处,而不是在Transformer层边界,从而在后向传递中移除了注意力模块的重计算,并且只需要存储其输出。LoongTrain【索引174,Loongtrain: Efficient training of long-sequence llms with head-context parallelism,2024,arXiv】引入了selective-checkpoint++,进一步优化了检查点过程,特别是在长序列训练中,通过将注意力模块添加到白名单中。该方法保存注意力输出和softmax统计信息(softmax_lse)。在前向传递期间,它保存白名单中模块的输出,在后向传递期间,它检索这些存储的输出而不是重新计算它们,继续计算图,从而减少了重新计算注意力的需要。
平衡内存与计算成本的静态策略。与最近主要关注LLM训练中注意力模块的手工检查点策略不同,Yuan等人【索引301,Accelerating the Training of Large Language Models using Efficient Activation Rematerialization and Optimal Hybrid Parallelism,2024,USENIX ATC 24】仔细测量了在模型训练期间重建每个激活张量所需的最小计算成本。他们通过枚举所有可能的检查点方法,得出了内存和计算成本的帕累托前沿。从这个帕累托前沿中,他们选择了一个在计算和内存成本之间实现最佳平衡的解决方案。
6.1.2 动态驱逐
动态驱逐策略。动态驱逐根据训练过程的当前状态实时决定丢弃和重新计算哪些激活。DTR【索引297,Dynamic tensor rematerialization,2020,arXiv】提出了一种贪婪的在线算法,用于在运行时为静态和动态模型启发式地驱逐和重物质化张量。MegTaiChi【索引298,Megtaichi: Dynamic tensor-based memory management for high-performance deep learning,2022,ICS】引入了一种动态张量驱逐,利用了在运行时跟踪的张量访问模式。Coop【索引299,Coop: Memory is not a commodity,2024,Advances in Neural Information Processing Systems】提出缓解由激活重计算方法引起的内存碎片问题,这些方法因驱逐张量时未考虑其连续性而导致。Coop采用一种高效的滑动窗口算法,确保只驱逐连续的内存块,从而最小化内存碎片。
6.2 冗余减少
冗余减少的动机。传统的数据并行方法在所有GPU上复制整个模型状态,导致大量的冗余内存使用。提出了冗余减少技术,通过消除或减少每个设备上的内存冗余来优化内存使用。这些技术通常寻求在内存效率和由此产生的通信开销之间取得平衡,从而以可接受的成本促进训练更大规模或批次大小。
6.2.1 完全分片
Zero Redundancy Optimizer (ZeRO)。ZeRO【索引145,Zero: Memory optimizations toward training trillion parameter models,2020,SC20】通过三个阶段(ZeRO-1、ZeRO-2和ZeRO-3)在所有GPU上完全分片模型状态来优化内存冗余。ZeRO-1在全球所有GPU上分布优化器状态。在训练期间,每个GPU独立进行前向和后向传播以计算梯度,然后使用ReduceScatter操作在数据并行组内的所有GPU之间同步。每个GPU负责更新模型参数的特定分片。之后,使用AllGather操作从其他GPU收集更新后的模型参数分片,确保所有GPU都拥有最新的模型参数。ZeRO-1将优化器状态的内存消耗从12Φ减少到12Φ/N,其中N是数据并行的大小。在ZeRO-1的基础上,ZeRO-2进一步在所有GPU上分片梯度,每个GPU只更新其参数分片,将持有梯度所需的内存从2Φ减少到2Φ/N。ZeRO-3除了优化器状态和梯度外,还对参数进行分区。每个GPU只持有一部分参数。当需要来自远程GPU的参数进行即将到来的计算时,通过AllGather操作收集它们,之后丢弃。在ZeRO-3中,每个GPU只持有与其特定参数分区对应的权重、梯度和优化器状态,将总内存消耗从16Φ减少到16Φ/N。ZeRO被许多框架广泛采用,如DeepSpeed【索引183,Deepspeed-moe: Advancing mixture-of-experts inference and training to power nextgeneration ai scale,2022,ICML】、PyTorch-FSDP【索引146,Pytorch fsdp: Experiences on scaling fully sharded data parallel,2023,VLDB Endowment】和ColossalAI【索引253,Colossal-ai: A unified deep learning system for largescale parallel training,2023,ICPP】。
6.2.2 部分分片
部分分片的权衡与实现。ZeRO面临通信挑战,因为集体通信操作的延迟随着通信规模的增加而增加。在分布式LLM训练中,内存利用率和通信成本之间存在权衡。可以通过在较小的GPU组(即大型GPU集群内的较小GPU集)中分片模型状态来优化通信开销。这种方法减少了节点间通信和通信规模,但由于模型状态的冗余增加,可能导致更高的内存使用。关键是在通信规模和内存利用率之间取得平衡【索引305,Amsp: Reducing communication overhead of zero for efficient llm training,2024】。
基于ZeRO的改进。为解决通信效率低下同时提高内存利用率,已提出了几种基于ZeRO框架的方法。ZeRO++【索引302,ZeRO++: Extremely Efficient Collective Communication for Giant Model Training,2023,arXiv】遵循ZeRO-3在全球所有设备上划分所有模型状态,并进一步在GPU子组内引入参数的次级分片。在前向阶段,它利用跨所有GPU的主分片收集参数,并在子组内(通常在同一节点内)维护一个次级参数分片。在后向阶段,它从此级分片收集参数,减少了通信规模和节点间通信。此外,ZeRO++使用量化来压缩参数和梯度,有效地减少了通信量,但以准确性为代价。MiCS【索引147,Mics: Near-linear scaling for training gigantic model on public,2022,VLDB Endowment】和FSDP【索引146,Pytorch fsdp: Experiences on scaling fully sharded data parallel,2023,VLDB Endowment】在子组内划分所有模型状态组件,并在子组间复制它们,从而减少通信规模和通信延迟,提高训练性能。AMSP【索引305,Amsp: Reducing communication overhead of zero for efficient llm training,2024】和PaRO【索引303,Rethinking memory and communication cost for efficient large language model training,2023,arXiv】集成了三种灵活的分片策略,包括完全复制、完全分片和部分分片,允许模型状态中的每个组件独立选择分片策略。AMSP构建了一个优化问题,以找到在内存约束下最小化通信成本的最优分片策略。此外,AMSP提出了一种定制的通信和计算重叠策略,结合这些灵活的分片策略以实现优化的训练效率。RTP(旋转张量并行)【索引304,Rtp: Rethinking tensor parallelism with memory deduplication,2023,arXiv】通过有策略地分片激活和旋转权重/梯度来最小化内存重复。
6.3 碎片整理
GPU内存碎片问题。GPU内存碎片指的是相邻张量之间产生的零散、不可用的GPU内存块。在LLM训练期间,由于不同张量的生命周期不同以及通用深度学习框架(如PyTorch【索引240,Pytorch: An imperative style, high-performance deep learning library,2019,Advances in neural information processing systems】和TensorFlow【索引326,Tensorflow: a system for large-scale machine learning,2016,OSDI 16】)的内存分配和释放方案效率低下,这个问题尤为突出。此外,像重计算和卸载这样的内存优化技术通过引入更频繁和不规则的内存分配和释放请求,加剧了这个问题【索引299,Coop: Memory is not a commodity,2024,Advances in Neural Information Processing Systems;索引306,Roam: memory-efficient large dnn training via optimized operator ordering and memory layout,2023,arXiv;索引308,Gmlake: Efficient and transparent gpu memory defragmentation for large-scale dnn training with virtual memory stitching,2024,ASPLOS】。碎片问题可能导致高峰值内存和内存不足(OOM)错误,限制了批次大小和整体训练效率。提出了碎片整理的努力,通过内存管理技术来缓解这些问题。
6.3.1 基于张量的碎片整理
基于张量分配和释放方案的碎片整理。深度学习框架通常使用带有内存池的缓存分配器,以实现快速的内存分配和释放,而无需设备同步。已提出几种基于缓存分配器中张量分配和释放方案来减少内存碎片的方法。ROAM【索引306,Roam: memory-efficient large dnn training via optimized operator ordering and memory layout,2023,arXiv】通过考虑张量的生命周期和大小,协同优化算子的执行顺序和张量分配。它引入了一种高效的基于树的算法来搜索最大化张量重用并减少数据碎片的执行计划。ROAM已在单GPU场景中进行了评估,具体是最大的模型为1.5B的GPT-2 XL【索引5,Language models are unsupervised multitask learners,2019,OpenAI blog】,但尚未在具有更大模型的分布式训练场景中进行测试,其中计算图可能变得更大。Imanishi等人【索引307,A Heuristic for Periodic Memory Allocation with Little Fragmentation to Train Neural Networks,2024,ISMM】提出了一种离线优化方法,将张量分配建模为二维装箱问题。在此模型中,每个张量分配表示为一个可垂直移动的矩形,反映了模型训练期间的周期性分配模式。他们提出了一种使用模拟退火的启发式算法来优化分配的拓扑排序,旨在最小化碎片。虽然有效,但当应用于LLM时,由于分配数量多且模式复杂,此方法可能存在可扩展性问题。MegTaiChi【索引298,Megtaichi: Dynamic tensor-based memory management for high-performance deep learning,2022,ICS】和Coop【索引299,Coop: Memory is not a commodity,2024,Advances in Neural Information Processing Systems】在为减少内存消耗而驱逐激活张量时考虑了内存碎片问题。
6.3.2 基于虚拟内存管理的碎片整理
利用VMM进行碎片整理。GMLake【索引308,Gmlake: Efficient and transparent gpu memory defragmentation for large-scale dnn training with virtual memory stitching,2024,ASPLOS】和PyTorch可扩展段【索引309,Pytorch expandable segments,2023,github.com】提出通过利用低级CUDA驱动程序应用程序编程接口的虚拟内存管理(VMM)功能来缓解碎片。这个低级API为开发人员提供了对GPU虚拟内存操作的直接控制,如保留、映射和管理虚拟内存地址。在此基础上,GMLake【索引308,Gmlake: Efficient and transparent gpu memory defragmentation for large-scale dnn training with virtual memory stitching,2024,ASPLOS】引入了一种虚拟内存拼接机制,通过虚拟内存地址映射将非连续的内存块整合成更大的块,最小化了数据移动和复制。类似地,PyTorch的可扩展段【索引309,Pytorch expandable segments,2023,github.com】使分配的内存段能够扩展到更大的尺寸以供重用。这两种方法对不同的模型和内存高效训练技术都是透明的,可以无缝集成到现有的深度学习框架中。此外,GMLake在多GPU上展示了出色的可扩展性,开销很小,并且不需要修改用户代码。PyTorch-v2.1也集成了可扩展段。
6.4 卸载
卸载技术的目的。为了在较少的GPU上高效训练LLM,已经提出了各种利用交换和卸载方法的工作。这些技术将部分计算和数据从GPU内存转移到外部资源,这些资源虽然便宜且速度较慢,但容量巨大。
6.4.1 CPU卸载
CPU卸载的静态和动态方法。许多研究提出了有效利用CPU内存来增强分布式LLM训练的方法。这些技术可大致分为两种主要方法:静态卸载和动态卸载。
静态卸载。静态卸载方法涉及在GPU和CPU内存之间预先确定模型组件的分配。L2L【索引310,Training large neural networks with constant memory using a new execution algorithm,2020,arXiv】逐层管理和移动张量。L2L同步地将即将进行的计算层所需的张量取入GPU内存,同时将其余层的张量存储在CPU内存中。L2L允许将模型扩展到任意深度,但无法跨多GPU扩展。相比之下,ZeRO-Offload【索引311,Zero-offload: Democratizing billionscale model training,2021,USENIX ATC 21】专注于多GPU训练。它将模型参数保存在GPU上,并将优化器状态和梯度存储在CPU内存中。此外,它将优化器更新计算卸载到CPU。这种方法使得在16个V100上可以训练多达70B的模型。然而,ZeRO-Offload可能会留下一些未使用的GPU内存,并且受到缓慢的CPU优化器更新的影响【索引312,Elixir: Train a large language model on a small gpu cluster,2022,arXiv】。
改进的静态卸载。为解决这个问题,Elixir【索引312,Elixir: Train a large language model on a small gpu cluster,2022,arXiv】使用一个搜索引擎,通过利用运行前模型分析来找到内存分区和卸载的最佳组合。与ZeRO-Offload不同,Elixir通过在GPU和CPU之间划分模型状态和优化器块,有效地利用了所有可用的GPU内存。Mobius【索引315,Mobius: Fine tuning large-scale models on commodity gpu servers,2023,ASPLOS】通过引入一种流水线并行方案,解决了在GPU间带宽有限且通信竞争激烈的商用服务器上的多GPU训练问题。该方案为每个GPU分配多个阶段,并在GPU和CPU内存之间动态交换它们。此外,Mobius通过预取和交叉映射来优化通信,以减少开销和竞争。Yuan等人【索引301,Accelerating the Training of Large Language Models using Efficient Activation Rematerialization and Optimal Hybrid Parallelism,2024,USENIX ATC 24】提出通过在流水线阶段的粒度上卸载和重载激活来缓解激活瓶颈,同时最大化激活传输与计算的重叠,从而避免减慢训练过程。与其他卸载工作相比,这项工作更侧重于改善计算和内存利用率之间的平衡,而不是在极其紧张的内存预算下进行训练。
6.4.2 动态卸载
动态卸载方法。动态卸载方法根据内存利用率和数据传输的实时优化,自适应地在GPU和CPU内存之间分配模型或张量的分区。STRONGHOLD【索引318,Stronghold: fast and affordable billion-scale deep learning model training,2022,SC22】提出在GPU和CPU内存之间动态卸载模型状态,并维持一个合适的工作窗口大小,以最小化卸载期间的GPU停顿。Harmony【索引316,Harmony: overcoming the hurdles of gpu memory capacity to train massive dnn models on commodity servers,2022,VLDB Endowment】使用基于启发式的调度器将计算和模型状态映射到物理设备。Harmony通过减少交换和快速的点对点交换来降低卸载的开销。TMOF【索引317,Tensor movement orchestration in multi-gpu training systems,2023,HPCA】引入了不相交的交换和双向重叠协调机制,以防止交换和卸载中的PCIe通道竞争。对于MoE模型,MPipeMoE【索引327,Mpmoe: Memory efficient moe for pre-trained models with adaptive pipeline parallelism,2024,IEEE Transactions on Parallel and Distributed Systems】设计了一种自适应且内存高效的流水线并行算法。具体来说,MPipeMoE通过消除内存冗余来采用高效的内存重用策略,并使用一个自适应选择组件来决定是卸载还是重新计算所需的张量,以减少内存需求。
细粒度动态内存管理。为促进更好的内存管理,一些研究提出了将张量分解为更细粒度单元的系统。TSPLIT【索引313,Tsplit: Fine-grained gpu memory management for efficient dnn training via tensor splitting,2022,ICDE】和PatrickStar【索引314,Parallel training of pre-trained models via chunk-based dynamic memory management,2022,IEEE Transactions on Parallel and Distributed Systems】是两个优化峰值GPU内存使用的动态内存管理系统。TSPLIT将张量拆分为微张量,并在微张量级别执行操作,实现了精确和动态的内存操作。PatrickStar将模型数据组织成内存块,这些块在CPU和GPU内存之间动态分布,并优化了CPU-GPU数据传输以及带宽利用率。此外,TSPLIT使用模型引导的规划算法为每个张量找到最优的内存配置,而PatrickStar则使用运行时内存跟踪、块驱逐策略和设备感知的算子放置来进一步最小化CPU和GPU之间的数据移动。
6.4.3 SSD卸载
利用SSD进行更大规模模型的训练。为了促进万亿级LLM的训练,仅依赖CPU卸载的方法是不够的,一些工作提出了在训练期间将数据卸载到CPU内存和NVMe SSDs。ZeRO-Infinity【索引319,Zeroinfinity: Breaking the gpu memory wall for extreme scale deep learning,2021,SC】将所有分区的模型状态卸载到CPU或NVMe内存,并将激活仅卸载到CPU内存。这种方法支持在32个节点(总共512个V100)上训练多达32T参数的模型。然而,激活的CPU卸载仍然需要大量的CPU内存。例如,训练10T模型需要约0.76 TB的CPU内存来存储激活检查点,而100T模型则需要约4 TB。Fuyou【索引322,Adding nvme ssds to enable and accelerate 100b model finetuning on a single gpu,2024,arXiv】专注于在CPU内存容量有限且只有一个GPU的商用服务器上训练LLM。与ZeRO-Infinity相比,Fuyou进一步将激活卸载到SSD,并将SSD-CPU通信作为额外的优化维度。它还提出了一种与后向传播阶段重叠的同步核外CPU优化器,并引入了一种自动激活交换机制,从而最大化GPU利用率。Smart-Infinity【索引321,Smartinfinity: Fast large language model training using near-storage processing on a real system,2024,HPCA】提出通过使用近存储处理设备进行参数更新来减少二级存储的带宽需求。MoESys【索引323,Moesys: A distributed and efficient mixture-of-experts training and inference system for internet services,2024,IEEE Transactions on Services Computing;索引324,Se-moe: A scalable and efficient mixture-of-experts distributed training and inference system,2022,arXiv】结合了各种存储设备(GPU、CPU内存和SSD)来保存稀疏参数状态和密集参数状态,并为MoE训练提出了一种2D预取调度策略,以便参数的计算可以与调度重叠。
7 通信优化
不同并行机制的通信模式。不同的并行机制引入了不同模式的网络通信流量。例如,张量并行需要在张量并行等级之间进行AllReduce操作。而数据并行则需要在每次迭代结束时,在数据并行等级之间进行梯度同步的AllReduce操作。流水线并行涉及在每个阶段结束时将激活值传递到下一阶段。通常,训练框架将需要高带宽的张量或序列并行通信组放置在高带宽域内(例如,同一节点内),而将带宽要求较低的数据并行或流水线并行通信组放置在高带宽域之间。图13显示了实践中LLM训练的通信热图,很好地反映了不同并行策略带来的数据流量。从这个热图中可以看出,LLM训练的通信流量呈现出清晰的模式和层次结构,大部分通信发生在较小的范围内,只有一小部分流量跨越整个集群。这一洞见启发了像rail-optimized拓扑【索引62,Doubling all2all performance with nvidia collective communication library 2.12,2018,developer.nvidia.com】这样的方法,它减少了不必要的核心交换机以降低成本。
Fig. 13: InternLM-2 102B预训练在单次迭代中使用128个GPU的通信流量热图,张量并行(TP)大小为8,流水线并行(PP)大小为4,数据并行(DP)大小为4,ZeRO阶段1(ZeRO-1)大小为4。拓扑排列的优先级是TP > DP/ZeRO-1 > PP。有四种不同的数据流量负载:1 TP的AllReduce;2 3 DP/ZeRO-1的ReduceScatter/AllGather;4 PP的Send/Recv。TP的通信利用了NVSwitch的全连接拓扑,在图中沿对角线形成了16个密集的方形流量模式,每个模式代表一个节点。DP和ZeRO-1的跨节点通信流量在图中显示为四个32x32矩形拓扑内的六条对称对角线。需要注意的是,DP/ZeRO-1也涉及节点内通信流量,这些流量与TP累积在同一个热图网格中。由于其通信量相对较小,PP在热图上坐标((32, 0), (128, 96))和((0, 32), (96, 128))处形成两条黄线。(在此图中,所有通信都使用基于环的集合算法)
本节内容概述。本节介绍优化分布式LLM训练集体通信性能的系统和技术。如图14所示,我们首先讨论利用预定义和合成算法的集体通信库。接下来,我们探讨旨在重组通信操作以与计算重叠的通信调度技术,从而减少延迟并加速训练过程。最后,我们深入研究网络内聚合(INA),它利用网络设备的计算能力来执行聚合操作,例如对深度学习模型的梯度求和。
Fig. 14: 分布式LLM训练的通信优化研究。
关于通信压缩的说明。压缩模型参数和梯度能有效减少分布式LLM训练期间的通信开销。各种研究探索了稀疏通信和量化方法。例如,ZeRO++【索引302,ZeRO++: Extremely Efficient Collective Communication for Giant Model Training,2023,arXiv】在通信前对权重采用量化,将每个模型参数从FP16缩小到INT8数据类型。然而,这些工作通常涉及有损的稀疏化或量化技术。本节我们不调查有损数据压缩技术,因为它们超出了本工作的范围。
7.1 集体通信
集体通信库的作用。消息传递接口(MPI)是在并行计算架构上广泛采用的大规模科学应用程序编程模型。MPI有多种实现,包括OpenMPI【索引328,Open mpi: Goals, concept, and design of a next generation mpi implementation,2004,Proceedings, 11th European PVM/MPI Users’ Group Meeting】、MPICH2【索引329,Mpich2: A new start for mpi implementations,2002,Recent Advances in Parallel Virtual Machine and Message Passing Interface】和MVAPICH【索引330,The mvapich project: Evolution and sustainability of an open source production quality mpi library for hpc,2013,WSSPE】。这些库提供了多种CUDA感知的原语,如AllReduce、AllGather和ReduceScatter,这些对于分布式LLM训练至关重要。在实践中,当前的训练框架更倾向于使用针对特定AI加速器、采用预定义或合成算法的集体通信。
7.1.1 预定义的集体通信算法
NCCL和RCCL。NVIDIA的NCCL【索引331,Nccl library,2016,github. com】和AMD的RCCL【索引332,Rccl library,2018,github. com】是高度优化的库,在其各自的AI加速器上通常性能优于基于MPI的集体通信库。这些库通常根据网络拓扑和输入张量大小等条件选择预定义的算法来执行集合操作。
环形算法 (Ring Algorithm)。环形算法用于像AllReduce这样的集体通信,以在所有GPU之间移动数据。使用该算法,输入张量被分成多个块,并在操作期间逐一传输。这种流水线减少了每个设备等待数据所花费的空闲时间。百度使用带宽最优的环形AllReduce算法【索引333,Bandwidth optimal all-reduce algorithms for clusters of workstations,2009,Journal of Parallel and Distributed Computing】进行分布式深度学习模型训练。Horovod【索引143,Horovod: fast and easy distributed deep learning in tensorflow,2018,arXiv】用NCCL替换了百度的环形AllReduce实现,并为分布式训练设计了一个用户友好的接口。
树形算法 (Tree Algorithm)。环形算法的延迟随着GPU设备数量的增加而增加【索引346,Gradientflow: Optimizing network performance for large-scale distributed dnn training,2019,IEEE Transactions on Big Data】。双二叉树算法【索引334,Massively scale your deep learning training with nccl 2.4,2019,developer.nvidia.com】被提出来解决这个问题。双二叉树依赖于一个事实,即在一个二叉树中,一半或更少的等级是节点,一半或更多的等级是叶子。因此,可以为每个二叉树构建第二棵树,使用叶子作为节点,反之亦然。该算法在基于MPI的库、NCCL和RCCL中都有实现。
混合算法 (Hybrid Algorithm)。一些方法提出使用混合算法来处理具有异构节点内和节点间通信带宽的训练集群上的集体通信任务。两级AllReduce【索引335,Highly scalable deep learning training system with mixed-precision: Training imagenet in four minutes,2018,arXiv】将单个AllReduce操作分为三个步骤:利用PCIe/NVLINK的节点内Reduce,利用网络的节点间AllReduce,以及节点内Broadcast。2D-Torus AllReduce【索引336,Massively distributed sgd: Imagenet/resnet-50 training in a flash,2018,arXiv】和ACCL【索引337,Accl: Architecting highly scalable disnication library,2021,IEEE micro】将单个AllReduce操作分解为三个阶段:节点内基于环的ReduceScatter,节点间基于树的AllReduce,以及节点内基于环的AllGather。BlueConnect【索引338,Blueconnect: Decomposing all-reduce for deep learning on heterogeneous network hierarchy,2019,MLS】将单个AllReduce操作分解为许多可并行的ReduceScatter和AllGather操作。每个操作可以映射到不同的网络结构,利用每个特定结构性能最佳的预定义实现。Plink【索引339,Plink: Discovering and exploiting datacenter network locality for efficient cloud-based distributed training,2020,MLSys】可以探测网络拓扑并高效生成两级混合通信计划,利用数据中心网络中的局部性。
7.1.2 合成的集体通信算法
为特定硬件定制的集体通信算法。一些方法出现了,它们合成了专门针对硬件拓扑的集体通信算法和内核,旨在在许多情况下超越通用的预定义算法。GC3【索引340,Gc3: An optimizing compiler for gpu collective communication,2022】引入了一种面向数据领域特定语言(DSL),用于设计自定义集体通信算法。它包括一个优化编译器,将这些算法转换为针对特定硬件配置优化的可执行形式。SCCL【索引341,Synthesizing optimal collective algorithms,2021,PPoPP】将集体通信合成问题编码为SMT(可满足性模理论)公式。这种方法旨在为帕累托最优算法推导出精确的调度,优化延迟和带宽利用率。TACCL【索引342,{TACCL}: Guiding collective algorithm synthesis using communication sketches,2023,NSDI 23】将寻找最优通信算法的问题公式化为一个混合整数线性规划(MILP)。它利用通信草图抽象来高效地收集基本信息并减少搜索空间,目标是最小化总执行时间。Blink【索引343,Blink: Fast and generic collectives for distributed ml,2020】在运行时通过探测每个作业可用的链路集来动态构建具有合适链路容量的拓扑。利用这个拓扑,它通过创建数据包生成树来优化通信速率,并生成CUDA代码。P2【索引344,Synthesizing optimal parallelism placement and reduction strategies on hierarchical systems for deep learning,2022,MLS】利用并行矩阵在系统级别划分并行轴,从而生成拓扑感知的并行放置和归约策略。通过模拟和预测通信成本,该方法减少了所需的实际评估次数。
7.2 通信调度
通信调度的目标。分布式训练中的通信调度重组通信操作以与计算重叠,从而减少延迟并加速训练过程。通信调度的关键概念是根据并行训练的数据依赖性重新排序通信操作。混合并行LLM训练需要多维通信调度方案来管理由数据、流水线、张量和序列并行以及它们的组合产生的通信。
7.2.1 基于FIFO的调度
先入先出(FIFO)调度。在后向阶段,与其等待所有梯度计算完成后再启动通信,不如在每个梯度准备好后立即开始通信。这种无等待反向传播方法利用依赖性有向无环图来高效管理任务。Poseidon【索引345,Poseidon: An efficient communication architecture for distributed deep learning on {GPU} clusters,2017,USENIX ATC 17】使用一个先入先出(FIFO)队列来调度AllReduce算子,确保每一层在其梯度生成后立即开始通信。受大张量上集体通信效率的启发,GradientFlow【索引346,Gradientflow: Optimizing network performance for large-scale distributed dnn training,2019,IEEE Transactions on Big Data】和Pytorch DDP【索引142,Pytorch distributed: Experiences on accelerating data parallel training,2020,arXiv】将多个顺序的AllReduce通信操作融合成一个单一操作。这种方法通过在后向阶段等待一小段时间然后将多个梯度组合成一个AllReduce操作,避免了在网络上传输大量小张量。
7.2.2 基于优先级的调度
优先级调度。FIFO方案通常是次优的,因为在后向阶段生成的通信序列与前向阶段的计算序列不同。这种不匹配可能导致通信阻塞计算,即使在启用重叠的情况下也是如此。因此,许多方法采用优先级队列来高效地调度通信算子。P3【索引347,Priority-based parameter propagation for distributed dnn training,2019,MLS】以更细的粒度调度AllReduce操作,将当前层的梯度通信与下一层的前向计算重叠。与基于FIFO队列的调度不同,该方法将层划分为固定大小的片,并根据它们在前向传播中处理的顺序优先同步这些片。因此,第一层获得最高优先级,优先级向末端递减。在使用参数服务器架构进行分布式模型训练时,TicTac【索引348,Tictac: Accelerating distributed deep learning with communication scheduling,2019,MLS】优先处理那些能加速底层计算图中关键路径的传输。
通用优先级调度框架。ByteScheduler【索引349,A generic communication scheduler for distributed dnn training acceleration,2019,SOSP】和PACE【索引350,Preemptive all-reduce scheduling for expediting distributed dnn training,2020,IEEE INFOCOM】被提出来,以在各种训练框架中推广基于优先级的通信调度。具体来说,ByteScheduler【索引349,A generic communication scheduler for distributed dnn training acceleration,2019,SOSP】引入了一个统一的抽象,以便在不破坏框架引擎内原始依赖关系的情况下进行通信调度。ByteScheduler通过使用贝叶斯优化来自动调整两个关键参数:分区大小和信用大小,从而获得良好性能。PACE【索引350,Preemptive all-reduce scheduling for expediting distributed dnn training,2020,IEEE INFOCOM】通过将原始AllReduce操作分割成更小的片段来实现抢占式通信。被抢占的AllReduce算子可以在稍后恢复。这种抢占策略防止了大通信张量的队头阻塞。此外,PACE使用动态规划方法来融合小通信张量,以减少处理大量小张量所带来的开销,从而实现更高效的带宽利用。
MoE系统中的优先级调度。为了提高MoE系统中的带宽效率,Lina【索引188,Accelerating distributed {MoE} training and inference with lina,2023,USENIX ATC 23】优先处理All-to-All操作而不是AllReduce。通常,专家并行(All-to-All)和数据并行(AllReduce)进程使用独立的CUDA流,这可能导致在没有协调的情况下发生重叠和带宽共享。Lina将张量分解成更小的块,确保All-to-All操作获得全部带宽,同时允许AllReduce微操作在空闲时间运行。此外,微操作使得All-to-All操作能够与专家计算重叠。
7.2.3 基于分解的调度
分解调度的思想。一些进展集中于将通信和计算操作分解为细粒度的任务,以更大的灵活性重新排序这些操作,从而最大化重叠并优化执行效率。
流水线阶段分解。当使用传统的流水线并行时,每个GPU存储一个连续的层段。Breadth-First【索引159,Breadth-first pipeline parallelism,2023,MLS】进一步将这些连续的阶段分解为分布在不同GPU上的更细粒度的阶段,通过连接第一个和最后一个GPU形成一个循环,这样每个GPU被分配多个阶段。这使得给定的微批次能够更早地到达流水线的末端,减少了流水线气泡。Breadth-First使用广度优先的调度策略来实现更大的计算-通信重叠。Fold3D【索引351,Fold3d: Rethinking and parallelizing computational and communicational tasks in the training of large dnn models,2023,IEEE Transactions on Parallel and Distributed Systems】采用“全进全出”的调度策略来将流水线的梯度同步过程与计算重叠。这涉及在流水线内进一步折叠模型片段,其中每个设备包含两个模型片段,允许一个片段的梯度同步与另一个片段的前向或后向计算重叠。
异步流水线并行中的调度。异步流水线并行放宽了梯度和参数更新之间的数据依赖关系。利用这一特性,TriRace【索引352,A multidimensional communication scheduling method for hybrid parallel dnn training,2024,IEEE Transactions on Parallel and Distributed Systems】推迟参数更新,以最大化计算与梯度通信的重叠。此外,TriRace将流水线阶段之间的双向P2P通信分解为两个独立的单向操作,并根据关键路径分析对它们进行优先级排序。
通信分解。通信原语可以分解为具有高调度灵活性的细粒度操作。Wang等人【索引353,Overlap communication with dependent computation via decomposition in large deep learning models,2022,ASPLOS】将通信操作(如AllGather和ReduceScatter)分解为一系列细粒度的点对点集合。此外,计算操作(如爱因斯坦求和)被划分为细粒度的任务,每个任务执行一部分计算。这种分解为重叠通信与计算创造了更多机会。SYNDICATE【索引354,Better together: Jointly optimizing {ML} collective scheduling and execution planning using {SYNDICATE},2023,NSDI 23】将通信操作分割成更小的子操作,称为Motifs,并使用一个采用马尔可夫链蒙特卡罗搜索的中央优化器来实现最优的重叠执行计划。Centauri【索引355,Centauri: Enabling efficient scheduling for communicationcomputation overlap in large model training via communication partitioning,2024,ASPLOS】采用不同的方法,通过使用原语分区、组分区和工作负载分区将通信操作分解为细粒度的原子操作。然后使用工作负载感知调度、后向调度和弹性调度来调度这些操作,以最大化重叠效率。DeAR【索引356,Dear: Accelerating distributed deep learning with fine-grained all-reduce pipelining,2023,ICDCS】也分解了通信原语,特别是将AllReduce分解为AllGather和ReduceScatter。这种分解允许后续操作与模型的前向传播过程重叠,从而消除了等待两个通信步骤完成的需要。
计算分解。当使用张量并行时,需要一个AllReduce通信来同步前向阶段的矩阵乘法输出。CoCoNet【索引357,Breaking the computation and communication abstraction barrier in distributed machine learning workloads,2022】通过将输出划分为更小的块,并在计算矩阵乘法内核内的每个结果块后立即启动AllReduce内核,来促进矩阵乘法和AllReduce的重叠。为最小化AllReduce内核的等待时间,数据块以精心安排的顺序送入矩阵乘法内核。T3【索引358,T3: Transparent tracking & triggering for fine-grained overlap of compute & collectives,2024,ASPLOS】采用硬件-软件协同设计方法,透明地将矩阵乘法与通信重叠,同时最小化资源争用。在硬件层面,T3引入了一个“跟踪-触发”机制来协调生产者的计算和通信活动。此外,它还使用计算增强型内存来处理通信过程所需的伴随计算操作。
乱序反向传播。反向传播会产生两种类型的梯度:用于计算前一层梯度的输出梯度,以及用于更新该层权重参数的权重梯度。这些权重梯度需要使用AllReduce与其他等级同步。传统的框架同时执行权重和输出的梯度计算。乱序反向传播(ooo-backprop)【索引361,Out-of-order backprop: An effective scheduling technique for deep learning,2022,EuroSys】将权重和输出的梯度计算解耦,灵活地将权重梯度计算安排在其原始顺序之外。这允许更关键的计算被优先安排。因此,ooo-backprop通过基于这种乱序计算策略调度通信来优化整体性能。Zero Bubble【索引156,Zero bubble pipeline parallelism,2023,ICLR】也使用该方案来降低流水线并行的气泡率。
重计算与通信的重叠。启用激活检查点后,训练框架需要在后向阶段重新计算激活。当使用张量并行时,这种重计算也涉及AllReduce通信。Oases【索引359,Automated tensor model parallelism with overlapped communication for efficient foundation model training,2023】通过始终将AllReduce通信作为重计算单元的最后一个前向通信操作,来减少重计算中的冗余通信,并进一步将批次拆分为更小的子批次,允许两个批次的通信和计算重叠。Lynx【索引360,Optimizing large model training through overlapped activation recomputation,2024】也利用了重计算和通信重叠的潜力,使用两种重计算调度算法OPT和HEU来搜索最优或近优的重计算调度策略,实现最佳的重叠和训练性能。
7.3 网络内聚合
网络内聚合(INA)的原理。网络内聚合(INA)利用网络设备的计算能力来执行聚合操作,如对深度学习模型的梯度求和。这项技术之前已被提出来加速大数据处理。值得注意的是,像NetAgg【索引369,Netagg: Using middleboxes for application-specific on-path aggregation in data centres,2014,CoNEXT】、SwitchAgg【索引370,Switchagg: A further step towards in-network computation,2019,FPGA】和CamDoop【索引371,Camdoop: Exploiting in-network aggregation for big data applications,2012,NSDI 12】等框架通过在直接连接拓扑中,在交换机附加的高性能中间件或服务器上执行数据聚合,展示了显著的性能优势。许多方法已被提出来将网络内聚合应用于深度学习模型训练,旨在减少在后向阶段梯度AllReduce操作期间节点间交换的数据量【索引372,When in-network computing meets distributed machine learning,2024,IEEE Network】。
7.3.1 基于以太网的聚合
基于可编程交换机的INA系统。许多基于以太网的网络内聚合系统依赖于可编程交换机,并且可以用于分布式LLM训练。SwitchML【索引362,Scaling distributed machine learning with {In-Network} aggregation,2021,NSDI 21】支持在分布式训练的后向阶段将集体通信操作卸载到可编程网络交换机。由于完整的模型更新可能超过交换机的存储容量,SwitchML通过交换机流式传输聚合,一次处理有限数量的向量元素的聚合功能。SwitchML有两个限制。首先,在处理浮点操作时,SwitchML不能直接为浮点张量执行集体通信(如AllReduce)。相反,它使用类似块浮点的方法将浮点值转换为32位整数。其次,SwitchML主要在DPDK上实现,虽然有支持RDMA的实现,但很难与训练框架集成。
改进的以太网INA系统。为了更好地促进分布式模型训练,FPISA【索引363,Unlocking the power of inline {Floating-Point} operations on programmable switches,2022,NSDI 22】将浮点计算实现为直接在可编程交换机上运行的P4【索引373,P4: Programming protocol-independent packet processors,2014,ACM SIGCOMM Computer Communication Review】程序。因此,训练框架可以将FP16张量上的集体通信操作卸载到交换机,而无需将其转换为32位整数。NetReduce【索引364,Netreduce: Rdma-compatible in-network reduction for distributed dnn training acceleration,2020,arXiv】支持与RoCE兼容的网络内聚合,充分利用RoCE的拥塞控制和可靠性设计,而无需在交换机中使用昂贵的网络协议处理栈。NetReduce的原型是使用连接到以太网交换机的FPGA板实现的。AllReduce-Switch【索引365,Scalable fully pipelined hardware architecture for in-network aggregated allreduce communication,2021,IEEE Transactions on Circuits and Systems I: Regular Papers】与NetReduce密切相关,并与其网络协议兼容。它引入了一种专为网络内聚合任务量身定制的新型交换机架构,并使用FPGA硬件实现了原型。PANAMA【索引366,In-network aggregation for shared machine learning clusters,2021,MLS】和ATP【索引367,{ATP}: In-network aggregation for multi-tenant learning,2021,NSDI 21】也通过为共享环境设计的网络内聚合框架为该领域做出了贡献。PANAMA专注于通过管理多个并发活动训练作业之间的带宽分配来优化网络负载。它解决了传统拥塞控制可能无法充分支持同时训练操作的挑战。另一方面,ATP支持多个并发租户同时运行多个作业,强调在共享环境中支持多样化工作负载。
特定工作负载的INA系统。某些工作是为特定训练工作负载量身定制的,因此不适用于分布式LLM训练。例如,Libra【索引374,Libra: An efficient and scalable in-network aggregation scheme for sparse gradients in distributed learning,2022,arXiv】是为使用参数服务器架构的稀疏模型训练设计的。它将频繁更新参数的聚合卸载到可编程交换机,而将不常更新的参数留给服务器处理。这种方法有效地减少了服务器负载。另一方面,iSwitch【索引375,Accelerating distributed reinforcement learning with in-switch computing,2019,ISCA】是为强化学习训练任务中的参数聚合设计的。虽然其基于FPGA的实现支持原生浮点操作,但其运行带宽显著较低。此外,iSwitch在聚合期间存储整个梯度向量,这对于强化学习工作负载是可行的,但对于大规模模型,尤其是LLM,则无法很好地扩展。
7.3.2 基于Infiniband的聚合
NVIDIA SHARP。NVIDIA Mellanox的可扩展分层聚合协议(SHARP)【索引368,Scalable hierarchical aggregation protocol (sharp): A hardware architecture for efficient data reduction,2016,COMHPC】是一种专有的网络内聚合方案,可在某些InfiniBand交换机和NVIDIA GPU中使用。SHARP建立在InfiniBand之上,利用链路层流控制和无损保证,并使用专用的片上FPU进行集体卸载。SHARPv1在InfiniBand EDR交换机上引入,SHARPv2在InfiniBand HDR交换机上得到增强,功能包括支持集体通信(如Barrier、Reduce、AllReduce和Broadcast)、整数和浮点操作(16/32/64位)以及GPUDirect RDMA。SHARPv2还使用流式聚合进行大向量的线速归约,与NCCL集成,并且易于被现有训练框架使用。在最新的InfiniBand NDR交换机上启用后,SHARP已为分布式LLM训练做好了生产准备,并已部署在许多训练集群中。除了Infiniband,NVIDIA的NVSwitch-v3【索引46,3.2 the a100 datacenter gpu and ampere architecture,2021,ISSCC】也集成了SHARP,以加速基于GPU的集群中的集体操作。
Fig. 15: 分布式LLM训练的容错技术研究。
8 容错
LLM训练对容错的需求。LLM训练通常涉及从数周到数月的延长训练期,使用数万个GPU的集群。所涉及的组件范围广泛,从底层基础设施到训练系统优化,都需要强大的容错机制来确保训练过程的可靠性。这是因为由于训练的同步性,系统中任何部分的单点故障都可能导致训练过程暂停。在本节中,我们首先对LLM训练中的故障进行分析,然后研究快速故障检测和恢复的方法。
8.1 LLM故障分析
LLM训练中的故障频率。来自各种来源的经验证据强调了LLM训练中故障的频率。例如,Bloom的训练在拥有384个GPU的集群上平均每周会遇到1-2次GPU故障【索引292,Bloom: A 176b-parameter open-access multilingual language model,2022,arXiv】。Meta关于在992个A100 GPU上训练175B OPT模型的全面训练记录【索引398,Metaseq: Opt-175 logbook,2022,github.com】记载了在两周内因硬件、基础设施和其他外部因素导致的40多次中断。最近的研究进一步凸显了这个问题。Acme【索引23,Characterization of large language model development in the datacenter,2024,NSDI 24】报告称,在使用超过1000个A100 GPU的训练过程中,平均每1-2天发生一次故障。字节跳动的MegaScale项目【索引71,{MegaScale}: Scaling large language model training to more than 10,000 {GPUs},2024,NSDI 24】,使用12,288个Ampere GPU,在几周内经历了超过100次故障。Meta的LLaMA3在拥有16,384个H100 GPU的集群上进行为期54天的预训练期间,经历了466次作业中断【索引9,The llama 3 herd of models,2024,ai.meta.com】。频繁的故障主要归因于这些系统的巨大复杂性和规模以及延长的训练期。整个训练系统涵盖了我们在前面章节中研究过的各种组件。此外,同步训练进一步加剧了这个问题,因为任何单个节点的错误都可能导致整个作业失败,使得系统特别容易受到即使是孤立的硬件或软件故障的影响。即使是在阿里巴巴的集群中观察到的单个节点1.5%的日故障率【索引383,Dlrover: An automatic distributed deep learning system,2024,github.com】,当扩展到拥有1000个GPU的系统时,也转化为惊人的84.8%的日故障率。然而,扩大训练系统的趋势仍在继续增长,强调了容错机制在维持系统可靠性方面的伴随挑战。
故障原因。这些故障背后的原因是多方面的,源于LLM训练系统的各个组件。根据Acme【索引23,Characterization of large language model development in the datacenter,2024,NSDI 24】,最严重的影响来自硬件故障,例如GPU(如CUDA-Error、ECC-Error)、NVLink和网络系统(如NCCL-Timeout-Error、Connection-Error)的问题。在阿里巴巴C4【索引377,Boosting large-scale parallel training efficiency with c4: A communication-driven approach,2024,arXiv】中也观察到了类似的现象。C4进一步观察到,大多数错误(约82.5%)局限于特定节点甚至单个设备,尽管用户观察到的大多数错误是NCCL错误。LLaMA3预训练【索引9,The llama 3 herd of models,2024,ai.meta.com】也报告称,78%的故障是硬件问题。此外,最新一代的GPU(A100和H100)往往表现出较高的错误率,这可能是由于快速开发、仓促交付和功耗增加所致【索引377,Boosting large-scale parallel training efficiency with c4: A communication-driven approach,2024,arXiv;索引399,An empirical study on quality issues of deep learning platform,2023,ICSE-SEIP】。除了硬件,分布式训练框架、数据预处理管道或库依赖关系中的软件相关问题也可能导致崩溃或意外行为【索引23,Characterization of large language model development in the datacenter,2024,NSDI 24;索引378,The infrastructure powering ibm’s gen ai model development,2024,arXiv;索引399,An empirical study on quality issues of deep learning platform,2023,ICSE-SEIP】。模型本身的复杂性可能引入不稳定性,如损失尖峰、数值上溢或下溢、梯度爆炸或优化困难【索引398,Metaseq: Opt-175 logbook,2022,github.com;索引400,Palm: Scaling language modeling with pathways,2023,JMLR】。数据中心的外部因素,如停电或冷却系统故障,进一步导致系统不稳定。例如,集群服务器机房的高温也容易导致GPU过热,这可能引起NVLink-Error或ECC-Error【索引23,Characterization of large language model development in the datacenter,2024,NSDI 24】或不稳定的训练速度【索引9,The llama 3 herd of models,2024,ai.meta.com】。
故障的影响。这些高频和多方面的LLM故障导致了GPU的巨大浪费。这种低效率主要体现在两个方面:故障恢复和性能下降。首先,LLM训练作业在运行时定期保存检查点以保持进度。发生故障时,系统维护人员必须首先定位和诊断问题,然后通过回滚到先前的检查点来重新启动训练。然而,一些硬件故障很难主动检测,并且通常需要相当长的时间来诊断和恢复,导致LLM训练长时间停滞。其次,集群中的掉队者(stragglers),由网络链路故障【索引377,Boosting large-scale parallel training efficiency with c4: A communication-driven approach,2024,arXiv】或异常的计算减速【索引71,{MegaScale}: Scaling large language model training to more than 10,000 {GPUs},2024,NSDI 24】引起,会显著降低MFU,进一步加剧了整体训练效率的低下。Meta的175B OPT模型的训练 exemplifies了这些低效率【索引398,Metaseq: Opt-175 logbook,2022,github.com】。虽然根据MFU估算的理想训练时间约为25天,但实际训练持续了大约57天。这意味着总时间的惊人的56%被浪费在处理各种故障上,凸显了系统不稳定性对LLM训练中资源利用和训练效率的严重影响。
8.2 异常检测
异常检测的重要性。快速检测和诊断LLM故障对于维持训练稳定性和效率至关重要。这个过程,称为异常检测,主要采用两种方法:统计监控和主动验证。
8.2.1 统计监控
统计监控的原理。统计监控是一种系统性的方法,用于在整个LLM训练过程中观察和分析各种指标。该方法涉及收集、处理和解释数据,以识别异常或偏离预期行为的情况。在一个典型的设置中,每个GPU都分配有一个专用的监控进程,负责收集基本信息和运行时统计数据【索引71,{MegaScale}: Scaling large language model training to more than 10,000 {GPUs},2024,NSDI 24;索引378,The infrastructure powering ibm’s gen ai model development,2024,arXiv;索引379,Unicron: Economizing Self-Healing LLM Training at Scale,2023,arXiv】。这些统计数据然后作为心跳消息传输到中央监控节点进行进一步分析。未能发送心跳消息的节点被认为已发生故障。该监控系统的主要目标是及时检测异常,以便快速恢复,从而最小化训练中断并维持整体效率。
监控的指标和工具。在LLM训练中监控的大多数运行时统计数据都与硬件相关,包括GPU和网络指标。最近的工作【索引71,{MegaScale}: Scaling large language model training to more than 10,000 {GPUs},2024,NSDI 24;索引378,The infrastructure powering ibm’s gen ai model development,2024,arXiv;索引379,Unicron: Economizing Self-Healing LLM Training at Scale,2023,arXiv】使用NVIDIA DCGM【索引401,Nvidia data center gpu manager (dcgm),2017,developer.nvidia.com】收集与GPU相关的统计数据,包括SM块利用率、SM占用率、SM管道利用率、PCIe流量速率、NVLink流量速率等。一个常见的问题是GPU内存行重映射,它在硬件中无缝地用稀疏单元替换已知的退化内存单元。Vela【索引378,The infrastructure powering ibm’s gen ai model development,2024,arXiv】通过利用DCGM的DCGM_FI_DEV_ROW_REMAP_PENDING统计数据来检测这一点。Megascale【索引71,{MegaScale}: Scaling large language model training to more than 10,000 {GPUs},2024,NSDI 24】和Transom【索引380,TRANSOM: An Efficient Fault-Tolerant System for Training LLMs,2023,arXiv】也通过分析训练日志中发生的错误来检测错误。
网络和通信监控。除了GPU指标,网络统计数据对于监控分布式训练性能也至关重要。MegaScale【索引71,{MegaScale}: Scaling large language model training to more than 10,000 {GPUs},2024,NSDI 24】跟踪RDMA流量指标以检测潜在的异常。它还开发了可视化系统来手动识别低效的GPU。Unicron【索引379,Unicron: Economizing Self-Healing LLM Training at Scale,2023,arXiv】在训练期间通过延迟通知检测NCCL超时、TCP超时和任务挂起等错误。C4【索引377,Boosting large-scale parallel training efficiency with c4: A communication-driven approach,2024,arXiv】收集连接细节,如RDMA IP和QP号,以及消息统计数据,包括传输的计数、大小和持续时间,以在传输层检测训练减速和挂起。集体通信活动也可以通过PyTorch内置的NCCL飞行记录器【索引273,Pytorch 2: Faster machine learning through dynamic python bytecode transformation and graph compilation,2024,ASPLOS】进行监控,该记录器将集体元数据和堆栈跟踪捕获到环形缓冲区中以供后续诊断。Meta进一步与PyTorch协同设计了NCCLX【索引9,The llama 3 herd of models,2024,ai.meta.com】,允许PyTorch访问其内部状态以进行快速准确的故障检测。NCCLX跟踪每个NCCLX通信的内核和网络活动,这有助于诊断通信问题。Vela【索引378,The infrastructure powering ibm’s gen ai model development,2024,arXiv】实现了一个增强的多NIC健康检查器,收集每个端口上所有2节点对的节点网络带宽数据。这些信息可用于检测RoCE/GDR性能下降的节点。Transom【索引380,TRANSOM: An Efficient Fault-Tolerant System for Training LLMs,2023,arXiv】利用LLM训练的关键特性作为先验知识,开发机器学习算法进行异常检测。
TPUv4中的统计监控。统计监控也使Google的TPUv4超级计算机【索引376,Resiliency at Scale: Managing Google’s TPUv4 Machine Learning Supercomputer,2024,NSDI 24】具有弹性。每个TPUv4机器都配备了一个healthd守护进程,对ICI(芯片间互连)链路、PCIe链路和TPU ASIC进行实时监控。检测到的严重症状会通知集群调度器采取适当的行动,例如驱逐受影响的作业或重新调度它们。
8.2.2 主动验证
主动验证方法。主动验证提供了一种替代基于在线统计监控的被动故障排除方法,旨在在故障发生前验证训练系统。然而,在验证测试时间和准确性之间存在权衡,因为全面的验证会显著影响有效的训练时间。MegaScale【索引71,{MegaScale}: Scaling large language model training to more than 10,000 {GPUs},2024,NSDI 24】引入了一套轻量级测试,包括网络内主机和NCCL测试,以诊断广泛的潜在故障。Vela【索引378,The infrastructure powering ibm’s gen ai model development,2024,arXiv】采用两层策略,轻量级测试定期在每个节点上运行,而更具侵入性的测试仅在节点空闲时执行。Google的TPUv4超级计算机在用户作业前实施预检【索引376,Resiliency at Scale: Managing Google’s TPUv4 Machine Learning Supercomputer,2024,NSDI 24】,包括一个端到端检查和一个用于硬件健康的意图驱动检查器。SuperBench【索引381,Superbench: Improving cloud ai infrastructure reliability with proactive validation,2024,USENIX ATC 24】提出了一个全面的基准套件,用于评估单个硬件组件,并集成了一个选择器来平衡验证时间与潜在问题相关的惩罚。
8.3 基于检查点的恢复
检查点恢复的困境。定期保存模型状态(即检查点),并在故障发生后从最新的检查点恢复计算是容错LLM训练的常见做法。然而,这带来了一个困境:频繁的检查点会产生高I/O开销,而不频繁的检查点则会在故障发生时导致大量的进度损失。为解决这一困境,设计了快速的持久化和内存中检查点方法。
8.3.1 持久化检查点
持久化检查点的原理和挑战。持久化检查点涉及将模型状态保存到非易失性存储(如SSD和远程云存储),确保数据在系统故障后仍然持久。该过程通常包括两个阶段:首先,快照阶段将模型状态从GPU复制到CPU内存;其次,持久化阶段将快照写入持久存储设备。尽管由于存储设备的低带宽导致I/O开销相当大,但由于其易用性和可靠性,持久化检查点仍然是一种广泛使用的容错方法。已提出了先进的持久化检查点方法来减少训练停顿,从而能够在没有显著性能损失的情况下更频繁地进行检查点。
同步检查点。为了保持模型参数的一致性,DeepSpeed的默认同步检查点【索引22,Deepspeed: System optimizations enable training deep learning models with over 100 billion parameters,2020,KDD】和Varuna【索引167,Varuna: scalable, low-cost training of massive deep learning models,2022,EuroSys】会定期暂停训练过程,在数据并行等级0上同步地将检查点保存到持久存储。这种方法在快照和持久化阶段都会导致GPU空闲,从而造成资源利用不足。认识到大多数故障可归因于单个GPU或网络设备,JIT-Checkpointing【索引382,Just-in-time checkpointing: Low cost error recovery from deep learning training failures,2024,EuroSys】提出了一种替代策略。它在故障发生后立即进行即时检查点,允许训练从这些JIT检查点恢复。这种方法显著减少了浪费的GPU时间成本,将其限制在最多一个mini-batch迭代的工作量。DLRover Flash-Checkpoint【索引383,Dlrover: An automatic distributed deep learning system,2024,github.com】利用分布式缓存服务加速了迁移效率。Universal Checkpointing【索引384,Universal checkpointing: Efficient and flexible checkpointing for large scale distributed training,2024,arXiv】引入了一种通用的检查点表示,以将分布式检查点存储与并行技术解耦。Universal Checkpointing可以根据需求无缝地将检查点从一种并行策略转换为另一种。
快照-停顿检查点。为了在保持检查点一致性的同时减少LLM训练的停顿,Check-N-Run【索引385,{Check-N-Run}: A checkpointing system for training deep learning recommendation models,2022,NSDI 22】将快照和持久化阶段解耦。它通过仅在快照阶段暂停训练来实现原子检查点,并使用专用的后台CPU进程异步地持久化快照。TorchSnapshot【索引386,Torchsnapshor: A performant, memory-efficient checkpointing library for pytorch applications, designed with large, complex distributed workloads in mind,,github.com】通过张量分块和多线程磁盘写入进一步优化了这一过程。通过创建分块快照,TorchSnapshot允许持久化阶段通过并行写入更早开始,从而减少了整体训练停顿时间。MegaScale【索引71,{MegaScale}: Scaling large language model training to more than 10,000 {GPUs},2024,NSDI 24】和InternEvo【索引18,Internlm2 technical report,2024,arXiv】也采用快照-停顿方法进行快速检查点和恢复。快照阶段暂停训练几秒钟以捕获模型状态,而持久化阶段则异步地将检查点从CPU内存传输到分布式文件系统。MegaScale通过指定数据并行组内的一个工作节点从分布式文件系统读取来优化恢复过程,从而缓解了低带宽瓶颈。该工作节点然后将检查点数据广播给其他GPU,从而在整个系统上实现更快、更高效的恢复。为了节省存储空间,InternEvo还异步地将检查点从昂贵的热存储移动到更便宜的冷存储。
异步检查点。异步检查点旨在通过与训练并发执行快照和持久化阶段来最小化训练停顿。DeepFreeze【索引387,Deepfreeze: Towards scalable asynchronous checkpointing of deep learning models,2020,CCGRID】在后台应用轻量级(快照)和重量级(持久化)两种持久化策略,并在数据并行GPU之间分片检查点以分散I/O工作负载。CheckFreq【索引388,{CheckFreq}: Frequent,{Fine-Grained}{DNN} checkpointing,2021,FAST 21】仔细地将快照和持久化阶段与后续迭代的前向和后向传递流水线化,确保在下一次参数更新前完成快照。它还动态调整检查点频率以平衡恢复成本和运行时开销。LightCheck【索引389,A cost-efficient failuretolerant scheme for distributed dnn training,2023,ICCD】利用迭代间的数据依赖性,引入了逐层检查点流水线以减少停顿。DataStates-LLM【索引390,Datastates-llm: Lazy asynchronous checkpointing for large language models,2024,arXiv】通过为快照预分配固定的主机内存来解决主机内存分配缓慢的问题,并利用高效的计算、快照和持久化逐层流水线。FastPersist【索引391,Fastpersist: Accelerating model checkpointing in deep learning,2024,arXiv】识别了完全异步持久化阶段中的风险,并将其与下一次迭代的参数更新同步。它通过双缓冲固定内存来提高SSD带宽利用率,并通过使用数据并行等级的子集进行检查点写入来减少硬件争用。
8.3.2 内存中检查点
内存中检查点的优势和实现。远程持久存储的低带宽严重限制了检查点的频率,内存中检查点通过将检查点存储在其他计算节点或专用内存存储系统的内存中来解决这些限制,显著减少了I/O开销并实现了更高的检查点频率。Gemini【索引392,Gemini: Fast failure recovery in distributed training with inmemory checkpoints,2023,SOSP】提出将检查点保存到CPU内存以实现更快的故障恢复,并提出了一种检查点放置策略以最小化检查点损失,以及一种流量调度算法以减少对训练的干扰。REFT【索引393,Reliable and Efficient In-Memory Fault Tolerance of Large Language Model Pretraining,2023,arXiv】异步地将模型状态缓存到主机内存和像Redis这样的内存存储中,绕过了检查点I/O并实现了高检查点频率。它还利用纠删码来实现RAIM5(受RAID5启发,将“磁盘”替换为“内存”),以保护数据免受节点故障的影响。虽然这些方法通过实现更频繁的检查点而无性能损失,显著推进了LLM训练的容错性,但它们可能无法提供与传统基于存储的方法相同的长期数据持久性。因此,结合内存中和持久化检查点的混合方法对于全面的容错策略是必要的。
8.4 无检查点恢复
无检查点恢复的原理。无检查点恢复方法旨在通过消除在发生故障时需要重启和回滚到先前检查点的需求来最小化训练停顿。这些技术依赖于自动故障检测机制来及时发现问题。当检测到故障时,无检查点方法会自动解决问题,并允许训练过程在没有中断的情况下继续。通过避免从检查点加载和重复计算的耗时过程,这些方法可以显著减少停机时间并提高整体训练效率。无检查点恢复策略可大致分为两种主要方法:实时迁移和模块冗余。
8.4.1 实时迁移
实时迁移的实现。实时迁移利用分布式LLM训练设置中固有的冗余,特别是跨不同数据并行流水线的模型副本,来在发生故障时恢复模型状态。当检测到故障时,实时迁移方法使用剩余的健康实例或通过将新实例并入训练集群来动态重新配置并行策略。然后,当前的模型状态被传输到这些重新配置的节点,允许训练过程以最小的中断继续。Parcae【索引394,Parcae: Proactive, Liveput-Optimized DNN Training on Preemptible Instances,2024,arXiv】提出了三种不同的迁移机制,每种机制具有不同的通信开销,以在不同的并行策略之间高效地传输模型状态。Oobleck【索引395,Oobleck: Resilient distributed training of large models using pipeline templates,2023,SOSP】采用基于流水线模板的方法进行实时迁移。它维护一组预定义的流水线模板,并在检测到故障时,根据这些模板迅速实例化新的异构流水线。
8.4.2 模块冗余
模块冗余的实现。模块冗余,像实时迁移一样,也利用了模型状态的冗余性。然而,这种方法不是在不同GPU之间恢复最新的模型状态,而是通过将计算路由到冗余模块来继续训练。Bamboo【索引396,Bamboo: Making preemptible instances resilient for affordable training of large {DNNs},2023,NSDI 23】在持有同一流水线内相邻流水线阶段的GPU中放置一个冗余的流水线阶段。这个冗余阶段在训练期间利用流水线气泡执行冗余计算,并在发生故障时被激活为正常阶段。SlipStream【索引397,Slipstream: Adapting pipelines for distributed training of large dnns amid failures,2024,arXiv】利用跨模型副本流水线的冗余,将失败节点的计算路由到不同数据并行流水线中的节点。SWARM【索引231,Swarm parallelism: Training large models can be surprisingly communication-efficient,2023,ICML】提出了类似的解决方案,但更侧重于连接不良、异构且不可靠的设备。除了冗余计算,SWARM还结合了实例迁移来重新平衡流水线,结合了冗余计算和实时迁移两方面的特点。
A5 结论
对LLM训练系统和基础设施进展的总结。LLM的兴起改变了人工智能,催生了个人助理、代码生成和科学研究等应用。像GPT、LLaMA和Gemini这样的模型设立了新的标准,但训练这些大规模模型,如LLaMA-3在16,384个GPU上耗时54天的过程,带来了可扩展性、效率和可靠性方面的挑战。管理庞大的GPU集群需要创新的硬件和网络解决方案。高效的训练要求优化计算、通信和内存使用。可靠性则涉及在长时间的训练期间建立强大的机制来检测和从故障中恢复。本综述回顾了LLM训练系统和基础设施的最新进展,重点介绍了增强可扩展性、效率和可靠性的方法。
未来展望:光子计算。由摩尔定律和登纳德缩放定律指导的传统基于数字电路的计算系统,在满足训练和部署LLM的计算需求方面正面临着重大的物理和经济限制。因此,AI行业需要创新的解决方案。一个有前途的方法是大规模光电集成技术,它用集成的硅光子学取代传统数字电路,以增强计算和通信能力【索引402,Taking stock of new data center computing paradigm,2023,eetimes.com】。这种光电混合数据中心技术将光计算与光网络相结合,提高了单节点计算能力和大规模分布式计算的效率。一些工作已经提出利用光网络进行LLM训练。例如,TopoOpt【索引67,{TopoOpt}: Co-optimizing network topology and parallelization strategy for distributed training jobs,2023,NSDI 23】在分布式训练中同时优化光网络拓扑和并行策略,增强了计算和通信效率。TPUv4【索引42,Tpu v4: An optically reconfigurable supercomputer for machine learning with hardware support for embeddings,2023,ISCA】使用光路交换机动态重构其3D环面互连拓扑,为LLM训练中密集的通信模式改善了数据流。此外,Taichi【索引403,Large-scale photonic chiplet taichi empowers 160-tops/w artificial general intelligence,2024,Science】探索了一种分布式衍射-干涉混合光子计算架构,以每瓦160万亿次操作(TOPS/W)的能效,有效地将光神经网络扩展到百万神经元级别。未来可能需要将LLM的训练和推理范式转向硅光子学。然而,这一转变将需要在系统设计和实现方面进行广泛的创新。
💬 评论讨论
欢迎在这里分享您的想法和见解!