如何在 ExperienceOperator 中使用空闲 GPU 训练辅助模型？

[SYSUzhouting]

背景与动机

Trinity-RFT 开发者您好！

我们正在使用 Trinity-RFT 框架进行一些关于RLHF的前沿研究。在我们的探索中，一个核心需求是能够基于 Experience 数据流，动态地训练和更新一个辅助小模型，用以实现更复杂的经验筛选或奖励修正策略。这对于扩展框架的研究能力、满足更多样的学术探索非常有价值。

目前，我们尝试在 trinity/buffer/operators 中实现一个自定义算子，该算子会初始化一个小模型，并在 process 方法中周期性地利用历史经验数据对其进行训练。

问题描述

我们发现，在这个自定义算子中，无法将待训练的小模型的训练任务加载到指定的 GPU 上。如果将模型和数据强制 to(device)，即使该设备在物理上是空闲的，程序也会立即报错并中断。这导致我们只能在 CPU 上进行训练，对于需要频繁更新的模型来说，效率非常低下。

复现概述

1. 我们新建了一个 OutlierRewardFilter，一个 ExperienceOperator，并尝试放到 GPU 上

training_device: str = "cuda:5"
self.preference_classifier = PreferenceClassifier(
            input_dim=self.input_hidden_state_dim, 
            hidden_layer_size=self.output_hidden_dim
        ).to(training_device)

2. 我们使用以下命令启动训练：

CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 trinity run --config xxx.yaml

cluster:
  node_num: 1
  gpu_per_node: 6

explorer:
  rollout_model:
    engine_num: 2
    tensor_parallel_size: 1

3. 报错信息

当 ExperienceOperator 被初始化时，程序崩溃，并抛出以下错误，指出对于该进程没有可用的 CUDA 设备：

可能原因和需求

我们猜测，这个问题可能与 Trinity-RFT（及其底层的 Ray 框架）的分布式资源管理机制有关；例如在 trinity run 启动时，框架可能会为每个工作进程（Ray Actor）设置独立的 CUDA_VISIBLE_DEVICES 环境变量，导致所有GPU都被锁了，无法再次进行使用。

希望能请教一下这个训练问题的原因，以及是否有比较合适的解决方案，谢谢！

以下是报错内容附图

[pan-x-c]

由于 vLLM 和 verl 都有自己独立的 GPU 调度机制，所以严格来说被分配给 vLLM 和 verl 的 GPU 都无法被其他进程共享，否则可能会导致 vLLM 和 verl 自身出现 OOM 或是其他异常现象。

一种可能的补救方法是将待训练的小模型封装为 restful API 服务，在另一台机器上部署，然后在 Operator 中通过网络发送参数并启动训练。

另一种可行方法是通过合理配置 cluster.gpu_per_node, explorer.rollout_model.engine_num, explorer.rollout_model.tensor_parallel_size，使得ray集群中能够空出 1 卡用于部署其他模型。

例如在 8 卡机器上设置 cluster.gpu_per_node 为 7, explorer.rollout_model.engine_num 为 3, explorer.rollout_model.tensor_parallel_size 为 1，即可实现 trinity 仅占用 7 卡，剩下一卡能够部署其他模型。但该方式想要实际生效还需要对 trinity.explorer.explorer.py 中的 _init_experience_pipeline 做硬编码式修改，为其设置 num_gpus=1

        return (
            ray.remote(ExperiencePipeline)
            .options(
                name=f"{self.config.explorer.name}_pipeline",
                namespace=ray.get_runtime_context().namespace,
                num_gpus=1,
                scheduling_strategy=NodeAffinitySchedulingStrategy(
                    node_id=node_id,
                    soft=False,
                ),
            )
            .remote(self.config)
        )