LLM Infra 你需要知道的一切

GPT 3

GPT 模型是 LLM 模型的起点，你如果说 Transformer 是，那也没错。既然我们要讨论 LLM Infra，那么我们需要知道我们的系统架构以及设计是给谁服务的。如果不了解模型架构，你对于后面介绍的 Infra 设计，比如推理阶段的 Prefill-Decode Disaggregation，Context Cache Migration，Traffic-Adaptive Request Routing 等，就会感到困惑。知其然，知其所以然。我们从 GPT3 开始，因为 GPT1 以及 GPT2 的参数量都太小了，GPT2 参数只有 1.5B, 而 GPT3 参数有 175B。如果对 Transformer 架构不熟悉，推荐这篇文章 How Transformers Work;中文用户推荐知乎文章 Transformer模型详解（图解最完整版）。

首先确定 GPT3 的输入输出：

• 输入：一段文本，长度为 2048 个 token
• 输出：输入 2048个 token 下一个 token 的概率分布

这里的 token 是使用 Byte Pair Encoding (BPE) tokenizer 进行标记的。比如我使用 OpenAI Tokenizer 对 I am xieydd, I love GPT3. 进行标记后的结果如下, token ids 为 [40, 716, 2124, 494, 88, 1860, 11, 314, 1842, 402, 11571, 18, 13]：

Encoding

Tokenizer 后的 token ids 还只是一串数字，并非向量。我们需要将其向量化，每一个 token 转换成 50257 维向量，向量 token id 位置为 1， 其他位置为 0。这样我们就得到了输入向量，长度为 2048 * 50257 = 102760896。

Embedding

这个向量太稀疏了，我们需要将其降维，将其降维到 2048 * 12288。

Positional Encoding

由于 Transformer 没有像 RNN 那样的递归机制，因此使用位置编码来提供有关序列中每个 token 的位置的信息。这使他们能够理解句子中每个单词的位置。Transformer 使用 sin 以及 cos 结合的方式，让 position encoding 可以编码任意长的句子。

Attention

Vector 以及 Positional Encoding 相加后，得到了输入向量，接下来我们就要进行 Attention 了。这里将 2048 * 12288 的输入简化成 3*512。

Image from https://dugas.ch/artificial_curiosity/GPT_architecture.html

• 输入向量，通过权重 Wq, Wk, Wv 得到 Query, Key, Value。
• Query 与 Key 进行点积，然后除以根号 d
• 进行 softmax 得到 attention score
• attention score 与 Value 进行点乘

Sparse Attention

GPT3 并没有使用 Attention, 而是使用 Sparse Attention。sparse attention 除了相对距离不超过 k 以及相对距离为 k，2k，3k，… 的 token，其他所有 token 的注意力都设为 0，如下图所示：

好处有两点：

1. 减少注意力层的计算复杂度，节约显存和耗时，从而能够处理更长的输入序列；
2. 具有“局部紧密相关和远程稀疏相关”的特性，对于距离较近的上下文关注更多，对于距离较远的上下文关注较少；

Multi-Head Attention

GPT3 将上述过程重复执行 96 次，有96个头,最后将输出在列行进行 concat。

Feed Forward

Feed Forward 层是一个全连接层，输入是 Multi-Head Attention 的输出，输出是 2048 * 12288。

Add & Norm

将进入 Multi-Head Attention 的前的输入和 Feed Forward 的输出相加，然后进行 Layer Normalization。

Decoding

我们需要将 Add & Norm 的输出 2048 * 12288 转换回 2048 * 50257，这样才能知道 2048 个 token 下一个 token 的概率分布。我们复用了 Embedding 层的权重，将其转置后与 Add & Norm 的输出相乘，得到 2048 * 50257 的输出。随后经过 Softmax 得到概率分布, GPT3 取 topk 个 token 作为输出。

了解了 GPT3 的架构，我们就可以更好的理解 LLM Infra 的设计。下面我们进入如何训练大模型的话题。这里我们使用目前 SOTA 的开源模型 DeepSeek-V3 来进行讲解。

Infra

GPU

集大成者 GB200 NVL72:

1. Blackwell Architecture
2. NVIDIA Grace CPU
3. Fifth-Generation NVIDIA NVLink
4. NVIDIA Networking

• NVIDIA Quantum-X800 InfiniBand
• NVIDIA Spectrum™-X800 以太网
• NVIDIA BlueField®-3 DPU

CUDA

Network

NVLink

机器内通信可以通过：

1. PCIe
2. 内存
3. NVLink

简单来说就是目前的 PCIe 5.0 的带宽无法满足 AI 分布式训练中参数同步的时间需求，需要带宽更大的机器内 device 通信的方式。NVLink 是给单机多 GPU 之间提供高速互联的技术，替换 PCIe。

随着 NVLink 的代次的发展，速度也在不断地提升，到目前第五代，单个 NVIDIA Blackwell Tensor Core GPU 支持多达 18 个 NVLink 100 GB/s 连接，总带宽可达 1.8 TB/s，比上一代产品提高了两倍，是 PCIe 5.0 带宽的 14 倍之多。

注意只有 SXM 口才能真正的使用 NVLink P2P 通信，而不是 NVLink Bridge 连接。

NVLink 最新的技术 NVLink C2C interconnect 技术：将两个高性能 NVIDIA Blackwell Tensor Core GPU 和一个 NVIDIA Grace CPU 连接到两个 Blackwell GPU，最终的目的是使 NVIDIA GPU、DPU 和 CPU 能够与定制芯片实现一致互连, 速度达到 900 GB/s，下图是 NVIDIA GH200 Grace Hopper Superchip 的架构图：

NVSwitch

当 GPU 通过 NVLink 进行 P2P 连接时，最高总带宽可以达到 1.8TB/s, 但是每个 GPU 必须将相同的 1.8TB/s 连接分成七个专用的 256 GB/s 点对点连接，也就是 GPU 和 GPU之间的通信最大只能达到 1.8TB/s。如果想突破这个限制就需要 NVLink Switch, NVSwitch 其实是个 ASIC 芯片。 NVLink Switch 可以保证 GPU 之间的通信带宽达到 1.8TB/s 的完整带宽，大幅压缩通信时间，在推理过程中实时响应 tok/s/user 相对于没有 NVSwitch 的情况提升 1.4倍，而且batch 越大，提升越明显。

InfiniBand

参考 知乎-AI 集群基础设施 InfiniBand 详解。

机器间通信可以通过：

1. TCP/IP 协议
2. RDMA(Remote Direct Memory Access) 协议

• InfiniBand
• RoCE (RDMA over Converged Ethernet)
• iWARP (Internet Wide Area RDMA Protocol)

RDMA 协议可以不像 TCP/IP 协议需要 CPU 参与，数据需要经过数据链路层，网络层，传输层，应用层，CPU 参与数据包的接受，解包等操作。RDMA的内核旁路机制，允许应用与网卡之间的直接数据读写，同时RDMA的内存零拷贝机制，允许接收端直接从发送端的内存读取数据，极大的减少了CPU的负担，提升CPU的效率。

RDMA网络，分为三类分别是Infiniband、RoCE、iWARP。其中，Infiniband 是一种专为 RDMA 设计的网络，从硬件级别保证可靠传输，而 RoCE 和 iWARP 都是基于以太网的 RDMA 技术，支持相应的 verbs 接口。

IB（InfiniBand）： 搭建基于 IB 技术的 RDMA 网络需要专用的 IB 网卡和 IB 交换机。

iWARP（Internet Wide Area RDMA Protocal）： 基于 TCP/IP 协议的 RDMA 技术，由 IETF 标 准定义。iWARP 支持在标准以太网基础设施上使用 RDMA 技术，但服务器需要使用支持 iWARP 的网卡。

RoCE（RDMA over Converged Ethernet）： 基于以太网的 RDMA 技术，也是由 IBTA 提出。RoCE 支持在标准以太网基础设施上使用RDMA技术，但是需要交换机支持无损以太网传输，需要服务器使用 RoCE 网卡。

IB 链路速率随着网络带宽的升级越来越快，如下图：

这里的 x4, x12 是线缆中 Link 数量，x4 代表 4 个 Link，x12 代表 12 个 Link，一个 Link 一来一回两条线。

InfiniBand 的线缆有别于以太网线缆以及光纤线缆，InfiniBand网络互连产品包括：DAC高速铜缆、AOC有源线缆以及光模块。

Mellanox 在全球 InfiniBand 市场的占有率基本上无敌的存在，在英伟达收购 Mellanox 之后，也于2021年推出了自己的第七代 NVIDIA InfiniBand 架构：NVIDIA Quantum-2。

NVIDIA Quantum-2 平台包括：NVIDIA Quantum-2 系列交换机、NVIDIA ConnectX-7 InfiniBand 适配器、BlueField-3 InfiniBand DPU以及电缆。 目前最新的是：

• NVIDIA Quantum-X800 InfiniBand
  • NVIDIA® ConnectX® InfiniBand HCAs(Host Channel Adapters),ConnectX-8 InfiniBand SuperNIC 可提供每秒 800 Gb/s 的数据吞吐量
  • InfiniBand Switches 最高 144 个口，每个口 800 Gb/s 一共 115Tb/s
• NVIDIA Spectrum™-X800 以太网
  • 通过 SuperNICs 以及 NVIDIA Spectrum-X800 SN5600 Ethernet switch 提供 800 Gb/s 的吞吐量
  • Switch 最高 64 个口，每个口 800 Gb/s 一共 51.2Tb/s
• NVIDIA BlueField®-3 DPU
  • 支持以太网以及 InfiniBand 连接， 400Gb/s 速度
  • NVIDIA DOCA 软件框架进行 SDN
  • 网络、安全、存储、管理、AI 加速等功能
• BlueField SuperNICs: Provides best-in-class remote direct-memory access over converged Ethernet (RoCE) network connectivity between GPU servers at up to 400Gb/s。

GPUDirect 是 NVIDIA 开发的一项技术，可实现 GPU 与其他设备（例如网络接口卡 (NIC) 和存储设备）之间的直接通信和数据传输，而不涉及 CPU。 GPU Direct 技术包括 GPUDirect Storage、GPUDirect RDMA、GPUDirect P2P 和 GPUDirect Video。

• GPUDirect P2P 主要是单机多卡场景下，让 GPU可以通过PCI Express直接访问目标GPU的显存
• GPUDirect RDMA 主要是在多机多卡场景下，计算机1的GPU可以直接访问计算机2的GPU内存，而不需要 CPU -> CPU 再由计算机 2 的 CPU -> GPU 多次拷贝

Communication

MPI (Message Passing Interface) 是一种并行计算中常用的通信协议，MPI 有多种实现，比如 OpenMPI, Gloo, NCCL 等。

NCCL（Nvidia Collective multi-GPU Communication Library）实现了针对 NVIDIA GPU 性能优化的多 GPU 和多节点集合通信原语。NCCL 提供了诸如 All Gather，All Reduce，Broadcast，Reduce，Reduce-Scatter 等实现，这些实现优化后可以通过 PCIe、 NVLink、InfiniBand 等高速互联，从而实现高带宽和低延迟。

介绍几种通信原语：

• Reduce：从多个sender那里接收数据，最终combine到一个节点上。
  
• AllReduce：Reduce 的变种，将所有节点的数据都reduce到所有节点上。
  
• Broadcast：一个节点将数据发送到所有节点。

不同的网络拓扑对 NCCL 通信的速度有很大影响，如下图所示：

• 第一种是两个GPU通过CPU然后通过QPI和另一个CPU上的两块卡相连，因此速度最慢，但也能达到>5GB/s。
• 第二种是两个GPU通过PCIe switch相连后再经过CPU连接，速度会稍微低一点。
• 第三种是四张卡都在一个PCIe switch上，所以带宽较高，能达到>10GB/s PCIe的带宽大小。
• 第四种是DGX-1架构，通过 NVLink 直接连接，速度最快，能达到 60GB/s。

• 左图2机8卡，机内PCIe，机间InfiniBand能达到>10GB/s的速度，InfiniBand基本上能达到机内的通信速度；
• 右图4机32卡，机内NVLink，机间InfiniBand，带宽能达到>40GB/s。

NCCL P2P 的几种 Level:

• LOC：从不使用P2P（始终禁用）
• NVL：当 GPU 通过 NVLink 连接时使用 P2P
• PIX：当 GPU 位于同一 PCI 交换机上时使用 P2P。
• PXB：当 GPU 通过 PCI 交换机（可能是多跳）连接时使用 P2P。
• PHB：当 GPU 位于同一 NUMA 节点上时使用 P2P。 流量将通过 CPU。
• SYS：在 NUMA 节点之间使用 P2P，可能跨越 SMP 互连（例如：QPI/UPI）。

Storage

Operation and Maintenance

Failure

Hardware Exception：

1. GPU ECC Error：通常需要重启机器或 Reset GPU。128 个节点每 1-2 天就有一个。
2. IB（Infiniband）/NCCL问题：通常是硬件问题，比如网卡坏了，或者网络抖动等。

• 由于网络问题导致训练降速，比如降速 20% 左右，需要通过二分法来判断异常的节点进行相应替换。
• 也可能导致任务直接异常，比如 ：p2p_plugin.c:141 NCCL WARN NET/IB : Got async event : port error。

3. 任务 Hang 住：通常和 IB/NCCL 问题相关，也需要人工检测。OPT-175B 训练中多次出现任务 Hang 住的情况。
4. GPU 掉卡：此时一般会触发 CUDA Error 或程序异常退出，比如：“RuntimeError: Caught RuntimeError in pin memory thread for device 4.”。
5. 机器异常：GPU 之外的硬件异常，比如硬盘、CPU 等异常，甚至机器直接挂掉。
6. 机器配置异常：比如发现某个机器开启了 MIG。
7. 集群维护：通常一个集群里不仅仅是支持一个任务，偶尔会需要维护升级，可能需要停止任务。
8. 存储挂掉：比如 Lustre 存储挂掉

Monitoring & Fault Tolerance

需要自动故障检测和快速恢复机制：

1. 当用户提交训练任务后，除了在每个 GPU 上创建一个训练 Executor 外，还会创建一个训练守护进程，定期向 Driver 发送心跳。心跳信号包含各种信息，以实现实时的异常检测和发送告警信息。Driver 在特定时间内未收到心跳信号时，也会触发故障恢复流程，包括：

• 暂停所有训练的 Executor，并执行一系列的自检诊断。
• 一旦识别到异常机器，将其驱逐，并将通过测试的同等数量的健康机器加入集群。除此之外，还提供了用户接口，允许用户手动识别并驱逐异常机器。
• 机器恢复完成之后，将从最近的 Checkpoint 恢复训练。对 Checkpoint 的保存和恢复过程也需要进行了优化，以最大限度减少对训练进度的影响。引入了两阶段方案，第一阶段每个 GPU 训练进程将显存中的状态信息写入主机内存中，由于高速 PCIe 带宽，这个过程通常只要几秒钟，这样训练进程可以马上继续训练。第二阶段中，一个后台进程异步的将状态信息同步到分布式文件系统，MegaScale 使用的是 HDFS。
• 其心跳信号包括 IP 地址，Pod 名字，硬件信息等，还包含当前的训练进度信息。除此之外，训练进程的 stdout/stderr 日志也会被收集，以便进行实时的聚合、过滤和分析。如果识别到 warning 或 error 等关键字，Driver 将实时的上报这些诊断信息。最后，RAMA 的流量指标也会包含在内，以便更好的识别网络利用率和通信效率

日常巡检（需要精确到毫秒级的监控系统，避免 Spike）：

• 启动大型矩阵乘法任务，捕获这些故障
• 还会启动一些简单的通信测试，以保证网络的正常
• 机内测试：回环测试（Loopback test）：测量机器内所有的 RAMA 网卡到各种机内端点（包括内存、GPU）的回环带宽。会进行 full-mesh 测试，覆盖所有可能的链路组合，以便推断 PCIe 配置中潜在的链路问题。
• RNIC-to-RNIC 测试：主要测试机内不同 RNIC 之间的连通性和带宽性能。以便发现 RNIC 是否符合速度规格或存在路由问题。
• NCCL 测试：在机内进行 All-to-All 测试，也会在同一个 TOR 交换机的机器之间进行 All Reduce 测试，以发现潜在的硬件故障和性能问题

Training

大模型训练，这里我们主要参考 Andrej Karpathy 在 2023 年 MSBuild 的分享，分为以下四个阶段：

1. Pre-training
2. Supervised Fine-Tuning
3. Reward Modeling
4. Reforcement Learning

在介绍这四个阶段之前，首先普及下 Hardware 的只是，因为大模型训练需要大量的计算资源，比如 GPU，此外数据的存储和网络通信都是需要考虑的。

Hardware

GPU

一提到 GPU （Graphics Processing Unit）我们自然而然将其和 NVIDIA 挂钩，这个是有一定的道理的，因为 NVIDIA 的 GPU 在深度学习领域占有很大的市场份额。不过这个市场巨大，有 AMD 以及 Brodcom 的 Application-specific Integrated Circuits (ASICs) 芯片也在不断地吞噬市场份额。

虽然 GPU 带了个 Graph，但是在 AI 中我们更多的是使用 GPGPU (General-purpose computing on graphics processing units)。如果你不了解 GPU 的术语，推荐你看看 Modal 维护的 GPU 术语表。GPU 相对于 CPU

Network

Storage System

Pre-training

Pretaining model 是大模型训练最消耗算力和数据的阶段，这个阶段的目标是让模型学习到语言的基本规律，比如语法，语义等。这个阶段的数据量很大，通常是 TB 级别的数据，通常是几周到数月。这个阶段的训练通常是无监督的，也就是说我们不需要标注数据。这个阶段的训练通常是在大规模的 GPU 集群上进行的，比如 OpenAI 的 GPT3 数千张 V100 上训练一个月天。

作为最消耗计算资源的预训练，我们在这个章节来分享 LLM 训练的一些技术，比如说分布式训练并行技术，训练优化技术。

Distributed Training

参考 OpenAI 的文章 Techniques for training large neural networks , 主要分为以下四个部分：

1. Data Parallelism
2. Pipeline Parallelism
3. Tensor Parallelism
4. Expert Parallelism

1. 这里不同颜色的块代表模型不同的层
2. 虚线代表分割到不同的 GPU
3. 箭头表示模型的前向传播以及反向传播

Data Parallelism

Data Parallelism 将数据分成不同的 subset, 然后分发到不同的 GPU 上。对于每一个 GPU, 模型是完整的，所以每一个 GPU 都需要存储模型完整的参数。所以当 GPU 显存无法存储模型参数时，Data Parallelism 就不适用了，不过现在有些 Offload Memory 的技术或者 KV Cache 可以缓解这个问题。

正是因为每个 GPU 需要完整的参数，所以当更新的时候，需要保证工作线程（一般和 GPU 保持一致）的参数保持一致。这样在工作线程上需要引入阻塞通信：

1. 每个 worker 计算梯度
2. 平均每个 worker 的梯度
3. 每个 worker 更新参数

当参数量越大，线程数越多，这个阻塞通信的开销就越大，当然可以使用异步的方式进行通信，但是这样可能会损害学习率，得不偿失，尤其对于 LLM 模型，投入资源大，不容有失。GPU 之间的通信通过 NCCL (NVIDIA Collective Communication Library) 这个库来实现。上述通信过程通常的叫法是 AllReduce 如下图所示：

我们注意到如果 AllReduce 在一个单独的进程进行处理，那么这个进程需要 Gather All 以及 Scatter All, 这样单进程的通信开销和节点数以及参数量成正比。为了消除节点数以及参数量带来的阻塞等待浪费，我们可以使用 Ring AllReduce，如下图所示：

首先将 N 个 worker 构建一个 ring, 每个进程将数据划分成 N 个 chunk。

第 N 个 worker 将第 N 个 chunk 发送给第 worker[N+1]，同时从 worker[N-1] 接受 chunk。

第 N 个 worker 将接受到的 N-1 chunk 和自己的 N-1 chunk 进行 reduce，然后发送到 worker[N+1]。

重复接受，reduce，发送的过程，重复 N-1 次，这样每个 worker 获得结果的一部分。最后再进行一轮的 Send 操作而不做 reduce 操作，这样每个 worker 就获得了最终的结果。

Pipeline Parallelism (Model Parallelism)

Pipeline Parallelism 将模型分成不同的 stage, 然后分发到不同的 GPU 上。对于大模型来说，比如 LLaMA 65B, 如果参数是 FP32 那么总显存需要 260GB 但是一般训练都采用半精度也就是 FP16，那么也需要 130GB 显存，目前最前进的 H200 显存也就是 80GB。目前来自 DeepSeek 团队训练的 DeepSeek V3 已经首次在超大规模模型上验证了FP8训练的可行性和有效性。但是 DeepSeek V3 是 MOE (Mixtures of Experts) 模型，参数达到 671B, 这个我们后面再聊。PP 的通信量相对比较小，因此常常会放在不同的机器上。

1. 模型每一层的输入和输出都有顺序，F 代表 Forward, B 代表 Backward, U 代表 Update
2. 每个 worker (GPU) 在同一时间只负责处理一个 stage
3. 时间序列上，资源利用出现大量的气泡

为了优化这个问题，我们可以将刚才 ring allreduce 的思想拿过来，也就是将一个 batch 数据进行切分，将计算时间和等待时间 overlap，如下图所示：

上图为 GPipe 的示意图，GPipe 将模型 layer 的 chunk 数据的激活值连续传递给下一个 worker，也连续进行向后传递。然后同步聚合梯度，更新参数。

而 PipeDream 则让 worker 交替处理向前和向后传递。相对 GPiPe，PipeDream 有更好的资源利用率，但是需要更多的通信开销。

Tensor Parallelism

Tensor Parallelism 将模型的 operator 分成不同的 subset, 然后分发到不同的 GPU 上, 比如说矩阵乘法。Pipeline parallelism 是将模型的层分到不同的 GPU, 而 Tensor Parallelism 是将模型层内的 operator 分到不同的 GPU。对于现代模型比如 Transformer, 将激活值和大的权重进行点积计算是计算的瓶颈。比如 MegatronLM 在 Transformer 的Self-Attention 和 MLP 层进行了并行化矩阵乘法。PTD-P 使用tensor，pipeline，以及data parallelism， pipeline scheduler 为每个设备分配多个非连续层，以网络通信为代价减少气泡开销。其通信量更大，因此往往会将 TP 放在单个机器内部，可以充分利用机内的 NVLink 高速带宽。

下图是采用的是 8DP 12PP 4TP 方案，因此需要的 GPU 为 8x12x4=384：

Expert Parallelism

Mixture-of-Experts (MoE) 是一种将模型每一层的某些部分放在同一个 GPU 上执行，这意味着你可以通过门控制器控制模型中哪些部分会应用到本次输入和输出。每一组门控制的一组参数都是一个 Expert，不同的 Expert 在不同的 GPU 上。

很明显 MoE 架构可以显著的提升模型的大小，这来源于一个经验，在有限的计算资源预算下，用更少的训练步数训练一个更大的模型，往往比用更多的步数训练一个较小的模型效果更佳。特别是在预训练阶段，与稠密模型相比，混合专家模型通常能够更快地达到相同的质量水平。MoE 是基于 Transformer 架构的，由下面两个部分组成：

1. 稀疏 MoE 层：MoE 层有若干个专家，如下图为 8 个专家，这些专家可以是 FFN 层，也可以是 MoE 层，像俄罗斯套娃。
2. 门控网络或路由：这个路由决定 token 会被发送到哪个专家，路由器的参数会在训练中进行训练。

上面讲到 MoE 的优点，但是 MoE 也有一些缺点：

1. 微调过程中泛化能力不足，容易过拟合
2. 推理过程中需要加载所有的参数到内存，即使只有一个专家被激活，比如 Mixtral 8x7B ，VRAM 需要放的下 47B 的参数，这里不是 56B 是因为 MoE 只有 FFN 是独立的，其他层会共享参数。
3. Token 被路由到不同的专家，导致批量大小分布不均，导致计算资源浪费。

对于门控网络，除了最典型的带有 Softmax 的网络，还有其他网络，比如包括带噪声的 TopK 门控 (Noisy Top-K Gating)，引入可调节的随机噪声并保留前 k 值。

• 随机噪声是为了让专家间的负载更加均衡，在训练过程中才避免只有受欢迎的专家得到训练，在 transformers 库中，可以通过 aux_loss 参数来控制辅助损失。
• topk 可以选择性的保留前 k 个值，加速训练和推理

列举两个典型的 MoE 模型或者说方法：

1. Gshard

为了保证负载平衡和训练效率，Gshard 除了使用上述的辅助损失外，还进行了以下的优化：

• 随机路由: 在 Top-2 设置中，我们始终选择排名最高的专家，但第二个专家是根据其权重比例随机选择的。
• 定义一个专家能处理多少令牌。如果两个专家的容量都达到上限，令牌就会溢出，并通过残差连接传递到下一层，或在某些情况下被完全丢弃。

2. Switch Transformer

Switch Transformer 使用接受两个输入，每个输入4个专家，而并非单输入，至少两个专家的设定。这样做的优势是：

1. 减少门控网络 (路由) 计算负担
2. 每个专家的输入 batch 至少可以减半
3. 降低通信成本，保持模型质量

于此同时 Switch Transformer 也对专家容量进行了研究，通过大于 1 的容量因子为 token分配不均匀时提供缓冲。提高容量因子 (Capacity Factor, CF) 可以增强模型的性能，但这也意味着更高的通信成本和对保存激活值的显存的需求。 在训练期间，对于每个 Switch 层的辅助损失被添加到总模型损失中。这种损失鼓励均匀路由，并可以使用超参数进行加权。 Switch Transformer 还探索了混合精度，也就是专家使用原始精度，其他部分使用 Bfloat16 精度，在保证训练稳定性的同时，增加了训练速度。

ST-MoE 的研究者们发现，编码器中不同的专家倾向于专注于特定类型的令牌或浅层概念。例如，某些专家可能专门处理标点符号，而其他专家则专注于专有名词等。与此相反，解码器中的专家通常具有较低的专业化程度。

训练 MoE 模型有哪些 tricks:

1. 稀疏模型更易于出现过拟合现象，因此在处理这些模型时，尝试更强的内部正则化措施是有益的，比如使用更高比例的 dropout。例如，我们可以为稠密层设定一个较低的 dropout 率，而为稀疏层设置一个更高的 dropout 率，以此来优化模型性能；
2. MoE 模型在小任务比如 SuperGLUE 容易过拟合，不如在大任务比如 TriviaQA 效果好；
3. 冻结 MoE 参数进行微调，可以显著加速并且对模型效果影响很小；
4. MoE 适合更小的 batch 和更高的学习率；
5. 在稀疏模型上进行微调效果好用稠密模型。

对于稀疏模型和稠密模型如何选择呢？

1. 稀疏混合专家模型 (MoE) 适用于拥有多台机器且要求高吞吐量的场景。在固定的预训练计算资源下，稀疏模型往往能够实现更优的效果。相反，在显存较少且吞吐量要求不高的场景，稠密模型则是更合适的选择。
2. 直接比较稀疏模型和稠密模型的参数数量是不恰当的，因为这两类模型基于的概念和参数量的计算方法完全不同。

我们刚才在讲分布式训练的优化手段，讲到了 Expert Parallelism，从而引申到 MoE 模型，那么为什么 MoE 模型需要 Expert Parallelism 呢？

1. 在 Expert Parallelism 中，专家被放置在不同的设备上，每个设备处理不同批次的训练样本；
2. 当单个专家参数可以在一个设备上时候，采用数据并行和专家并行的方式训练
3. 如果单个专家模型不能再一个设备上，需要加入模型并行，将单个专家模型分布到不同的设备上，同时配合数据进行并行训练加速

部署专家模型优化方式：

1. 对专家模型进行蒸馏
2. 路由器被修改为将整个句子或任务直接路由到一个专家。这样做可以提取出一个用于服务的子网络，有助于简化模型的结构。
3. 合并各个专家的权重，在推理时减少了所需的参数数量

Zero-DP

Sequence Parallel

Context Parallelism

Prepare Data

LLaMA 做 Pre-training 时的训练数据如下所示：

收集完原始数据后，需要进行数据处理，比如 Tokenization。

Model

我们看下 2020 年 GPT3 以及 2023 年 LLaMA 模型的一些参数对比：

解释下上图中的参数：

1. Vocabulary Size: 50257。这是指的是 Tokenizer 中词汇表的数量，也就是输入向量的维度。
2. Context Length: 2048。这是指的是生成模型在生成下一个 token 之前会看到的窗口的大小。
3. 虽然 GPT3 训练的参数量是 175B，而 LLaMA 只有65B 但是 LLaMA 训练的 token 1.5T > GPT3 300B。
4. 一些训练的超参数，比如 batch size, learning rate, 头数量以及模型层数等。

训练的输入如果是以下文本, 这里会由特殊 token <|endoftext|> 进行标记

• Row 1: Here is an example document 1 showing some tokens.
• Row 2: Example document 2<|endoftext|> Example document 3<|endoftext|> Example document
• Row 3: This is some random text just for example<|endoftext|> This
• Row 4: 1,2,3,4,5

Tokenization 后的结果如下：

这里 Batch 取 4，Context Length 取 10。其中每个单元格只能看见同行的单元格，绿色的是当前 highlight 的 token 所能看到的 context, 红色是它的目标。

Training

Supervised Fine-Tuning

在这个阶段需要小的高质量的数据，一般是人类标注的数据，比如 prompt 以及相应理想的回复， 一般需要 10~100k。

这一阶段会将 Pre-training 的模型加载进来，然后在这个数据集上进行训练，得到的数据就是 SFT(Supervised Fine-Tuning) 模型。目前进行 SFT 的模型有很多，比如 Instruct, Coder, Math 以及 Reasoning 分别针对指令性的 QA、编程、数学以及推理任务场景。

Reward Modeling

RLHF(Reward Learning from Human Feedback) 阶段分为两个部分，一个是 Reward Model，一个是 Reinforcement Learning。

Reward Model 会将数据收集变成比较的形式，举个例子

这里，人类需要输入相同的指令，在不同的输出中进行 Rank，得到 pair 数据集，大概 100k~1M。

训练中:

1. 每一行的蓝色 prompt 是一样的
2. 黄色是 SFT 模型输出
3. 绿色是 reward token，也就是 SFT 评价模型输出的质量，和人类评价的质量进行比较
4. 损失函数衡量 reward token 与人类标记的 ground truth 的一致性

Reinforcement Learning

准备 10k~100k 的 Prompt 数据, Reward Model 对这些模型进行训练, 100 GPUs天级别。

1. 每一行的蓝色 prompt 是一样的
2. 黄色是 SFT 模型输出, 作为初始化值，作为 Reward Model 的训练数据
3. 绿色是 reward token，这个 token 会将 sampled token, 也就是黄色部分进行评价，如果高，则黄色部分 token 在后续的采样中会被采样的概率会增加。

PPO 算法就是 RLHF 模型, 为什么要使用 RLHF, 参见下图，RLHF 可以显著减少 predict 的熵，也就是预测的更加稳定。

参考知乎姜富春​博士总结的 聊聊Reasoning Model的精巧实现（ReFT, Kimi K1.5, DeepSeek R1） 介绍了 ReFT, Kimi K1.5, DeepSeek R1 这三个 Reasoning Model 的实现,其中 RL 是关键。

RLHF 可以参考 OpenRLHF.

Serving

Model Optimization

模型这块本来想放在 Model 里面，但是这里的模型优化主要是让推理更快，所以放在 Serving 里面。

MHA vs MQA vs GQA vs MLA

这里首先解释下这几个名词：

• MHA: Multi-Head Attention
• MQA: Multi-Query Attention
• GQA: Global Query Attention
• MLA: Multi-Latent Attention

参看苏大的博客 缓存与效果的极限拉扯：从MHA、MQA、GQA到MLA。

MHA

在 Transformer 解码器中，由于 token 的注意力依赖于前面的 token，因此不会重新计算前面的上下文，而是缓存其 Key 和 Value, 为什么没有缓存 Query 呢？

如下面这个 不带 Causal Mask 的 Attention 公式，从上到下，在 T 阶段，只有 Q 参与了每一个地方的计算，但是 K 和 V 是已经算好的，不需要再次计算。但是引入 KV Cache 后虽然计算量减少了，但是常驻内存占用增加了，而且随着 Context 长度的增长，对显存的占用也越大，单卡可能无法放得下。如果拓展到多卡或者多机，那么将会引入通信带宽的瓶颈限制。所以这一章节就是为了优化这个 KV Cache 的问题。

MQA

MQA（Multi-Query Attention） 思想比较简单，就是所有 header 的 Q 共用一个 K V，所以 KV Cache 就减少到原来的 1/h。其缺点就是对模型性能会带来一定的影响。

GQA

GQA（Grouped-Query Attention）为了缓解 MQA 对模型性能带来的影响，讲 h 个 Head 划分为 g 个 Group, 这样 KV Cache 就压缩原先的 h/g 倍。llama2/3-70B中，GQA的g=8，h=32。这里的 g=8 是考虑到单机 8 卡，这样在做计算时尽可能保证K、V多样性的同时最大程度上减少卡间通信。

MLA

MLA（Multi-Latent Attention）很多的数学概念，对于我有点复杂了，可以参看上面苏大的博客。我总结下思路吧：

1. 在 GQA 的基础上，将 K, V 拼在一起看做由输入和一个参数矩阵的乘积结果[K,V] = x.W
2. 通过不同的投影矩阵来增强了GQA的能力，并且推理时可以保持同样大小的KV Cache
3. 那么反过来，如果我们只需要跟GQA相近的能力，那么 dimension W 取更小的值（DeepSeek-V2取了512），从而进一步压缩KV Cache。

Serving Pattern

随着 Chat，Code Assistant 等 LLM 应用的蓬勃发展，LLM Serving 由单卡逐步拓展到多卡，甚至是多个实例。由此衍生出 Orchestration Pattern, 例如 prefill-decode disaggregation, context cache migration, traffic-adaptive request routing 等。

Prefill-Decode Disaggregation

为什么要将 P(Prefill) 以及 D(Decode) 分开呢？参考prefill 和 decode 该分离到不同的卡上么？ - Chayenne Zhao的文章 - 知乎

P 和 D 分离的本质还是 P 和 D 阶段的计算模型不同，P 阶段是 Compute bound，D 阶段是 Memory bound。

• Compute bound: 输入所有 tokens Attention 的计算，增大 batch 会被这个所限制，没有更多算力可算。
• Memory bound: 每一个 token 生成都需要 KV Cache 从显存中频繁的读取，增大 batch 会被这个所限制。

优势：

• 显著的降低了显存的；P 阶段计算时中间激活值参数保留在显存中，而 D 阶段只需要 P 阶段的 KV Cache 即可，如果 P 和 D 不分离，在 D 阶段激活值占据的显存就是浪费，在模型较大以及 context 较长的情况下，显存的浪费会更加明显。
• 提高吞吐量：多卡系统上分离 prefill 卡和 decode 卡可以实现更高效的 pineline 并行，提高吞吐量。

缺点：

1. 通讯开销：跨卡传输 KV cache 无疑会带来新的通讯压力。
2. 隐藏的开销：为了加强通讯，networking 硬件的成本不可忽略。此外，大规模的 data center 部署很多不同机型造成的 fragmentation 也有不小代价。

P/D 分离带来了架构上的创新，例如：

1. MoonCake

• 在满足 SLO(TTFT（time to first token），TBT(token between token)) 情况下最大化 KV Cache 的缓存命中以及最大化 MFU（Model FLOPs Utilization）。因为复用 KV Cache 涉及到远程通信，例如从其他介质（CPU,DRAM,SSD）中读取到显存，所以会增大 TTFT；而更大的batch 增加 MFU 的同时也会增加 TBT 的时间。
• Mooncake 设计了 global scheduler（conductor）。对于每个 request，Conductor 需要为其选择一组用于 prefiil 和 decode 的设备，然后进行调度。首先迁移尽可能多的 KV cache 到 prefill 设备上；接着在 prefill 设备上通过 chunked and layer-wise prefill 的方法连续地 stream prefill 所得的 KV cache 到 decode 设备上；最后在 decode 设备上加载 KV cache，将此 request 加入到 continuous batching 中，完成 decode。
• Mooncake 将 GPU 集群内的 CPU DRAM SSD RDMA 资源分别建立了资源池。
• 在 GPU 和 GPU 之间转移 KV cache 由单独的 GPU direct RDMA 设备完成，被称为 messager。原文中基于哈希的前缀存储能够为上层用户提供 context caching API。
• 如果某个请求中未被 cache 的 tokens 数目超过了一个特定的 threshold（prefill_chunk）。这个请求将会被拆为多个 chunks，然后流水执行。通常而言，prefill_chunk 的大小大于 1k。
• 在 prefill 阶段进行预估而没有被拒绝的请求可能在 decode 阶段出于 SLO 的缘故而被拒绝，这也是一个显著的 staleness 问题。

chunked prefill 可以参考 基于 chunked prefill 理解 prefill 和 decode 的计算特性 - Chayenne Zhao的文章 - 知乎.

MLC 的 MicroServing API 通过细粒度的 API编排推理引擎实现：

• Data Parallel
• Prefill-Decode Disaggregation
• Balanced prefill-decode disaggregation： prefill 和 decode 工作负载可能会不平衡。当处理长 prompts 时，prefill引擎可能会被过度利用，而decode引擎则以低利用率运行，甚至保持空闲状态。将部分prefill计算动态卸载到decode引擎中。为了实现这一点，路由器需要决定“decode_start”（decode引擎开始prefill的位置）并将其传递给所有 API.

Context Cache Migration

在为 QA 工作负载提供服务时，开发人员倾向于将不同类别的上下文缓存放入不同的引擎中，并根据其匹配的上下文类别来调度传入流量。考虑有多个引擎，其中一些专门用于历史背景，另一些则专门用于科学背景。如果科学请求多于历史请求，我们可能希望通过上下文迁移将一些历史引擎切换到科学引擎，反之亦然。上面的 MicroServing 细粒度控制推理引擎 API,在引擎之间实现高效的 KV 传输，而无需中断服务。

Traffic-Adaptive Request Routing

根据工作负载特征动态重新配置编排模式，而无需更改底层引擎。可编程路由器允许不同编排策略之间的无缝切换：

• 当prefill：decode比率（prefill输入token的时间/decode所有输出token的总时间）增加时，可以通过balanced prefill-decode disaggregation 将更多prefill计算转移到 decode引擎。
• 如果大部分 prompt 数据在 decode 引擎的Context Cache中找到，系统可以完全绕过prefill引擎，直接调用start_generate来处理 prompt 的非 cache 部分。

Serving Optimization

• KV Cache
• Paged Attention
• Quantization
• Speculate Decode
• Constrainted Decoding
• Chunked prefill
• Prompt Cache
• KV Compression

KV Cache

参看大模型推理加速：看图学KV Cache - 看图学的文章 - 知乎。

大模型推理性能优化的一个常用技术是KV Cache，该技术可以在不影响任何计算精度的前提下，通过空间换时间思想，提高推理性能。具体可以参考👆的参考，总结下：

• Transformer Encode 阶段是自回归的
• 自回归下 Attention 的第 k 次计算只依赖于第 k 个 Q,其他部分是重复计算，可以通过 KV Cache 来缓存

Paged Attention

参看 图解大模型计算加速系列之：vLLM核心技术PagedAttention原理 - 猛猿的文章 - 知乎。

Paged Attention 主要解决直接给输入请求预分配显存导致的显存利用率低的问题；通过借鉴现代操作系统的内存分页技术，提高推理的显存利用率。

<\div align="center">