ray-ascend:Ascend-native hardware plugin for ray

Ray Ascend 插件

| 关于昇腾 | 文档 |

概述

ray-ascend 是一个社区维护的硬件插件，支持在昇腾 NPU 加速器上运行高级 Ray 功能。

默认情况下，Ray 原生支持将昇腾 NPU 作为预定义资源类型，用于绑定 actor 和任务（详见 Ray 加速器支持）。作为增强工具，ray-ascend 提供了 Ray 上的昇腾原生功能，例如通过 华为集合通信库（HCCL） 实现的集合通信、Ray 直接传输（RDT） 等。

性能基准测试详见 性能基准测试报告。

前提条件

• 架构：aarch64、x86
• 操作系统内核：Linux
• Python 依赖：
  • python >= 3.10，<= 3.11
  • CANN >= 8.2.rc1
  • torch >= 2.7.1；torch-npu >= 2.7.1.post2
  • torch 与 torch-npu 版本需相互兼容。
  • ray >= 2.55.0

快速开始

安装

pip install "ray-ascend[yr]"

Ray Actor 间的 HCCL 集合通信

import ray
from ray.util import collective
from ray_ascend import register_hccl_collective_backend

register_hccl_collective_backend()

@ray.remote(resources={"NPU": 1})
class RayActor:
    def __init__(self):
        register_hccl_collective_backend()

collective.create_collective_group(
    actors,
    len(actors),
    list(range(0, len(actors))),
    backend="HCCL",
    group_name="my_group",
)

# Each actor broadcasts in SPMD manner
collective.broadcast(tensor, src_rank=0, group_name="my_group")

通过HCCS传输昇腾NPU张量

import ray
import torch
from ray.util.collective import create_collective_group
from ray_ascend import register_hccl_tensor_transport

register_hccl_tensor_transport()

@ray.remote(resources={"NPU": 1})
class RayActor:
    def __init__(self):
        register_hccl_tensor_transport()

    @ray.method(tensor_transport="HCCL")
    def random_tensor(self):
        return torch.zeros(1024, device="npu")

    def sum(self, tensor: torch.Tensor):
        return torch.sum(tensor)


sender, receiver = RayActor.remote(), RayActor.remote()
group = create_collective_group([sender, receiver], backend="HCCL")

tensor = sender.random_tensor.remote()
result = receiver.sum.remote(tensor)
ray.get(result)

通过 HCCS 传输昇腾 NPU 张量，通过 RDMA 传输 CPU 张量

开放元融数据系统 （YR）允许用户使用 Ray 对象传输 NPU 张量（通过 HCCS）和 CPU 张量（若提供，则通过 RDMA）。

import ray
from ray_ascend import register_yr_tensor_transport

register_yr_tensor_transport(["npu", "cpu"])

@ray.remote(resources={"NPU": 1})
class RayActor:
    def __init__(self):
        register_yr_tensor_transport(["npu", "cpu"])

    @ray.method(tensor_transport="YR")
    def transfer_npu_tensor_via_hccs(self):
        return torch.zeros(1024, device="npu")

    @ray.method(tensor_transport="YR")
    def transfer_cpu_tensor_via_rdma(self):
        return torch.zeros(1024)

sender = RayActor.remote()
npu_tensor = ray.get(sender.transfer_npu_tensor_via_hccs.remote())
cpu_tensor = ray.get(sender.transfer_cpu_tensor_via_rdma.remote())

Ray 版本兼容性

Ray 版本	YR 传输	HCCL 集合通信	HCCL 张量传输（RDT）
>=2.55, <2.56	✅	❌	❌
>= 2.56	✅	✅	✅

贡献指南

有关详细信息，请参见 CONTRIBUTING 和 开发者指南——这是一份逐步指南，可帮助您设置开发环境、构建和测试。如果您发现错误或需要新功能，请通过 提交 issue 告知我们。

许可证

Apache License 2.0。详见 LICENSE 文件。