`device.py` — 设备抽象层¶

文件路径: verl/utils/device.py

文件概述¶

device.py 是 verl 框架的设备抽象层，让框架可以在不同硬件加速器（NVIDIA GPU / 华为 Ascend NPU / CPU）上统一运行。它是几乎所有其他模块的底层依赖。

背景知识¶

在深度学习中，模型训练通常在 GPU（NVIDIA CUDA）上进行。但近年来华为的 Ascend NPU 也逐渐被使用。verl 通过这个文件抽象了硬件差异，使得上层代码不需要关心具体使用的是哪种设备。

核心函数详解¶

1. 设备可用性检测¶

def is_torch_npu_available(check_device=True) -> bool:
    try:
        if not hasattr(torch, "npu"):
            return False
        if check_device:
            return torch.npu.is_available()
        else:
            return True
    except ImportError:
        return False

is_cuda_available = torch.cuda.is_available()
is_npu_available = is_torch_npu_available()

框架启动时，首先检查 CUDA 和 NPU 的可用性，结果缓存为模块级变量。check_device=False 表示只检查是否安装了 torch_npu 包，不要求设备真正存在。

2. 获取设备名称和模块¶

def get_device_name() -> str:
    if is_cuda_available:
        device = "cuda"
    elif is_npu_available:
        device = "npu"
    else:
        device = "cpu"
    return device

def get_torch_device():
    device_name = get_device_name()
    try:
        return getattr(torch, device_name)
    except AttributeError:
        return torch.cuda

get_device_name() 返回字符串如 "cuda"，而 get_torch_device() 返回对应的 torch 模块（如 torch.cuda），可以用来调用 memory_allocated()、current_device() 等方法。

3. 分布式通信后端¶

def get_nccl_backend() -> str:
    if is_npu_available:
        return "hccl"   # 华为通信库
    else:
        return "nccl"   # NVIDIA 通信库

分布式训练需要 GPU 间通信，NVIDIA 用 NCCL，华为 NPU 用 HCCL。

4. 显存管理¶

def set_expandable_segments(enable: bool) -> None:
    if is_cuda_available:
        torch.cuda.memory._set_allocator_settings(f"expandable_segments:{enable}")

开启 expandable segments 可以让 CUDA 内存分配器扩展已有的内存段，而不是分配新的，有助于避免 OOM。

5. IPC 支持检测（NPU 专用）¶

def is_support_ipc() -> bool:
    if is_cuda_available:
        return True
    if is_npu_available:
        software_version, cann_version = get_npu_versions()
        return check_ipc_version_support(software_version, cann_version)
    return False

IPC (Inter-Process Communication) 在 NPU 上需要特定版本的驱动支持。GPU 始终支持。

核心函数/类列表¶

函数	说明
`is_torch_npu_available()`	检查 NPU 是否可用
`get_device_name()`	返回设备名（cuda/npu/cpu）
`get_torch_device()`	返回 torch 设备模块
`get_device_id()`	返回当前设备索引
`get_nccl_backend()`	返回通信后端名
`get_resource_name()`	返回 Ray 资源名（GPU/NPU）
`get_visible_devices_keyword()`	返回可见设备环境变量名
`set_expandable_segments()`	设置内存分配器选项
`auto_set_device()`	自动配置设备
`get_device_capability()`	获取 GPU 算力版本
`is_support_ipc()`	检查 IPC 支持

与其他模块的关系¶

被几乎所有 verl/utils/ 下的模块依赖
profiler/performance.py 用它获取设备信息
megatron/memory.py 用它创建设备上的缓冲区
activation_offload.py 用它获取 stream 对象
distributed.py 用它选择通信后端

小结¶

device.py 是 verl 多设备支持的基石。它通过简洁的抽象函数，让上层代码以统一方式操作 CUDA GPU、华为 NPU 或 CPU，极大降低了跨平台适配成本。

device.py — 设备抽象层¶