四个关键模型的角色与交互

画图并说明 PPO 训练时，Actor、Reference、Critic、Reward Model 这四个模型之间的数据流。¶

下图展示了在单次 PPO 更新中，四个模型的交互关系与数据流向。

数据流详解：

阶段1：在线采样（Rollout）

用户 prompt x 输入 Actor（策略 π_θ），Actor 自回归地生成完整回复 y，并记录下每个 token 的旧对数概率 log π_old。此时只有 Actor 在工作。

阶段2：奖励与 KL 惩罚计算

阶段3：PPO 训练更新

整个数据流中，RM 和 Reference 始终冻结，不参与梯度更新；Actor 和 Critic 在训练阶段交替更新，两者的梯度通常不共享（共享底座时除外），构成了典型的 Actor-Critic 架构。

为什么 Reference 模型必须保持冻结？其参数更新会导致什么问题？¶

Reference 模型必须保持冻结，其核心作用是提供一个固定不变的参考分布锚点。如果它参与更新，RLHF 的训练框架将彻底崩塌。

Reference 模型的角色：

它充当了一个行为正则器，防止 Actor 为了追求奖励模型的虚高分数而生成怪异的、不符合语法的文本。

如果 Reference 模型参数更新会怎样？

移动靶问题：KL 惩罚原本是为了将 Actor “拴”在一个固定的、安全的语言分布上。如果 Reference 也在变，Actor 试图靠近一个不断移动的目标，这就变成了“追着自己尾巴咬”的困境。Actor 可能漂移，Reference 也跟着漂移，两者之间的 KL 散度即使很小，整体语言能力却已经共同退化到了一个离原始 SFT 分布很远的荒原。
自我强化的奖励黑客：更危险的是，Reference 模型可能会被 Actor “带坏”。如果 Actor 发现了一种奖励黑客的捷径（例如不断重复一个高分单词），并生成了大量这样的畸形文本，这些文本如果被用来更新 Reference 模型，Reference 也会逐渐认为这种畸形模式是“正常”的。这样一来，KL 惩罚就完全失效了，因为 Actor 做那些畸形动作时，Reference 也会认为这是高概率的，KL 散度将保持很低，RLHF 彻底失去语言质量的约束。
优化目标的矛盾：Reference 模型本身没有明确的优化目标。如果它的更新只是单纯地跟随 Actor，则相当于弱化了正则化；如果给它一个相反的目标（例如想让它保持原始分布），则又需要引入另一个复杂的对抗训练过程。因此，冻结 Reference 是最简洁且数学上最干净的方案。

结论：冻结 Reference 模型确保了 KL 散度始终以最初那个可靠、流畅的 SFT 策略为基准，为整个 RL 优化提供了一个稳固的“地面参考系”。它的稳定性是 RLHF 不崩溃的最后一道防线。

如果 Critic 模型从 Reward Model 初始化，会有哪些好处？有什么需要注意的？¶

好处：

语义空间的继承与加速收敛：RM 已经在大量人类偏好数据上训练过，其内部表征已经学会了区分回答的“好坏”。Critic 的任务是估计给定状态（部分生成的文本）的未来回报。从 RM 初始化，意味着 Critic 在训练初期就拥有了一个质量不错的“语言理解与价值判断”底座，不需要从零开始学习语言。这使得 Critic 能够更快地提供准确的价值估计，从而稳定训练初期的优势计算，加速整体收敛。
目标对齐：RM 的最终输出是对完整回答的评分，而 Critic 需要预测这个评分在生成过程中的期望。两者对“好”的认知在底层应该是高度一致的。共享初始化的参数有助于保持这种一致性，避免因不同初始化引入的概念分歧。

需要注意的事项：

灾难性遗忘与过拟合：RM 是为了精确区分完整回答的细微质量差异而训练的，其顶层参数可能对某些表面特征（如长度、特定格式）已经产生了过拟合。如果直接用它初始化 Critic，并在 PPO 过程中继续训练，Critic 可能会延续 RM 的这些偏见。更危险的是，如果 Critic 在某个 batch 中过度拟合到了一些噪声样本，它可能会“忘记” RM 教给它的一些基本判断力。
任务的不匹配性：RM 的训练数据是完整、高质量的回复，而 Critic 在 RL 训练中需要评估各种“半成品”甚至是错误序列。RM 的初始权重在面对这些分布外的烂尾文本时，可能会给出错误的价值估计（例如，对一个开篇流畅但后面崩溃的文本仍然给高分），导致 Actor 被误导。因此，需要让 Critic 尽快接触到 RL 过程中产生的各种中间状态样本，来适应这个“脏乱差”的环境。
更新时的谨慎调整：由于初始权重已经很好，Critic 的学习率可能需要设置得比 Actor 更小，以避免剧烈更新破坏其已经具有的评判能力。同时，价值函数裁剪（value clipping）等技术对于保护继承自 RM 的权重尤其重要。

总结：这是一种高效的“热启动”策略，能极大缩短 Critic 的冷启动期。但必须时刻警惕 Critic 对 RM 偏见的继承，并用充分的 RL 探索样本和保守的学习率来帮助它适应策略优化过程中的各种非完美状态。

在显存紧张时，能否让 Actor 和 Critic 共享底座？共享有什么优缺点？¶

完全可以，这是在显存受限下训练大语言模型 RLHF 时的常见妥协方案。

共享底座的做法：

Actor 和 Critic 使用同一个 Transformer 编码器作为主干网络，只是在顶层有所区别：Actor 接语言模型头（词表大小线性层），Critic 接标量输出头。在前向传播时，对于同一个序列，底座的计算只需要执行一次，同时产出 Actor 的 logits 和 Critic 的 value。

优点：

显存的极大节省：去掉了 Critic 的独立底座，节省了近一半的模型参数显存。这使得在相同硬件上可以使用更大的基础模型或处理更长的序列，效果立竿见影。
表示的同源性：Actor 和 Critic 共享完全相同的世界知识和语言理解，Critic 的价值估计和 Actor 的策略决策建立在完全相同的特征之上。这天然避免了因不同底座导致的认知偏差，可能使优势函数的估计在特征层面更加一致。

缺点：

梯度干扰：Actor 的 PPO 损失和 Critic 的 MSE 损失对共享参数可能会产生相互矛盾的梯度。例如，为了获得高奖励，Actor 希望底座学习某种特征表示，但 Critic 为了准确预测回报，可能希望底座学习另一种特征。这两种梯度在共享层上角力，可能导致两者都无法充分学习，收敛速度变慢甚至陷入局部最优。
灾难性遗忘加剧：Actor 的策略在持续变化，导致输入分布不断漂移。这迫使共享底座去适应 Actor 的变化，而 Critic 原本基于旧分布建立的价值认知会被迅速颠覆。Critic 需要不断追着 Actor 的屁股跑，价值估计的方差增大。
学习率匹配困难：Actor 和 Critic 通常需要不同的学习率来稳定训练（Critic 一般需要更小或更大的学习率，取决于具体环境）。共享底座使得两者无法独立调节学习率，通常需要折中，导致至少一方的训练效率受损。
价值函数裁剪的复杂性：如果底座更新过于剧烈，可能直接导致 Critic 的价值预测发生跳变，而价值裁剪通常只作用于最后的标量头，难以约束底座的变化。这增加了训练的不稳定性。

实践中的折中：如果非共享不可，建议不完全共享，而是保留少数几层独立的 Transformer 层（如最后 2-4 层分别独立），仅共享大部分底层。同时，对共享参数部分使用更小的学习率，并对来自 Actor 和 Critic 的梯度进行缩放或动态平衡。

四个模型在单次 PPO 更新中，分别进行多少次前向和反向传播？分析计算量分布。¶

我们以标准的一次 PPO 迭代（采样一个 batch，然后进行 K 个 epoch 的 mini-batch 更新）为例来分析。

采样阶段
Actor 模型：进行一次前向传播（自回归生成），无反向传播。它需要生成序列中的每个 token，计算量巨大，且受序列长度影响。
Reference 模型：在奖励计算时，对生成序列执行一次前向传播，计算每个 token 的对数概率。这是为了计算 KL 惩罚，无反向传播。
Reward Model：在奖励计算时，对完整序列执行一次前向传播，输出一个标量奖励。无反向传播。
Critic 模型：此阶段无计算。

计算量分布（采样阶段）：计算量主要集中在 Actor 的自回归生成上，占大头。Reference 和 RM 的推理是一次性编码整个序列，计算量相对小得多，但仍然是不可忽略的固定开销。

PPO 训练阶段
我们将采样得到的经验数据（prompt, response, logπ_old, reward 等）重复用于 K 个 epoch 的训练。
在每个 epoch 的每个 mini-batch 中：
Actor 模型：进行一次前向传播（计算新策略的对数概率）和一次反向传播（计算 PPO 损失的梯度，更新参数）。前向传播不需要生成，只需要计算已有 token 的概率，因此比采样阶段快很多。
Critic 模型：进行一次前向传播（计算每个 token 的价值 V(s_t)）和一次反向传播（计算价值损失的梯度，更新参数）。

计算量分布（训练阶段）：Actor 和 Critic 的前向和反向计算量大致相当（取决于网络相对大小），都是计算瓶颈。由于要进行 K 次重复，这个阶段的总计算量通常远大于采样阶段。

总计：在单次完整的 PPO 迭代中，Actor 进行了 1 (gen) + K * 1 (fw) + K * 1 (bw) 次计算；Critic 进行了 K * 1 (fw) + K * 1 (bw) 次；Reference 和 RM 各进行了 1 次 fw。计算成本从高到低大约是：Actor（训练+采样）> Critic >> Reference ≈ RM。

如果 Reward Model 的输出不是一个值而是一个向量（如分段奖励），Critic 应如何设计？¶

设计一：线性加权后再接标量 Critic 这是最直接的做法。在训练 PPO 时，我们用一个固定的权重向量 w 将多维度奖励融合成一个标量 r=w⊤r。然后，Critic 仍然保持原样，即输出一个标量价值 V(s)，去拟合这个标量化的综合回报。这允许我们在推理时通过调整权重来动态改变模型的偏好（例如儿童模式更看重安全性）。

设计二：向量价值 Critic (Vectorized Critic)

设计三：多头共享底座 Critic

Critic 主体仍然是一个 Transformer，但顶部有多个标量头，每个头对应 RM 的一个维度。它们共享底层表示，但各自独立预测各维度的价值。这兼具了参数效率和多维评估的精细度。

关键考虑：

维度相关性：不同维度的奖励可能具有相关性（如流畅的回答通常也有用）。向量价值 Critic 可以通过参数共享来捕捉这些关系。
训练难度：多维输出的训练目标更复杂，可能需要进行维度间的平衡，防止某个维度的损失主导训练。同时，需要引入多维度的价值裁剪。
信用分配：多维价值可以让我们更细粒度地分析模型的行为。例如，Actor 的哪些动作提高了安全性，却可能损害了有用性，这在单标量 Critic 下是无法区分的。

为什么 Reference 模型有时也被称为“先验策略”？它在 KL 散度计算中充当什么角色？¶

在贝叶斯统计和强化学习的交汇处，“先验”一词有着深刻的含义，Reference 模型正是这样一个角色。

Reference 模型作为先验策略

在 RLHF 中，我们使用 SFT 模型作为 RL 微调的起点，它不仅提供了初始参数，更代表了我们在见到任何用户偏好数据之前，对“好的生成策略应该是什么样”的初始信念。这个初始信念源自海量的互联网文本和高质量的指令数据，它编码了语法、常识、世界知识等一切让语言“像人话”的要素。因此，Reference 模型就是我们的“先验策略”。

在 KL 散度中的角色——正则化先验

PPO 的目标是最大化奖励，同时最小化与 Reference 模型的 KL 散度。这等价于在贝叶斯框架下求解最大后验概率（MAP）估计。奖励模型的对数概率构成了“似然项”（数据告诉我们应该朝哪个方向优化），而 KL 散度则构成了“先验项”（初始信念告诉我们应该保持在什么样的分布附近）。

动态视角：如果 Reference 模型随着训练更新，就相当于在一个动态变化的先验上进行优化，这违背了贝叶斯推断的一致性前提。一个固定的先验确保了我们在整个 RL 过程中始终朝向一个稳定、可靠的“正常语言锚点”靠拢。

在线采样时，Actor 生成回答后，需要立即调用 Reward Model 和 Reference Model 计算奖励和 KL，这部分计算的延迟如何影响训练效率？¶

在线采样后的奖励和 KL 计算是一个典型的串行瓶颈，对训练效率有显著影响，因为它打断了 Actor 的采样流水线。

延迟的具体体现：

Actor 的生成是逐 token 自回归的，速度较慢，但完成后是一口气生成一个完整的序列。紧接着，系统必须阻塞等待 Reference 和 RM 的前向传播结果，才能组装出完整的训练数据。

RM 的推理：RM 需要对完整的序列进行一次深度 Transformer 编码，对于长序列，这需要不小的计算和显存。
Reference 的推理：同样是一次完整的编码，并且需要逐 token 地输出对数概率，计算量更大。这两个模型的计算都是必须的，且通常是在相同的 GPU 上以串行方式执行（除非用多 GPU 流水线）。此时，Actor 是空闲的，GPU 计算资源在 RM 和 Reference 之间切换，导致了计算核心的空泡和利用率下降。

对训练效率的冲击：

墙钟时间增加：在千亿参数模型下，RM 和 Reference 的一次前向传播可能耗时零点几秒到数秒。如果 batch size 较大，累积的推理时间会显著拖慢单次迭代的周期。
通信开销：如果使用分布式训练，生成的序列数据可能需要在 GPU 间复制或广播，增加了通信延迟。
优化手段：
管道并行与异步：将采样和奖励计算放在不同 GPU 上，形成流水线。例如，Actor 在 GPU A 上生成，生成完毕后数据立刻发送到 GPU B 上的 RM/Ref 进行计算，同时 GPU A 立即开始下一个 batch 的采样。这可以隐藏计算延迟。
共享底座与批量推理：将 RM 和 Reference 的部分层与 Actor 共享（如果架构允许），或者将多个序列打包成更大的 batch 进行一次前向传播，提高 GPU 的利用率，摊销延迟。
模型量化与优化：对 RM 和 Reference 进行 INT8 量化或使用 TensorRT 推理加速，减小其计算时间。

是否可以将 Reference Model 替换为一个固定的均匀分布或其他先验？会有什么效果？¶

理论上可以替换，但效果将与原版 RLHF 南辕北辙，导致生成质量灾难性地退化。

替换为均匀分布：

效果：Actor 会极力避免任何 token 的概率变得太高，因为任何集中都会增加 KL 散度。结果将是，模型生成完全随机、混乱的文本，毫无语义可言。RL 优化会变成“如何在乱码中寻宝”，完全失败。SFT 阶段所学的一切语言能力都将被均匀分布的惩罚项冲刷殆尽。

替换为另一固定先验（如旧的 SFT 版本）：

如果替换为一个更早的、能力更弱的 SFT 模型，那么 KL 惩罚会将策略向那个较弱的模型拉扯。这等于是在训练过程中故意遗忘新学的知识，迫使策略倒退。优化过程会陷入奖励提升与能力倒退的拉锯战，最终得到一个“对齐但弱智”的模型。

原版 Reference（当前 SFT 模型）的不可替代性：

当前 SFT 模型是当前策略的“最优初始点”。以它为 Reference，意味着我们只惩罚策略在当前基础上的过激偏移，而不是逼它退回一个遥远的、能力更差的分布。它既保留了最强的语言能力基线，又允许策略在这个基线的坚实基础上进行局部的、安全的探索。用任何其他先验替代，都是在用一套错误的坐标系来导航，轻则让模型退化为平庸，重则直接导致语言崩溃。

如果 Critic 模型欠拟合，对 Actor 的训练会产生什么影响？会导致高方差还是偏差？¶

更危险的是系统性偏差

欠拟合更常见的表现是 Critic 学到了一些简单但错误的模式，形成了系统性偏差。例如，Critic 可能只是学会了根据序列长度来给出价值估计（因为长的回答在训练数据中往往得分高），而忽略了内容本身的优劣。这种偏差会形成错误的方向感。它会持续给一些质量平庸但恰好符合其偏差模式的回答以过高的优势，引导 Actor 涌向这个错误的“洼地”，形成奖励黑客。此时，即使梯度方差不大，Actor 也在被平稳地、一步步地带入歧途，这比高方差但零均值的噪声更加危险。

后果：

策略优化目标扭曲，Actor 学习到的是“如何迎合有偏差的 Critic”，而不是“如何生成对人类更好的回答”。
信用分配完全失效。一个关键的优秀 token 可能因为处在 Critic 不看好的位置而被惩罚，一个无关紧要的 token 却可能被过度奖励。
最终，Actor 的性能上限被 Critic 的拟合能力锁死。一个弱的 Critic 无法为 Actor 提供攀登更高性能山峰的精细阶梯。

缓解措施：

保证 Critic 的学习率足够高，并且在其损失函数上使用足够的正则化。
动态监控 Critic 损失，如果持续高位且不下降，需要调整其架构或训练策略。
对优势进行批次归一化，可以减轻由于 Critic 输出的尺度漂移带来的部分方差。

在训练循环中，Critic 的更新频率和 Actor 的更新频率是否一致？为什么通常设置一致的步数？¶

在标准的 PPO 实现中，Critic 和 Actor 的更新频率是一致的，即每个 mini-batch 同时计算两者的损失，并执行一次参数更新。

为什么通常设置一致？

架构的耦合性：在大多数实现中，Actor 和 Critic 要么共享底座，要么独立但结构类似。同步更新在工程上最简单、最直接。每次新数据进来，我们都需要同时获得最新的策略（来自 Actor）和最新的价值估计（来自 Critic）。如果两者更新不同步，那么价值估计将基于旧策略，而策略更新又基于旧的价值估计，这会引入额外的 off-policy 误差，使训练不稳定。
PPO 的 on-policy 特性：PPO 使用同一批由当前策略（旧策略快照）生成的数据进行更新。在重复利用这批数据的几个 epoch 内，Actor 和 Critic 都在基于相同的、最近一次采样得到的分布进行优化。保持同步更新，确保了两者始终在同一个“数据页”上。

异步更新的特例及风险：

有些研究中会尝试让 Critic 更新得更频繁（例如 Critic:Actor 更新步数比为 5:1），目的是为了让价值估计更快地收敛，减少策略梯度方差。然而在 RLHF 中，由于：

采样代价极高，我们不想浪费宝贵的数据。
语言模型对价值估计误差非常敏感。过快地更新 Critic 可能会导致它对当前策略分布过拟合，反而在策略稍微变化后就变得不准。因此，标准的 1:1 同步更新是最稳健的选择，它避免了因更新频率不匹配而引入的复杂调参和潜在不稳定。

PPO 更新时，Critic 的梯度是否会反传到共享底座，从而影响 Actor 的特征？如果共享，如何平衡？¶

是的，当 Actor 和 Critic 共享底座时，Critic 的梯度会直接反传到共享底座，深刻地影响 Actor 使用的特征表示。这种影响既可能是正面的促进作用，也可能是毁灭性的干扰。

影响的本质：

共享底座意味着特征空间同时用于两个目标：生成下一个 token（Actor）和评估当前状态的价值（Critic）。Critic 的损失优化会“拉扯”底座参数，使其学习到的特征对 Critic 的价值回归任务更有利。如果这些特征对 Actor 的决策也有帮助，那就形成了双赢。但现实中，二者经常冲突：Critic 可能需要一些与内容质量间接相关但易于回归的特征（如长度、某些标识词的出现频率），而 Actor 为了生成有创造性、精确的回答，可能需要更精细的语义特征。Critic 的梯度会不断将底座朝其偏好方向修正，可能“洗掉”对 Actor 至关重要的特征。

如何平衡这种冲突？

差异化学习率：对共享底座使用一个极小的学习率（例如是独立头学习率的 1/10），减缓 Critic 对底座的修改速度。让底座主要追随 Actor 的更新，而 Critic 主要调整自己的独立价值头来适应底座。
梯度缩放与掩码：在反向传播时，对来自 Critic 的梯度乘以一个小于 1 的系数（如 0.5），削弱其对共享参数的牵引力。极端情况下，可以让 Critic 的梯度完全不更新共享底座，只更新其独立的价值头。这牺牲了 Critic 的表征适应性，但保护了 Actor 的特征稳定性。
交替训练：不进行联合更新，而是在几个 mini-batch 中只更新 Critic（冻结共享底座），接着在几个 mini-batch 中只更新 Actor（冻结共享底座）。这避免了梯度在同一更新步中发生直接冲突，但实现复杂，且可能降低数据利用效率。
正则化解耦：在损失函数中加入鼓励底座特征正交或解耦的正则项，促使 Actor 和 Critic 使用底座特征的不同子空间，但这在实践中较难实现且效果不一定理想。

经验法则：在显存受限必须共享时，最常用且有效的平衡手段是给共享底座设置一个远小于独立头部学习率的学习率，并可以适当缩放来自 Critic 的梯度。

如何通过观察 Critic 的 loss 曲线来判断价值估计的准确性？¶

正常健康的曲线：

初期快速下降：在训练刚开始时，由于 Critic 从 RM 或 SFT 初始化，其价值估计并不准，损失会迅速下降。这表明模型在学习基本的状态价值。
缓慢下降并趋于稳定：在大部分训练过程中，Critic 损失应该呈缓慢下降趋势，并最终在一个较低的水平上波动。这说明价值估计在持续、渐进地改善。
无剧烈尖峰：整体曲线应该比较平滑，没有突然的、巨大的向上跳变。跳变通常意味着遇到了某个高度不可预测的状态分布或异常回报样本。

异常曲线与对应问题：

损失不下降或过早平台化：可能的原因是 Critic 学习率过低、模型容量不足、或者输入特征（共享底座）发生了剧烈变化（被 Actor 带偏）导致其无法拟合。这预示着价值估计存在系统性偏差，无法为 Actor 提供准确优势信号。
损失剧烈震荡，且幅度较大：这表明价值估计的方差很高。可能是训练数据中的回报本身波动过大（需要增大 batch size），或者价值网络对输入变化过于敏感（需要调整学习率或增加正则化）。这会导致 Actor 收到的优势信号充满噪声。
损失突然飙升并导致 NaN：这是典型的价值估计崩溃。可能是遇到了极端离群值的回报样本（例如 RM 偶尔给出了极高或极低分），或者学习率过大导致参数更新跳入一个恶劣区域。
损失很低但价值估计的方差很大：这是过拟合的迹象。Critic 在训练集上完美拟合，但在新的采样上表现极差（泛化能力弱）。这需要增加正则化（如 Dropout、L2 正则）或使用更大、更多样的训练数据。

高级监控：除了 loss，还可以定期记录 Critic 预测值与真实回报的 Spearman 相关系数。这个指标比 MSE 更能反映价值估计对状态好坏排序的正确性，是判断 Critic 是否真正“理解”价值的关键。

四个模型是否可以放在不同的 GPU 上？这会引入什么通信开销？¶

完全可以，并且这是大规模 RLHF 训练的标准配置，即所谓的“模型并行”和“流水线并行”。尤其在千亿参数模型下，单个 GPU 根本放不下四个模型。

典型的放置策略：

Actor 和 Critic：由于它们需要频繁地进行梯度同步，且是训练的计算核心，通常被放置在多个 GPU 上，使用张量并行或流水线并行。
Reference 模型和 Reward Model：它们是冻结的推理模型，计算量相对小。常常被单独放置在 1-2 个 GPU 上，或者与 Actor/Critic 共享部分 GPU 但逻辑隔离。

通信开销分析：

主要的通信瓶颈发生在从采样 GPU 到推理 GPU 的数据传输，以及分布式训练时模型间的梯度同步。

序列数据传输：Actor 生成完一批回复后，需要将这些文本序列（及其 token IDs）发送给 Reference 和 RM 所在的 GPU。这部分传输的原始数据量（通常是 token IDs 和 attention masks）相对模型参数来说较小，但 batch size 大或序列极长时，也可能成为不可忽视的延迟。通常通过 NCCL 高速互联进行。
奖励与 KL 结果回传：RM 和 Reference 计算出奖励分数和 token 级对数概率后，需要将这些小规模张量（标量或与序列长度相当的小张量）传回 Actor 的训练 GPU。这部分通信开销极小。
分布式训练同步：Actor 和 Critic 在各自的多 GPU 组内进行前向和反向传播时，需要进行梯度的 AllReduce 通信。这才是真正的通信大头。RM 和 Reference 不参与训练，因此不产生这部分开销。

优化策略：

异步流水线：Actor 采样完一批数据后，立刻开始下一批的采样，同时将当前批数据异步发送给 RM/Ref 进行计算。当下一批采样完成时，上一批的奖励和 KL 已经计算完毕，可以几乎无延迟地开始训练。这完全隐藏了推理模型的通信和计算延迟。
本地化部署：尽可能让 Reference 和 RM 模型与 Actor 共享一个节点内的 GPU，利用 NVLink 等高速互联降低传输延迟。

如果只能保留两个模型在 GPU 上，你会如何选择并设计 offload 策略？¶

这是一个极致的显存优化问题，意味着我们只能在 GPU 显存中驻留两个最核心的模型，其余模型必须通过 Offload 到 CPU 内存甚至 NVMe 硬盘上，以时间换空间。

选择保留的模型：Actor 和 Critic。

理由：只有这两个模型需要频繁地进行前向和反向传播以更新参数，是训练的计算核心。将它们驻留在 GPU 上，保证了训练循环（占整个 RLHF 流程 90% 以上的计算量）不受 Offload 的严重影响，保持了尽可能高的训练吞吐。Reference 和 RM 在整个训练过程中只进行少量的推理，计算频率低，可以被 Offload。

Offload 策略设计：

Reference 模型和 Reward Model 的 Offload：
在采样阶段，当 Actor 生成完毕一个 batch 的数据后，异步地将数据（token IDs）传输到 CPU 内存。
CPU 上部署着完整的 Reference 和 RM 模型（或它们的量化版本）。利用成熟的 CPU 推理框架（如 llama.cpp、ONNX Runtime），让它们在 CPU 上执行前向传播，计算所需的奖励和 KL 惩罚。
计算完成后，将得到的标量和 token 级奖励等小规模数据传回 GPU。
这个过程可以与下一个 batch 的采样重叠，完全隐藏 CPU 推理的延迟。只要 CPU 推理时间小于一个 batch 的生成时间，就不会对总吞吐造成影响。如果 CPU 推理成为瓶颈，可以对 RM 和 Reference 进行 INT8 量化加速。
优化器状态的 Offload (ZeRO-Offload)：
即使只保留了两个模型，它们的参数、梯度和优化器状态（如 Adam 的动量）仍然可能超出显存。此时，需要采用 ZeRO-Offload 的思想。
将优化器状态和 FP32 的主权重副本 Offload 到 CPU 内存。GPU 上仅保存 FP16（或 BF16）的模型参数和梯度。
在更新时，将 GPU 上的梯度传回 CPU，在 CPU 上完成优化器更新（如 Adam 的动量更新、权重衰减），然后将更新后的 FP16 权重传回 GPU。
这同样可以与 GPU 计算流水线化，用通信和 CPU 计算掩盖延迟。
检查点与重计算：
为了节省训练时激活占用的显存，可以对 Actor 和 Critic 启用激活检查点（gradient checkpointing）。这会增加部分前向计算时间，但能大幅降低显存占用。

风险与权衡：

大量依赖 CPU Offload 和异步通信，使得系统架构变得复杂，调试困难。CPU 与 GPU 之间的 PCIe 带宽是核心瓶颈。如果模型极大，即使 Offload 也无法实时处理，训练吞吐将显著下降。
我们假设 CPU 推理足够快，对于百亿参数级别的 RM 和 Reference，CPU 推理可能成为瓶颈，需要部署在多个 CPU 核上并进行模型并行，这将增加工程复杂度。

总结：保留 Actor 和 Critic 在 GPU，Offload Reference 和 RM 到 CPU，并配合 ZeRO-Offload 与异步流水线，是极端显存限制下最可行、对训练吞吐影响最小的方案。