跳转至

constants_ppo.py — 该文件定义了 PPO 训练中 Ray 运行时所需的环境变量常量

文件概述

模块路径: verl.trainer.constants_ppo

该文件定义了 PPO 训练中 Ray 运行时所需的环境变量常量。这些环境变量控制了 tokenizer 并行、NCCL 通信日志级别、vLLM 推理引擎配置等底层行为。在分布式训练启动时,这些环境变量会被注入到每个 Ray Worker 的运行时环境中。

在训练流程中的位置

这个文件在训练的最开始被使用:main_ppo.py 中的 run_ppo() 函数调用 get_ppo_ray_runtime_env() 来获取运行时环境变量,然后传给 ray.init() 完成 Ray 集群初始化。

关键代码讲解

1. 环境变量字典定义

PPO_RAY_RUNTIME_ENV = {
    "env_vars": {
        "TOKENIZERS_PARALLELISM": "true",        # 启用 tokenizer 并行
        "NCCL_DEBUG": "WARN",                    # NCCL 只打印警告级别日志
        "VLLM_LOGGING_LEVEL": "WARN",            # vLLM 只打印警告级别日志
        "VLLM_ALLOW_RUNTIME_LORA_UPDATING": "true",  # 允许运行时更新 LoRA 权重
        "CUDA_DEVICE_MAX_CONNECTIONS": "1",       # 限制 CUDA 设备最大连接数
        "NCCL_CUMEM_ENABLE": "0",                # 禁用 NCCL 统一内存,防止挂起
        "VLLM_DISABLE_COMPILE_CACHE": "1",       # 禁用 vLLM 编译缓存(避免缓存损坏)
        "HCCL_HOST_SOCKET_PORT_RANGE": "auto",   # 华为 NPU 相关配置
        "HCCL_NPU_SOCKET_PORT_RANGE": "auto",
    },
}

这些环境变量非常重要: - NCCL_CUMEM_ENABLE=0 防止在 Actor 和 Rollout 之间同步权重时出现挂起或崩溃 - CUDA_DEVICE_MAX_CONNECTIONS=1 是分布式训练的常见最佳实践

2. 环境变量过滤函数

def get_ppo_ray_runtime_env():
    """
    过滤函数:返回 PPO Ray 运行时环境。
    避免重复设置已存在的环境变量。
    """
    working_dir = (
        json.loads(os.environ.get(RAY_JOB_CONFIG_JSON_ENV_VAR, "{}"))
        .get("runtime_env", {})
        .get("working_dir", None)
    )

    runtime_env = {
        "env_vars": PPO_RAY_RUNTIME_ENV["env_vars"].copy(),
        **({"working_dir": None} if working_dir is None else {}),
    }
    # 移除已经在系统环境中存在的变量,避免覆盖用户的自定义设置
    for key in list(runtime_env["env_vars"].keys()):
        if os.environ.get(key) is not None:
            runtime_env["env_vars"].pop(key, None)
    return runtime_env

这个函数的巧妙之处在于:如果用户已经通过系统环境变量设置了某个值(比如 NCCL_DEBUG=INFO),它不会被默认值覆盖。这提供了灵活性。

核心类/函数列表

名称 类型 作用
PPO_RAY_RUNTIME_ENV dict 预定义的 Ray 运行时环境变量字典
get_ppo_ray_runtime_env() function 获取过滤后的运行时环境变量(避免覆盖已有设置)

数据流和调用关系

main_ppo.py: run_ppo()
    |
    +-- get_ppo_ray_runtime_env()  --> 获取环境变量
    |
    +-- ray.init(runtime_env=...)  --> 初始化 Ray 集群,注入环境变量

小结

这个文件虽然简短,但定义了分布式训练环境中至关重要的配置。理解这些环境变量有助于排查分布式训练中的通信问题、挂起问题和性能问题。