PyTorch FSDP 后端
======================

最后更新：2025/12/01。

我们通过实现用于 actor、critic、reference、rollout 和 reward 模型的各种 worker 来提供对 PyTorch FSDP 后端的支持。

**优势**

- 随时支持各种模型。

  - 用户只需要针对 FSDP 和 vLLM 之间的权重同步实现相应的
    ``dtensor_weight_loader``。而对于 ``hf_weight_loader``，用户可以直接应用任何同时在 HF 和 vLLM 中受支持的模型，而无需任何代码更改。

- 易于组织每个模型的前向和反向计算。

**劣势**

- 当处理大规模模型（例如 Llama 70B 和 405B）时，可扩展性较差。
- actor 和 rollout 之间的重新分片开销可能比 Megatron-LM 后端更大。

鉴于其简单性，我们推荐使用 FSDP 后端进行算法研究和原型验证。

FSDP Worker
--------------

ActorRolloutRefWorker
^^^^^^^^^^^^^^^^^^^^^

Actor/Rollout HybridEngine
''''''''''''''''''''''''''

1. HybridEngine、Actor 和 Rollout 初始化 API。

.. code:: python

   @register(dispatch_mode=Dispatch.ONE_TO_ALL)
   def init_model(self):

``ONE_TO_ALL``：当从驱动进程调用 ``init_model`` 函数时，每个 worker（位于 GPU 上）将执行以下模型初始化过程。

HybridEngine、Actor 和 Rollout 的初始化细节如下：

1. ``DataParallelPPOActor`` 实现了当模型基于 FSDP 构建时简单的 PPO 计算逻辑，包括计算对数概率、模型更新。
2. ``vLLMRollout`` 支持使用 vLLM 进行生成。我们修改了 vLLM Engine，使其在 SPMD 下执行，以适应我们的
   ``WorkerGroup`` 设计。

请参考 `源代码 <https://github.com/volcengine/verl/blob/main/verl/workers/fsdpw_workers.py>`_ 了解更多信息。

1. 生成序列并重新计算对数概率

.. code:: python

   @register(dispatch_mode=Dispatch.DP_COMPUTE_PROTO)
   def generate_sequences(self, prompts: DataProto):

- ``Dispatch.DP_COMPUTE_PROTO``：数据将沿着 DP 维度进行分发和收集

- 在此函数中，rollout 模型将执行自回归生成，而 actor 模型将重新计算生成响应的旧对数概率。

3. 更新 actor 模型

.. code:: python

   @register(dispatch_mode=Dispatch.DP_COMPUTE_PROTO)
   def update_actor(self, data: DataProto):

- 使用 PPO 和熵损失更新 actor 模型权重。

ReferenceModel
''''''''''''''

1. 参考模型初始化

参考模型使用与 actor 模型相同的函数进行初始化，但不初始化 HybridEngine 和 Optimizer。然后 actor 模型也被 ``DataParallelPPOActor`` 包装。

2. 计算参考对数概率

.. code:: python

   @register(dispatch_mode=Dispatch.DP_COMPUTE_PROTO)
   def compute_ref_log_prob(self, data: DataProto):

- 在此函数中，参考模型将调用 ``DataParallelPPOActor`` 中的计算对数概率函数来计算参考对数概率。

CriticWorker 和 RewardWorker
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

1. 模型初始化

与参考模型非常相似。CriticWorker 将为 Optimizer 执行额外的初始化。

2. 为 CriticWorker 计算值

.. code:: python

   @register(dispatch_mode=Dispatch.DP_COMPUTE_PROTO)
   def compute_values(self, data: DataProto):

3. 更新 Critic

.. code:: python

   @register(dispatch_mode=Dispatch.DP_COMPUTE_PROTO)
   def update_critic(self, data: DataProto):

4. 计算奖励

.. code:: python

   @register(dispatch_mode=Dispatch.DP_COMPUTE_PROTO)
   def compute_rm_score(self, data: DataProto):


HybridShard
------------

我们不支持 FSDP `HybridShard`。要支持此功能，我们可能需要构建一个 2D 设备网格，并为每个模型测试相应的
``dtensor_weight_loader`` 和 ``hf_weight_loader``。