尚能学否

较难无痛的各种PO眼保健操：VPG TRPO PPO DPO GRPO

Sat, 01 Nov 2025 00:00:00 GMT

Intro

本文旨在以或数学或物理或ML的「直觉」+ 少量关键位置的精确数学推导来描述我所理解的Policy Gradient系列常用于大模型后训练的RL算法。

在学习本文资料之前，本人没有任何的强化学习背景，有一些不厉害的统计学与计算机科学背景。对于本文所引用的资料学习难度，对本人来说属于不伤脑的极限状态，偶发一些完全超出知识体系的内容就当语文书背了算了。

本文的主要知识脉络来源于OAI此教程（https://spinningup.openai.com/en/latest/user/introduction.html）

其他引用paper

TRPO: https://arxiv.org/pdf/1502.05477

PPO: https://arxiv.org/pdf/1707.06347

DPO: https://arxiv.org/pdf/2305.18290

GRPO: https://arxiv.org/pdf/2402.03300

Basement Problem in RL

强化学习的基本问题假设与符号约定见OAI教程，本文不再重复建议先阅读此处Introduction的三个部分: https://spinningup.openai.com/en/latest/spinningup/rl_intro.html

强化学习讨论的是一个跟交互有关的问题。

以下围棋为例，Agent（棋手），Environment（对手）。从棋手本人视角，根据盘面（）思考下一个落子的地方（），对手会根据新的盘面落子；对手落子完毕，对棋手本人，会获得一个新的盘面（）。而如果这个时候有一个围棋上帝，他算尽了所有的盘面，告诉棋手刚刚的给他带来了多少胜率的增加（假设就是胜利结局除以总结局的比例，只是其中一种的表达方式）；那么经过足够多的训练数据后，理论上就拥有了一个输入，能够输出正确的，以获得更高的的模型；输入，输出的模型就叫做策略模型policy或者；

Policy Gradient的方法就是通过建模这个学习过程，然后通过找出梯度向量来提升目标函数。

其中，有一些可以前置思考的点

（1）这样的值从样本数据中获得的跟「ground truth」也许非常遥远。比如战争游戏等带有迷雾/信息不全特征的interaction loop，我们只能获得一个观察（observation），是的一个在低维空间上的投影。而，在真实世界里，甚至是不可能获得的，只能通过一类统计函数，想方设法表达出来。当然，如果有围棋上帝，那么他可以通过minimax树的计算来表达棋手的落子跟最佳点的落子之间的胜率差，这可以认为是一种「ground truth」。

（2）没想好

VPG：Vanilla Policy Gradient

J函数的定义

用语文话来说就是，在Policy Gradient的框架下，我们求解的目标是获得一个在样本数据下的最优策略，是环境的状态空间，是action空间。训练必须要有监督信号，有一个标量函数定义的好坏，这个函数就是。

类比下围棋，一盘棋谱是有一个轨迹（）生成的，轨迹的元素是，即状态→行动→新状态loop的过程。这个轨迹可以视为对的一次采样。对于某一个特定的轨迹来说，从采样出来的概率是；同时由每一次action决策与环境状态转移概率的连乘共同决定。即：

每一条完整的轨迹，我们都可以通过一些不同的统计函数对其定义reward，也就是。

从整体样本空间的角度，评价具体的的好坏，就是对依概率分布的期望，即。

下棋来说，就是AI下棋，一盘下得很臭，两盘下得很好，如何评估？首先，随机性来源于自身的策略存在采样（模型输出分布，随机采样根据分布决定action），以及也有随机性。下得臭的棋局对AI来说，他的出现概率是可计算的，而下得好的棋局同样。因此，同样是两胜一负，如果一个策略函数生成臭棋的概率小，另一个生成臭棋的概率大；在同样的三个reward计算中，就是远好于对方的策略。公式来说就是

再换一句话，我们可以认为，环境是客观存在的一种概率分布，这种客观包括状态转移，以及对状态的奖励。Policy Gradient的行为是在客观存在的奖励体系中，找出一个贴合奖罚分布的决策机制，使得采用高奖励的动作概率更高，低奖励的动作概率更低。

J函数对的梯度

原始Policy Gradient的精华就是以上的推导过程。

其中我认为最重要的是第三步，使用Log-derivative trick，将「概率分布的梯度」转换为「概率和log概率分布的梯度的乘积」。

首先，这一步是统计、机器学习中的常用技巧，来源于求导的基本公式

是我们的老朋友了，几乎所有的机器学习模型都是在学习log likelihood的梯度。

除了数值计算上的便利，这一变换更深远的意义在于，在原始公式中的对的求导，是一个混合环境概率的连乘公式。而环境概率这一值，我们大多数情况下是无法获得的。而取了log之后，将连乘转变为连加，将与无关的项变成了单独的项，不再是必须得系数。再取梯度后，这些讨人厌的不知道怎么找的数据都消失了。

从这个公式出发，就是log策略概率的梯度的和，可以通过样本计算到得的。

降低variance的过程

从统计性质上说，上文得到的是对的梯度的无偏估计的一种，评价梯度的好坏除了无偏是一个必要条件外，稳定性/收敛速度也是很重要的。这就需要降低梯度的variance来提高迭代效率。

第一个显然的insight是对于，原始定义是整条链路的reward之和，也就是以下的公式。这显然不合理，我们不能用过去发生的奖励来“鼓励”之后执行的决策。举个例子，围棋手布局下得很好，前序的state让他有了很大的优势；中期打了个大勺让对手差点反败为胜。棋局最终是赢的，的绝对值，不会特别地小，但这不能反应这局棋在前后期决策质量上的分布差异。前序的action我们应该给它很高的reward，后序的臭棋我们应该给它很低的reward，而不是被前序“平均”掉，认为后序手其实“还行”。这样模型就不会在这局棋中学坏。换言之，我们要学习的是单纯地根据state做出action的平均后果，而不是如何复现整个的轨迹。

修改后的梯度公式，赋予了奖励系数时序性，避免了奖励的后向“泄露”。同时，在数学上可以证明，该梯度的期望依然是无偏的，而且variance更小。

第二点，沿着这个路径，我们会发现，梯度的variance的来源就是奖励系数，这个跟无关的常量。前文说到，奖励本身存在一定的概率性和可能的由计算定义带来的方差。减去一个跟无关的，只跟有关的函数，就可以减少方差。而可以证明的最优选择是，即后序的奖励和的期望。这个证明我不会，但结论是需要引入一个模型预估，也就是PPO里常说的critic Model；而函数就是reward model。

TRPO: Trust Region Policy Optimization

TRPO的标准形式是长这样的：

TRPO中的TR（Trust Region）来自公式中对的KL散度的约束。文中定义这样一个约束，通过一系列证明和推导（暂时不大看得懂）以确保对的更新，能够使得performace function（文中是类似J函数的一个目标）的值是增加的。

而从新定义的目标函数，是在梯度上和原始形式等价的，称为surrogate objective function的形式。也是通过一系列看不太懂的过程中推导出来的。

其实第一次看到的形式是在deepseek的paper中关于GRPO的论述。当时硬看公式的一个理解是跟“学习率”相似的一个值。而这个学习率是一个分布，通过对优势函数的提升比例的分布来学习到更好的策略。这样相比原始形式，对了一个“对比”的含义，确保了“总的来说，在的邻域”内，比老策略好。

TRPO的求解是又一个比较复杂的问题。这里面运用了三个基本的数学技巧（1）泰勒展开（2）拉格朗日对偶法，求强对偶问题的解（3）共轭梯度法求海森矩阵。最终求得了一个更新的近似解析解。

（1）泰勒展开

是，

KL divergence的展开第一项为0，应该展开到第二项，因此引入了KL divergence对的二阶梯度信息，即Hessian Matrix。

近似展开后，问题形式变换成以下：

（2）拉格朗日对偶

在解凸问题时的常用手段，得到的解析解

（3）共轭梯度法

由于直接求解是参数量的平方。因此引入共轭梯度法来迭代求解替代值。

最终算法

最终的TRPO Algorithm如上。

是从样本中计算的策略梯度，与VPG是一样的。其中；Q值从样本的中来，V函数从样本的中来。

是通过共轭梯度法求解到的替代求解的近似值。

的更新中，相比解析解的公式，增加了一个调节超参，是因为解析解是近似地从泰勒展开而来，需要可调节的参数调整方差。

即的参数，是一个MSE问题。

PPO：Proximal Policy Optimization

PPO在我看来是VPG和TRPO两者结合的一个经验公式。没有太多的数学道理，简单实用。重新定义了目标函数，看似复杂其实非常简洁，就是在TRPO目标函数的基础上，加上了clip + min的裁剪。

TRPO的理论看似高深美妙，但实用有以下缺陷：实现复杂；需要计算二阶梯度信息。

从过往经验来看，漂亮的理论都胜不过Adam这样的一阶方法。可能就是一阶方法小步快跑比较“有机会”鲁棒，而最优化方法对数据/分布的假设太强，容易在实践中被打破。

PPO从公式来说，就是将新旧策略的概率比，控制在 ]之间。物理含义是，如果新策略距离旧策略太远，他“提升过多”或“降低过多”的奖励，就会被减少权重。

KL Divergence是一种偏硬的正则措施，且为了求解KLD，二阶泰勒展开会引入二阶导。clip并不在更新结果上限制新旧的分布距离，而是在奖励的监督信号上对「过分的更新」做惩罚。同时，PPO的相比 TRPO的，提供了更容易理解的超参物理量。

另外，一阶SGD方法也提供一个便利。在TRPO中，我们一个网络只能解决一个问题，需要通过解析解更新Policy Network，需要通过的MSE估计作为降方差的baseline。一阶方法可以通过混合的多头loss，同时给Policy Network和Reward Model提供梯度。即：

下图为PPO-Clip Algorithm

DPO: Direct Preference Optimization

DPO的Direct体现在，他没有采用Policy Optimization这一套强化学习的建模；这个PO其实不是PO。而是在大模型后训练上引入了类似对比学习的框架。

DPO的公式非常直观，也有明确的物理含义

忽略paper中一大堆论证过程，从「直觉」上来说。

公式中的第一项相当于

A和B的含义是win/lose数据在待优化策略上的比特信息，两者相减是一种类KL散度的度量；最大化A-B或最小化B-A，就是梯度的来源。

为什么不直接优化比特信息的二重差分，而是还要套一层sigmoid变成概率？

其中一个角度是，我们拿到的label是一个偏好分布坍缩后的真值，并不是比特信息的能量差。第二个角度上sigmoid函数的梯度能够防止数值优化上的问题。二者从底层来说似乎是一致的，指向了人类偏好的选择服从统计物理中的玻尔兹曼分布。或者GPT告诉我这叫做「二能级系统的占据概率公式」。

而从数据收集/数据质量的维度考虑，相比Reward数据，DPO所需的数据在单人评估上，潜在标准更加一致；监督信号的噪声更小，数据更易收集。

GRPO: Group Relative Policy Optimization

回到PPO的算法问题，我们收集到了一个带人类评分的偏好数据集，如：

第一，数据集制作需要时间，第二数据集的是在同一个Policy（LLM）上求解的，算法又要考虑新策略与旧策略的距离度量不宜过大。

GRPO提出了一个在一类任务上的无需打标的更加on-policy的策略梯度方法。

其中，总的来说GRPO的surrogate object形式与PPO是一致的。在样本数量上，对单个Prompt上做了扩展。PPO是在原始策略上，用一个prompt( )，生成response（），交给人类打分（）；GRPO是，用一个prompt( )，采样生成个response( )，每个输出交给规则打分（）。同时参与训练的reward，不再是PPO中的绝对值，而是相对同源response组的归一化相对值。

从两个「直觉」角度解析。第一，组内相对优势就是「提拔」模型求解空间中，「做对」的部分，提高其相对优势的概率分布；第二点，相比类PPO的的优势函数，计算的是「某一个动作」发生时的「样本观测到的动作-状态奖励」 - 「预测状态奖励」；而GRPO计算的优势函数是「G个动作」分别发生时「样本观测到的动作-状态奖励」- 「样本观测到的状态奖励」。从learning的可解释上，就是样本变多+预估环节变少，非常有利于学习。

此类方法只能运用在有「解」的数据中，答案形式如选择题，简单填空等可以work。

GPU分布式训练缩写眼保健操：PP、DP、DDP、MP、ZeRO、FSDP

Tue, 28 Oct 2025 00:00:00 GMT

学习的第一步是让眼睛先熟悉

分布式计算的几个原生问题

在谈论分布式计算或大规模服务系统时，可能会谈论这几个问题

扩展的方向：纵向或横向

纵向扩展：比如一个复杂的业务系统，将其按功能和开发形式，一分为二，这是纵向扩展；纵向扩展后，微服务系统A和系统B，可以分开地考虑自己的迭代，部署方式，资源增长形式等问题。又比如TP数据库，做MYSQL分库分表的垂直拆分，将多写的表与多读的表，或按业务的性质划分不同的实例也是一种纵向扩展。纵向扩展的思想是分而治之，不同的问题分开看待，消耗了一定的团队/硬件资源，但清晰了权责，缩小了问题的边界。
横向扩展：把同一个问题切分成identity的若干个变小的问题。比如上述微服务系统A，是计算密集型应用，计算紧张时，可以部署若干套。输入的请求划分成若干份；若干份资源对应若干份请求，问题和解决方案被切分成更小等价的局部。再说TP数据库，MYSQL做分表分表的水平拆分，假设用户表的规模对单个实例来说压力过大，则可以通过路由的形式，拆分成N个等价的实例。AP系统更不用说，横向扩展是基本结构；计算100w个数字的avg，拆分成100个计算实例，分别计算1w个数的avg，将100个结果做合并，mapreduce的过程即横向扩展。

shard与replica

Shard是横向扩展的概念，表明横向扩展成几份，称呼shard数，做sharding，sharding集群等名词
Replica是在横向扩展基础上衍生出的高可用需求。在无状态系统里做sharding，不需要考虑这个问题，节点之间是完全一致的。在有状态系统里做sharding需要考虑冗余，典型地如分布式存储系统。在有状态系统里，看似横向扩展的shardA和shardB，在支持的计算和存储的类型这些元信息上一致，但存储的内容本身就是一种有状态：对用户A的统计和用户B的统计计算过程一致，但用户A的结果绝不可能替代用户B。因此系统需要冗余来提供可用性，冗余本身也是一种并行度。计算100w个用户的avg，这100w个用户shard成10份，先拆分成10个负责10w用户的子任务；每个shard有3个replica，每个shard的计算任务，理论上“有机会”纵向扩展成3个计算3w用户的子任务。做不做这件事看系统的特性去论证收益。

通信模式与一致性：P2P or Master-Slave

做了sharding，做了replica，replica之间的数据就会产生一致性要求，即N份理论上要求一致的数据，如何保持一致？
一种模式是Master-Slave，典型如Raft/Paxos等分布式协议，通过选举投票机制确定“话事人”/Master，话事人执政期间，数据一切以Master为主。代价是大多数工作需要由master承接，承担“业务”的开销与“管理”的开销大头。优点是对于一致性来说，可靠且高效。
一种是P2P模型，replica之间没有一个明确的master角色，靠互通有无来决定是否有新数据。如果“外部人”不等他们同步，则只能取到不一定可靠的数据，数据还可能被回滚。但如果“外部人”愿意等待，数据最终会一致，此为最终一致性。代价就是不可靠，低效，好处是对每个节点的性能要求比较平均。

纵向系统里的流与批

spark或者starrocks或者flink的分布式引擎做的是通过分析任务，先纵向拆分，再在纵向切分好的子任务里确定如何横向扩展。区别在于spark一类的批系统是按阶段（stage）在时间轴上串行执行的。而flink一类的流系统是按数据的流动在时间轴上并行执行的。
打个比方，spark系统的数据是快照性质的，集合当前数据的一个计算；就像修100km长的路，建设局是一段一段去修，有的路段难度小，用了100人的团队，有的地方难度大，需要用1000人，但不会在修下一段时，上一段的团队，资源还占着不放。而flink更像一个维护公路运作的团队，100km的路，每10km分为一个服务小组，不知道什么时候有多少车车会来，不可能等车来了再去组建团队，更不可能去现修马路。
微批：在flink大行其道之前，spark生态利用其完善的计算生态去完成流的计算，批的拆分就是流，从宏观层面不能说错误。

存算分离 vs 存算一体

TP系统存算一体较多，AP系统存算分离较多。存算一体：性能好适合实时计算；存算分离：性能较差，但存储可以单独考虑成本与规模，计算有机会弹性扩容，如spark与hdfs，redshift等。
在GPU计算里，”存“可以看做显存或内存，而不是磁盘存储。

PP: Pipeline Parallelism & GPipe

Gpipe: https://arxiv.org/pdf/1811.06965

GPipe的研究大的方向是把模型纵向切分，第二个是运用了微批，第三个小优化是运用了re-materialization的技巧。

其概念，如论文中最重要的图所示，没有太多需要解释的地方，即把模型按串行的层序切分成K个部分，每个部分放进一个理论上的“加速设备”上。层间通过通信，交换必要的信息；如激活值，如梯度等。

Naive的串行方式由此引入了一个问题，在figure(b)部分所示，”Device i” 计算时，其他K-1个部分是空闲的。因此引入了micro-batch的概念，将整体的mini-batch，分成M份，“流动地”推进计算。这样有两个好处，一是计算设备的空闲时间减少了，另一个是单位时间内，mini-batch的逻辑大小可以成倍增加，有助于训练的收敛和稳定性。

第三点的优化是re-materialization。在单个设备内，假设存在若干层网络（平均是#L/#K），计算梯度所需的input，可以缓存下来，也可以不缓存。大多数情况下GPU分布式系统是memory-bound与network-bound的，计算能力是溢出的。因此不缓存下来，节省了内存消耗同时不增加通信开销，是一种很好的策略。类比spark，就是RDD的lineage机制，但那是为了节省因不可靠设备导致的开销，出发点不同。PagedAttention中也提及了内存淘汰的re-compuation策略，一种普遍的时间换空间的策略。

一个遗留的问题是，如何处理“跨样本的，依赖batch内全局信息”的计算，如batch norm。文中是这么描述的。

也就是说，使用了micro-batch内的数据来做mini-batch的统计量的一个预估。从逻辑上是会影响计算结果的，从统计意义上，我认为是无损的。因为文中给的实验，batch_size少说在十万级别。

另一个问题：Pipeline Parallelism能否像spark计算一个，分stage（layer）计算完成，再加载下一层的模型推理。当然，实现上是完全可行的。但与spark task相比，加载模型的权重对比加载next stage这样的计算图程序，需要的内存加载大得多，中间结果（activations, optimizer states）也大得多。

TP: Tensor Parallelism via Megatron-lm

纵向扩展面临的几个问题：(1) 纵向切分网络需要做“规划” (2) 切分的均衡性有一定考验，且最坏情况时，单层也许超过单卡的内存，使得进行。

Megatron-lm提出的是横向扩展的方案，即我们应当如何横向切分权重矩阵为等价的若干个小矩阵，什么时候做聚合数据等问题。

如图所述，(a) 描述的是一个两层MLP + GeLU + Dropout，应当如何切分矩阵。第一层权重A按列切分，每台设备获取全量的X与sharding分片的矩阵做乘法。分片的乘法结果，与第二层权重B，按行切分后继续相乘，随后聚合结果。至于为什么A和B两层的权重切分方式不同，以及三层MLP应该怎么做，在另一篇文章中有说明。blog.zjfwork.cc

blog.zjfwork.cc

（b）描述了在原始transformer self-attention中是怎么横向扩展的。考虑self-attention，有head_size的概念，head_size影响softmax的范围。如果按naive的思路，将attention head横向切分，则会导致softmax的计算需要跨节点，v的计算也需要跨节点。因此合理的策略是，将同一个attention head放在一台机器上，使其有能力完成一个local的self-attention计算。这也是横向扩展，只是sharding by的维度不同，(a)中描述的mlp是sharding by column or row，(b)中描述的是将attention head视为一个整体，sharding by attention head。

以上方法称为tensor parallelism，是model parallelism的一种。pipeline parallelism也是model parallelism的一种。

DP， DDP: Data Parallelism and Distributed Data Parallelism

DP，如《PyTorch Distributed: Experiences on Accelerating Data Parallel Training》一文中所说：

DP的思想很直观，一个模型的版本训练速度是确定的，那么理论上多部署一些模型就可以加速。但加速不是免费的，是边际效用递减的，因为最终若干个模型replica之间需要通信来保持参数的一致。

思考几个场景：（1）单机单卡（2）单机多卡（3）多机多卡。前两者在本机内通信，后者通过网络通信，需要考虑的网络开销不是一个数量级。DDP主要的工作是在确保分布式正确性的前提下，加速梯度通信。

pytorch DDP主要通过NVDIA的NCCL提供的”collective communication primitives”实现多模型实例之间的梯度同步。文中主要提及有以下几个特性

（1）通过NCCL的ring all-reduce tree-all reduce，从通信模型上减少了通信压力。相比master-slave结构的，更类似p2p的通信协议。

（2）bucketing gradients

通过实验证明，对于tensor数据的通信，nccl等通信库，少量多次 vs 大量少次，批量的总体性能表现更好。因此通信时机策略上会偏好梯度计算的过程中，累积足够的数据后再通信。

（3）overlapping communication with computation

文中论述了在GPU/CPU上完整计算一次反向传播的消耗时间，与all reduce一次梯度的时间在同一个数量级，意味着可以设计一种并行策略，将梯度ready的单位划分成更小的粒度，在计算梯度时，已经全部ready的梯度单位则可以先行开启all-reduce。这一部分的时间收益几乎是免费的。

（4）skipping synchronization

是一种防止死锁的策略，与性能无关

ZeRO:

ZeRO paper中声明了一个关于Megatron-lm的性能实验，以表述了一个问题，Tensor Parallelism的方法在单机多卡中可行，而多机多卡需要网络通信时，则硬件利用率很低。

ZeRO的insights

简单类比CPU分布式系统里的存算一体或存算分离，DP和MP都属于存算一体的特征；参数所在的机器就是负责其前向计算，后向计算和参数更新的机器。

（1）单纯DP中，参数被复制N份，每一完整复制品所在的机器负责一部分的数据，进行前向计算，后向计算，all-reduce后更新参数。单纯MP中，参数被切割成M份，每一份独特的分片所在的机器，负责所有数据对当前分片矩阵的前向计算，后向计算，结果同步与更新参数。

（2）DP的存算一体，带来了成倍的固定显存消耗，同时也限制了模型能塞进单卡的显存尺寸；MP的存算一体，虽然解除了显卡尺寸的限制，但MP alone的训练策略，无法对数据并行训练；DP + MP的模式，模型尺寸的限制因素从单卡变成了集群的显存总和，但要除以DP的并行度，因为集群的整体显存随着DP并行度线性增长。

（3）从存算分离的角度，模型/优化器状态/梯度这三者主要的显存占用（存），从全局而言，只需要一份。而计算则多多益善。这样，存算分离后，只要整个集群的整体显存塞得下单份模型的模型/优化器状态/梯度这三者，理论上就能达到了当前总显存大小下，能训练的最大模型（不考虑offload策略）。

（4）那么代价是什么，代价就是通信。ZeRO需要论证，在存算分离之后带来的计算/通信过程的改变，会让“算”与“存”的交互的额外时间开销不仅不会增加总体物理训练时间，反而会成倍加速。

ZeRO的方法与结果

1. ZeRO-DP

DP是很好的在充分利用Flops的情况下加速训练的手段，DP的问题就是memory的利用效率与单卡的限制。通过显存占用的分析，将其分为三个对象与三个ZeRO的stage:

（1） Optimizer States: 如Adam中的Momentum与Variance；而文中讨论Mixed-Precision设置下，优化器需要存储三个值，fp32版本的参数，M与V。因此在文中K = 4 * 3 = 12。Stage1是将这一部分分布在集群的所有GPU上，即每个分片存储。是数据并行度，在单纯ZeRO-DP里可以认为是GPU的数量。（2）stage2/stage3继续扩展，将fp16版本的“负责存储”梯度与参数的状态分片。

（3）在ZeRO-DP里，分片的是“状态”，是“存储”，计算则是每块GPU都会计算完整的模型过程，因此就是GPU的数量，就是一个mini-batch分多少个micro-batch的数量。

实验表明，在单机内，ZeRO与Megatron-lm的方法差距并不十分大（实现ZeRO的复杂度比Megatron-lm大很多）。当TP方法需要跨节点时，ZeRO在计算利用率上有数倍的提升。

2. ZeRO-R

除优化器外，剩余的内存占用问题主要是(1) 激活值的存储（activations） (2) 算子的中间值等（temporary buffers）以及由此引发的external memory fragmentation（如Partitioned Attention中讨论的问题）。

（1）激活值

激活值的理论集群显存占用在分布式训练中跟三个要素有关，①sequence length ②模型大小，层数与hidden size ③batch_size。前两者取决于模型训练的需求，后者主要影响训练速度。在单纯ZeRO-DP中，一份计算对应一份GPU，因此一份GPU对应一份micro-batch的激活值显存占用，并没有什么sharding的空间。因此ZeRO-R中需要引入MP的概念，在基于ZeRO-DP的一个计算组内（一整份的计算和一个分片的“状态”的逻辑单位），再次横向拆分。接着一整份计算对应的激活值，就可以被MP degree降低占用。举个例子，同样16卡的双机集群，可以设置DP=16，也可以设置DP=8, MP=2或DP=2, MP=8，取决于内存与通信占用分布与训练速度之间的取舍。

同时，在一个DP组内，activation也不会layer by layer存储，而是使用activation checkpoint策略，采样sqrt(L) 层存储，其他层在需要时recompute，以此减少模型配置带来的显存增长。

（2）temporary buffers与external memory fragmentation

文中没有细说，大意是通过连续内存管理等工程方式降低了浪费。

ZeRO的通信量分析

设模型参数量为

ZeRO-DP

以单纯DP为baseline，单纯DP在计算出local的梯度后，对分布式的梯度做all-reduce。ring all-reduce为其通信量最佳实现，分为ring reduce-scatter + all gather两阶段。两阶段分别的单机通信量为，和为。具体过程说明可见https://andrew.gibiansky.com/blog/machine-learning/baidu-allreduce/
与，其过程均为先对partition的元素做reduce-scatter，后进行all gather，因此通信量是一致的。
对于，额外的过程是需要拉取一个的参数量，因此通信量是

ZeRO-R

以单纯MP为baseline，对比activation partition带来的额外通信量
结论是10%左右，计算需要考虑多种因素较为复杂。

FSDP： Fully Sharded Data Parallel

https://arxiv.org/abs/2304.11277 fsdp是zero stage3在pytorch中的一个实现 The FSDP algorithm is motivated by the ZeroRedundancyOptimizer technique from DeepSpeed but with a revised design and implementation that is aligned with the other components of PyTorch.

1. FSDP Unit

ZeRO paper中并没有提到，应该如何partition states，以什么方式partition。 Naive地，我想有几种可能性（1）对参数全局编号，分段partition；显然这不是很好的方式，会造成forward时layer by layer加载时的分片网络热点。（2）全局编号，hash partition；分散了layer by layer的热点通讯。但在元数据管理上显得比较单薄，比如在加载策略上，我们在无法做更多的工作，按矩阵加载 - 计算 - 释放的方式显然会容易存在computation和communication的bubble。其他还要考虑self-attention/MLP等计算内容上的不同。

如果我们能在使用时表达一种元数据，表明A和B两个参数的单位是一组，我希望这一组同时拉取/同时释放，帮助训练引擎“掌握”更多的对训练中模型的理解。 FSDP提供了一个新的概念，unit，同一unit中的参数拥有相同的生命周期。这样我们可以将较小的layer的通信fusion在一起，从而摊销通信的overhead。

2. Hybrid Sharding

类比分布式数据库，ZeRO只提供了一种sharding=N, replica=1的策略；而FSDP提出了Hybrid Sharding的概念，sharding states允许存在replica，以适配hardware的locality。提供了一种在模型size扩展和通信量增加之间中间参数调节方案。即假设模型的大小没有大到超过一台物理机的显存size，但又希望有多机并行训练，显然通过全局的ZeRO satge3 shard会增加网络延迟。因此，在计算上依然是以单卡做一个DP，在存储分片上，一台物理机保留一个完整副本，物理机之间为replica；可以将参数通信保持在单个硬件内完成。

Megatron-LM中GEMM的推导

Wed, 08 Oct 2025 00:00:00 GMT

问题定义：

有

寻找一种适当的切分的方式与数据流程，可以得出等价的以及

前向传播的讨论

Feature1: 按列拆分

权重矩阵按列拆分时，权重矩阵的列方向数据是完整的，要在一个block上计算完整结果矩阵的数据，则要求矩阵在行方向上是”完整“的。因此对矩阵同样做列拆分是无意义的。这一点从矩阵维度的代数表达上也成立，即我们要求跟计算时，列维度必须是embed_size。
若将按行拆分，即, 则将分布在一块设备上时，只能计算出1/4的数据。
因此，必须将通信到所有的设备上，使用完整的与按列拆分的进行矩阵乘法。即：

Feature2: 按行拆分

同样地，对按行拆分或全复制（不改变列维度），不符合矩阵乘法的要求，不具备意义。
将按列拆分时，即 , 可得

可见，行拆分的权重特点是，输入可拆分，输出再聚合计算；而列拆分权重则是输入全复制，输出无计算

因此，矩阵乘法分片的两种基本形式是

权重按列分片，输入广播，输出同步拼接

权重按行分片，输入按列分片，输出同步相加

Feature3: 按列拆分后，输出内容的设备分布为按列分布，利用这个特性，也可以进行第二次乘法运算（也可以穿插若干有row locality的运算，如softmax/layer norm等）。简单过程如下

重新定义
第一次运用列拆分分布式乘法后可得，
第二次运用行拆分分布式乘法后可得，

反向传播

按列拆分时，

同理对按行拆分

二者梯度在分布设备上分别计算即可，无需聚合

Paged Attention&Orca Paper笔记与解读

Fri, 12 Sep 2025 00:00:00 GMT

序

正如PagedAttention Paper开篇第一句话所说

ORCA: A Distributed Serving System for Transformer-Based Generative Models

Paper地址：https://www.usenix.org/system/files/osdi22-yu.pdf 会议视频：https://www.youtube.com/watch?v=Ob9PPLxETYU&t=1s

Orca的目标问题

在CPU计算系统中对计算资源的管理，以OLTP，OLAP系统分类讨论。OLTP系统，典型以http/rpc等服务来说，一般以线程池为硬件计算资源在操作系统中的映射，随即典型以golang、python的asyncio、erlang otp等M:N类协程调度机制组织。典型的OLTP的服务应用中一般不存在直接利用CPU SIMD的场景，协程/线程之间的计算是不做合并的。在OLAP系统中，使用SIMD就比较常见，如Spark/Doris等“矢量化”计算的场景；因为OLAP的计算一般任务比较复杂且容易并行化，调度系统的职责就是将其切成“小块”的同质化运算，如对一批同类型的数据做同一聚合函数。但OLAP的计算隔离性一般也以task/query为维度隔离，不同的任务之间是争抢的关系。这类OLAP系统可以使用如k8s分布式虚拟化CPU，来将计算资源切分成小块，提供最大化调度的机会。

在GPU计算系统中，最大的区别是，并行化本身这件事，是运行在硬件上的。本质差别是什么？本质是单个CPU的执行是线性的，因此CPU单位的所谓并行，是用软件层的调度塞满CPU的执行流水线。任何形式的CPU调度，无论是线程池/虚拟化，都是为了安排好空闲等待的任务不去占用CPU。而GPU的执行任何时候都是批量的。因此如果GPU计算的最小“请求”单位，无法在单次计算中填满GPU，那当前批次的计算就是浪费的。于是，GPU调度的系统需要去跨请求，寻求在整体系统层面的最大化利用。

总的来说，Orca或其他llm serving system的课题之一就是如何合并请求来最大化GPU的单次计算，提升吞吐量/降低延时。

Naive地来说，假设我们的baseline是一个非常朴素的，等待一批，处理一批的系统，他存在哪些问题以及有什么优化手段？这个Orca和vLLM讨论的核心命题。

Orca的Solution

batching策略上的几个挑战：（1）以请求为最小颗粒调度时，调度的粒度不对齐；存在早结束（ealy-finished）和晚到等待（late-joining）的请求。

早结束；意味着一个batch的请求时长跟随max_output_length的请求一起结束，延长了所有请求的完成时间

晚到等待：意味着下一个batch的请求，其等待时间也是上一个batch中最长请求的完成时间。

因此假设是naive的batching思路，“集合一批，处理一批”，latency取每一批中的最长请求为基本单位，且存在计算资源的闲置和内存的空占。

iterative-level的思路也很简单，自回归本身即循环调度f(x)，因此不必以request为最小颗粒，而是可以以一批对f(x)的调度为batching；那就可以将输出长度为n的请求和长度为m的请求视作(m+n)次调度，当某请求完成时，即退出调度；新的请求可以进入调度。

（2）在解决1的基础上，使用iterative-level，即遇到第二个问题，tensor对齐

在训练或request-granularity的推理中，计算之所以可以batching，是因为self-attention计算框架约定了[batch_size, max_token_length, hidden_size]这样的tensor shape。此处可以分别做一个讨论，在问题框架上，将推理prefill和decode by step。

prefill阶段的naive batching可以认为是将token长度简单的padding到max_token_length可以解决
decode by step阶段，应用简易的kv cache思想，q是当前新增token的一维矩阵，多个请求batching，q的数量可以放置在batch维度中，设置为[n, 1, hidden_size]；K, V为前序token缓存二维矩阵，其shape为 [n, ?, hidden_size] —— 在iterative-level的调度中，K, V的长度不尽相同；若推理时，仍然按照潜在可能的max_token_length对齐，内存空间的使用依然是按照潜在最长输出占用。

当然latency上有优化，允许已完成的请求退出，允许等待的请求开始。理想的状况是，一个bacth中的多个请求，内存占用与真实的序列长度线性增长。

selective-batching：思考batch_size维度在计算上存在的意义，先不考虑multi-head这个因素。同一个batch内，除了scaled-dot attention的QK计算需要行列交互，其他的计算如MLP，softmax，layer norm——更不用说element-wise，都是在行内的操作。因此若假设两个长度不一的请求（两段prompt），假设长度为s1, s2；在原始计算框架下需要合并成两个batch，[2, max(s1, s2), hidden_size]；而在selective-batching方法中，可以合并成[1, s1 + s2, hidden_size]。在non-request-aware operation中，可以当做一个请求处理；在执行attention时，再拆成两个请求（也许是[2, max(s1, s2), hidden_size]）。

selective-batching不仅可以合并同样在prefill阶段的不同长度请求，在另一个方面也让prefill阶段和decode by step阶段可以合并计算；prefill阶段的输入是长度为任意s的用户promot；decode by step阶段输入是长度为1的last token配合有状态的kv cache manager。在计算时，系统只需要在attention时，能够区分(1) 不同请求的划分（2）不同请求的阶段 - 决定是否需要查询kv cache参与计算。

🤔思考：在selective-batching的框架下，是否需要考虑[X, hidden_size]的形状对GPU计算效率的影响。X可能会变得远大于hidden_size。

（3）Orca scheduling algorithm: iterative-level和selective-batching解决了“颗粒”和“颗粒“细化后的合并问题；下一个Orca讨论的命题就是，面对N个request产生的M个iteration全集里，在时刻t的由（requtest_i, iteration_j）组成的request pool里，每一个调度计算时，如何选择的出一个集合执行计算的问题。

Orca遵循先进先出（文中称FCFS，first come first serve）的原则；对请求到达时间排序，并以内存资源为优先的瓶颈，填充资源

n_rsrv可以理解为KV cache使用的连续内存空间；每个request的prefill被执行时，会扣除request.max_token这么大的资源预占（这也是paged attention要解决的问题）；而因为按到达序排序，因此prefill之后的请求会稳定地出现在之后批次的serve batch里，直到请求结束，占用的KV cache slots被释放。
PagedAttention主要优化的问题，即KV cache的生命周期管理与分布方式 & 具体地与decoding algo结合的使用场景。

（4）paper中讨论的分布式并行方式（intra/inter-layer），我感觉不在重点讨论范围。本文更关键的思路还是如何通过调度原语（split or merge），分配计算资源。

Efficient Memory Management for Large Language Model Serving with PagedAttention

paper地址：https://arxiv.org/pdf/2309.06180

PagedAttention的目标问题

基于Orca这样解决了调度Batching问题的系统基础上，还有哪些地方的优化可以提升系统整体的吞吐量和延时？

Insight：Orca系统中最显眼的“浪费”是在处理batching请求时，batch中的每个request都占用了request.max_token大小的显存用于存放KV cache；且该占用的生命周期是整个request存在的期间。直观上，显存的利用率是比较低的。

若引入一种机制，使得request占用的显存尽量贴近当前真实使用的大小，那么同样大小的GPU memory就大概率会允许更多的request进入当前batch（假设batch大小对计算不是瓶颈），从而在memory-bound假设的系统中吞吐量会上升，延时会降低。

此外Orca中还有一类“隐藏问题”没有被提及，LLM的decoding算法不总是基于贪心的单一线性路径，单一线性路径下token generation path是唯一的；若使用beam search这样的多分支搜索，如何管理KV cache的临时存储与释放？Naive的方法是在beam search时，每一个可能的N条路径都预占一份len == request.max_tokens的KV cache且需要在“路径”被淘汰时释放，在“路径”被扩展时copy，这显然是不可行的。因此执行beam search时，在KV cache管理上需要允许引用公共前缀与独立分叉的申请与释放。这样势必导致KV cache在至少一个路径上是不连续。于是是否KV cache一定要存放在连续显存中，成为了一个疑问。

PagedAttention的方法

引用文中Abstract的描述

On top of it, we build vLLM, an LLM serving system that achieves (1) near-zero waste in KV cache memory and (2) flexible sharing of KV cache within and across requests to further reduce memory usage.

一句话描述：模拟操作系统的虚拟内存，在KV cache需要时才实际占用显存。

随之而来的几个问题

如何定义一页/一块（page or block）

OS VMM模块管理的是一一映射的对等概念，虚拟内存与物理内存单位都是byte，只是按一定的page size划分；无论OS层面运行什么样的软件，one page就是4kb/8kb/16kb等整数。而vLLM的管理对象是KV cache，是一个有具体的计算含义的变长实体。具体来说，一个Transformer-based的模型有attention layer count, hidden_size, head_size，其他还有一些算法变体等参数会影响单个token的KV Cache到底多大。此外无论如何，单个token的所有KV cache放在一起过于巨大（100kb~10MB之间），理应有一个“切割”的思考，按类（K or V）切 or 按层切 or 两个维度都切分，也是一个因素。

如何trade off非连续的KV cache存放机制与实际计算时内存连续的要求

当KV cache以物理上非连续的方式在逻辑数据结构中被索引时，前序token的KV矩阵就不再是可以直接被GPU矩阵乘法计算的对象；不像Orca中申请max_tokens为基数的连续内存方便。

HBM不足时，evict策略是什么

Orca的连续内存策略能够“悲观”地保证进入bacth的请求，有足够的内存资源能够完成选中的请求；但虚拟内存的方式，只在申请物理内存时，“乐观”地事后知晓物理内存分配不足；（1）若没有evict策略，直接请求失败，则会导致类似惊群或死锁的问题，请求集体失败，集体重试，再集体失败（2）若有evict策略，调度系统逐出一些潜在低效的占用，那么有若干个问题：evict的时机是什么，慢介质是什么，选择策略是什么，数量怎么决定

如何与类似beam search的搜索路径配合，复用的策略

问题1：以token为单位，定义one block的token size(block size)；K V分开存放；异构模型size不同如何解决，文中并未提到。

问题2：分块计算当前q向量与K的乘积，并开发相应的cuda kernel进行优化；且vllm是一个分布式系统，以local compute为计算策略是显然的。这里会涉及到较多细节且复杂的模型并行通信原语。

问题3：此处文中提出了两种逐出介质，并在消融实验中声称各有优劣

引入sequence group的概念，意即完成一个请求所带来的需要KV cache，根据解码算法不同，一个request会产生一到多条推理sequence; 本文采用了all or nothing的策略，在sequence group的维度上，要么逐出sequence group的所有的KV cache，要么不逐出任何一块。因为vLLM假设不存在一种有收益的场景叫做，部分KV cache存在，可以先进行部分的计算。这种比iteration-level更为复杂的计算机制，超出了Orca/vLLM等最细粒度的框架假设。

解法一：逐出到CPU内存中；解法二：清除内容，后续重新计算，重新计算的速度是不慢的，重新获得kv cache就是经历一次新的prefill阶段；；需要考量内存显存带宽 pk 计算能力

请求调度的额外策略：当逐出发生时，系统就不再接受新的请求，直到先前被抢占的任务完成。这也是一个FCFS的原则（那么在被抢占任务完成前，系统是否有一段时间是跑不满的？）；文中讲到这获得了一个额外的特性，CPU RAM作为缓冲区，容量需求不会超过GPU HBM大小。

问题4：引入reference count + copy on write复用与释放内存

并行decoding采样。并行decoding采样定义：N个互不相关的随机采样序列以特定的随机参数运行。在这个场景下，request prompt是确定的N个sequence的prefix，token长度为L的prompt会产生至少向下取整(L/block_size)个完整的block，这些block在request生命周期中被每一个sequence的每一次generation复用；最后 < block_size个prompt的KV cache填不满一页，会被copy N份用于后续采样的一部分。

beam search采样过程中，更复杂的点在于，路径会动态地生成与销毁。prompt也是确定可以被多个sequence复用的；进行过程中对token产生的block进行引用计数；beam search with k的设定下，t时刻下，m（m≤k）条路径依赖来自t-1时刻产生的前缀，那么reference count记为m。在t+1时刻，m条路径会派生出m*k条新路径参与下一轮概率选择，其中有[0, k]个路径会生出，由此也会决定，在t时刻，reference count是[0, k]中的一个数字。若某block的reference count降为0，则释放。

这两类采样算法，在PagedAttention中的原理完全是一致的。通过reference count + copy on write实现了复用，区别只是fork和reference count的操作时机不同。前者的fork只存在第一个token生成时（若token length与block size是整数关系，则没有fork过程）；后者的fork在行进中随时会分叉。前者无需管理reference count；后者需要通过reference count管理释放时机。

总结

通过阅读Orca与Paged Attention Paper，大致了解了LLM serving中的一些基本问题。

但偏insight，还需要在实现中寻找更多的技术细节。

比如，vllm架构设计上是怎么样的，分为哪些模块；vllm管理分布式GPU是如何计算和通信的；哪些是kernel；block table是什么结构等