执行摘要

• 一句话：修复 ROCm Ray 权重传输 GPU 可见性问题
• 推荐动作：该 PR 属于 ROCm 平台特定 bugfix，值得相关平台开发者阅读，特别是 _weight_transfer_ray_env_vars 函数的设计可作为处理 ROCm+Ray GPU 可见性问题的通用模式。

功能与动机

在 ROCm 平台上，Ray 权重传输测试使用 NCCL/RCCL 进行跨进程通信。Ray 会设置 HI_VISIBLE_DEVICES/ROCR_VISIBLE_DEVICES，而 vLLM 通过 CUDA 兼容的 API 推理设备信息，导致 RCCL 在错误逻辑设备上初始化，引发未处理的 CUDA/HIP 错误。PR body 明确描述了此问题机制。

实现拆解

1. 新增环境变量设置函数 (_weight_transfer_ray_env_vars)：仅在 ROCm 平台返回三个 RAY_EXPERIMENTAL_NOSET_* 环境变量，阻止 Ray 修改 GPU 可见性；非 ROCm 平台返回空字典。
2. 初始化 Ray 辅助函数 (_init_ray_for_weight_transfer)：封装 ray.init 调用，传入 runtime_env 包含上述环境变量，避免重复初始化。
3. 设备发现与设置函数：_get_ray_assigned_device 通过 ray.get_gpu_ids() 获取实际分配的 GPU ID，兜底为 cuda:0；_set_ray_assigned_device 调用 torch.accelerator.set_device 确保 PyTorch 使用正确的设备。
4. 修改集成测试用例：在 trainer_broadcast_tensor、inference_receive_tensor、trainer_broadcast_sparse_tensor 等 Ray remote 函数中，使用 _set_ray_assigned_device() 获取设备，并替换所有硬编码 device=0 或 "cuda:0" 为动态值。同时将测试入口的 ray.init 替换为 _init_ray_for_weight_transfer。
5. 仅有测试文件变更，未触及任何生产代码或配置。

关键文件：

• tests/distributed/test_weight_transfer.py（模块 权重传输；类别 test；类型 test-coverage；符号 _weight_transfer_ray_env_vars, _init_ray_for_weight_transfer, _get_ray_assigned_device, _set_ray_assigned_device）: 唯一变更文件，包含所有新增的辅助函数和修改的集成测试逻辑，是 PR 的核心。

关键符号：_weight_transfer_ray_env_vars, _init_ray_for_weight_transfer, _get_ray_assigned_device, _set_ray_assigned_device

评论区精华

审核人 tjtanaa 在 _weight_transfer_ray_env_vars 函数处留下评论，建议将来将该函数移至公共工具文件，供其他 ROCm Ray 测试复用。开发者回应表示同意，但暂未在本 PR 中抽取为公共模块，仅作为测试内部辅助。

• 辅助函数抽取为公共工具 (design): 开发者同意但暂未抽取，保持当前测试文件内私有。

风险与影响

• 风险：本 PR 仅修改测试文件，不涉及生产代码，回归风险低。但需注意：torch.accelerator.set_device 在不同 PyTorch 版本上行为可能不一致，将来若 Ray 更新改变 RAY_EXPERIMENTAL_NOSET_* 行为，可能需要适配。
• 影响：仅影响 ROCm 平台上 Ray 权重传输集成测试的稳定性和正确性。对其他平台（NVIDIA、CPU、Intel GPU 等）无影响；对生产运行时无影响。
• 风险标记：仅测试文件变更, 平台特定逻辑（ROCm）, 依赖 Ray 实验性环境变量

关联脉络

• 暂无明显关联 PR