学习日志 on Fain的Blog

LLM 学习日志 #5 后训练框架：veRL、OpenRLHF、TRL和NeMo

Tue, 21 Apr 2026 18:24:07 +0800

上一篇博客系统性的整理和介绍了“预训练框架”，那么接下来就是“后训练框架”了。如果读者了解过的话，可以知道这个领域一个很有意思的现象：在“后训练框架”中，中国的社区、开发者和公司所占据的影响力和份额比预训练框架要高得多的多。这很大程度上要归功于Deepseek这一家公司：DeepSeek-R1几乎以一己之力引爆了RLVR的热潮，而GRPO也是DeepSeek提出的，这两个贡献直接定义了当前后训练的主流方向。随着Qwen、Seed的快速追赶，中国的相关实体在这一领域的早期贡献产生了明显的奠基者效应。在接下来的介绍中，将会更深刻的体会到这一点。我们将从一些先前遗漏的（因为较为不重要，或者较为陈旧了）较为小的后训练专用技术开始，随后扩展到当前的主流框架都有哪些。

⚠️注意事项：这部分我也是初学，许多内容原本不甚了解，因此大量参考了网络资料和其他在线内容进行学习后整理，供读者参考。如有问题，欢迎指出。

一些后训练技术

PEFT（Parameter-Efficient Fine-Tuning）

Prefix/Prompt-Tuning：在模型的输入或隐层添加$k$个额外可训练的前缀tokens（这些前缀是连续的伪tokens，是在Embedding空间内可微梯度下降训练找到的，不对应真实的tokens），只训练这些前缀参数；

Adapter-Tuning：将较小的神经网络层或模块插入预训练模型的每一层，这些新插入的神经模块称为adapter（适配器），下游任务微调时也只训练这些适配器参数；

LoRA：通过学习小参数的低秩矩阵来近似模型权重矩阵$W$的参数更新，训练时只优化低秩矩阵参数。

历史上这三种方法几乎同期出现（2021-2022 年），但LoRA最终成为事实标准，原因是几个优势的叠加：推理零开销（合并后消失）、实现简单（包一层wrapper就行）、效果最好（直接改权重比改输入或加模块更有效）、和量化技术的配合更自然（QLoRA）。

OpenRLHF（开源，Jian Hu等）

以"简洁 + 高性能"的平衡为核心的开源 RLHF 框架，代码量最小但性能不输重量级方案。

定位：面向学术研究者和想深入理解 RLHF 工程的开发者。核心卖点是代码精简易读（约 8,500 行），同时通过 Ray + vLLM + DeepSpeed 的组合实现了工业级性能。已被多所大学课程采用为教学框架。适合的场景是从单机到多机的 RLHF/RLVR 训练，模型规模支持到 70B+。不适合的场景是需要 Megatron 后端做超大规模训练的场景（只支持 DeepSpeed）。

组成：Python + PyTorch，核心依赖三个组件——Ray（分布式编排和 GPU 资源调度）、vLLM（高效 rollout generation，支持 AutoTP 和 PP）、DeepSpeed ZeRO-3（训练端显存优化）。架构上为每个模型（Actor/Critic/Reward/Ref）分配独立的 worker group，通过 Ray 调度。

有什么：

算法方面：PPO、REINFORCE++、REINFORCE++-baseline、GRPO、RLOO、DAPO、SFT、DPO、Reward Model 训练。

推理引擎：集成 vLLM，支持 AutoTP 和 PP，rollout 吞吐量很高。

分布式方面：基于 Ray 的灵活 GPU 分配，Hybrid Engine 支持所有模型和 vLLM 引擎共享 GPU 资源。支持 DeepSpeed ZeRO-3、deepcompile、AutoTP、RingAttention。

高级特性：异步 RL 训练（async_train）、异步 agent RL、统一的 token-in-token-out agent 执行范式（SingleTurn 和 MultiTurn 统一接口）、VLM RLHF（支持 Qwen3.5 等视觉语言模型）、多 agent RL（MARTI 基于 OpenRLHF 构建）。

没有什么/不足：

训练后端只支持 DeepSpeed，不支持 FSDP 和 Megatron。多轮 agent RL 的支持相对较新，成熟度不如单轮场景。没有专门的环境接口（需要通过 agent_func_path 自定义）。

veRL（火山引擎，ByteDance）

为性能和灵活性而生的 RLHF 框架，支持最多的训练后端和最高的显存效率。

定位：面向需要大规模 RLHF/RLVR 训练的工业团队和高级研究者。核心卖点是灵活的多后端架构和 resharding 机制。适合的场景是大规模训练（支持到数千 GPU）、需要 Megatron 后端的超大模型训练、需要精细控制训练-推理资源分配的场景。不适合的场景是快速原型验证（代码复杂度高、上手门槛大）。

组成：Python + PyTorch，核心依赖 Ray（编排）+ vLLM（推理）。训练后端灵活——同时支持 FSDP、DeepSpeed 和 Megatron Core 三种选择。核心创新是 WorkerDict（把多个模型放在同一组 worker 上共享资源）和 HybridEngine（同一组 GPU 上在训练和推理模式之间动态切换）。代码量约 32,000 行。

有什么：

算法方面：PPO、GRPO、REINFORCE、SFT、DPO、Reward Model 训练。

训练后端：FSDP + DeepSpeed + Megatron Core 三选一（其他框架最多支持一到两个）。

推理引擎：集成 vLLM。

核心创新——Resharding：在训练和推理之间动态转换参数的并行分布方式（比如训练时用 FSDP 的 sharding，推理时转成 TP），实现最高的显存效率。In-place transformation 避免了多余的参数拷贝。

分布式方面：基于 Ray 的 WorkerDict 架构，多个模型共享同一组 GPU，资源利用率高于 OpenRLHF 的独立 worker group 方案。

高级特性：支持 agentic mode（多轮 RL）。

特别值得一提的贡献：HybridEngine + resharding 的设计思想影响了后续很多框架，veRL 的论文（HybridFlow）发表在 EuroSys'25。

没有什么/不足：

代码复杂度最高（32,000 行），上手门槛大。resharding 引入了额外的延迟（虽然显存效率最高）。文档和教程不如 TRL 和 OpenRLHF 完善。异步 RL 的支持不如 OpenRLHF 成熟。VLM RLHF 支持相对较新。

TRL（Transformers Reinforcement Learning，HuggingFace）

HuggingFace 官方的后训练库，以易用性和生态集成为核心卖点。

定位：面向不想深入分布式细节、只想快速跑通后训练流程的研究者和应用开发者。最适合的场景是在 HuggingFace 生态内做 SFT/DPO/GRPO 实验，模型规模在 7B-70B，硬件在 1-8 卡。不适合的场景是需要大规模分布式 RLHF（多模型协调）、需要高吞吐 rollout generation、或者需要灵活的多轮 agent RL 的场景。

组成：Python + PyTorch，深度集成 HuggingFace Transformers + PEFT（LoRA/QLoRA）+ Accelerate（封装了 FSDP/DeepSpeed 做分布式）+ Datasets。不直接集成 vLLM 等高效推理引擎。代码量约 19,000 行。

有什么：

算法方面：SFT、DPO、SimPO、KTO、PPO、GRPO、RLOO、Reward Model 训练、Online DPO。

PEFT 集成：原生支持 LoRA、QLoRA（通过 PEFT 库），可以一行代码启用。

分布式方面：通过 Accelerate 支持 FSDP 和 DeepSpeed 后端，但用户不需要直接接触底层。

易用性方面：和 HuggingFace Hub 无缝对接（数据集加载、模型推送），Trainer 风格的 API，配置简单。

没有什么/不足：

没有集成 vLLM 等高效推理引擎，PPO rollout 阶段用的是 HuggingFace 的原生 generate，吞吐量远低于 vLLM。没有 Ray 等分布式编排，多模型（Actor/Critic/Reward/Ref）的 GPU 分配不灵活。不原生支持多轮 RL 和任意环境交互。不支持 Megatron 后端。不支持异步 RL。整体大规模训练效率不如 OpenRLHF 和 veRL。

其他值得注意的后训练框架

NeMo RL

NVIDIA 官方的后训练引擎，NeMo-Aligner 的下一代替代品，面向企业级部署。

定位：面向使用 NVIDIA 基础设施的企业团队，需要一站式、有官方支持的后训练方案。适合的场景是在 DGX 集群 / SLURM / Kubernetes 上做大规模后训练，需要和 NeMo 生态（Curator、Evaluator、Guardrails、NIM 部署）打通。不适合的场景是个人研究者或小团队快速实验（太重、依赖太多）。

组成：Python + PyTorch，底层训练用 Megatron Core，编排用 Ray，和 HuggingFace 有一定程度的集成。是 NeMo 全家桶的一部分。

有什么：

高性能 Megatron Core 后端，支持各种并行策略。支持端到端 FP8 训练。支持异步 RL。和 NeMo Gym（RL 训练环境库）配合。支持和 TRL、veRL 的互操作。被用于训练 Nemotron 3 模型家族。结构化数据流设计、清晰的接口定义。

没有什么/不足：

生态相对封闭（深度绑定 NVIDIA 全家桶）。社区活跃度和开源影响力不如 OpenRLHF 和 veRL。安装和配置复杂（推荐用 NGC 容器）。文档在 NeMo 重组过程中仍在完善。

DeepSpeed-Chat

最早的完整 RLHF 开源方案之一，但维护停滞，逐渐被取代。

定位：历史上是第一个把 SFT → Reward Model → PPO 全链路打通的开源方案，对 RLHF 工程的普及有重要贡献。但现在已经不推荐新项目使用——社区维护活跃度低，bug 多，性能不如 OpenRLHF 和 veRL。

组成：Python + PyTorch + DeepSpeed ZeRO。不使用 Ray 做编排（用 DeepSpeed 自己的 engine），不集成 vLLM（rollout 用自己的推理实现）。

有什么：

完整的 SFT → Reward Model → PPO 三阶段 pipeline。和 DeepSpeed ZeRO 深度集成。Hybrid Engine（在同一 GPU 上切换训练和推理模式）。

没有什么/不足：

不支持 GRPO、DPO、RLVR 等较新的算法。不集成 vLLM，rollout 吞吐量低。不使用 Ray，GPU 资源调度不灵活。不支持异步 RL、多轮 agent RL。社区广泛反映的问题包括：配置复杂、调试困难、某些场景下训练不收敛。

ChatLearn

阿里内部的后训练框架，开源但社区影响力有限。

定位：面向阿里内部及其云服务用户。支持 RLHF 的多模型编排和训练。在阿里内部有一定使用，但在开源社区中的采用率和讨论度远低于 OpenRLHF 和 veRL。

后训练框架核心特性汇总表

特性 TRL OpenRLHF veRL NeMo RL DS-Chat
─────────────────────────────────────────────────────────────────────────────────────
维护方 HuggingFace 社区(中国) ByteDance NVIDIA Microsoft
代码量 ~19K ~8.5K ~32K - -
定位 易用生态 简洁高性能 灵活极致性能 企业级 已过时

SFT ✓ ✓ ✓ ✓ ✓
DPO/SimPO/KTO ✓ ✓ ✓ ✓ ✗
PPO ✓ ✓ ✓ ✓ ✓
GRPO ✓ ✓ ✓ ✓ ✗
REINFORCE++ ✗ ✓ ✗ ✗ ✗
RLOO ✓ ✓ ✗ ✗ ✗
DAPO ✗ ✓ ✓ ✗ ✗
Reward Model训练 ✓ ✓ ✓ ✓ ✓

推理引擎集成
 vLLM ✗ ✓ ✓ ✗ ✗
 自有推理 HF generate - - Megatron推理 DS推理

分布式编排
 Ray ✗ ✓ ✓ ✓ ✗
 自有编排 Accelerate - - - DS Engine

训练后端
 FSDP ✓(Accelerate)✗ ✓ ✗ ✗
 DeepSpeed ✓(Accelerate)✓(ZeRO-3) ✓ ✗ ✓
 Megatron Core ✗ ✗ ✓ ✓ ✗

LoRA/QLoRA ✓(PEFT) ✓ ✓ ✓(PEFT) ✗
异步RL ✗ ✓ 部分 ✓ ✗
多轮Agent RL ✗ ✓ ✓(agentic) ✓(NeMo Gym) ✗
VLM RLHF ✗ ✓ 部分 ✓ ✗
FP8训练 ✗ ✗ ✗ ✓ ✗

GPU资源共享
 (Hybrid Engine) ✗ ✓ ✓(WorkerDict) ✓ ✓
Weight Sync方式 - NCCL/IPC Resharding Megatron内部 共享内存

易用性 最高 高 中 低 低
学习/教学价值 中 最高 中 低 低(已过时)
社区活跃度 高 很高 很高 中 低
─────────────────────────────────────────────────────────────────────────────────────

LLM 学习日志 #4 预训练框架：Megatron-LM、DeepSpeed和FSDP

Tue, 21 Apr 2026 15:27:53 +0800

这一篇接续上文，将会介绍一下值得一提的当下主流训练框架。

Megatron-LM（NVIDIA）

这是最早的祖师爷级别的预训练框架，也是预训练框架中最大而全的框架。它最早的系统性提出了TP + PP组合方案的预训练框架，定义了整个领域的标准做法。

定位

Megatron-LM的定位是明确的面向大规模的、工业级别的场景，其核心卖点是极致的大规模训练性能，面向的是需要千卡级训练的工业团队和研究机构。代码量大、学习曲线陡峭，但它定义了整个领域的"标准做法"。几乎所有其他框架的TP实现都在模仿Megatron的Column/Row Parallel模式。此外，Megatron-LM同时也是学术界研究大规模并行策略的首选参考实现，它的几篇论文（Megatron-LM 2019、Efficient Large-Scale Training 2021、Reducing Activation Recomputation 2022）本身就是分布式训练领域的经典文献。

它不适合的场景是所有相反的场景：小团队快速实验、硬件资源较少的实验室环境、单机微调，复杂的部署和源码量带来的陡峭学习曲线使其得不偿失，且性能优势无法发挥。

组成

基于Python+PyTorch，依赖NCCL+TransformerEngine（NVIDIA的混合精度库）+CUDA/C++开发和运行。

核心分为两部分：Megatron-LM是包含训练脚本的参考实现，Megatron Core是可复用的组件库。前者包括了预训练脚本、模型定义、数据处理流程，后者提供transformer building blocks、并行策略的封装、混合精度管理等等功能。也因此，许多（实际上是大部分）第三方框架可以（且选择）直接依赖Megatron Core而不用整个 Megatron-LM，例子包括了NeMo、veRL等著名后训练框架。

特性

字面意义上的在预训练这块大而全。

并行策略方面：TP（Column/Row Parallel，定义了行业标准）、PP（1F1B+Interleaved 1F1B）、DP、SP、CP、EP（包括Parallel Folding）。

混合精度方面：FP16、BF16、FP8（通过TransformerEngine）、FP4。

显存优化方面：Activation Checkpointing（full+selective recomputation）、Distributed Optimizer（ZeRO-1级别）。

工程基建方面：分布式Checkpoint（支持不同并行拓扑间的resharding）、和HuggingFace的checkpoint互相转换（Megatron Bridge）。

缺点在于Zero Bubble调度不是原生内置的。后训练能力本身不提供（SFT/RLHF由NeMo-Aligner/NeMo-RL等后训练框架承接）。ZeRO-2/3级别的参数和梯度sharding不是Megatron原生的（其原生Distributed Optimizer只做ZeRO-1，更深的sharding需要和 DeepSpeed 结合）。对非NVIDIA硬件基本不可用。文档不够完善，很多高级功能需要读源码才能理解。

DeepSpeed（Microsoft）

以ZeRO系列优化为核心，通过最小代码改动让PyTorch模型实现大规模训练的库。

定位

面向不想大幅改动模型代码但需要训练大模型的团队。核心卖点是低侵入性，通过修改配置而不是代码本身就能启用ZeRO，不需要重写模型。不适合的场景是追求极致MFU的超大规模训练，因为TP/PP的实现质量和性能不如Megatron。此外也不太适合需要精细控制并行策略的场景。

组成

同样以Python+PyTorch为基础平台，在此基础上有自定义CUDA kernel（fused Adam、fused kernels等），通信层封装了PyTorch的 torch.distributed，支持NCCL后端，也支持ROCm/AMD。核心组件包括ZeRO Optimizer（Stage 1/2/3）、DeepSpeed Engine（训练循环实现的主要组件）、Pipeline Engine、配置系统（JSON config，不需要改代码）。

特性

ZeRO系列是DeepSpeed系列的核心贡献，包括Stage 1（shard 优化器状态）、Stage 2（shard 优化器状态+梯度）和Stage 3（shard 优化器状态+梯度+参数）。ZeRO-Offload（优化器状态offload到CPU）、ZeRO-Infinity（扩展到NVMe）。ZeRO++是后续优化，包含 hierarchical partitioning、quantized all-gather等。

并行策略方面：3D并行（DP(ZeRO)+PP+TP）。值得一提的是社区认为其TP的实现不如Megatron原生。PP支持1F1B调度。Sequence Parallelism采用Ulysses方案，和Megatron的SP不同。Automatic Tensor Parallelism（对HuggingFace模型自动做TP，但仅限推理）。

混合精度方面：FP16、BF16支持。

通信优化方面：1-bit Adam / 0/1 Adam / 1-bit LAMB。这部分梯度通信压缩可以减少通信量最高26x，是其区别于FSDP的重要feature之一。

此外，其也包括了后训练需要的feature，以DeepSpeed-Chat（RLHF pipeline，包括SFT和PPO流程）为主。除此之外，还包括了Activation Checkpointing，Gradient Accumulation，Sparse Attention。

特别值得一提的是，ZeRO论文本身是分布式训练领域引用最高的论文之一。ZeRO的三级sharding思想直接影响了PyTorch FSDP的设计。"改配置而不是改代码"的设计哲学影响了后续很多框架。

DeepSpeed并不包括Context Parallel和Ring Attention。此外，也没有FP8训练支持。有MoE支持但不如Megatron完善。社区维护活跃度近年有所下降，DeepSpeed-Chat被广泛反映问题较多。

Megatron-DeepSpeed（混合方案）

取Megatron的TP/PP+DeepSpeed的ZeRO，组合成工业界实际使用最多的大规模训练方案。

定位

面向需要同时使用高效model parallelism和memory-efficient data parallelism的超大规模训练场景。BLOOM-176B、MT-NLG-530B等知名模型的训练都使用了这个组合。

组成

其实就是Megatron-LM和DeepSpeed的集成方案。TP和PP由Megatron提供，ZeRO（通常是Stage 1，和PP配合时Stage 2/3有性能问题）由DeepSpeed提供，通信压缩等额外优化也来自DeepSpeed。

ZeRO和PP配合时通常只用Stage 1。Stage 2在PP下需要额外的reduce-scatter通信，开销大。Stage 3理论上可以但当时性能不够好。这主要是历史遗留问题，DeepSpeed官方后来承认如果重新评估可能会选Stage 3。

PyTorch FSDP/FSDP2（PyTorch）

PyTorch官方的ZeRO实现，在演化过程中从独立库演变为框架原生组件，是当前PyTorch生态中data parallel的标准方案。准确的说，它不是一个独立框架，而是一个组件，在演化过程中以特殊功能被嵌合在生态内。

定位

不是独立的预训练框架，而是一个并行策略组件。被TorchTitan、TRL、OpenRLHF等框架作为底层使用。面向所有使用PyTorch的训练场景。适合需要ZeRO级别显存优化但不想引入DeepSpeed外部依赖的场景。不适合需要TP/PP的场景，FSDP本身不提供这些，需要配合其他组件。

组成

纯PyTorch原生，不依赖任何外部库。FSDP1基于FlatParameter实现，FSDP2基于per-parameter DTensor sharding实现。

特性

FSDP1等价于ZeRO-3，也支持ZeRO-2和不shard（纯DDP）。支持CPU offload、mixed precision和activation checkpointing。

FSDP2在此基础上额外支持了per-parameter sharding（不再把参数flatten成一大块），和 torch.compile 兼容，显存效率比FSDP1高约7%，性能平均提升约1.5%。

本身不提供TP、PP、SP、CP等model parallelism。不提供训练循环、数据加载、checkpoint管理等框架级功能。不是一个完整的训练解决方案，需要搭配其他组件使用。

此外，值得一提的是FSDP的演化历史。它其实经历了三个阶段，FairScale（Meta的早期实现。这个时候是独立库）→PyTorch FSDP1（被PyTorch吸收进入主库）→FSDP2（用DTensor重写，和TorchTitan共同开发）。作为数学原理和DeepSpeed完全相同（都是 ZeRO 的三级 sharding），但起源和演化历史完全不同的对位存在，其实现方式和API设计与DeepSpeed的异同之处是一件很值得研究的有意思的事情。

实际上，它的生态位和前面的DeepSpeed高度重叠：前者本来就是以最小侵入为卖点，后者索性直接在PyTorch内原生内置，两者的数学原理又相同。FSDP在逐步蚕食DeepSpeed的市场份额。在2026年，对于大多数团队在2-8卡上微调7B-70B模型的场景，PyTorch FSDP已经成为默认推荐。这个转变在过去两年内发生得很快。

其他预训练框架

Nanotron（HuggingFace）

以简洁和教学价值为核心的minimalistic 3D并行预训练库。

定位：面向想理解3D并行实现原理的学习者和研究者，以及需要一个轻量级预训练工具的小团队。配套的Ultrascale Playbook是目前最好的分布式训练实践指南之一。不适合追求极致性能的工业场景，且feature覆盖不如Megatron。

组成：Python + PyTorch，使用torchrun启动，依赖Flash Attention和Triton。核心组件包括DistributedTrainer（训练编排）、ParallelContext（管理DP/TP/PP的process group）、PipelineBlock（PP的stage封装）、NanotronModel（模型基类）。

特性：

具有的特性：TP（模仿Megatron的Column/Row Parallel）、PP（1F1B调度）、DP、ZeRO-1、Activation Checkpointing、支持Llama/Mistral/Qwen等主流架构、支持和HuggingFace checkpoint的互相转换、支持Tied Parameters。

不具有的特性：ZeRO系列没有Stage 2/3的实现、没有EP/CP/SP、没有FP8训练、PP调度只有基础1F1B（没有Interleaved或Zero Bubble）、没有分布式Checkpoint resharding。

TorchTitan（Meta+Pytorch官方）

PyTorch官方的预训练参考实现，是为了展示PyTorch最新分布式特性而开发的showcase实现，同时也是一个clean-room implementation。

定位：面向研究团队和框架开发者，作为PyTorch分布式训练能力的展示平台和最佳实践参考，同时也是可用于生产的预训练系统。不适合的场景是需要MoE/EP支持的场景，以及需要完整后训练pipeline的场景。

组成：完全基于PyTorch原生API，核心抽象是DTensor（分布式tensor表示）和DeviceMesh（设备拓扑抽象）。使用FSDP2作为data parallel层。不依赖Megatron或DeepSpeed。

特性：

具有的特性：ZeRO-1/2/3（通过FSDP2这一原生实现获得）、TP、PP、CP（Context Parallel）、Float8/MXFP8训练（也即对Blackwell GPU硬件特性的支持）、torch.compile集成、Activation Checkpointing（layer级和operator级selective）、异步分布式Checkpoint（DCP）、SymmetricMemory（节点内高效的peer-to-peer通信实现）、AsyncTP（通信计算重叠）、SFT 支持、Flight Recorder（调试工具）、Gradient Accumulation。

不具有的特性：没有MoE/EP支持、PP调度没有Zero Bubble等高级方案、整体feature数量不如Megatron、没有通信压缩。

此外，其重要贡献和卖点之一在于DTensor+DeviceMesh的统一抽象（和FSDP2一样），即让不同并行维度可以作为独立的、可组合的模块叠加，模型代码几乎不需要改动。这个设计思想对整个PyTorch分布式生态有深远影响。

各框架核心特性汇总表

此表有ai辅助排版。感谢ai。如下：

特性 Megatron-LM DeepSpeed FSDP/FSDP2 TorchTitan Nanotron
─────────────────────────────────────────────────────────────────────────────────────────
维护方 NVIDIA Microsoft PyTorch Meta/PyTorch HuggingFace
定位 工业极致性能 低侵入易用 并行原语组件 官方参考实现 教学/轻量

DP ✓ ✓ ✓ ✓(FSDP2) ✓
TP ✓ ✓(弱) ✗ ✓ ✓
PP ✓ ✓ ✗ ✓ ✓
SP ✓ ✓(Ulysses) ✗ ✗ ✗
CP ✓(含Dynamic) ✗ ✗ ✓ ✗
EP (MoE) ✓ ✓ ✗ ✗ ✗

ZeRO-1 ✓(Dist.Opt.) ✓ ✓ ✓(FSDP2) ✓
ZeRO-2 ✗(需DS) ✓ ✓ ✓(FSDP2) ✗
ZeRO-3 ✗(需DS) ✓ ✓ ✓(FSDP2) ✗
ZeRO-Offload ✗ ✓ ✓(CPU) ✗ ✗
ZeRO-Infinity ✗ ✓ ✗ ✗ ✗

PP调度: 1F1B ✓ ✓ ✗ ✓ ✓
PP调度: Interleaved ✓ ✗ ✗ ✗ ✗
PP调度: Zero Bubble ✗ ✗ ✗ ✗ ✗

Activation Ckpt ✓(+selective) ✓ ✓ ✓(+selective) ✓
Mixed Precision-BF16 ✓ ✓ ✓ ✓ ✓
FP8 Train ✓(TE) ✗ ✗ ✓(Float8) ✗
Distr-Checkpoint ✓ ✓ ✗ ✓(DCP,Async) ✗
torch.compile Partial ✗ ✓ ✓ ✗
通信压缩 ✗ ✓(1-bit Adam) ✗ ✗ ✗
SFT支持 ✗(NeMo) ✓(DS-Chat) ✗ ✓ ✗

非NVIDIA硬件 ✗ ✓(ROCm) ✓ ✓(ROCm) ✓
代码简洁度 低 中 高 高 最高
学习曲线 陡峭 中等 平缓 平缓 最平缓

可以看出，即使是最简单的教学类预训练框架，其feature的复杂性也不比一般的推理框架少。对于后训练框架，这个问题还会更明显。也许，训练框架的开发在未来可能比推理框架的开发更具有可持续性。也就是说，Agent Infra的后续发展和完善也许会进一步成为后续几年的热点。不过，究竟会是以什么形式、挑战会是什么，这仍然是一个值得思考的问题。

LLM 学习日志 #3 训练、预训练和后训练基础

Mon, 20 Apr 2026 01:28:45 +0800

今天进行简单的训练框架的调研。和我们在日常的博客、知乎文章和其他零零碎碎的地方看到的开源资料不同，其实从Megatron-LM出世到现在，训练框架已经发生了长足的发展，乃至于接近发生了范式级别的变化，不过，就我个人的观察而言，目前的开源资料针对这些新变化的普及其实还不够充分。具体而言，至少在最近1~2年的时间跨度内，工业界的前沿热点已经逐渐从单纯的实现高性能、大规模的预训练，转向中等规模的微调、后训练等等面向不同具体领域和需求的后训练相关内容。

这个演化的方向和近1~2年中Agent的崛起密切相关，同时也和参数规模的膨胀相关。随着LLM的scaling law发力，在5年间，LLM的参数量已经从单卡可以推理的约100M增加到了令人发指的约1000B，而后者光是训练就需要以十万计数量的GPU，进行持续的为期数月甚至将近年为单位的后训练。许多更小的模型则随后从这个教师模型中提取知识进行蒸馏。这意味着一件事情：基模本身的训练越来越依赖于规模效应，而规模效应和硬件资源挂钩，硬件资源和财力上限挂钩。因此，越来越多的无法在这个维度上竞争的企业（包括企业内部的团队）转向后训练，即通过更小卡需求的方式和可控规模的训练数据量，实现基于基模的垂直领域性能表现提升。也因此，这一热点需求的改变反过来催生了训练框架范式的改变：从规整的、追求规模可扩展性、强调分布式设计和并行效率的大规模预训练框架，变成强调灵活协调性、追求复杂训练流程和支持不同训练算法的后训练和强化学习训练框架，同时可能还需要和推理框架有效整合。

值得注意的是，所有的框架几乎全部基于Python+Pytorch，底层通信依赖NCCL，通过PyTorch的torch.distributed封装调用。具体来说，PyTorch提供了 ProcessGroup、all_reduce、all_gather、reduce_scatter 等原语，各框架在此之上构建自己的并行策略。部分框架（Megatron-LM、DeepSpeed）还包含自定义的C++/CUDA扩展来加速特定算子。Triton kernel在较新的框架（TorchTitan、Nanotron）中也有使用。这再次提醒了Python语言（实际上是Pytorch平台）在LLM时代的统治力。即使它并不是性能最优化的语言，也依然在工业场景里没有被C/C++淘汰，很大程度上和早期的奠基者效应以及完善的生态依赖有关系。

什么是训练

准确的说，能够让模型性能提升的操作都叫训练。在实践中，它又分为两个领域，预训练和后训练。就像在前面提到的那样，前者追求极致性能和规模的scaling，而后者追求训练调度复杂性的处理。

预训练

预训练部分的技术重点是维度切分和集合通信。这一部分在先前的学习日志#1和#2当中已经比较详细的整理了，在此不做水字数的冗余讲解。就现在而言，在当下已经落地和可预见的大模型架构下，可切分的维度基本已经被完全发掘，数学的理论意义上已经几乎没有数量级层面的突破可能性。因此，这部分的热点目前在于“大规模”、“高性能”和“基于MoE的异构”。很明显，这些都远远超过了个人、小型团队甚至许多中型企业团队的能力范畴，这从近几年国内的有竞争力的基模团队和相应人才的演化趋势当中可见一斑：豆包（字节跳动-Seed）、Qwen（阿里巴巴）、混元（腾讯）、GLM（智谱华章）、Kimi（月之暗面）、Deepseek（幻方量化）、文心一言（百度），无一不是具有雄厚财力支撑下的结果（也许Kimi除外）。

后训练

监督微调（SFT）

监督微调虽然被归类为后训练，但其实其在训练的特性上和预训练没有过多的差别，甚至可以说很相似。它的通俗解释就是“对只会续写文字的预训练基模进行训练，使得其可以按照对话的形式，根据提问和上下文给出回答”。它的技术要点在于loss和梯度更新的mask：模型是不需要对固定的格式部分，例如<|assistant|>，只需要在assistant回复的内容部分上算loss。除此之外，没有特别需要注意的地方了，大部分工作都在于数据本身（system/user/assistant的多轮对话格式），而非训练的实现。

基于偏好的对齐（Preference Alignment）

在这个领域内，可以说产生的算法和算法家族是最多的，也是最复杂的。笔者本人在学习的时候，也很长时间没能整理出一个具体的分类关系和所以然来。经过反复的调研，个人认为以下的心智模型是最好的，记录在这里，供读者参考。

后训练中的对齐/强化学习方法，可以沿两个维度来分类：

第一个维度：用于评价好坏的数据来自哪里？是1、事先收集好的（offline），还是2、训练过程中模型实时生成的（online）。

第二个维度：reward 信号来自哪里？是1、人类标注的偏好，2、训练好的 reward model，还是说3、可验证的规则。

把这两个维度交叉起来，就可以对所有目前的主流方法进行无重复无遗漏的归类。

数据来源的两种模式

模式一：离线方法。这种模式意味着在训练开始之前，已经完整拥有了一个完整的偏好数据集，其中每条数据是一个prompt加上一个 chosen response和一个rejected response（正负样本对），这些数据是提前收集好的，其来源不一，可能由人类标注，也可能由更强的模型生成。这些数据接下来以其原始文本的形式直接被用于新模型的训练，而不是通过任何中间层学习（再重新生成）之后再用于训练，这是和在线方法的区别之处。

训练过程中，模型不需要生成任何东西。它直接读取这些固定的数据对，通过特定的loss函数（鼓励policy相对于reference更偏好 chosen）来更新参数。整个训练循环和 SFT 几乎一样，也和预训练更类似，其工作流大致是读batch、forward、算loss、backward、update。没有rollout，也就是模型不新生成这些固定数据对之外的数据以获得评价（在这个工作流下无法获得评价），不存在一个可以对模型本身生成内容的评估也就是没有reward model的在线打分。

实现简单且对数据量要求低，但这个模式存在distribution shift问题：偏好数据是由某个旧policy（或人类）生成的，但随着训练进行，模型被鼓励避免旧输出模式之后，当前policy的行为分布会偏离训练数据的分布。模型无法在自己新的输出模式上得到反馈。

模式二：在线方法。训练过程中，模型自己生成回复，然后在这些自己生成的回复上获得反馈并学习。一个online RLHF的step大概是这样的：

1. 从 prompt 池中采样一批 prompt
2. 当前 policy 对这些 prompt 生成回复（rollout）
3. 获取 reward 信号
4. 用 RL 算法更新 policy
5. 回到第 1 步

复杂很多，但模型始终在自己当前的分布上生成数据并获得反馈，不存在 distribution shift问题。不过，它本身有比较严重的reward hacking问题，因为reward model在泛化之后通常质量不如人类直接进行偏好标注。

Reward信号来源的三种模式

模式一：人类直接标注。最原始的RLHF，人类看模型的输出，给偏好排序。好处是reward质量高（人类直接进行选择），坏处是成本极高、速度很慢，无法规模化。早期InstructGPT/ChatGPT用的是这种模式，但现在已经很少直接用了。

模式二：Reward Model打分。先用人类标注的偏好数据训练一个reward model（本质上是一个回归模型，输入prompt+response，输出一个标量分数）。之后的RL训练中，用这个reward model代替人类做在线打分。PPO和GRPO传统上都用这种模式。好处是可以规模化，reward model速度比人类标注快几个数量级。坏处是reward model本身可能不准确，而且policy可能学会hack reward model，即找到一个方法，可以使得reward model系统性的给高分，但实际回复的质量对人类来说并不高。

模式三：可验证的规则，基于有ground truth的reward信号进行后训练，这就是RLVR。这种reward信号有几个非常好的性质。第一，它是ground truth，有客观评价指标，没有模糊空间，因此可以避免reward hacking问题。第二，它是完全免费的，不需要人类标注和reward model，一个Python脚本就可以直接进行验证并且输出+1或者-1的reward。第三，它是无限可扩展的，通过自动方法可以生成无限多的数学题和编程题作为训练数据。不过，并不是所有问题都是可验证的，许多语言风格、道德倾向等等的训练目标完全无法通过自动形式验证，因此这个模式可以做的训练目标本身受限。

Offline方法

DPO

虽然写在前面，但其实这个方法的诞生晚于后面提到的PPO。其核心思路是绕开reward model，直接用偏好数据训练policy，从而简化训练流程和在小样本下提升训练的稳定性和可行性。每条数据（数据对）包含三个部分，一个prompt + 一个chosen response + 一个rejected response。训练时需要：

用一个冻结的reference model分别对chosen和rejected计算log probabilities
用当前policy model同样分别计算log probabilities
Loss基于这四组log probabilities构造，鼓励policy相对于reference更偏好chosen

一个数据对的训练因此实际上要跑四次forward过程（ref×chosen、ref×rejected、policy×chosen、policy×rejected）。实际上，它很难真的和其他的后训练方法被放在一起，因为它实际上处理的是整个sequence level的优化，没有一个RL的过程，换句话说根本没有进入RL的框架。

此外，也存在SimPO和KTO这样的DPO变体，分别处理不需要reference model的情况和无数据对只有答案的好坏性质（不成对，单独存在）的情况，在此不逐一介绍，感兴趣的读者可以自行了解。它们的关系如下：

 数据格式 reference model 长度处理
─────────────────────────────────────────────────────────
DPO 偏好对 需要 sum（有长度偏差）
 (chosen, rejected)

SimPO 偏好对 不需要 average（无偏差）
 (chosen, rejected) + margin

KTO 逐条标注 需要 sum
 (desirable 或
 undesirable)

Online RLHF

PPO

这是最经典和最复杂的后训练算法，其思想直接借鉴了早期的RL相关的理论和实践。其完整流程涉及四个模型：

Rollout 阶段（推理）：policy model对prompt做自回归的generation，生成回复
Reward 阶段（推理）：reward model对生成的回复打分
Advantage 估计：value model估计每个token的 value，结合reward算GAE
Policy update（训练）：用PPO clip loss和相对reference model的KL散度更新policy，同时更新value model

区别在于多出了一个value model和一个reward model。它们是协作关系：reward model需要对每一个policy model的Prompt整体给出一个标量估计，而value model负责对于每一步状态给出一个估计，即从这个状态继续出发按照现有模型的策略和分布，最终期望能拿到什么程度的reward model，因此实际上是一个把reward从sequence的整体level分解到任意中间状态的辅助评估者。

实际上，从这个视角观察的话，reward model和value model在思路上有相似之处，都是基于模仿的扩展者。reward model从标注的数据对中学习，把少量的sequence泛化到对任意整体sequence都有一个标量可供评估，而value model从reward model中学习，把对任意整体sequence都有一个连续分布的标量泛化到不完整的任意中间状态sequence都有一个连续分布的标量可供评估。我认为这个心智模型是有助于理解和记忆的。

值得一提的是，技术上的关键难点之一在于rollout阶段是自回归推理，需要KV cache、高效 batching、高效通信，以及很多推理相关的技术。这就是为什么veRL/OpenRLHF要集成vLLM。

GRPO

DeepSeek 提出的简化版PPO，核心改进是只使用三个模型，也就是说去掉了value model。对每个prompt生成一组（比如8个）回复，用reward model打分后，在组内做归一化得到advantage。也就是说，它和DPO一样，是一个sequence level赋予reward而不是token level进行评估的方法，其消除噪声的方法在于通过一组的密集输出和reward的归一化来将sequence level的不精确性部分抵消。它的最终性能上限不一定有PPO好，但简单的多。这样就只需要policy + reference + reward三个模型，不需要value model去做critic。实现上比PPO简单不少，效果也被验证过，目前非常热门。

RLVR（RL with Verifiable Rewards）

这是最近几年被Deepseek带起来的时兴架构，也就是DeepSeek-R1、Qwen等模型采用的路线。它和传统RLHF的区别在于reward不来自训练好的reward model，而来自可验证的规则（数学题的答案是否正确、代码是否通过测试用例）。工程上它依然是RL训练，因此需要rollout、需要reward scoring、需要policy update，但reward的来源从一个神经网络或者一个人类的标注结果变成了一个确定性的验证器。

RLVR的产生是LLM本身能力目标发生范式转移的标志之一。随着DeepSeek-R1的成功，它展示了纯RL（不经过传统的SFT中间步骤）就能让模型学会复杂的推理行为：长链chain-of-thought、自我纠错、"wait let me reconsider"这种反思模式。从此，推理能力成为新战场。2024-2025年的LLM竞争焦点从"对话能力"转向"推理能力"（数学、代码、逻辑），而推理任务恰好是最容易做verifiable reward的领域。这让RLVR有了天然的应用场景。

训练会用到的其他技术特性

除了1F1B这类调度特性、TP-PP-DP-CP-SP-EP的并行切分之外，其实还有许多无法简单以谱系方法归类，但仍然值得一提的技术特性。

FP8精度训练

FP8精度训练是一种对计算进行优化的方法。在这个方法下，权重和激活以FP8格式直接喂给tensor core做矩阵乘法，GPU的计算单元直接使用8-bit精度做运算。H100/B200的tensor core原生支持FP8 GEMM，而且FP8的吞吐量是BF16的两倍（因为同样的硬件单元一个周期能处理两倍数量的8-bit元素）。因此，使用FP8训练既省显存（tensor变小了），又提升速度（FLOPS翻倍）。

具体来说，一个FP8训练的GEMM工作流如下：1、输入的X和W都是BF16格式，2、将X和W量化为X_fp8 = quantize(X / scale_X), W_fp8 = quantize(W / scale_W)，均为FP8格式，3、计算：Y_fp32_acc = X_fp8 @ W_fp8，其中Y_fp32_acc是FP32格式，4、输出Y_bf16 = Y_fp32_acc * scale_X * scale_W，其中Y_bf16和输入一样，是BF16格式。

在更细节的层面，FP8有两种格式，针对不同场景：

E4M3（4-bit指数+3-bit尾数）：动态范围较小但精度较高，用于forward pass的权重和激活。

E5M2（5-bit指数+2-bit尾数）：动态范围较大但精度较低，用于backward pass的梯度。这是因为梯度的数值范围波动更大，需要更大的动态范围。

LoRA和Q-LoRA

LoRA的思想相对简单：用低维度矩阵及其乘法结果的地址矩阵作为可训练量参与计算，即：$y=Wx+ABx$，其中A、B是可训练的小矩阵，$A: [r, d_{in}]$，$B: [d_{out}, r]$，其中r是rank，通常取8、16、32、64这样的值，远小于d_in和d_out（例如，~4096）。这里不多展开。

QLoRA的出发点是：LoRA虽然省了可训练参数的显存，但冻结的基座权重W仍然占着14GB（7B×BF16，即2）。能不能把这个也压缩呢？对此，QLoRA的做法是：把冻结的基座权重量化到4-bit存储，forward时动态反量化回BF16做计算，LoRA部分仍然以BF16训练。

具体技术有三个关键点。

NF4（NormalFloat 4-bit）。这是QLoRA论文提出的一种4-bit量化格式。它的设计思路是：预训练模型的权重分布近似正态分布，那么 4-bit 的16个量化档位应该按正态分布的分位点来划分，而不是均匀划分。这样信息论意义上最优，即每个量化档位承载的信息量相等。和均匀量化相比，NF4在同样的bit数下能更好地保留权重信息。

Double Quantization。量化时每个block（比如64个权重）需要一个scale factor（FP32，4bytes）。如果block size=64，那么scale factor的额外开销是4/64=0.0625bytes per parameter。虽然看起来不多但7B参数下也有~0.4GB。Double quantization是对这些scale factor本身再做一次量化（量化到FP8），进一步压缩这个开销。

Paged Optimizers。QLoRA 还用了NVIDIA的unified memory（CUDA managed memory）来管理优化器状态，当GPU显存不够时自动page到CPU内存，避免OOM，这其实和后面ZeRO-Offload的思路类似。

方法 基座权重 优化器状态 总计
─────────────────────────────────────────────────
全量微调 BF16 14 GB 84 GB ~120 GB（需要多卡）
LoRA BF16 14 GB 0.5 GB ~30 GB（单卡 A100）
QLoRA 4-bit 3.5 GB 0.5 GB ~10 GB（单卡 RTX 4090）

QLoRA让单张24GB消费级GPU就能微调7B模型，这对学术界和个人开发者的意义巨大。但它也有其局限性。第一，LoRA的低秩约束意味着表达能力上界。第二，QLoRA的激进4bit量化会带来不可忽略的精度损失，而且反量化的计算引入了额外延迟，训练速度通常比纯BF16 LoRA慢 20-30%。第三，LoRA需求一个训练好的预训练基座，因此不适合预训练。

Zero-1/2/3/Offload/Infinity

前三者不再赘述，可以参见之前的博客，这里讲后两个，他们都是关于存储层次利用的工作。先回顾ZeRO-3的状态：模型参数、梯度、优化器状态全部被shard到各张GPU上。每张GPU只持有1/N的完整状态，需要某个参数时通过all-gather临时拼回来。这已经把GPU显存的利用率压到很低了，但如果模型再大，即使shard之后每张卡分到的那1/N也放不下呢？

ZeRO-Offload的思路是：GPU显存不够，就往CPU内存搬。CPU内存通常比GPU显存大一个数量级（比如一台机器GPU显存总共640GB，但CPU内存可能有 1-2TB），而且便宜得多。

具体来说，ZeRO-Offload把训练过程中的不同计算和数据按照"计算密度"来划分——计算密集的操作留在GPU，内存密集但计算量不大的操作搬到CPU：

留在GPU上的：forward和backward的计算（矩阵乘法、attention等），因为这些是计算密集型，GPU的算力优势在这里是不可替代的。

搬到CPU上的：优化器状态（Adam的exp_avg和exp_avg_sq）和优化器更新计算（parameter update）。Adam的更新本质上是若干个逐元素运算，读四个 tensor、做加减乘除、写回去。计算量不大，但占显存的大头（fp32 master weights+两个矩估计=参数量的12倍字节）。把这些搬到CPU，用CPU做Adam update，然后把更新后的参数传回GPU。

关键的性能瓶颈在PCIe带宽。GPU和CPU之间的数据传输走PCIe，带宽大约32-64GB/s（PCIe Gen4/Gen5），比GPU内部的HBM带宽（几TB/s）和节点内NVLink（几百GB/s）差了一到两个数量级。所以ZeRO-Offload需要精心设计通信与计算的overlap，在GPU做当前层的backward的同时，把上一层的梯度通过PCIe传到CPU；在CPU做optimizer update的同时，GPU继续做下一个microbatch的forward。

ZeRO-Offload的收益场景是：GPU数量少但模型大。同时，ZeRO-Infinity是ZeRO-Offload的进一步延伸：CPU内存也不够的话，就往NVMe SSD搬。这个在工业上通常过慢，不具体展开。不过，其中的通算overlap和prefetch思想值得借鉴和考虑。

Activation Checkpointing

是一个很有意思的细粒度技术。它的主要思想是：保留部分中间计算值，丢弃部分中间计算值，以减少缓存压力。在没有保存中间值的部分，需要计算梯度的时候，则重新通过最近的激活值通过前向传播重新回到这个地方，进行计算，在计算完成后再丢弃。此外，和这个相关伴生的是反向传播情况下的Kernel设计和fusion相关的工作。几个著名的例子是Flash Attention的backward版本，以及optimizer fusion相关的Kernel，即通过合并Adam的四个需要更新的Tensor（param+grad+exp_avg+exp_avg_sq）的更新和合并逐元素操作来显著减少memory bandwidth压力；还有fused grad accumulation，即backward产生的梯度直接累加到梯度buffer里而不是先写出来再加，在gradient accumulation场景下省一次读写。

总的来说，推理的融合目标主要是减少kernel launch开销和memory bandwidth，训练的融合目标除了这些之外，还多了一个"减少中间 tensor的materialize以配合activation checkpointing的显存策略"。

Fault tolerance

这具体来说是指一类计数，即容错，也是典型的分布式场景的技术应用。这个概念其实远早于，但在LLM的训练场景下由于不同的workload特性被赋予了。在这一领域，目前经常采用的技术就是Checkpoint（存盘恢复），还有自动重启和弹性训练。此外，梯度和Loss的异常检测、静默故障检测、心跳检测和快速故障发现等技术也有所应用。这个领域目前的主要故障类型是硬件故障（ECC错误）、软件故障（NCCL通信超时和OOM）、静默错误，热点挑战之一是弹性的扩缩容和并行度的变化。例如，如果用TP=4、PP=2训练配置保存了checkpoint，故障恢复时某些机器挂了，随后想换成TP=2、PP=4重新开始，这时候就需要checkpoint格式支持不同并行拓扑之间的resharding，而这其实是一个non-trivial的课题。此外，网络层面的容错也需要进行相应的处理（慢节点、链路故障等等）。

Ray

准确的说这是一个分布式框架而不是一个计数特性，更准确的说它其实本来和LLM毫无关系：它本质上是一个分布式任务调度框架，其原始定位是通用的分布式Python程序执行框架。它的核心概念包括了两个基本类型：

Task：一个无状态的远程函数调用。写一个普通Python函数，加上 @ray.remote 装饰器，Ray会把它调度到集群中某台机器的某个进程上执行，返回一个future（异步句柄）。

Actor：一个有状态的远程对象。写一个普通Python类，加上 @ray.remote 装饰器，Ray会在某台机器上实例化它，之后就可以远程调用它的方法。每次调用都在同一个实例上执行，所以它有持久状态。

它最经典的应用就是在OpenRLHF，连接训练端的DeepSpeed ZeRO和推理端的vLLM。

LLM 学习日志 #2 集合通信拓扑、原语和实现

Sun, 05 Apr 2026 16:39:11 +0800

这篇博客主要是重温一下之前学过的集合通信的基本模式，包括拓扑、原语和实现。说是“重温”是因为这部分在之前学习NCCL的时候其实就都已经粗略的学习过了，但之后没有用到过，因此掌握不熟练、印象不深。现在趁着Pico-vLLM的开发机会，把这部分知识彻底的掌握起来。

集合通信的原语

点对点（P2P）

Broadcast

Broadcast意味着一对多，一个GPU的数据广播给所有人。

GPU 0: [A B C D] → GPU 0: [A B C D]
GPU 1: [ ] GPU 1: [A B C D]
GPU 2: [ ] GPU 2: [A B C D]
GPU 3: [ ] GPU 3: [A B C D]

Reduce

Reduce意味着多对一，所有GPU的数据归约到一个GPU。要注意的是它不是拼接而是运算，其最终得到的数据的数据量和原本每个GPU自己的数据量是相同的。Reduce本身并不是一个具体的操作，而是一系列具有特殊性质的规约算子的集合，其满足的性质是为可结合（associative）且通常可交换（commutative）的二元运算。Reduce的常见算子包括了：

Sum（求和）
Prod（乘积）
Max / Min（极值）
Avg（平均）
BAND / BOR / BXOR（按位运算）
MinLoc / MaxLoc（极值+极值所在位置）

LLM训练中99%的Reduce操作都是求和，其他算子出现的极少。NCCL甚至不支持MinLoc / MaxLoc操作。

GPU 0: [A₀] → GPU 0: [A₀+A₁+A₂+A₃]
GPU 1: [A₁] GPU 1: [ ]
GPU 2: [A₂] GPU 2: [ ]
GPU 3: [A₃] GPU 3: [ ]

All-Reduce

All-Reduce意味着将所有GPU的数据归约，结果每个GPU都有一份完整拷贝。从逻辑上可以理解为先完成一个Reduce再完成一个Broadcast——实际上的操作在概念上类似但具体实现并不相同，同样是分解为两个阶段，但具体分解成的是Reduce-Scatter和All-Gather。

为什么不是一个Reduce+一个Broadcast？实际上，在早期的参数服务器模式下，确实就是如此。但它有一个问题：可扩展性。Reduce 阶段所有数据汇聚到一个节点（CPU），Broadcast 阶段再从这一个节点（GPU）发出去。这个节点（GPU）的带宽成为整个操作的瓶颈，其他 GPU 在等待时空闲。Reduce-Scatter+All-Gather模式实际上是通过一个线性的偏移把每一个节点（GPU）变成了一个$1/N$的参数服务器，每个只负责和自己编号相同的第$K$块的数据的Reduce和Broadcast。在集合通信的视角下，这恰好就是Reduce-Scatter+All-Gather的模式。这样讲，可能在概念上是最容易理解的。

GPU 0: [A₀] → GPU 0: [A₀+A₁+A₂+A₃]
GPU 1: [A₁] GPU 1: [A₀+A₁+A₂+A₃]
GPU 2: [A₂] GPU 2: [A₀+A₁+A₂+A₃]
GPU 3: [A₃] GPU 3: [A₀+A₁+A₂+A₃]

Reduce-Scatter

Reduce-Scatter操作意味着归约后，结果被切分，每个 GPU 只拿到一部分。但这种简化的描述方式实际上不完全准确，因为理解它在过程中的副产物，比理解它的开始状态和结束状态更有助于理解它是如何运作的。以最经典的ring为例子讲解，一张图会更容易帮助读者理解到底发生了什么：

每步：每个 GPU 向下一个 GPU 发送一块数据，同时从上一个 GPU 接收一块并求和。
规则：第 1 步发送自己编号对应的块，之后每步发送刚更新过的块。

初始状态:
 GPU 0: [A₀ B₀ C₀ D₀]
 GPU 1: [A₁ B₁ C₁ D₁]
 GPU 2: [A₂ B₂ C₂ D₂]
 GPU 3: [A₃ B₃ C₃ D₃]

第 1 步: GPU 0 ──A₀──→ GPU 1
 GPU 1 ──B₁──→ GPU 2
 GPU 2 ──C₂──→ GPU 3
 GPU 3 ──D₃──→ GPU 0

 GPU 0: [A₀ B₀ C₀ D₀+D₃ ]
 GPU 1: [A₁+A₀ B₁ C₁ D₁ ]
 GPU 2: [A₂ B₂+B₁ C₂ D₂ ]
 GPU 3: [A₃ B₃ C₃+C₂ D₃ ]

第 2 步: GPU 0 ──D₀₃──→ GPU 1 （发送刚更新的 D 块）
 GPU 1 ──A₀₁──→ GPU 2 （发送刚更新的 A 块）
 GPU 2 ──B₁₂──→ GPU 3 （发送刚更新的 B 块）
 GPU 3 ──C₂₃──→ GPU 0 （发送刚更新的 C 块）

 GPU 0: [A₀ B₀ C₀+C₂₃ D₀₃ ]
 GPU 1: [A₀₁ B₁ C₁ D₁+D₀₃ ]
 GPU 2: [A₂+A₀₁ B₁₂ C₂ D₂ ]
 GPU 3: [A₃ B₃+B₁₂ C₂₃ D₃ ]

第 3 步: GPU 0 ──C₀₂₃──→ GPU 1 （发送刚更新的 C 块）
 GPU 1 ──D₀₁₃──→ GPU 2 （发送刚更新的 D 块）
 GPU 2 ──A₀₁₂──→ GPU 3 （发送刚更新的 A 块）
 GPU 3 ──B₁₂₃──→ GPU 0 （发送刚更新的 B 块）

 GPU 0: [A₀ B₀₁₂₃=B_all C₀₂₃ D₀₃ ]
 GPU 1: [A₀₁ B₁ C₀₁₂₃=C_all D₀₁₃ ]
 GPU 2: [A₀₁₂ B₁₂ C₂ D₀₁₂₃=D_all]
 GPU 3: [A₀₁₂₃=A_all B₁₂₃ C₂₃ D₃ ]
规约完成。

因此，它的最终结果就是下图。但需要注意的是，它并不是一个没有副产物的操作。如果是原位操作的话，它实际上会修改每个GPU没有得到最终规约结果的相应数据位置，使得它不再是原本值，而是一个和自身rank决定的偏移量、数据偏移量都相关的部分累加和，也可以和某种意义上的前缀和相类比。

GPU 0: [A₀ B₀ C₀ D₀] → GPU 0: [A₀+A₁+A₂+A₃]
GPU 1: [A₁ B₁ C₁ D₁] GPU 1: [B₀+B₁+B₂+B₃]
GPU 2: [A₂ B₂ C₂ D₂] GPU 2: [C₀+C₁+C₂+C₃]
GPU 3: [A₃ B₃ C₃ D₃] GPU 3: [D₀+D₁+D₂+D₃]

All-Gather

每个GPU有一部分数据，收集后每个GPU都有完整数据。Ring All-Gather中没有任何归约（求和）操作，每个GPU拿到的块直接原样转发给下一个人就行。每步每个GPU把自己最近收到的完整块沿环传递，N-1步之后每个GPU就收集齐了所有块。

GPU 0: [A] → GPU 0: [A B C D]
GPU 1: [B] GPU 1: [A B C D]
GPU 2: [C] GPU 2: [A B C D]
GPU 3: [D] GPU 3: [A B C D]

All-to-All

每个 GPU 向每个其他 GPU 发送不同的数据。由于其$O(n^2)$的通信连接数（有时候还有通信量）和复杂度，其被称为是“集群性能的试金石和考验”。它几乎总是可能引起拥塞。另一方面，大部分运行在集群上的应用都会尽量避免产生频繁的All-to-All通信，因为它很可能严重拉低相较于规整的集合通信模式下的集群性能。对All-to-All的通信优化算法和工程实践思路以处理其产生的拥塞控制为主。此外，就像上一篇的博客中对于MoE的介绍一样，许多All-to-All是动态的而非静态的，从而进一步加剧了这个问题。

GPU 0: [A₀ A₁ A₂ A₃] → GPU 0: [A₀ B₀ C₀ D₀]
GPU 1: [B₀ B₁ B₂ B₃] GPU 1: [A₁ B₁ C₁ D₁]
GPU 2: [C₀ C₁ C₂ C₃] GPU 2: [A₂ B₂ C₂ D₂]
GPU 3: [D₀ D₁ D₂ D₃] GPU 3: [A₃ B₃ C₃ D₃]

All-to-All-v

每个 GPU 向每个其他 GPU 发送不同的数据，但收发的总数据量不一定每个都相同。

集合通信的常见拓扑

Ring（环形）

最经典的集合通信的拓扑模式。

优点是带宽最优——每个 GPU 的发送和接收带宽被充分利用，总通信量和 GPU 数无关（都是 2×(N-1)/N × 数据量）。

缺点是延迟随 GPU 数线性增长——步数是 2(N-1)，GPU 越多步数越多。

Tree（树形）

优点是延迟最优——只需 O(log N) 步。

缺点是带宽利用率不如 Ring——叶节点只参与部分通信，根节点成为瓶颈。

Double Binary Tree

两棵互补的二叉树同时工作，结合了Tree的低延迟和更好的带宽利用率。NCCL在某些场景下会用这种策略。但其复杂且难以理解，而且对集群的节点数量有更严格的要求。

Recursive Halving-Doubling

递归地把GPU分成两半，先在小组内归约，再逐步扩大。兼顾延迟和带宽。

集合通信的常见实现

XCCL：NCCL、HCCL、ACCL...许多初创的CCL通信库实际上都是直接魔改的NCCL。

NCCL 的 kernel 是沿三个维度预编译的模板实例：

算法（Algorithm）：决定数据流动的拓扑模式

Ring: 环形，带宽最优
Tree: 树形（双二叉树），延迟最优
CollnetDirect / CollnetChain: 利用 SHARP 等网内计算

协议（Protocol）：决定数据传输的同步和打包方式

Simple: 大块传输，高带宽，适合大消息
LL: Low Latency，8 字节粒度，用 flag 做同步，适合小消息
LL128: 128 字节粒度，利用 NVLink 的原子操作，延迟和带宽的折中

归约操作 × 数据类型：

Sum/Prod/Max/Min × FP16/BF16/FP32/FP64/INT8/...

这些维度的组合在编译时就被实例化成了大量的 kernel 变体。例如 Ring + Simple + Sum + FP16 是一个 kernel，Tree + LL128 + Max + FP32 是另一个 kernel。这也是为什么 NCCL 的编译时间很长。

Cost Model：运行时选择

NCCL 的 cost model 是默认调优决策的核心。这个模型以时间为指标评估集合操作的开销，用于选择正确的协议和算法。cost model 考虑多种因素，包括数据量、GPU 架构、拓扑结构、网络和算法属性。

MSCCL：微软用的一种DSL，基于NCCL自定义通信算法。

MPI：太经典不讲。

硬件互联的数感建立

这一节列举常见网络和网卡设备的所有规格和速度，以帮助读者，包括写博客的我自己，建立对于网络的数字感知。

一、GPU 节点内互联：NVLink

代际 首发GPU 每条链路带宽(双向) 每GPU链路数 每GPU总带宽(双向)
NVLink 1.0 P100 40 GB/s 4 160 GB/s
NVLink 2.0 V100 50 GB/s 6 300 GB/s
NVLink 3.0 A100 50 GB/s 12 600 GB/s
NVLink 4.0 H100 100 GB/s (注1) 18 900 GB/s
NVLink 5.0 B200 100 GB/s 18 1,800 GB/s
NVLink 6.0 Rubin 200 GB/s 18 3,600 GB/s

注1: H100 的 NVLink 4.0 每条链路实际是 50 GB/s 双向，但 NVIDIA 在产品规格中
 按 sub-link 口径标注为 900 GB/s（18 sub-links × 50 GB/s）。
 B200 的 NVLink 5.0 SerDes 速率翻倍，同样 18 条但每条 100 GB/s，所以总带宽 1.8 TB/s。

二、GPU 节点内互联：NVSwitch

代际 配套GPU 每芯片端口数 每芯片交换带宽 单节点GPU数 节点总NVLink带宽
NVSwitch 1.0 V100 18 900 GB/s 8 (DGX-2) ~2.4 TB/s
NVSwitch 2.0 A100 36 (注2) ~3.2 TB/s 8 (DGX A100) ~4.8 TB/s
NVSwitch 3.0 H100 64 25.6 Tbps 8 (DGX H100) ~3.6 TB/s
NVSwitch 4.0 B200 72 NVLink5 14.4 TB/s 72 (NVL72) ~130 TB/s (机柜级)

注2: DGX A100 有 6 片 NVSwitch 2.0，每个 A100 的 12 条 NVLink 分别连到 6 片 NVSwitch
 （每片 2 条），实现 8 GPU 全连接。

NVSwitch 3.0 (Hopper) 开始引入 SHARP（Scalable Hierarchical Aggregation and Reduction Protocol），支持在交换芯片内部直接做归约计算（如 All-Reduce 的 sum），数据不需要绕回 GPU。

NVSwitch 4.0 (Blackwell) 把 NVLink 域从 8 GPU 扩展到了 72 GPU（整个机柜），是一个质变——以前跨节点才需要 InfiniBand，现在机柜内全部走 NVLink。

三、CPU-GPU 互联：PCIe

代际 发布年份 每通道速率 x16 单向带宽 x16 双向带宽
PCIe 3.0 2010 8 GT/s ~16 GB/s ~32 GB/s
PCIe 4.0 2017 16 GT/s ~32 GB/s ~64 GB/s
PCIe 5.0 2019 32 GT/s ~64 GB/s ~128 GB/s
PCIe 6.0 2022 64 GT/s ~128 GB/s ~256 GB/s
PCIe 7.0 2025(规范) 128 GT/s ~256 GB/s ~512 GB/s

对比：H100 的 NVLink 总带宽 900 GB/s，而 PCIe 5.0 x16 双向才 128 GB/s，相差 7 倍。这就是为什么 TP 必须走 NVLink 而不能走 PCIe。

四、节点间互联：InfiniBand

代际 发布年份 每通道速率 4x 链路带宽 单/双向 典型网卡 典型交换机
SDR 2001 2.5 Gbps 10 Gbps 单向 - -
DDR 2005 5 Gbps 20 Gbps 单向 - -
QDR 2007 10 Gbps 40 Gbps 单向 ConnectX-2/3 -
FDR 2011 14 Gbps 56 Gbps 单向 ConnectX-3 SwitchX
EDR 2014 25 Gbps 100 Gbps 单向 ConnectX-4/5 SwitchIB-2
HDR 2019 50 Gbps 200 Gbps 单向 ConnectX-6 Quantum
NDR 2022 100 Gbps 400 Gbps 单向 ConnectX-7 Quantum-2
XDR 2025 200 Gbps 800 Gbps 单向 ConnectX-8 Quantum-X800
GDR ~2028 400 Gbps 1.6 Tbps 单向 (规划中) (规划中)

换算成更直观的单位（双向）：

EDR: ~25 GB/s 双向
HDR: ~50 GB/s 双向
NDR: ~100 GB/s 双向 ← 当前主流 AI 集群配置
XDR: ~200 GB/s 双向 ← 2025 年开始部署

一台服务器通常配备多块网卡。例如 DGX H100 配 8 张 ConnectX-7 NDR 400G 网卡，每 GPU 一张，节点间总带宽 = 8 × 100 GB/s = 800 GB/s 双向。即便如此，仍然只有 NVLink (900 GB/s) 的约 89%。

五、节点间互联：以太网（RoCE）

速率 双向带宽 延迟（典型） 备注
10 GbE ~2.5 GB/s ~10-50 μs 传统数据中心
25 GbE ~6.25 GB/s ~5-20 μs 
40 GbE ~10 GB/s ~5-20 μs 
100 GbE ~25 GB/s ~2-10 μs RoCEv2 常见配置
200 GbE ~50 GB/s ~2-5 μs 
400 GbE ~100 GB/s ~1-3 μs 正在部署
800 GbE ~200 GB/s ~1-2 μs 2025 年开始

以太网 + RoCE（RDMA over Converged Ethernet）是 InfiniBand 的主要替代方案。带宽可以做到类似，但延迟通常高于 InfiniBand，且需要额外的拥塞控制（PFC/ECN）来模拟 IB 的无损特性。

NVIDIA 的 Spectrum-X 平台就是基于以太网的 AI 网络方案，面向不想用 InfiniBand 的客户。

六、带宽层级

带宽 (GB/s, 双向) 技术 用途
─────────────────────────────────────────────────────────
8,000 HBM3e (B200) GPU 内部显存带宽
3,600 NVLink 6.0 (Rubin) 机柜内 GPU-GPU (未来)
1,800 NVLink 5.0 (B200) 机柜内 GPU-GPU
 900 NVLink 4.0 (H100) 节点内 GPU-GPU
 800 8×NDR IB (DGX H100) 节点间总带宽
 600 NVLink 3.0 (A100) 节点内 GPU-GPU
 200 XDR IB (单卡) 节点间单卡带宽
 128 PCIe 5.0 x16 CPU-GPU
 100 NDR IB (单卡) 节点间单卡带宽
 50 HDR IB (单卡) 节点间单卡带宽
 25 EDR IB (单卡) 节点间单卡带宽 / 100GbE

数量级关系:
 HBM 带宽 > NVLink > 多卡 IB 总带宽 > 单卡 IB >> PCIe

七、延迟层级

带宽之外，延迟同样重要：

操作 典型延迟
GPU 内部 SRAM (共享内存) 访问 ~几十 ns
GPU HBM 访问 ~100-300 ns
NVLink GPU-GPU ~1-2 μs
PCIe GPU-CPU ~2-5 μs
InfiniBand RDMA (节点间) ~1-3 μs
RoCE (节点间) ~2-10 μs
TCP/IP 以太网 (节点间) ~10-100 μs

LLM 学习日志 #1 并行方案和通信模式

Thu, 02 Apr 2026 17:20:54 +0800

在开始下一个阶段的开发之前，要先打好理论地基。虽然这部分之前零零散散的学过很多，但在通常的技术博客中通常只讲解每一种并行模式的技术原理和思想，将其和通信模式逐一详细联系起来并且做通信次数、数据量整合，甚至和反向传播的训练过程整合起来的并不多。

这篇文章打算稍微弥补一下这个问题，把所有零散的、常见或者不常见的并行切分方案，并且逐一评估其通信的需求、消耗、适用情况，在我的单人用户场景下，是否值得实现、实现难度如何。下面是总结。如果后面遇到了（甚至产生了）新的并行方案，也会一并记录在这里。

此外，在后面一篇博客当中，也会介绍一下每一种通信模式，以及通信模式对应的集群拓扑结构和相应常见的具体实现有哪些。

1、Tensor Parallelism（TP）— 层内切分

这个应该是最经典的并行方式了。它的切分方式就是把每一层的权重矩阵沿着head或者hidden维度切到多卡上，每卡算一部分，合并结果。不同的卡持有不同的head、或者FFN的权重矩阵的不同部分。

特别值得一提的是，Megatron-LM中一种经典的做法是，FFN中有两个线性层，可以把第一个线性层（W_gate和W_up）按列切、第二个按行切。这样的好处是，在中间的计算过程中，中间结果不需要通信。

它的流程如下：先用被切分的，只有自己这部分head的W_Q、W_K、W_V矩阵乘以输入得到自己head的Q、K、V，然后用Attention算出自己head的输出，然后把自己的head送入按照同样模式切分的、自己本地的相应的W_o部分做乘法，得到一个部分和，然后进行all reduce。值得注意的是，由于切分是相同的，因此绕过了单卡情况下，对head的全局cat操作因此无需通信，只需要cat自己本地的head就可以了。all reduce之后，送入自己本地的第一个FFN权重，然后逐元素激活，送入第二个权重，算出本地部分和，然后进行第二次all reduce。

通信模式：All Reduce

通信次数：每个Transformer Block里有两次All-Reduce（Attention之后一次，FFN之后一次）。总共次数为2*模型层数。

通信量：每一层中，涉及到的需要通信的，相关的每层数据量 = batch_size × seq_len × hidden_size × dtype_bytes。在整个模型中，总数据量 = 每层数据量 × 2（每层两次） × layer_num（层数）。对于ring All-Reduce的情况，总通信量为总数据量 × TP 度 = 总数据量 × 2 × (n-1)/n。

涉及的模块：Attention、FFN

优点： 1、可以达到很细的切分粒度，因为不同Head、FFN权重都可以切分，每张卡的显存占用可以按 TP 度线性减少。 2、实现很成熟，是最早的并行模式之一，有现实的工业级别落地。 3、天然是完全负载均衡的。

缺点： 1、通信频率很高，对卡间带宽要求很高，产生的网络延迟在非NVlink的机内和机间基本不可接受 2、切分过细的矩阵可能导致GEMM的算术强度降低，从而降低效率

训练中反向传播的情况

Megatron-LM给出了一个很漂亮的结论：对于所有前向的All-Reduce情况，在反向中的算子为Identity（不作数学操作，仅把同一个输入x复制到每张卡上）；对于所有的前向的Identity情况，在反向中的算子为。因此，在反向传播中，通信产生的次数和前向传播一模一样、数据量也相同，但出现的位置和前向传播相反。前向传播在结尾处做all reduce，而反向传播在开始处（当然同时是反向传播自己的结尾处）做all reduce。

这其实很符合直觉，简单推导即可得到，在此按下不表。

需要注意的是，反向传播的一个stage，在一般论文里大约占据两个slot，也就是两个时间步。换言之，被建模为前向传播的时间的两倍。

2、Pipeline Parallelism（PP）— 层间切分

Layer 分离 — 切分不同的Transformer Block

最简单直接的并行模式。也就是直接把模型的不同层分配到不同卡上。例如，对于一个32层模型进行对4卡的PP，那么每一张卡恰好持有 8层。数据从第一卡计算完成后送入下一张卡，以此类推直到流向最后一卡，像流水线一样。

通信模式：点对点（Point-to-Point）

通信次数：整个前向传播过程中总通信次数N-1次，N为卡数。

通信量：每次切换卡，单次卡间通信量 = batch_size × seq_len × hidden_size × dtype_bytes，一次前向传播的总通信量 = num_gpus × 单次卡间通信量。

涉及的模块：不涉及具体模块，按卡的stage进行划分，粒度为layers。

优点： 1、通信量小、频率低，可以其通信量和频率要求允许跨节点进行通信。 2、实现简单、计算规整，可以保证层内计算和单卡完全相同，不产生任何浪费和算术强度降低。 3、反向传播计算流和单卡情况相同，不需要对反向传播的组织模式进行特别的处理。

缺点： 1、存在流水线空泡（Pipeline bubble）问题，流水线的启动。 2、对于训练场景有额外的问题：反向传播需要反向流水线，交叉进行在简单的调度策略下会造成更多空泡，需要micro-batch、1F1B调度等等更复杂的策略。 3、对于层数不能被卡数整除的场景，会产生负载不均衡。此外，首尾层的embedding/lm_head可能产生lagging。

训练中反向传播的情况

Pipeline Parallelism的反向传播在数学上和单卡完全一致，每个stage各自做标准的反向传播就行。但难点在于调度：怎样安排各stage的前向和反向的执行顺序，使得GPU尽量少空闲。

需要了解的调度策略如下：

1、GPipe 朴素方案。先把所有 micro-batch 的前向做完，再统一做反向。会产生很多空泡，而且需要保存大量的micro-batch的激活值。值得注意的是：重计算想法（只存储部分中间激活值以节省显存）也是在这部分研究当中一同产生的。

2、1F1B（One Forward One Backward）在第一个forward做完之后立刻开始第一个backward的反向过程，和后面的forward过程穿插。通过组织调度数学上可以保证每个stage在任意时刻最多只需缓存PP个micro-batch的激活值。它的核心是：在稳定阶段，每个

3、Interleaved 1F1B 每个GPU不再持有连续的若干层而是持有总数相同但是在模型前向传播过程中不连续的层。这等价于增加了stage的数量，只不过可能有多个stage在同一张卡上。比如PP=4的情况下，8层模型、interleave=2；那么第一个stage（gpu）持有0、4层，第二个stage持有1、5层，第三个stage持有2、6层，第四个stage持有3、7层。它的好处是可以通过很简单的方式增加流水线深度，从而允许更少的流水线空泡，坏处则很明显：大大增加了需要的通信量，进而增加了设备需求和延迟。

4、Zero-Bubble Pipeline 这是一个很巧妙的想法。它把反向传播的过程中的成分做了更细粒度的划分。具体来说，反向传播其实包含两个独立的计算：

B（backward for input）：计算 ∂L/∂x，即对输入激活的梯度
W（backward for weight）：计算 ∂L/∂W，即对权重的梯度

前者被前一个stage依赖的，前一个stage只有获得了B才能开始自己的反向传播计算过程。而后一个实际上是悬挂在依赖链上的分支，它依赖后一个stage的结果但本身并不被其他人依赖。W的结果只是累加到本stage的grad buffer里，等最后optimizer.step() 时用。它什么时候算都行，只要在optimizer.step()之前算完就行。那么一个自然的想法就是：能否先算有依赖的B，让后面的先开始，在有空的时候自己慢慢算W？答案是可以的。

通过手工调度，或者基于**ILP（整数线性规划）**求解的自动调度，可以最小化bubble的数量，甚至于接近0bubble的比例。当然，它实际上会引入一个问题：由于W延后，实际上一个step的总Latency变长了。但通过约束，可以完全保证跨step重叠至多在一个step以内，对训练收敛几乎没有可观测的影响。

Prefill-Decode 分离（PD Disaggregation）— 切分请求的不同阶段

Prefill 和 Decode 的计算特性完全不同：

Prefill: 大矩阵乘法，compute-bound，高算术强度，适合高算力卡
Decode: 矩阵-向量乘，memory-bound，低算术强度，适合高带宽卡

混在一起跑会互相干扰。Decode的延迟敏感请求被Prefill的大计算堵住，Prefill的高吞吐被Decode的频繁小请求打断。

更进一步的说，两者的基本的算术强度就不同：Decode的算术强度一般可以估算为，FFN部分为engine中的batch数量的两倍（fp16或bf16意义下。fp8则乘以2）、Attention部分固定为1。实测数据显示，总体强度一般是40-90 FLOPs/byte左右，远低于H100 BF16的ridge point 295。而Prefill的算术强度随着seq_len同步线性增长，一般远远超过任何gpu的ridge point。

Decode 阶段

线性层（FFN、QKV 投影、Output 投影）

Decode 每步只生成 1 个 token，所以是矩阵-向量乘法（GEMV）：

权重 W: [d, d]
输入 x: [batch_size, d] （每个请求只有 1 个 token）

FLOPs: 2 × batch_size × d²
数据搬运: d² × bytes_per_param（权重）+ batch_size × d × bytes（输入，通常可忽略）

权重远大于输入，所以：

算术强度 ≈ 2 × batch_size / bytes_per_param

精度	batch=1	batch=8	batch=32	batch=128
BF16 (2B)	1	8	32	128
FP8 (1B)	2	16	64	256

Attention（decode 的 score 计算）

每个请求要用 1 个 query token 去和自己的整个 KV cache 做 attention：

Q: [batch_size, 1, d]
K cache: [batch_size, seq_len, d] ← 每个请求有独立的 KV cache

FLOPs: 2 × batch_size × seq_len × d
数据搬运: batch_size × seq_len × d × bytes（KV cache，每个请求各自的）

注意 KV cache 不能跨请求共享（每个请求的上下文不同），所以 batch 增大时数据搬运也等比增大：

算术强度 ≈ 2 / bytes_per_element ≈ 1（BF16）

不随 batch size 增长。

而Prefill的算术强度可以如此计算：

Prefill 阶段

线性层

Prefill一次处理整个 prompt，所以是矩阵-矩阵乘法（GEMM）：

权重 W: [d, d]
输入 X: [batch_size × seq_len, d]

FLOPs: 2 × batch_size × seq_len × d²
数据搬运: d² × bytes_per_param（权重）
 + batch_size × seq_len × d × bytes（输入激活）

当 batch_size × seq_len 足够大时，权重搬运的成本被大量计算分摊掉：

算术强度 ≈ 2 × batch_size × seq_len / bytes_per_param

精度	seq=1024, batch=1	seq=4096, batch=1	seq=4096, batch=4
BF16 (2B)	512	2048	8192
FP8 (1B)	1024	4096	16384

这些数字远远超过任何 GPU 的 ridge point。所以 Prefill 的线性层基本都是 compute-bound。

Attention（prefill 的 self-attention）

Prefill 中 Q、K、V 的长度都是 seq_len：

Q: [batch_size, seq_len, d]
K: [batch_size, seq_len, d]

Q×K^T 的 FLOPs: 2 × batch_size × seq_len² × d_head × num_heads
数据搬运: batch_size × seq_len × d × bytes × 3（Q, K, V）

算术强度大约是：

算术强度 ≈ seq_len / (3 × bytes_per_element)

这些差异是如此的明显，以至于这个disaggreation是最先被提出的。因此，分离之后各自可以针对性地优化硬件配置和调度策略。

Attention-FFN分离（AF Disaggregation）— 切分Transformer的不同阶段

AF分离是一个相对来说更新的分离模式，而其和MoE的诞生关系密切，是一个很细粒度的分离。这是因为Expert的“专家”指的就是不同FFN的专家，它们共享Attention的所有参数，但各自具有各自的FFN层和相应参数。更进一步的说，两者占据显存的东西并不相同，Attention部分主要是因为KV Cache占据了大量显存，KV cache在不同卡之间的来回拷贝可能造成Attention性能的严重下降；而FFN中，权重占据了显存的绝大多数（甚至可以达到80%以上）。对于一个MoE来说，显存的膨胀同时也允许（或者说需求了）更细粒度的分离模式。此外，MoE的通信模式本身就需要大量的、不规则的、不可预知的点对点通信，Attention→FFN的传输可以和All-to-All的dispatch合并成一步，这意味着AF分离的边际代价相较于dense模型的情况减少了。反正都是无法避免要付出的，不如一口气全做完。

总的来说，Attention更倾向于小TP度（保持GEMM效率，减少KV cache副本和通信，而且Attention并没有非常庞大的参数量导致的需求），而FFN/Expert想要大EP度（容纳更多expert）。可以说MoE天然适合AF分离——既在动机上需要它（Attention和Expert的并行度不匹配），又在代价上容忍它（All-to-All吸收了额外通信）。

3、Data Parallelism（DP）— 请求间切分

这应该是最朴素的并行了，朴素到不需要怎么讲。它的复杂度不在于推理（推理就完全是不同的实例在独自工作，无通信），而是在于训练的冗余参数存储的优化，这里诞生了许多大模型训练重要的早期成果——Google Brain、ImageNet、GPT-2、BERT，一直到单个模型在单卡上彻底放不下为止。

通信模式：推理无（训练则是All Reduce（经典训练DP）/All-Gather（ZeRO-1）/Reduce-Scatter（ZeRO-2）/两者（ZeRO-3））

通信次数：推理0，训练则是每个训练步 2 次集合通信（1 次 Reduce-Scatter + 1 次 All-Gather）（ZeRO-1）/ L+1 次集合通信（L 次 Reduce-Scatter（反向传播中逐层做）+ 1 次 All-Gather（更新后同步权重））（ZeRO-2）/ 3L 次集合通信（前向 1 次 All-Gather + 反向 1 次 All-Gather + 反向 1 次 Reduce-Scatter）（ZeRO-3）。

通信量：推理0，训练则是2 × Φ（可训练参数量） × dtype_bytes（经典DP和ZeRO-1和ZeRO-2）/ 3 × Φ（可训练参数量） × dtype_bytes（ZeRO-3）。

涉及的模块：推理无，训练涉及到优化器状态（ZeRO-1）、梯度状态（ZeRO-2）和权重（ZeRO-3）。

优点： 1、无敌简单和容易理解。 2、推理完全无通信开销。 3、天然是完全负载均衡的。

缺点： 1、对大模型的一个卡装不下的情况毫无帮助。 2、训练的反向传播过程依然需要通信。

训练中反向传播的情况

这是ZeRO-1、2、3和FSDP大显身手的地方。

ZeRO-1：切分优化器状态

切分方式：优化器状态（Adam 的 m 和 v）按参数均分到 N 张卡。权重和梯度每卡完整保存。

通信模式：Reduce-Scatter + All-Gather

通信过程：

反向传播结束后:
 1. Reduce-Scatter 梯度
 每卡得到自己负责的那 1/N 参数的聚合梯度
 通信量: Φ × dtype_bytes（每卡发出完整梯度，收到 1/N）

 2. 每卡用聚合梯度更新自己那 1/N 的参数（本地操作，无通信）

 3. All-Gather 更新后的权重
 每卡把自己更新的 1/N 权重广播出去，收集到完整权重
 通信量: Φ × dtype_bytes

通信次数：每个训练步 2 次集合通信（1 次 Reduce-Scatter + 1 次 All-Gather）

每步总通信量：2Φ × dtype_bytes（和经典 DP 的 All-Reduce 相同）

ZeRO-2：切分优化器状态 + 梯度

切分方式：优化器状态和梯度都按参数均分到 N 张卡。权重每卡完整保存。

通信模式：Reduce-Scatter + All-Gather

通信过程：

反向传播中（逐层）:
 1. 每层算完梯度后立刻做 Reduce-Scatter
 每卡只保留自己负责的 1/N 梯度，其余丢弃
 → 梯度显存从 2Φ 降到 2Φ/N

反向传播全部结束后:
 2. 每卡更新自己负责的 1/N 参数（本地操作）

 3. All-Gather 更新后的权重
 通信量: Φ × dtype_bytes

通信次数：每个训练步 = L 次 Reduce-Scatter（反向传播中逐层做）+ 1 次 All-Gather（更新后同步权重）

每步总通信量：2Φ × dtype_bytes（仍然和经典 DP 相同）

ZeRO-3 / FSDP：切分优化器状态 + 梯度 + 权重

切分方式：优化器状态、梯度、权重全部按参数均分到 N 张卡。每卡只存 1/N。

通信模式：All-Gather + Reduce-Scatter（前向和反向每层都要做）

通信过程：

前向传播（逐层）:
 对于第 i 层:
 1. All-Gather 收集完整权重 通信量: Φ_i × dtype_bytes
 2. 用完整权重做前向计算
 3. 丢弃非本卡的权重（释放显存）

反向传播（逐层，从后向前）:
 对于第 i 层:
 4. All-Gather 再次收集完整权重 通信量: Φ_i × dtype_bytes
 5. 计算梯度
 6. 丢弃非本卡的权重
 7. Reduce-Scatter 梯度 通信量: Φ_i × dtype_bytes
 每卡只保留自己负责的 1/N

全部反向结束后:
 8. 每卡更新自己的 1/N 参数（本地操作，无通信）
 （不需要额外 All-Gather 同步权重，因为下一步的前向会自动收集）

通信次数：每层 3 次集合通信（前向 1 次 All-Gather + 反向 1 次 All-Gather + 反向 1 次 Reduce-Scatter），总计 3L 次

每步总通信量：

前向 All-Gather: Φ × dtype_bytes （所有层加起来）
反向 All-Gather: Φ × dtype_bytes
反向 Reduce-Scatter: Φ × dtype_bytes
──────────────────────────────────────
总计: 3Φ × dtype_bytes

比经典 DP / ZeRO-1 / ZeRO-2 的 2Φ 多了 50%。多出来的就是前向传播中那一次 All-Gather——因为权重也切了，前向时也要收集。

4、Expert Parallelism（EP）— 专家参数间切分

这个同样是相对来说很容易理解的一种并行模式，也是MoE的必然选择，可以说这个并行模式和MoE架构是严格相互绑定的。这么做的主要原因有两个：第一，Expert的总量非常大，导致相同量级的MoE模型，计算量低得多，但所有Expert加起来的参数量远大于dense模型，导致模型根本不可能放在同一张卡上；第二，这些expert之间是独立的，不同token去不同的expert，中间过程不需要通信或者拼接，天然适合分布到不同卡上。

通信模式：All-to-All-v（数据量不规则）

通信次数：每个 MoE 层: 2 次 All-to-All（dispatch + combine），非 MoE 层（如 Attention）: 按 TP 的通信模式（All-Reduce）。

通信量：不可精确预测。每次All-to-All的通信量取决于token的分发模式，进一步地取决于router。理想均匀情况下:

每个 GPU 持有 batch_size × seq_len 个 token
每个 token 选 top-k 个 expert
总共需要分发的 token 数: batch_size × seq_len × top_k
其中留在本卡的比例: EP_degree 分之一
需要发送的: batch_size × seq_len × top_k × (1 - 1/EP) × hidden_size × dtype_bytes 

每个 MoE 层: 上述 × 2（dispatch + combine）

总通信量: 上述 × MoE层数。

但实际情况往往（高度）不均匀，热门 expert 收到的 token 多，冷门 expert 收到的少。

涉及的模块：FFN。

优点： 1、可以在不让计算量增加到难以承受程度的情况下，获得接近大参数量模型的性能。 2、切分粒度可以非常高。

缺点： 1、几乎总是一定有明显的负载不均衡，且无法在离线情况下静态的提前规划调度。 2、对显存的需求非常非常大。 3、训练比dense模型更加困难。

负载不均衡是MoE最难以处理的问题。Router是动态的，每个batch的token分发模式都不同。某些expert可能突然变"热门"，收到大量 token，而其他expert空闲。这会造成GPU的lagging情况，从而拖慢整个集群的性能。EP也是唯一一个通信模式动态变化的并行策略。这使得它的优化比其他并行困难得多，因为没法提前知道每一步的通信模式是什么样的。对它的调度策略比起前面的精密的线性规划模式，可能更接近于时序预测、控制论、排队论、资源调度等等范畴。可以想见，关于MoE模型的网络设计、集群设计、调度策略设计将会比前面的模式更有可以做的地方，可能也需要更长的时间才能在工业界收敛到一个相对来说足够优秀的统一策略，也可能根本不收敛。

训练中反向传播的情况

EP的训练问题并不在于技术上的问题——在DP、PP、TP等模式里已经基本上完全解决了。它面对的是专家坍缩和辅助负载均衡损失等等问题。它指的是一种Expert之间的马太效应：Router是可学习的，它在训练过程中会"偏心"——逐渐把越来越多的token路由给少数几个 expert，其他expert收到的token越来越少，最终几乎不参与计算。这是一个正反馈循环：某个expert因为偶然收到更多token，得到更多训练，变得更强，router就更倾向于选它，它就收到更多token……最终模型退化成一个dense模型，MoE的容量优势完全消失。更详细的介绍超出了这篇博客的范畴，因此就不多讲了。

此外，值得一提的是，Router的top-k选择是一个离散操作（不可微），但我们需要梯度回传。常用的做法是：

对被选中的expert，正常回传梯度
对没被选中的expert，梯度为零（它们根本没参与计算）
Router的梯度通过加权系数（gating weight）回传——这些系数是连续的、可微的

这意味着router的训练信号只来自被选中的expert，它对未选中expert的"好坏"一无所知。这也是expert collapse的一个根源。

5、Sequence Parallelism （SP） — 序列维度切分

SP是基于前文提到的TP的一种改进（同时也必须在TP的基础上做，无法单独存在。此外，它的名字实在是容易引起误解），它处理的是非FFN和非Attention的、没有被TP切割的地方，如LayerNorm的激活值、Dropout的激活值、残差连接的激活值等等，当然也是一个训练专门使用的优化手段。这是因为，激活值存储这个问题在推理里根本不存在。SP优美的地方在于，它不增加通信量，而是通过拆解通信原语的方式无代价的获得收益——也许可以被称为一种通信的Kernel Fusion？

回顾TP用到的通信模式：每个Attention/FFN模块结束时做一次All-Reduce。而在底层通信原语的逻辑上，All-Reduce其实一般是拆成两步实现的：

All-Reduce = Reduce-Scatter + All-Gather

SP 把这个拆分利用起来了：

不用 SP 时:
 TP 模块结束 → All-Reduce → 每卡得到完整结果 → LayerNorm(完整)

用 SP 时:
 TP 模块结束 → Reduce-Scatter → 每卡得到 1/N 的结果 → LayerNorm(1/N)
 → 下一个 TP 模块开始前 → All-Gather → 收集完整输入 → TP 计算

通信模式：Reduce-Scatter + All-Gather（替代 TP 原本的 All-Reduce）

通信次数：和 TP 相同，每个 Transformer Block 4 次集合通信（前向 2 次 + 反向 2 次），只是把 All-Reduce 拆成了 RS + AG

总通信量：Reduce-Scatter + All-Gather = All-Reduce。通过把All-Reduce拆分和还原，做到了在不增加额外开销的情况下减少大量激活值存储。

6、Context Parallelism （CP） — 长序列 Attention 切分

这是最近比较新的一种并行模式。CP应该是目前主流并行策略中最新加入的一维。Llama 3的训练就用了4D并行：TP + CP + PP + DP，CP 是第四维。

CP和Ring Attention技术密切相关。它是把输入序列沿token维度切分到多张卡上。每卡只持有序列的一段，以此实现利用并行性来加速的效果。

对于非Attention的操作（FFN、LayerNorm等），token之间没有交互，每卡独立计算自己的部分即可，完全不需要通信。唯一需要通信的是 Attention，因为每个token的Query需要和所有token的Key、Value交互。

Meta 在Attention模块中实现了两种CP变体，通常被称为Ring Attention。

Pass-KV（经典Ring Attention）

每卡保留自己的 Q 不动
K、V 块在 GPU 之间沿环传递
每轮：用本地 Q 和收到的 KV 块计算部分 Attention

标准Ring Attention，适合大多数场景。

Pass-Q（反向传递）

每卡保留自己的 K、V 不动
Q 块在 GPU 之间沿环传递
每轮：用收到的 Q 和本地 KV 计算部分 Attention

Pass-Q在某些特定场景下更优。当KV cache命中率低于5%时，Pass-KV更优。