worker.py`¶

文件路径¶

verl/single_controller/base/worker.py

文件概述¶

这个文件定义了 Worker 基类，代表分布式系统中的一个计算单元（通常对应一张 GPU）。Worker 负责：

初始化分布式环境变量（RANK、WORLD_SIZE、MASTER_ADDR 等）
配置 GPU 设备（CUDA_VISIBLE_DEVICES）
管理分布式通信所需的元数据（dp_rank 映射、collect 信息等）
提供几个预注册的方法（如 execute_with_func_generator）

关键代码讲解¶

1. 分布式信息数据类¶

@dataclass
class DistRankInfo:
    tp_rank: int
    dp_rank: int
    pp_rank: int
    cp_rank: int

@dataclass
class DistGlobalInfo:
    tp_size: int
    dp_size: int
    pp_size: int
    cp_size: int

这两个 dataclass 用于描述分布式训练中的多维并行信息：

缩写	含义	说明
`tp`	Tensor Parallelism	张量并行，把模型的某一层切分到多张卡上
`dp`	Data Parallelism	数据并行，每张卡处理不同的数据
`pp`	Pipeline Parallelism	流水线并行，把模型的不同层放到不同卡上
`cp`	Context Parallelism	上下文并行，用于处理超长序列

DistRankInfo 表示当前 Worker 在各维度上的排名，DistGlobalInfo 表示各维度的总大小。

2. WorkerHelper 辅助类¶

class WorkerHelper:
    @staticmethod
    def _get_node_ip():
        if os.getenv("WG_BACKEND", None) == "ray":
            return ray.util.get_node_ip_address()
        else:
            raise NotImplementedError("WG_BACKEND now just support ray mode.")

    @staticmethod
    def _get_free_port():
        with socket.socket() as sock:
            sock.bind(("", 0))
            return sock.getsockname()[1]

    def get_available_master_addr_port(self):
        return self._get_node_ip().strip("[]"), str(self._get_free_port())

辅助方法用于获取当前节点 IP 和空闲端口，供分布式通信（如 PyTorch 的 init_process_group）使用。

_get_node_ip()：通过 Ray API 获取当前节点 IP
_get_free_port()：绑定 port 0 让操作系统分配一个空闲端口
get_available_master_addr_port()：组合上面两个方法

3. Worker 类初始化¶

class Worker(WorkerHelper):
    fused_worker_attr_name = "fused_worker_dict"

    def __init__(self, cuda_visible_devices=None) -> None:
        import os
        self._setup_env_cuda_visible_devices()

        world_size = int(os.environ["WORLD_SIZE"])
        rank = int(os.environ["RANK"])
        self._rank = rank
        self._world_size = world_size

        master_addr = os.environ["MASTER_ADDR"]
        master_port = os.environ["MASTER_PORT"]

        local_world_size = int(os.getenv("LOCAL_WORLD_SIZE", "1"))
        local_rank = int(os.getenv("LOCAL_RANK", "0"))

        store = {
            "_world_size": world_size,
            "_rank": rank,
            "_local_world_size": local_world_size,
            "_local_rank": local_rank,
            "_master_addr": master_addr,
            "_master_port": master_port,
        }
        if cuda_visible_devices is not None:
            store[f"_{get_visible_devices_keyword()}".lower()] = cuda_visible_devices

        self._configure_with_store(store=store)

        self.fused_worker_dict = {}
        self.__dispatch_dp_rank = {}
        self.__collect_dp_rank = {}

初始化流程： 1. 配置 GPU 设备：调用 _setup_env_cuda_visible_devices() 处理 CUDA/HIP/ROCR 设备环境变量 2. 读取分布式环境变量：从环境变量获取 RANK、WORLD_SIZE 等信息（这些环境变量由 WorkerGroup 在创建 Worker 时设置） 3. 构建配置存储：将所有分布式配置打包成字典 4. 应用配置：通过 _configure_with_store 设置环境变量 5. 初始化内部状态： - fused_worker_dict：用于 FusedWorker 模式下存储共存的子 Worker - __dispatch_dp_rank：记录各 mesh 下的 dp_rank 分发映射 - __collect_dp_rank：记录各 mesh 下是否需要收集结果

4. GPU 设备环境变量处理¶

def _setup_env_cuda_visible_devices(self):
    from verl.utils.ray_utils import ray_noset_visible_devices
    is_ray_noset_visible_devices = ray_noset_visible_devices()

    rocr_val = os.environ.get("ROCR_VISIBLE_DEVICES", None)
    hip_val = os.environ.get("HIP_VISIBLE_DEVICES", None)
    cuda_val = os.environ.get("CUDA_VISIBLE_DEVICES", None)

    if hip_val:
        # 将 HIP_VISIBLE_DEVICES 统一转换为 CUDA_VISIBLE_DEVICES
        val = os.environ.pop("HIP_VISIBLE_DEVICES")
        ...

    if rocr_val:
        # ROCR_VISIBLE_DEVICES 和 CUDA_VISIBLE_DEVICES 不能同时设置
        if cuda_val:
            raise ValueError(...)
        cuda_val = os.environ.pop("ROCR_VISIBLE_DEVICES")
        os.environ["CUDA_VISIBLE_DEVICES"] = cuda_val

    if is_ray_noset_visible_devices:
        # 如果 Ray 没有自动设置设备，手动从 Ray 运行时获取
        local_rank = ray.get_runtime_context().get_accelerator_ids()[device_name][0]
        os.environ["LOCAL_RANK"] = local_rank
        get_torch_device().set_device(int(local_rank))

这段代码处理了三种 GPU 环境变量的兼容性问题： - CUDA_VISIBLE_DEVICES：NVIDIA GPU 标准 - HIP_VISIBLE_DEVICES：AMD GPU（ROCm） - ROCR_VISIBLE_DEVICES：AMD GPU 的底层 ROCr 运行时

核心目标是统一使用 CUDA_VISIBLE_DEVICES，避免不同环境变量之间的冲突。

5. 环境变量配置¶

def _configure_with_store(self, store: dict):
    store_env_dict = {f"_{key.lower()}": store.get(f"_{key.lower()}", None) for key in type(self).env_keys()}
    self.__dict__.update(store_env_dict)
    for key in type(self).env_keys():
        val = self.__dict__.get(f"_{key.lower()}", None)
        if val is not None:
            os.environ[key] = str(val)
    os.environ["REDIS_STORE_SERVER_HOST"] = (
        str(self._master_addr).replace("[", "").replace("]", "") if self._master_addr else ""
    )

这个方法将配置字典中的值同时写入： 1. Worker 实例的属性（self.__dict__） 2. 进程的环境变量（os.environ）

这样后续的 PyTorch 分布式初始化可以直接读取环境变量。

6. Dispatch 和 Collect 信息管理¶

def _register_dispatch_collect_info(self, mesh_name: str, dp_rank: int, is_collect: bool):
    if mesh_name in self.__dispatch_dp_rank or mesh_name in self.__collect_dp_rank:
        raise ValueError(f"mesh_name {mesh_name} has been registered")
    self.__dispatch_dp_rank[mesh_name] = dp_rank
    self.__collect_dp_rank[mesh_name] = is_collect

@register(dispatch_mode=Dispatch.ONE_TO_ALL)
def _query_dispatch_info(self, mesh_name: str):
    assert mesh_name in self.__dispatch_dp_rank
    return self.__dispatch_dp_rank[mesh_name]

在多模型共存场景（如 Actor 和 Critic 共享 GPU），不同模型可能有不同的 DP 划分。每个 Worker 需要知道在某个"mesh"（逻辑网格）中自己的 dp_rank 是什么。

_register_dispatch_collect_info：注册某个 mesh 下的 dp_rank 和是否收集
_query_dispatch_info：被 WorkerGroup 远程调用，查询本 Worker 在某个 mesh 下的 dp_rank
_query_collect_info：被 WorkerGroup 远程调用，查询本 Worker 在某个 mesh 下是否需要收集结果

注意 _query_dispatch_info 用了 @register(dispatch_mode=Dispatch.ONE_TO_ALL)，意味着 WorkerGroup 调用时会把同一个 mesh_name 广播给所有 Worker。

7. 预注册的执行方法¶

@register(dispatch_mode=Dispatch.DP_COMPUTE_PROTO_WITH_FUNC)
def execute_with_func_generator(self, func, *args, **kwargs):
    ret_proto = func(self, *args, **kwargs)
    return ret_proto

@register(dispatch_mode=Dispatch.ALL_TO_ALL, execute_mode=Execute.RANK_ZERO)
def execute_func_rank_zero(self, func, *args, **kwargs):
    result = func(*args, **kwargs)
    return result

execute_with_func_generator：接收一个函数和参数，在 Worker 上执行该函数。使用 DP_COMPUTE_PROTO_WITH_FUNC 模式，会自动切分数据参数但把函数广播给所有 Worker。
execute_func_rank_zero：只在 rank 0 的 Worker 上执行给定函数。

核心类/函数列表¶

名称	类型	说明
`DistRankInfo`	dataclass	当前 Worker 在各并行维度的排名
`DistGlobalInfo`	dataclass	各并行维度的总大小
`WorkerHelper`	类	获取 IP 和端口的辅助方法
`Worker`	类	Worker 基类，管理分布式环境和计算
`Worker.__init__()`	方法	初始化环境变量和设备配置
`Worker._setup_env_cuda_visible_devices()`	方法	处理 GPU 设备环境变量兼容性
`Worker._configure_with_store()`	方法	将配置写入属性和环境变量
`Worker._register_dispatch_collect_info()`	方法	注册 mesh 的 dp_rank 信息
`Worker.execute_with_func_generator()`	方法	在 Worker 上执行给定函数
`Worker.execute_func_rank_zero()`	方法	在 rank 0 上执行给定函数

与其他模块的关系¶

继承自 WorkerHelper：获取网络通信辅助功能
使用 decorator.py：导入 Dispatch、Execute、register
被 worker_group.py 中的 WorkerGroup 管理：WorkerGroup 持有多个 Worker 的引用
被 ray/base.py 使用：Ray 后端通过 ray.remote 将 Worker 变成 Ray Actor
子类化：用户定义的 Worker（如 ActorWorker、CriticWorker）继承此类

小结¶

Worker 类是分布式计算的最小单元。它的核心职责是： 1. 在初始化时正确设置分布式环境变量（RANK、WORLD_SIZE 等） 2. 通过 @register 装饰器声明方法的分发策略 3. 管理 mesh 维度的 dp_rank 映射信息

用户编写自己的 Worker 时，只需要继承 Worker，然后用 @register 装饰需要远程调用的方法即可。

worker.py — verl/single_controller/base/worker.py¶