一文解析 DeepSeek 大模型高效训练背后的极限 AI 工程优化

来源：时序人

国产之光 Deepseek 大模型成功出圈，在效果比肩非开源模型ChatGPT的同时，其运行成本很低。那么 Deepseek 是如何实现这么高效的训练呢？文本将简要介绍 Deepseek 母公司幻方开源的 HAI-platform 大模型训练工具，来一窥极限 AI 工程的秘密。

为了更好地发挥 GPU 集群的计算能力，训练具有惊人功能的、强大的万亿参数模型，一个高效简洁的大模型训练工具十分必要。幻方-深度求索研发了一个能对 GPU 资源进行统一高效利用的 AI 平台，同时实现了四种并行训练方式：ZeRO 支持的数据并行、流水线并行、张量切片模型并行和序列并行。这种方式可以极大优化集群的使用，提高模型训练的显存效率和计算效率。

【项目地址】

https://github.com/HFAiLab/hai-platform

显存与计算效率

首先，我们先了解训练巨大模型的显存和计算效率的挑战。

01

显存效率

训练具有万亿参数的模型对显存的需求极为庞大，远超过单个 GPU 的承载能力。举例来说，当采用 Adam 优化器进行混合精度训练时，仅保存模型的状态信息（包括参数、梯度及优化器状态）就需要大约12TB的显存空间。以英伟达A100 GPU为例，其显存容量为40GB，这意味着仅为了存储模型状态，就需要多达400张 A100 GPU 的显存资源。

在模型的前向传播过程中，中间激活值会被保留，直至反向传播完成并计算出损失函数后才会被释放。这部分额外的显存消耗与批量大小成正比。即便将批量大小设为1，训练一个万亿参数的模型也会产生超过 400GiB 的激活显存占用。尽管可以通过 Checkpoint 技术（以计算时间换取显存空间）来提前处理这些激活显存，但整体上，训练过程中的显存需求依然极为巨大。

因此，为了确保这类大规模模型能够在显存不溢出的情况下进行训练，必须将模型状态量和激活显存有效地分散到多个GPU设备上。这种分布式策略是实现大规模模型训练的关键。

02

计算效率

据估算，端到端训练一个拥有万亿参数的模型大约需要消耗 5000 Zflops 的计算量（即5乘以10的24次方；这一估算基于OpenAI的研究《Scaling Laws for Neural Language Models》）。这一庞大的计算需求意味着，若使用4000张英伟达 A100 GPU 并以50%的计算效率运行，大约需要100天的时间才能完成训练。

虽然大型超级计算 GPU 集群可能配备超过4000个 GPU，但由于批量大小（batch size）的限制，要在如此规模上实现高计算效率仍然面临巨大挑战。计算效率与计算时间对通信时间的比率成正比，而这一比率又与批量大小直接相关。然而，训练模型时批量大小的设置存在一个上限：超过该上限，模型的收敛性能将显著下降。

以目前最大的模型之一 GPT-3 为例，其训练批量大小约为 1500。如果使用约 4000 张 GPU 进行训练，即使将批量大小设置为 4000，每张 GPU 上的批量大小也仅为1，这将严重限制模型训练的扩展性。因此，如何在保证模型收敛性能的同时优化批量大小和计算效率，是实现大规模模型训练的关键难题。

AIGC 集群：

整合 GPU 资源

构建一个能管理大规模 GPU 集群资源的 AI 平台，主要会遇到如下几个痛点问题：

• 资源调度：算力规模不断扩大，而训练任务的计算需求又多种多样，如何处理任务和算力的关系以最大化集群资源使用？

• 使用效率：集群使用有峰谷差异，同时又要应对突发任务的需求，如何兼顾时效性和集群整体效率？

• 迭代适配：集群会有升级迭代，在调整算力规模、类型、使用规则等场景下，如何让平台快速适配，尽可能降低切换成本？

幻方提出以任务级分时调度共享 AI 算力的理念，将集群零散资源进行整合再分配，成功支持在上千计算节点上稳定运行深度学习训练和其他多类型任务，日常算力占用率 95% 以上，日常 GPU 使用率 75% 以上。

01

任务调度

任务是集群使用的基本单位，而非用户。所有用户提交的任务都由集群统一管理和调度。

用户独占 vs 分时调度

幻方提出分时调度理念对集群资源进行管理。用户提交任务，如运行 python / bash 代码，启动开发容器等，由平台根据当前资源需求、集群忙闲程度等进行任务的中断和加载。任务代码需要遵循平台编码规则以确保可以断点续跑，具体流程如下：

• 接受集群的打断信号；

• 保存 checkpoint（模型参数，优化器参数等）；

• 通知集群打断；

• 从 checkpoint 恢复，继续运行。

这里，集群不会将 GPU 资源池化，而是以计算节点为基本单位，根据资源类型、网络区域等条件进行分类标记。AIGC 集群鼓励用户一次性用满多张 GPU，进行并行训练。用户提交任务时需选定节点数量 n，则该任务可获得 n 的整数倍个 GPU，比如在8卡计算节点上提交使用4个节点的任务，则该任务会获得32张 GPU 进行并行训练。

02

用户管理

集群提供以配额的方式记录集群的各类资源，通过优先级管理用户的使用权限。例如：某用户在 NORMAL 优先级上拥有10个计算节点的配额，其在集群上提交若干个任务，则最多同时以 NORMAL 优先级调度10个计算节点运行。算力资源优先满足高优先级用户的训练需求，同优先级内则交替使用集群算力资源。

同时，每个用户都可以归属于某一个用户组。计算资源和优先级配额是按用户为单位进行分配的，而同组的用户可以共享自定义环境、私有数据存储。

03

数据管理

AIGC 集群会将训练数据存储在文件系统中。训练中的任务会实时从文件系统中读取数据，流转到计算节点。这里，优良的文件系统将是提高 GPU 利用率，进而提升集群整体效率的决定性因素之一。

为了最大化 GPU 利用率，降低数据读取的成本，幻方自研了高速文件系统 3FS，其专门针对深度学习训练中样本读取的场景进行优化，能支持计算和存储节点间 7TB/s 以上的数据吞吐。3FS 需要用户将原始数据中较多的小文件聚合成大文件进行存储，从而实现优秀的性能。

并行训练

01

3D并行

设计高效的并行方式对训练大模型至关重要。deepseek 结合集群的特性，可以实现三种并行训练方式的灵活组合：数据并行、流水并行、张量并行，并使用序列并行方式对 Transformer 进一步并行优化，极大地提升显存利用和计算效果。

彩色块显示有 32 张 GPU （八个计算节点，每个节点有四个 GPU）的系统上进行 3D 并行训练的 GPU 映射。同一颜色的 GPU 在同一节点上

02

数据并行

为了减少显存的占用，Deepseek 采用 ZeRO 数据并行的方式，把训练优化器的状态均分到 1/N 个 GPU 上 （N 是数据并行的数量），如下图所示：

在做前向和反向传播时，ZeRO 数据并行会先做 allgather 获得完整的参数，然后在前向和反向传播结束后释放掉，只保留 1/N 的参数和梯度。

03

流水并行

流水并行把模型切分成 M 个阶段，分别放在 M 个 GPU 上，能够大幅节省显存占用和提升扩展性能。训练的时候每个阶段接收上一阶段的输出作为输入，然后传给下一个阶段。如下图所示：

可以看到，这种并行方式会造成大量的 GPU 空闲（Idle）。为了减少浪费，Gpipe 和 PipeDream 的概念被提出：如果同时进行多个迭代，每个节点在同一时刻负责不同的迭代的计算，就可以避免数据依赖，不用在原地干等了：

在 DeepSeek 中，模型的切分是在 Builder 里由用户切分的，用户可以继承 并实现自己的 build_model 方法。 Builder 创建出来的模型包含一个 forward_backward 方法，用户可以传进输入、损失函数、标签，该方法会执行前向传播和反向传播并返回 Loss。

04

张量并行

张量并行把全连接层的参数和计算切分到 K 个 GPU 上，能够节省显存占用和提升扩展性能。假设 K=2，张量并行把矩阵乘法 X @ A @ B 变成 X @ A1 @ B1 + X @A2 @ B2 ，其中 A 沿着纵轴切成 A1 和 A2，B 沿着横轴切成 B1 和 B2；GPU 0 会存放参数 A1 和 B1 并负责 X @ A1 @ B1 的计算，GPU 1 则存放参数 A2 和 B2 并负责 X @ A2 @ B2 的计算。

05

序列并行

Megatron-LM 提出的序列并行方案旨在通过一些方式分摊张量并行中无法分摊的显存。

如上图所示，在张量并行的基础上，将 Transformer 核的 LayerNorm 以及 Dropout 层的输入按序列长度维度进行了切分，使得每个 GPU 上面只需要做一部分的 Dropout 和 LayerNorm 。这样做进一步降低了计算资源和显存的开销。