本帖最后由 中国计算机学会 于 2024-4-19 15:02 编辑
摘要—人工智能（AI）已经成为当今社会不可或缺的一部分。近年来，为了满足人工智能模型的计算资源需求，国产AI芯片迅速崛起。然而，由于配套软件的不足，目前国产芯片难以充分发挥其计算能力。本文探讨了将国产算力应用于大模型训练的挑战和机遇，同时结合我们的经验提出了潜在的解决方案。
马子轩（清华大学）
翟季冬（清华大学）
郑纬民（清华大学）
关键词 ：人工智能 大模型 高性能计算
引言 人工智能（AI）技术迅速发展，已广泛渗透到日常生活的各个方面，从智能语音助手[1]、自动驾驶汽车[2]，到个性化推荐系统[3]和智能医疗诊断[4]，AI都展现出了巨大的潜力和价值，成为当今社会中不可或缺的一部分。这一切的成功得益于模型规模的不断扩大及其效果的持续优化。随着算法和硬件技术的不断进步，AI模型正变得愈加复杂和庞大，能处理更加复杂的任务和数据，并取得更好的结果。大模型的兴起是这一发展趋势的重要体现。例如以GPT系列[5]和LLaMA[6]系列为代表的大规模神经网络，拥有数百亿甚至上千亿的参数量，为自然语言处理和计算机视觉任务提供了前所未有的性能，不断拓展人工智能的认知边界。 然而，大模型的成功和广泛应用也带来了一个巨大的挑战，即如何满足其对计算资源的巨大需求。大模型的“大”体现在两个方面。一方面，庞大的参数量需要大量的内存资源。一个千亿规模的模型仅存储其所有参数，就需要几百GB甚至TB级别的内存空间，再考虑其训练过程中需要存储的梯度、更新参数等数据，所需的内存资源就更多了。另一方面，庞大的计算量需要大量的计算资源。更大的模型本身就意味着更大的计算量，对于每个样本，模型参数量越大，计算量就越大。此外，训练更大规模的模型，通常需要更大规模的数据，也就进一步提高了模型训练需要的总计算量。对内存和计算资源的高需求，使得大模型面临巨大的资源挑战。 近年来，国产人工智能芯片发展迅速。以华为昇腾、海光DCU、寒武纪MLU等AI加速器为代表的国产芯片经过多年的优化设计和迭代后，在性能、功能、易用性和可靠性方面都有了显著提升，成为解决国内算力需求的重要支撑。 当前，各家厂商的技术路线不同，体系结构各异，百花齐放，中国AI芯片发展进入了黄金时代。虽然硬件架构发展迅速，但其配套软件发展却不充分，成为限制这些新兴硬件平台在大模型任务中使用的主要因素。同时，国产芯片的快速发展也导致其配套软件存在发展速度不匹配的问题。这也是限制在主流AI训练任务中使用国产算力的主要因素。 为了研究如何高效地将国产算力应用到大模型应用上，2020年我们在新一代神威超级计算机上移植了大模型训练应用。该工作探索了当前国产平台系统软件面临的主要问题，并尝试性地提出了一系列解决方案。本文将结合该工作的经验，探讨当前国产平台需要解决的问题、面临的主要挑战以及可能带来的机遇。 国产算力平台 目前，国内企业与高校正积极投入大量计算资源以应对大模型训练任务的需求。然而，为了有效发挥规模效应并支持更大规模的模型训练，必须将这些资源整合为一体，构建完整的算力系统。目前，大规模人工智能训练任务主要依赖英伟达等GPU厂商提供的硬件和相关软件栈。然而，受复杂多变的国际局势影响，国内面临着无法获取足够数量GPU的困难，并且已有的GPU使用分散，难以建立规模足够大的GPU算力系统，严重制约了国内人工智能领域的发展。此外，依赖进口GPU不仅带来了经济负担，还可能受到国际政治因素的干扰，进而对人工智能产业的可持续发展构成威胁。因此，开发国内的算力平台，将国内自主研发的硬件和技术应用于大模型的训练中，已经成为当务之急。这不仅是出于技术自主性和国家安全的考量，更是为了确保我国人工智能产业能够实现长期可持续发展并保持国际竞争力。
近年来，国内芯片领域发展迅猛，呈现出多元化的趋势。表1列出了一些国内代表性的芯片公司及其产品。显而易见，不同公司采用了不同的技术路线，产生了多样化的芯片架构。而且不同的芯片采用不同的编程语言，支持不同的数据类型，需要不同的优化策略，用户很难将自己的程序迁移到不同的芯片上。因此，如何更好地利用国产硬件资源对用户来说是一个重大挑战。另一方面，为了满足大模型的算力需求，需要进一步发展基于国产硬件的大型智算中心。然而，目前基于国产硬件的大规模系统存在软件支持相对薄弱和封闭性较强的问题，难以满足不同AI用户的个性化需求，这也是国内算力平台发展面临的另一个挑战。因此，国产硬件算力系统的持续发展，对满足日益增长的人工智能需求，确保我国在该领域的领先地位至关重要。 典型模型训练流程 典型的模型训练流程可以从计算的角度分为三个关键步骤：前向计算、反向计算和参数更新。这些步骤构成了模型在训练中的核心计算过程。 前向计算 在前向计算阶段，模型使用输入数据和当前的模型参数进行计算，生成预测的结果。系统会将这些预测的结果与标注后的数据进行比较，计算出当前预测的损失（loss）。 反向计算 在反向计算阶段，系统会使用前向计算得到的损失，计算所有参数的梯度（即损失对参数的导数）。这些梯度指示了如何调整模型参数，以减小损失。这个过程也被称为反向传播。 参数更新 所有梯度计算完成后，模型使用这些梯度更新参数，以提高模型端到端的正确率。通常使用梯度下降等优化算法实现参数的更新。这个过程在训练中不断迭代，直到模型收敛到满意的状态。 在此基础上，大模型的发展需要更多的计算资源与存储资源，随着模型规模的扩展，单个机器或单个计算卡很难满足大模型训练的需求。采用多机多卡并行已是大模型训练中必不可少的一环，以下是常见的四种并行策略：数据并行、张量并行、流水线并行和专家并行。 数据并行 该策略将输入数据分配到不同的计算节点，每个节点保存一份完整的模型。在前向和反向计算过程中，各个节点独立进行计算，无须通信。在参数更新时，需要进行全局通信（All-Reduce）来汇总不同节点的梯度。数据并行允许扩展模型的训练吞吐量，但是不能有效提高参数量。 张量并行 该策略将模型参数分配到不同的计算节点上。计算过程中通过通信操作（All-Reduce或All-Gather）协同使用多个计算节点计算同一组数据。张量并行允许扩展模型的参数数量，但不能有效提高训练吞吐量。 流水线并行 该策略将模型的不同层分配到不同的计算节点上。计算过程中各层的计算结果会被传输。流水线并行也可以扩展模型的参数数量，但不能有效提升吞吐量。张量并行和流水线并行也被统称为模型并行。 专家并行 为了满足大模型训练对参数量和吞吐量的要求，近年来出现了一种新的预训练模型，称为混合专家（MoE）模型[7]，MoE将大模型划分为多个小模型。在MoE模型中，数据划分和专家划分同时存在，导致全局All-to-All通信，但能有效满足大模型的参数规模和训练吞吐量需求，这种并行模式被称为专家并行。专家并行同时具有高性能和参数规模扩展的特点。MoE是一种适用于新一代大模型的并行策略，谷歌的Switch Transformer[8]、GPT-4均采用了MoE架构。 在应对大模型日益复杂的计算需求时，一方面，模型本身的复杂架构使开发人员难以直接编写人工智能应用，需要更多的抽象和支持；另一方面，单一的并行模式已无法满足不断扩展的模型结构、参数量以及训练吞吐量的需求，因此混合多种并行策略成为了常见的使用策略。这表明，为了支撑人工智能应用的发展，系统软件的重要性不可忽视，而这也给国产算力系统带来了更多挑战。 国产算力面临的挑战 国产算力面临的挑战主要集中在系统软件方面。学术界和工业界普遍认为，系统软件是影响国产硬件平台使用的关键因素。对于人工智能应用开发者来说，使用英伟达等GPU提供的软件栈非常方便，这些软件易于快速部署，可靠性高、可用性好，而且有大量用户群体。相比之下，国产硬件平台的软件系统不够全面，对有些模型的支持不完善，有些模型甚至需要额外的工作才能支持。此外，不同的硬件平台使用不同的系统软件，用户群体分散，学习资料少，这无疑增加了用户的学习使用成本。要成功训练大模型，需要一系列系统软件的协同支持，包括编程语言、编程框架、并行加速库、通信库、算子库、AI编译器等，如图1所示。这些软件为大模型的开发和训练提供了必要的基础。
编程语言 编程语言提供了在硬件上编写并行程序的接口。为了发挥硬件性能，编程语言一般要求覆盖硬件的全部功能。在此基础上，研究人员可利用编程语言编写人工智能模型的基本算子的高效实现。目前流行的编程语言包括英伟达的CUDA、英特尔的oneAPI、寒武纪的BANGC等。 编程框架 随着AI模型越来越复杂，使用编程语言直接开发AI模型也变得愈发困难。为了降低开发复杂度，出现了一系列编程框架，它们为开发人员提供了一种更方便、更高效的方式构建AI模型。通常，这些框架将AI程序描述为不同的算子组成的复杂程序，通过依次调用算子库的形式实现计算。常见的编程框架包括PyTorch[9]、TensorFlow[10]等。 并行加速系统 随着大模型的参数数量的增加，单个芯片的内存不足以满足模型训练的需求。因此，多机多卡协同分布式训练成为必要。并行加速库可以为多机多卡环境下的人工智能模型提供高效的并行训练能力，使其能充分利用计算资源。通过对并行训练方法进行抽象，并行加速库允许用户快速地将模型在一个集群上并行训练部署，从而降低开发难度。流行的并行加速系统包括微软的DeepSpeed[11]、英伟达的Megatron-LM[12]、清华的FastMoE[13]等。 上述系统软件直接面向用户提供计算能力，以及实现模型训练部署的能力。此外，为了支撑这些软件，算力系统还需要其他的系统软件。 通信库 为了支持多机多卡训练，机器之间需要时刻传输数据，导致通信成为训练过程中的重要性能瓶颈。通信库提供了跨机跨卡的通信能力，以及能够满足大模型训练所需的各种复杂通信模式。通信库需要针对系统的网络特征，充分利用通信带宽，实现高效的数据通信。常见的通信库有MPI[14]、英伟达的NCCL库[15]、华为HCCL等。 算子库 算子库提供了人工智能模型高性能实现所需的基本操作，如矩阵乘法、卷积、池化等。这些操作构建了复杂神经网络模型的基础。目前常见的算子库有英伟达的cuDNN[16]、cuBLAS[17]等。 AI编译器 AI编译器的主要目的是在异构处理器上提供高效的目标代码，以加速人工智能模型的执行。通常来说，AI编译器的设计目标是为当前算子库无法支持的计算负载生成代码。近年来，AI编译器得到了进一步发展，支持更大范围的代码生成以及多个算子实现更高性能的合并。代表性的AI编译器有TensorFlow XLA[18]、TVM[19]等。 其他支撑软件 除了上述软件外，大模型训练还需要其他一系列支撑软件，例如调度系统、内存管理系统、容错系统、存储系统等。这些软件在确保训练过程高效运行方面起着重要作用。 总的来说，大模型的训练是一个复杂任务，需要多个系统软件协同工作。更加完整、易用、高效的系统软件对国产平台至关重要。系统软件的发展和优化也有助于推动人工智能研究和应用领域的进步，为开发人员提供更多灵活性和性能优势。 新一代神威平台的探索 为了研究如何高效地将国产算力应用于大模型任务，我们在新一代神威超级计算机上探索了如何解决训练百万亿参数量大模型时面临的挑战。在这一工作中，我们在并行策略、模型划分、混合精度、负载均衡、内存性能等方面提出多项关键技术，主要针对新一代神威系统的计算框架、并行加速系统和算子库进行设计与实现，并基于此提出一系列解决方案。 最终，该工作成功地将大模型训练任务扩展到新一代神威超级计算机的全机规模，并被命名为“八卦炉”，支持最高174万亿参数量的模型训练任务。“八卦炉”在训练1.93万亿、14.5万亿参数量的模型时，持续性能超过1 EFLOPS，并保持了接近线性的可扩展性。“八卦炉”还支撑了M6[29]、悟道2.0、GLM[30]等一系列模型的训练。该工作是大规模算力系统与人工智能协同设计的一个典型例子，为将来打造针对其他国产算力平台的训练系统提供了宝贵的经验。 在当前大模型高速发展的大趋势下，国产算力平台的崛起成为我国人工智能领域发展的一大机遇，同时也给系统软件设计带来了挑战。我们必须认真应对，确保该领域的长期可持续发展，以保持我国在这一领域的领先地位。 ■
脚注： 参考文献： [1] Amodei D, Ananthanarayanan S, Anubhai R, et al. Deep speech 2: End-to-end speech recognition in english and mandarin[C]// ICML’16: Proceedings of the 33rd International Conference on International Conference on Machine Learning. 2016: 173-182. [2] Grigorescu S, Trasnea B, Cocias T, et al. A survey of deep learning techniques for autonomous driving[J]. Journal of Field Robotics, 2020, 37(3): 362-386. [3] Bobadilla J, Ortega F, Hernando A, et al. Recommender systems survey[J]. Knowledge-Based Systems, 2013, 46(7): 109-132. [4] 糜泽花, 钱爱兵. 智慧医疗发展现状及趋势研究文献综述[J]. 中国全科医学, 2019, 22(03): 366-370. [5] Brown T B, Mann B, Ryder N, et al. Language models are few-shot learners[C]// Proceedings of the 34th International Conference on Neural Information Processing Systems. 2020: 1877-1901. [7] Masoudnia S, Ebrahimpour R. Mixture of experts: a literature survey[J]. Artificial Intelligence Review, 2014, 42: 275-293. [8] Fedus W, Zoph B, Shazeer N. Switch Transformers: Scaling to trillion parameter models with simple and efficient sparsity[J]. Journal of Machine Learning Research, 2022, 23(120): 1-39. [9] Paszke A, Gross S, Massa F, et al. PyTorch: An imperative style, high-performance deep learning library[C]// Proceedings of the 33rd International Conference on Neural Information Processing Systems. 2019: 8026–8037. [10] Abadi M, Barham P, Chen J, et al. TensorFlow: A system for Large-Scale machine learning [C]// 12th USENIX Symposium on Operating Systems Design and Implementation (OSDI 16). 2016: 265-283. [11] Rajbhandari S, Rasley J, Ruwase O, et al. ZeRO: Memory optimizations toward training trillion parameter models[C]// Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis. 2020: 1-16. [12] Narayanan D, Shoeybi M, Casper J, et al. Efficient large-scale language model training on GPU lusters using Megatron-LM [C]// Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis. 2021: 1-15. [14] Gropp W, Lusk E, Skjellum A. Using MPI: Portable parallel programming with the message-passing interface[M]. The MIT Press, 2014. [16] Chetlur S, Woolley C, Vandermersch P, et al. cuDNN: Efficient primitives for deep learning[OL]. [2014-12-18]. arXiv preprint arXiv:1410.0759, https://arxiv.org/abs/1410.0759. [19] Chen T, Moreau T, Jiang Z, et al. TVM: An automated End-to-End optimizing compiler for Deep Learning[C]// 13th USENIX Symposium on Operating Systems Design and Implementation (OSDI 18). 2018: 578-594. [20] Ma Z, He J, Qiu J, et al. BaGuaLu: Targeting brain scale pretrained models with over 37 million cores[C]// Proceedings of the 27th ACM SIGPLAN Symposium on Principles and Practice of Parallel Programming. 2022: 192-204. [21] 马子轩, 翟季冬, 韩文弢, 等. 高效训练百万亿参数预训练模型的系统挑战和对策[J]. 中兴通讯技术, 2022, 28(2): 51-58. [23] He J, Zhai J, Antunes T, et al. FasterMoE: Modeling and optimizing training of large-scale dynamic pre-trained models[C]// Proceedings of the 27th ACM SIGPLAN Symposium on Principles and Practice of Parallel Programming. 2022: 120-134. [24] Zhai M, He J, Ma Z, et al. SmartMoE: Efficiently training Sparsely-Activated models through combining offline and online parallelization[C]// 2023 USENIX Annual Technical Conference (USENIX ATC 23). 2023: 961-975. [25] Wang H, Zhai J, Gao M, et al. PET: Optimizing tensor programs with partially equivalent transformations and automated corrections[C]// 15th USENIX Symposium on Operating Systems Design and Implementation (OSDI’21). 2021: 37-54. [26] Zheng L, Wang H, Zhai J, et al. EINNET: Optimizing tensor programs with Derivation-Based transformations[C]// 17th USENIX Symposium on Operating Systems Design and Implementation (OSDI 23). 2023: 739-755. [27] Ma Z, Wang H, Xing J, et al. PowerFusion: A Tensor Compiler with Explicit Data Movement Description and Instruction-level Graph IR[OL]. [2022-07-11] arXiv preprints, https://arxiv.org/abs/2307.04995. [28] Tang S, Zhai J, Wang H, et al. FreeTensor: A free-form DSL with holistic optimizations for irregular tensor programs[C]// Proceedings of the 43rd ACM SIGPLAN International Conference on Programming Language Design and Implementation. 2022: 872-887.
版权声明：中国计算机学会(CCF)拥有《中国计算机学会通讯》(CCCF)所刊登内容的所有版权，未经CCF允许，不得转载本刊文字及照片，否则被视为侵权。对于侵权行为，CCF将追究其法律责任。
| 
京公网安备 11010802032778号 )
发表于 2024-4-18 09:51:15