本文内容来源于蚂蚁集团 AI Infra部分担任人张科，在 GTC 2024 大会 China AI Day 线上专场的讲演。在讲演中张科同享了 AI 工程当时的现状和首要应战，以及蚂蚁集团在 AI 工程范畴的实践经验和开源项目，也欢迎 AI 工程范畴的同行们一起参与开源项目的共建与同享。
张科于 2009 年中国科学院硕士结业参加微软，首要作业有必应搜索引擎排序体系重构和广告主平台研发，开源深度学习模型标准 ONNX 以及推理引擎 ONNX Runtime 的首要作者之一。2020 年回国后参加蚂蚁，担任蚂蚁 AI 根底设施建造。
扫描上方图片二维码或点击【阅读原文】，观看视频，并可下载讲义。
AI工程现状与应战
下图趋势展现了当今 AI 工程的现状：模型参数规划呈指数级增加，动辄达万亿规划，甚至更高，与此相伴的是对算力的需求急剧攀升，其需求是万卡乃至几十万卡，EFLOPS 级别的核算才能。值得庆幸的是，在 AI 工程层面，开源生态蓬勃开展，有诸如 DeepSpeed、NVIDIA Megatron-LM, TensorRT-LLM 等开源结构可供业界借鉴与运用。
当时 AI 工程面临多重应战：
首要，本钱问题是首要痛点，无论是练习仍是推理本钱，尤其是 AI 练习的试验本钱亟需下降。正如咱们不希望核试验频繁操作，AI 模型练习相同希望下降试错本钱，追求试验效率的最大化。此外，鉴于练习是一个长期且杂乱的分布式体系运作进程，咱们迫切要求练习进程高效且稳定，分布式体系的可靠性至关重要。
另一方面，面临很多开源结构和技能选项，AI 工程师在技能选型上面临巨大应战，包含需求考虑开源技能与自研技能的权衡、怎么在琳琅满目的开源资源中挑选适合本身需求的技能途径等，一起，这也意味着企业在人力投入上需求深思熟虑，既要看准技能趋势，又要保证资源的有用运用。
蚂蚁 AI 工程实践
蚂蚁集团在 AI 根底设施建造中饯别全栈战略，包含使用层、平台产品层、引擎结构和异构 AI 集群。本文将详细介绍这四个层面：首要，聚集于咱们发布的智能分布式练习服务 DLRover，其次，探讨咱们在分布式推理场景中的实践经验，再者，介绍咱们在显存优化与数据传输方面的作业作用 GLake；最终会同享咱们在工程智能范畴的具体作业，特别是模型优化器的根底算法研发。
其间，DLRover、GLake 以及模型优化器均已开源，欢迎 Star 重视：
- DLRover：github.com/intelligent…
- GLake：github.com/intelligent…
大规划分布式智能练习服务 DLRover
DLRover 是咱们推出的专心于大规划分布式练习的智能服务。面临大模型分布式练习的反复迭代流程，DLRover 致力于解耦模型界说、练习逻辑与练习优化，以便算法工程师专心于中心的模型研发作业，一起赋予工程研发人员进行优化定制的才能。
DLRover 的规划初衷包含三个方面：
- 首要，完成主动并行，针对大规划分布式练习的并行调优，力求使各种并行战略高效协同，以到达最优效率。
- 其次，供给主动资源装备功用，减轻算法研发人员与硬件和工程参数交互的担负，让他们无需关心非专业硬件和工程参数设置，专心中心的算法研发。
- 最终是提高大规划分布式练习的稳定性，保证数千乃至数万个分布式节点在长达数月的运转进程中，始终坚持长期稳定且高效的运转状态。
总结来说，DLRover 的中心方针仍是在于提高整个体系的效率，包含资源、功用、人效。其关键特性体现为：
- 深度解耦：分离算法研发与练习结构的依赖关系，并减少练习角色及调度体系的耦合度，使算法工程师可以聚集于算法创新，而非底层技能细节；
- 资源调度：选用中心化决议计划机制，具有全局资源视角，实时把握大规划分布式练习使命以及集群全体资源运用情况，然后作出更优的资源调度决议计划；
- 引擎优化：DLRover 整合了多种分布式优化战略，并可以主动生成最优决议计划组合；
- AI for Sys：经过构建模型来智能决议计划最优战略组合，经过超参数调优练习功用，资源运用率，以及模型结构与作用，然后在练习进程达成综合性的优化使命。
DLRover 的中心架构包含三个部件：
- 完好的服务模块；
- 引擎优化扩展库；
- 面向 AI 工程规划的优化，从主动化到智能化进程的算法库，服务于 AI for Sys 部分；
在整个练习服务流程中，DLRover 统筹离线和在线练习需求，包含资源的动态弹性弹性、容错机制、资源装备办理，以及对动态数据的高效分发与康复。其间，动态数据的高效分发和康复，会和容错机制、大模型练习稳定性提高模块、异构加速器提效和稳定性模块，配合运用。在引擎优化层面，DLRover 重点重视练习结构优化，以及 TensorFlow 的优化扩展，因为 TensorFlow 在大模型年代前，其使用较为广泛，尤其是搜推场景。为此，咱们开发了 TensorFlow 优化扩展库 TFPlus，以及 PyTorch 优化扩展库 ATorch。这两个库均作为扩展组件，需与相应的 TensorFlow 或 PyTorch 版本配套运用，而非独立运转。至于 AI for Sys 部分，咱们专心于主动化到智能化的转化，首要体现在自研优化器和在线学习算法库上。该算法库包含了处理数据推迟反馈、处理在线学习中数据漂移等问题。
在蚂蚁集团内部，DLRover 智能练习服务的实践使用展现出显着的作用。
- 首要，其智能化程度已到达让算法用户无需手动装备任何与硬件相关的资源参数，诸如 CPU 数量、内存大小、GPU 数量及显存分配等，全部交由体系主动处理。
得益于 DLRover 集成的容错机制和弹性弹性功用，练习作业的成功率完成了大幅提高，下图数据显现已到达 96%，实践最新数据更为出色。尤为值得一提的是，DLRover 的智能化装备相较于人工手动装备，在练习速度上展现出了显着优势，有用地缩短了练习周期。
- 此外，如下图右下角第五点，在大模型练习方面，DLRover 的体现相同引人注目。
面临千亿参数模型及千卡规划的并行练习场景，DLRover 可以完成高达 60% 以上的高功用目标 (HFU)，而且这一数据仍在持续前进。而在稳定性方面，DLRover 保证了极高的有用练习时长占比，即仅核算模型真实起作用的时刻，除掉无效及反常停机时刻，目前可到达 97% 以上，最优情况下甚至能到达 99.4%，显现了杰出的练习稳定性。一起，咱们也与 NVIDIA 合作，在 FP8 低精度练习上取得了一定的进展。
如下图所示，DLRover 的 TFPlus 扩展库经过引入 GPU 同步练习机制，尤其在稀少模型练习场景下，功用提高作用十分显着。在TFPlus扩展库中，咱们特地开发了 hybrid embedding 优化技能，针对稀少 embedding 占用存储空间大且访问热度不均的特色，创造性地将 SSD 等存储介质融入其间，构建了一套分层的 embedding 存储计划。这套计划使得即便是处理万亿级别的大规划模型预练习使命，练习速度也能坚持高效而不受任何影响。一起，内存运用量大幅减少，节约率高达 80% 以上，堪称是一项极为出色的优化作业。
分布式推理引擎
在 AI 推理引擎部分，下图展现的是一张详细的推理工程栈架构图，其间的组件包含前端、图优化、运转时及算子优化等，本文重点阐述蚂蚁集团在推理引擎建造中的技能选型和实践作用。
首要，咱们选用 ONNX 路线，一致支撑LLM和搜推模型的极致推理优化。考虑到公司内部涉及多种练习结构如 TensorFlow 和 PyTorch，作为企业界的 AI 工程中台部分，需求统筹各类模型的推理需求。因而，分布式推理、流式推理以及 Continuous Batching 和 LookAhead 等中心特性被纳入推理引擎规划中。其间，Continuous Batching 和 LookAhead 就是预先做 decoding，可以根据实践情况动态组合，以完成最佳功用。例如，在较高 Batching 情况下，LookAhead 的重要性相对减弱。此外，咱们还实践了高功用算子，包含 FlashAttention、Flash Decoding 以及量化技能等，不断提高推理效率。就目前可以同享的作用显现，在百亿参数模型上，运用两张 NVIDIA A10 Tensor Core GPU 并行推理，咱们可以完成每秒处理超越 560 个 tokens 的吞吐量，首次呼应推迟控制在 200 毫秒以内。而且这个数据并非最新作用，但展现出令人满意的多卡加速比和推理功用。
练习推理显存和传输优化 GLake
在显存优化与传输优化方面，咱们开发了名为 GLake 的服务，并已有一篇论文《GMLake》被 ASPLOS24 收录。GLake 旨在完成显存池化、层次化办理和碎片优化，充分运用一台物理机内部多张 GPU 卡之间的高速互联，完成多路并发和提高数据传输速率。
经实践验证，GLake 在大模型练习场景下体现杰出，吞吐量提高了近四倍，使得 batch size 可大幅度增大，一起数据传输速率提高了 3 到 12 倍。相较于 PyTorch 和 DeepSpeed，选用 GLake 进行大模型做调优（fine-tuing）时，吞吐量也能提高大约一倍。值得重视的是，GLake 作为一个通用的加速库，能很好地与现有开源结构和工程项目协同，例如 DeepSpeed、NVIDIA Megatron-LM，vLLM 等。
模型优化器三部曲
最终，本文会介绍蚂蚁集团在工程智能范畴的探索，特别是针对大模型练习的优化战略。这项作业最初起源于对普通模型练习进程的深入研究，包含怎么经过工程与算法相结合的方式，优化资源运用、提高模型泛化才能和加速收敛进程。为此，咱们推出三种优化器，而且每种优化器都有一篇对应学术论文可供参考。
- 第一个优化器是 Group Optimizer，这是一种无损压缩技能，旨在节约资源。经过该优化器，模型可以完成 80% 以上的压缩率，且不影响模型功用，广泛适用于各类在线学习场景。
- 第二个优化器 WSAM 聚集于提高模型泛化才能。咱们知道，深度神经网络的泛化功用与其极值点的平整程度密切相关。WSAM 经过将平整度作为正则化项，以优化练习极值点的平整度，然后有用提高模型的泛化才能。在某个落地营销场景，经过选用 WSAM 优化器，咱们成功将全体 GMV 提高了近 30%。
- 最终侧重介绍用于加速收敛的优化器 AGD，之所以如此重视，是因为它与大模型练习的效率提高严密相关。这款优化器旨在缩短模型练习时长，快速收敛至理想状态，一起坚持良好的模型功用。其内在数学原理借鉴了拟二阶优化器的思想，以完成更高效的收敛速度。
在实践验证进程中，咱们经过对 GPT-2 模型进行试验，搜集数据来查验该优化器的作用。试验成果显现，在 GPT-2 模型上，该优化器能使收敛速度提高 1.5 倍。
开源共建&同享
在大模型技能敏捷演进的年代，技能创新与工程实践的前进可谓瞬息万变，令每一位从业者目不暇接。咱们以为，AI 范畴的中心之一在于同享精神与共赢心态。跟着算法的不断打破，必然伴跟着工程层面的优秀实践涌现。在这个行业中，特别是在工程范畴，很少有肯定保密或耐久抢先的秘诀。因而，倡导开源共建与同享显得尤为重要。咱们鼓舞所有从业人士，不论是产业界的实践者仍是学术界的科研人员，积极参与进来，一起推进AI工程和AI根底设施建造的开展。经过同享各自宝贵的实践经验与研究作用，大家可以一起探讨、相互启发，一起提高整个行业的技能水平。
目前蚂蚁 AI Infra 团队已经开源了 2 个项目，协助处理大模型练习进程中的工程问题，欢迎 Star 重视：
- DLRover，致力于经过提高深度学习练习进程的智能性，来处理整个体系的提效问题：
- GLake， 一个作业在底层（虚拟与物理显存办理）与体系层（包含多卡、多通道、多使命）的加速库以及相关工具集，旨在对 GPU 显存+传输进行一体优化：
