engine_workers.py — DetachActorWorker 扩展了 ActorRolloutRefWorker

文件路径: verl/experimental/separation/engine_workers.py

文件概述

DetachActorWorker 扩展了 ActorRolloutRefWorker，增加了将模型在 GPU 和 CPU 之间动态迁移的能力。这在"分离式训练"（Separation Training）中至关重要：同一组 GPU 在不同阶段承担不同角色（训练 vs 推理），需要在阶段切换时释放/恢复模型。

核心类：DetachActorWorker

继承关系

ActorRolloutRefWorker（基础 Actor Worker）
    └── DetachActorWorker（增加模型迁移能力）

三个核心方法

1. save_model_to_cpu - 将模型保存到 CPU

def save_model_to_cpu(self):
    """将 GPU 上的模型参数保存到 CPU 内存"""
    if self.strategy == "fsdp":
        # FSDP 策略：需要先聚合分片参数，再复制到 CPU
        with FSDP.summon_full_params(self.model):
            self.cpu_state_dict = {
                k: v.cpu().clone() for k, v in self.model.state_dict().items()
            }
    elif self.strategy == "fsdp2":
        # FSDP2 策略：类似但 API 不同
        ...
    elif self.strategy == "megatron":
        # Megatron 策略：需要处理模型并行
        ...

为什么需要这个？ 当 GPU 需要用于其他模型（如 Critic 或 Rollout）时，先把当前模型参数安全保存到 CPU，然后释放 GPU 显存。

2. restore_model_from_cpu - 从 CPU 恢复模型

def restore_model_from_cpu(self):
    """将 CPU 上保存的模型参数恢复到 GPU"""
    if self.cpu_state_dict is None:
        raise RuntimeError("No saved CPU state dict found!")

    # 加载保存的参数
    self.model.load_state_dict(self.cpu_state_dict)

    # 清理 CPU 副本以节省内存
    self.clear_cpu_model()

3. clear_cpu_model - 清理 CPU 副本

def clear_cpu_model(self):
    """释放 CPU 上保存的模型副本"""
    if self.cpu_state_dict is not None:
        del self.cpu_state_dict
        self.cpu_state_dict = None
        import gc
        gc.collect()

支持的分布式策略

策略	说明	特殊处理
FSDP	全分片数据并行	需要 summon_full_params 聚合分片
FSDP2	FSDP 第二版	API 略有不同
Megatron	Megatron-LM 并行	需要处理张量/流水线并行

使用场景

在分离式训练中，一组 GPU 可能需要：

阶段1: 运行 Actor 模型做推理（rollout）
        ↓ save_model_to_cpu()
阶段2: 运行 Critic 模型做训练
        ↓ restore_model_from_cpu()
阶段3: 运行 Actor 模型做训练

通过在阶段之间迁移模型，可以用更少的 GPU 完成原本需要更多 GPU 的训练任务。

核心类/函数列表

名称	类型	说明
DetachActorWorker	类	可分离的 Actor Worker
save_model_to_cpu	方法	GPU -> CPU 保存
restore_model_from_cpu	方法	CPU -> GPU 恢复
clear_cpu_model	方法	释放 CPU 副本

与其他模块的关系

• 继承自 verl.workers.actor.ActorRolloutRefWorker
• 被 SeparateRayPPOTrainer（ray_trainer.py）使用
• 与 FSDP/FSDP2/Megatron 分布式策略紧密相关

小结

DetachActorWorker 的核心价值是时分复用 GPU 资源。通过在 GPU 和 CPU 之间迁移模型，使得有限的 GPU 可以在不同训练阶段服务不同的模型，从而实现更高效的资源利用。