一文读懂什么是AI Infra和大模型算法工程

什么是AI Infra和大模型算法工程

大模型的三要素是算法、算力和数据，我们需要海量的数据搭配上千卡甚至万卡算力集群再加上优秀的算法才能训练出一个优质的大模型。但把这三者连接起来并不是一件容易的事情。通常搞算法的同学如果想独立地训一个10B以上的模型，会有一段非常痛苦的调试期。在来我司之前，我在硕士期间是做的AI算法相关，当时实验室拿到了8块NVIDIA 4080非常兴奋，觉得可以尝试玩玩小的大模型了，结果找了个开源代码仓花了几天时间终于把代码调通，发现立马就OOM(out of memory)了，又上网找了一些省显存的技巧，模型终于跑起来了，但巨慢无比，可谓是一波三折。

AI infra或者说大模型算法工程的角色，就是介于算法和算力芯片之间，搭起一道桥梁，能让算法工程师更快速地把算法能在算力集群上高效地运行起来。具体来说，有下面几件事要做：

维度

关键词

一句话定义

数据

Feature Store / 数据治理

把「原始日志」转成「模型可用特征」并持续保鲜。

训练

分布式训练 / 显存优化/ 性能优化

让 10B~100B 模型在千卡集群上高效收敛。

推理

Serving / KV Cache / 编译器

让大模型在 GPU/NPU上低延迟、高吞吐、低成本地跑起来。

平台化

MLOps / 监控 / CI-CD

让算法同学像写 Python 脚本一样把模型端到端上线。

小结：AI Infra = Data Infra + Training Infra + Inference Infra + MLOps Infra

说起来看似简单，但从最底层的加速芯片，如NVIDIA的A100/H100和我们的昇腾910，一直到最上层的算法，有不止一层软件栈。大致上可以分为：

1. 底层的硬件驱动和算子加速库（Nvidia的CUDA，昇腾的CANN）
2. AI编程框架（Pytorch/Tensorflow/MindSpore）
3. 分布式训练/推理加速库（Megatron-Core/TransformerEngine/DeepSpeed/FSDP/MindSpeed/vLLM/sglang/MindIE）
4. 上层模型与应用套件（veRL/nemo/Megatron-LM）

行业标杆NVIDIA在这方面的积累很深厚，我们为了做好昇腾生态，近几年也投入了大量的资源，要做好AI Infra其实是投入不菲的，这个“坑”巨大同时机会也很多。另外，这两年大模型的发展非常迅速，从一开始的chatgpt到llama再到deepseek，模型的规模和参数量越来越大，模型架构也在不断迭代升级，但硬件的发展速度没有上层算法的发展快，这就会导致现有的infra不能满足新的算法带来的需求。因此，只要大模型算法和应用在不断的发展，AI infra技术就需要不断的演进来匹配新算法带来的新需求，来更好的链接算法与算力。或者，从公司和商业角度来说，作为AI Infra人员，只要你的业务方算法能不断地有新的算法或者新的业务形态，你就会有新的场景需要优化。基于此，目前大部分的头部AI应用厂商例如阿里/字节/腾讯/百度以及算力厂商NVIDIA，昇腾等等都有AI Infra岗位的需求，仍然处于高速发展的阶段。另外，DeepSeek更是证明了AI Infra的重要性，杰出的Infra和算法协同设计可以突破算力的限制，并大幅减少大模型serving的成本，打开了算法和Infra Co-Design的新时代。

当前AI Infra发展的主力在开源社区，例如业界知名的Megatron/vLLM两个训练和推理加速框架都是开源软件，各个AI厂商都在持续贡献开源社区，一起在推动着这个行业的发展。生态的开放性更是使得Infra行业参与的开发者团体十分巨大。所有在社区的贡献都会被记录和认可，开源代码就是每个人的技术名片，这对于从业人员来说更是一项无形技术资产的积累。同时，大模型训练和推理优化涉及到计算机网络、内存，计算和负载均衡等方方面面的术问题，想要做好需要系统性地优化，并在多种约束条件下取得trade-off，通过设计一项好的加速特性，开发者可以体会到系统工程的魅力。

本文后续将会按照底层算子、AI框架、分布式训练框架、分布式推理框架、强化学习框架的顺序进行介绍，在此之前建议先阅读大模型的一些前置知识

底层算子

上层软件定义的计算逻辑最终实在AI Core硬件单元上执行的。在底层硬件上有非常多的硬件执行单元，为了保证计算的正确性和高效性，需要对这些单元进行精细地指令控制，是一个非常复杂的过程。算子层的作用就是将一系列计算动作，例如矩阵乘，rms norm等封装成可供上层软件直接调用的操作函数（OP），屏蔽上层用户对硬件底层实现的感知，把复杂性留给算子本身，让上层用户可以通过调用一系列封装好的OP执行自己需要的计算流程。

目前我大部分的工作仅停留在分布式训练/推理加速库，并未深入到算子本身的优化，因此本章节只能点到为止，更多的信息大家可以学习昇腾的CANN或者NVIDIA的cuda，二者在底层算子实现上有很多的积累和实践。

AI框架

在有了底层的算子实现能够让AI计算在GPU/NPU上高效地执行起来以后，上层开发者仍然需要一套框架能够方便地将多个OP组合成算法业务逻辑。这个过程需要同时具备易用性和高效性，尽量能让开发者聚焦在算法业务逻辑的开发，尽量屏蔽复杂的底层算子实现。AI编程框架对开发者提供高阶API，把矩阵运算、自动求导、分布式通信等底层细节封装起来，研究者只需写几十行代码就能训练大规模神经网络，当前最主流的AI编程框架是Pytorch和Tensorflow，其中Pytorch更是占据了大约90%以上的使用场景。

Pytorch框架支持动态图和即时算子下发模式（Eager Mode），写法上与普通Python代码无异，每执行一行python代码就下发一个算子，计算图可以不断动态拓展。这给深度学习算法开发带来了极大地便利性，写法与普通 Python 无异，打印中间变量、断点调试、即时修改网络结构都更直观，科研迭代极快。由于极佳的易用性，Pytorch已经成为了业界最主流的AI框架，已经成为了事实上的标准。

同时，Pytorch可支持多种硬件后端，为GPU/TPU/XPU/NPU等多种芯片提供了注册接口，例如torch-npu项目实现了插件化的npu适配，能支持大部分pytorch原生特性，同时无缝接入了NPU底层算子，从而能在NPU上支持torch生态项目。

我个人认为，AI Infra/大模型算法工程最好的学习切入点是Pytorch框架，它是上层大部分训练加速库以及AI软件的框架基础，同时也起到了算子下发连接异构加速硬件的作用，可以说是AI软件届的"linux操作系统"，通过学习torch框架的基本用法与原理至关重要。

分布式训练加速库

随着大模型的参数量与日俱增，当前参数量已经来到了万亿级别，早期的单机单卡训练模式已经无法满足百亿/千亿/万亿大模型的训练。当训练拓展到多机多卡时，就需要卡间和机间的通信来确保计算结果正确，分布式并行算法就成为了必需品。同时，大模型训练消耗海量的计算资源，通常需要千卡集群运行数天甚至数月，每一秒钟都是真金白银，优化计算效率和训练吞吐对于模型成本和迭代效率至关重要。

因此，在AI Infra层需要沉淀出一个大模型专用加速库，给大模型训练提供多种分布式并行算法，例如数据并行（DP），张量并行（TP），流水并行（PP）等，让大模型能在千卡集群上运行起来。同时也要在计算，通信，显存等多维度沉淀大模型加速能力。

业界主流加速库：Megatron/MindSpeed/FSDP/DeepSpeed

目前业界主流的分布式训练加速库主要有Megatron，pytorch社区原生的FSDP和DeepSpeed，另外针对昇腾设备上的大模型训练加速，昇腾推出了MindSpeed支持插件化适配Megatron，支持原生Megatron能力的同时提供亲和昇腾NPU的训练加速能力。

Megatron

Megatron是NVIDIA推出的基于Pytorch的大模型加速库，起初Megatron的核心主打特性是张量并行（tensor parallelism），通过模型参数的张量切分实现了模型并行。相比于之前工作聚焦于数据并行，Megatron-LM首次将模型并行工程化落地，也成为了后续3D并行范式的起点。

Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism

Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM

Reducing Activation Recomputation in Large Transformer Models

后续Megatron又陆续推出了流水线并行（Pipeline Parallel）、序列并行（sequence parallelism）和传统的数据并行（DP）一起构成了3D并行范式，即TP-PP-DP。3D并行能够以90%以上的集群线性度将任意规模的大模型训练扩展到千卡甚至万卡集群，同时Megatron后续引入了一系列计算、通信和显存优化，极致优化大模型训练效率，目前是开源大模型训练加速库中性能最好的。

MindSpeed

MindSpeed是专为昇腾设备设计的大模型加速解决方案，通过插件化适配的方式，使原生Megatron-LM框架支持昇腾软硬件环境。用户只需进行简单的适配，即可在昇腾设备上使用Megatron提供的大部分加速特性和并行。

同时，MindSpeed针对昇腾硬件特点，额外提供了昇腾亲和的加速算法，用户可以使能MindSpeed提供的各类加速、优化技术，从而获得在昇腾上大模型训练的极致性能，性能提升可达30%以上。

另外，对于使用自研的类Megatron-LM分布式训练框架的用户，MindSpeed同样提供了灵活的集成方案。通过集成MindSpeed中的关键特性，用户可以实现加速效果。MindSpeed中的各类特性相互独立、相互兼容，通过封装成独立的模块或API，为用户提供了方便的集成和优化途径。此外，MindSpeed采用开源开放策略，用户可直接获取MindSpeed源码，并根据自身需求进行源码修改和特性集成。MindSpeed的很多特性在客户界面得到了应用和借鉴，同时部分优秀特性也被Megatron吸收例如选择性重计算框架，是昇腾上训练大模型兼具易用性和高性能的解决方案。

DeepSpeed

DeepSpeed是一个由微软开发的开源深度学习优化库，旨在提高大规模模型训练的效率和可扩展性。它通过多种技术手段来加速训练，包括模型并行化、梯度累积、动态精度缩放、本地模式混合精度等。DeepSpeed还提供了一些辅助工具，如分布式训练管理、内存优化和模型压缩等，以帮助开发者更好地管理和优化大规模深度学习训练任务。此外，deepspeed基于pytorch构建，只需要简单修改即可迁移。

FSDP

Fully-Sharded Data Parallel(FSDP) 是pytorch官方推出的大模型分布式训练方案，可基于torch原生API支持，无需依赖上层加速库，在易用性方面备受推崇，受到了广大算法人员的喜爱。FSDP通过在DP域内对模型也同时进行分片，缓解了纯DP并行模型参数冗余的问题，但同时也引入了高昂的通信。总体而言，FSDP目前在易用性方面极佳，但在极致性能以及集群线性度方面和Megatron对比尚有差距，不过后续随着社区的不断优化，FSDP的性能也会得到不断提升，逐渐逼近Megatron，因此需要对FSDP保持关注。

TP/PP/SP/DP/CP/EP

由于目前Megatron是当前性能最好的分布式训练加速框架，头部厂商的内部训练框架大多基于Megatron进行二次开发，且Megatron在性能优化方面积累了很多优秀的特性，值得学习与分析。

一切大模型训练性能优化的起点是分布式并行策略，即模型并行和数据并行。模型并行层面主要包括张量并行（TP）和流水线并行（PP），分别在层内和层间对模型进行切分。数据并行层面主要包括上下文并行（CP）和数据并行（DP），分别在序列维度和批次维度对数据进行切分。