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 集群,注入环境变量
小结¶
这个文件虽然简短,但定义了分布式训练环境中至关重要的配置。理解这些环境变量有助于排查分布式训练中的通信问题、挂起问题和性能问题。