大模型的核心功能与技术架构解析

2. 通讯墙

大模型通过模型并行、流水线并行切分到 AI 集群后，通讯便成了主要的性能瓶颈。随着机器规模的扩大，基于同步的 All Reduce 通讯聚合方式，会因为大量的 AI 芯片和服务器之间频繁进行同步，出现水桶效应，也就是最慢的一路通讯，将会决定整个 AI 集群的通讯的高度。如果采用目前比较流行的 Ring-AllReduce 的通信聚合方式，当通讯的环越大，通讯的延迟将会不断地被扩大。另外网络协议的多次握手的方式，诸如此类的开销会导致训练无法有效利用带宽。

3. 性能墙

性能墙主要是指计算资源利用率的问题。随着大模型的提出，对算力需求更加迫切。理论上在 4K 的集群上每块卡快 1 分钟，总体就快了 68 个小时。大模型会增加对算力的需求，但是随着大模型引入各项分布式并行技术的同时，会降低计算资源的利用率。

• 算子层 (Operator Level)：小算子过多，可以通过算子融合进行优化；子实现不够高效，类似于卷积 CONV 算子针对 2x2 和 3x3 Kernel 大小的使用 Winograd 算法代替；内存局部性差，对算子和内存的开销进行分析，可以对计算时算子输出有相同的 shape 进行复用。
• 图层 (Graph Level)：如何搜索和切分出计算效率更高的子图，分配到不同的机器上进行计算；数据在通讯和内存之间如何增加 overlay 重叠部分，提高单卡整体的计算率。
• 任务层 (Task Level)：在任务层级，性能的瓶颈从计算转移到了通信，如何降低通信量，缩减通信域规模，尽可能把通信时延隐藏在计算中，是大规模训练的核心关注点。

4. 调优墙

所以在数千节点的集群上，需要考虑到提升算法工程师分布式调试调优的效率，另外还要考虑降低工程师对大模型进行并行切分的难度。除了对人的考虑，还要对硬件集群的管理，需要保证计算的正确性、性能、可用性。要是有一台机器坏了，如何快速恢复训练中的参数。

四、CV 和 NLP 场景实践

1. 分布式训练痛点

除了上述的数据量级大，不同场景下分布式训练的痛点对 CV 和 NLP 场景尤为明显。

• 模型复杂：单机性能要求高，比如卷积的计算时间在 CPU 上和 GPU 上相差十倍到几十倍。业界主要使用高性能的 GPU 进行计算，并采用 All-reduce 的通信拓扑进行参数的同步更新。
• 模型大 (DenseNet 部分)：比如 NLP 领域，GPT-3 这样的模型高达 1750 亿参数，显存占用高达 2.8 TB，单机内存无法容纳。而 Bert-Large 虽然只有 3.4 亿参数规模，但由于各种内存占用，在 16G V100 上，训练也仅能使用 batch Size=8。

当面对 GPT-3 这种 DenseNet 部分大的模型，Allreduce 单卡内存无法容纳，我们需要采用模型并行 (model parallelism) 的方式将计算图划分到不同的设备上构建有向无环图 (DAG) 进行分布式训练，其中 Gpipe, Megatron, Oneflow 和 Whale 都提出模型并行的相关解决方案。相比于数据并行每个 worker 只有一部分数据，模型并行下每个 node 使用所有数据。

2. 并行策略详解

• Intra-layer parallelism (Tensor Parallelism)：主要是将一层 Layer 中的矩阵计算分别拆分到不同的机器上进行运算。比如简单的 Y_1=W_1 X_1 这一次矩阵乘法中，我们将模型参数 W_1 或者输入数据 X_1，按某个维度分别拆分到不同设备上计算，比如 1D Megatron。
• Inter-layer parallelism (Pipeline Parallelism)：会将模型的 layers 拆分到不同的机器上，则一次 forward 就需要跨过不同的机器串行地进行运算。流行并行通过将 batch size 切分为更小的 micro batch，减少数据依赖，从而将整个计算过程异步起来，最大化资源利用率。举例来说，在一个简单的三层 MLP 中 (Y_i = W_i X_i, i=1,2,3) 会存在三次矩阵乘法 W_i X_i，流水线并行会把 W_i X_i 分别分配到三台机器上进行运算。

3. 并行方法选择

• 数据并行：假如整个模型设两个节点，一个模型节点 0、另一个模型做的节点 1，整个模型都做了数据并行，数据各一半要拿去训练学习。但是要注意训练完了以后不是最终的结果，因为只输入了一半的数据。因此这中间要 All-Reduce。
• 模型并行：将整个模型切成一半，一半模型节点 0，一半模型节点 1，数据是整个拿去训练。训练完了以后出来的结果也不是最终结果，因为只训练了一半的模型，出来还有 All-Gather，也是做通信的。

究竟用哪种并行方法呢？实际上这跟计算机结构有关。如果每台计算机之间通信都非常快，那么用数据并行就可以；如果你的通信比较慢，就要考虑模型并行。因此，这些模型如何跟数据、机器实际情况匹配？这就涉及到软硬件协同设计。

五、模型保存与部署

1. 小规模模型方案

在模型比较小的时候，比如 100G 以下，模型还有可能单机存储。这个时候的方案是 tensorflow 分布式训练 + savedmodel，分布式训练可以用多个 ps(tensorflow 自带的)，资源管理可以用 yarn。用分布式是由于样本数大，同时多 ps 也能异步加速训练。saved_model 一般由 chief worker 保存，但存下来以后，会抹掉 ps 的痕迹，保存的模型跟单机训练的一模一样。

2. 大规模模型方案

当模型比较大的时候，要求的样本数也更大，训练完 dump 出来的模型会很大，一个单机放不下，尤其是 embedding table。这个时候怎么办？一个自然的思路就是，把训练时候的 ps 拷贝同步一份给 serving ps，线上由该 ps 做 serving。注意后面谈到的 serving ps，都是自己开发或者根据某个开源软件修改而成 (比如 ps-lite)。

如果是天级模型，可以用 tensorflow 原生的 worker 和 train ps，但依然用 saved model 方式把模型存放到 hdfs，然后从 hdfs 读入另外一个 serving ps。如果是实时训练，则 serving ps 还得跟训练 ps 进行实时的网络连接，在内存就解决掉 weight 同步的处理，这个时候就不能用 tensorflow 原生的 ps 了，因为原生的 ps 没有实现同步接口。ps 变了，worker 也得跟着变，worker 大多数都是基于 tensorflow 的 c++ 底层接口开发，底层用到 tf 的 session 接口。

六、核心解决思路

针对上述的问题，各个大厂的训练框架进行很多相关优化，目前总结下来，核心的两点，一个在于分布式通信拓扑的设计，还有一个在于 Embedding Lookup 的性能优化。

1. 并行策略选择

只要单卡放的下，走数据并行，ps 或 allreduce 都行，allreduce 通信成本小一些。若规模变大：

• 稀疏模型：稀疏参数特殊处理。使用 ps，加上一些稀疏 tensor 的优化，且将 embedding 存储和更新负担转嫁到 ps。
• 稠密参数：allreduce，想办法解决稀疏 tensor 的存储、通信成本。比如 HybridBackend 架构中参数放在 worker 上：稠密参数 replication 存放，每个 worker 都有副本，梯度更新时进行 allreduce；稀疏参数 partition 存放，每个 worker 只有部分分片，梯度更新时进行 alltoall。allreduce 和 alltoall 都会使用 nccl 进行同步通信，效率较高。hb 进行 alltoall 时，通信的是稀疏梯度，而不是完整的参数，通信量上和 ps 是一致的，但是通信效率会更高。

2. 模型并行优化

稠密模型，单卡无论如何也放不下了，就只能采取模型并行及附属的一些优化方案。

• 算子拆分：单个矩阵乘法可以分到两个 device 上计算 Y = WX = [W1,W2]X = [W1X,W2X]。我们在工程上要做的就是：将切分到两个 device 上，将复制到两个 device 上，然后两个 device 分别做矩阵乘法即可。有的时候，切分会带来额外的通信，比如矩阵乘法切到了 reduction 维度上，为了保持语义正确，就必须再紧跟一个 AllReduce 通信。这里复杂之处在于，你不能无脑地将所有算子全部做拆分，因为拆分可能会引入一些额外通信，降低总体吞吐。所以你得做些分析，决定哪些算子被拆分，现在大部分框架都不支持这种全自动化策略，要么是半自动或纯手工，要么是针对某种模型把它的拆分方案写死。所以只能造轮子解决这个事。
• 流水并行：不切算子，而是将不同的 Layer 切分到不同的 Device 上，就可以形成 Pipeline 方案，GPipe 就是这样一种方案，提出了将一个 batch 拆分成若干个 micro-batch，依次推入到 Pipeline 系统中，即可消除 Bubble time。和算子拆分类似，全自动化方案工作量不小，比如 Pipeline 怎么切，才能让通信量小，计算还能均匀，这需要一定的算法和工程量。

3. 进阶优化技术

• 梯度累积 (Gradient Accumulation)：我们的模型可能会很大，或者数据量会很大。仅仅用一块 GPU 卡可能连模型都放不下，或者 batch size 只能设置的很小，但是我们知道有些情况下大的 batch size 往往会提供更好的效果。假设我们只有一个 GPU，我们的模型一次只能输入 batch size 为 8 的数据，那么我们怎么样实现 batch size 为 32 的更新呢？那就需要时间换空间了，即我们训练 32/8=4 步才去更新模型，也就是所谓的梯度累积。
• Gradient Checkpointing：那么如果你的 GPU 连 batch size 为 1 都跑不了怎么办？我们在训练深度学习模型的时候，需要先做前向传播，然后将中间得到的激活值存储在内存中，然后反向传播的时候再根据 loss 和激活值计算梯度。也就是说内存消耗其实跟模型的层数线性相关。那么怎么减少这个内存消耗呢？最简单的想法就是我不存这些中间信息，计算梯度的时候，到了某一层我重新计算它的激活值，这个方法虽然可以让内存消耗是个常量，但是运行时间会是 O(n^2)，这是没法接受的。那么就有一个折中的办法，我不存全部的中间数据，只存部分，那么我们在计算梯度的时候不需要从头计算了，只需要从最近的 checkpoint 点计算就好。
• 混合精度训练：我们训练模型一般都是用单精度 (FP32) 的参数，但是其实我们还使用半精度 (FP16)。半精度可以降低内存消耗，从而训练更大的模型或者使用更大的 batch size；同时运算时间受内存和算术带宽的限制，在有些 gpu(Tensor cores) 上可以为半精度提供更大的算术带宽，从而提高训练效率，减少 inference 用时。BF16 格式也在新一代硬件中得到支持，提供了更好的数值稳定性。

七、总结与展望

大模型技术的落地不仅依赖于算法的创新，更依赖于系统工程能力的提升。从分布式训练框架的选择，到硬件资源的调度，再到模型压缩与加速，每一个环节都需要精细化的设计与优化。未来，随着硬件算力的持续提升和软件生态的完善，大模型将在更多垂直领域发挥关键作用，推动人工智能从通用智能向行业专用智能迈进。开发者应重点关注分布式训练原理、显存优化策略以及模型部署流程，以应对日益复杂的工程挑战。