全文 - 05 part - NVIDIA 集合通信库(NCCL)文档

05 语言绑定

NCCL 官方语言绑定的 API 规范。

NCCL 核心 API 使用 C/C++ 实现(参见 NCCL API)。语言绑定在该实现之上提供封装,使其他编程语言也能使用 NCCL 的功能。

  • Python 绑定(NCCL4Py)

05.01. Python 绑定(NCCL4Py)

  • Communicator
    • Communicator 类
    • 创建与生命周期方法
    • 集合通信方法
    • 点对点与信号方法
    • 内存注册方法
    • 设备通信器设置
    • 状态与工具方法
  • 配置
    • NCCLConfig
    • CTAPolicy
    • NCCLDevCommRequirements
    • 需求项(Requirement entries)
  • 组操作
    • group
    • group_start
    • group_end
    • GroupSimInfo
  • 内存管理
    • mem_alloc
    • mem_free
  • 通信器资源
    • CommResource
    • RegisteredBufferHandle
    • RegisteredWindowHandle
    • CustomRedOp
    • DevCommResource
    • 设备资源句柄
  • 类型与常量
    • 数据类型
    • 归约操作符
    • Team
    • 类型别名
    • 异常
  • 参数
    • params
    • dump_params
  • 版本
    • show_versions
    • get_version
    • VersionInfo
    • LibraryInfo
  • 框架互操作
    • CuPy
    • PyTorch

05.01.01. Python 绑定(NCCL4Py)

  • Communicator
    • Communicator 类
    • 创建与生命周期方法
    • 集合通信方法
    • 点对点与信号方法
    • 内存注册方法
    • 设备通信器设置
    • 状态与工具方法
  • 配置
    • NCCLConfig
    • CTAPolicy
    • NCCLDevCommRequirements
    • 需求项(Requirement entries)
  • 组操作
    • group
    • group_start
    • group_end
    • GroupSimInfo
  • 内存管理
    • mem_alloc
    • mem_free
  • 通信器资源
    • CommResource
    • RegisteredBufferHandle
    • RegisteredWindowHandle
    • CustomRedOp
    • DevCommResource
    • 设备资源句柄
  • 类型与常量
    • 数据类型
    • 归约操作符
    • Team
    • 类型别名
    • 异常
  • 参数
    • params
    • dump_params
  • 版本
    • show_versions
    • get_version
    • VersionInfo
    • LibraryInfo
  • 框架互操作
    • CuPy
    • PyTorch

05.01.01.01. Communicator 类

class nccl.core.Communicator(ptr: int = 0 )

基类:object

用于集合操作和点对点操作的 NCCL 通信器。

通信器代表一组可以执行集合操作(例如 allreduce、broadcast)和点对点操作(send/recv)的 rank。每个 rank 在 [0, nranks) 范围内拥有唯一的 ID。

Communicator 实例公开了若干用于检查的属性(ptrnranksdevicerank,以及与设备 API 相关的属性,如 cuda_devnvml_devdevice_api_supportmultimem_supportgin_typen_lsa_teamshost_rma_supportrailed_gin_type);详情请参阅各属性的文档。

init(ptr: int = 0 ) → None

使用原始 NCCL 指针初始化通信器。

与 init() 类方法不同,此构造函数允许 ptr=0,用于创建空通信器(例如当 split() 排除某个 rank 时)。空通信器之后可以通过 initialize() 进行初始化,或者作为 grow() 的调用者加入已有的通信器。

参数:

ptr ------ 表示 NCCL 通信器指针的整数(0 表示空通信器)。默认为 0。

属性

标识

Communicator.ptr

指向底层 ncclComm_t 结构的原始指针(已销毁或为空时为 0)。

Communicator.is_valid

通信器是否有效(未销毁且非空)。

Communicator.nranks

通信器中的 rank 总数。

抛出:

NcclInvalid ------ 如果通信器未初始化。

Communicator.device

与此通信器关联的 CUDA 设备。

返回一个 [cuda.core.Device](https://nvidia.github.io/cuda-python/cuda-core/latest/generated/cuda.core.Device.html#cuda.core.Device),提供额外功能,例如用于获取 NVML 设备的 to_system_device、设备属性以及同步。

抛出:

NcclInvalid ------ 如果通信器未初始化。

Communicator.rank

调用者在该通信器中的 rank(0 到 nranks - 1)。

抛出:

NcclInvalid ------ 如果通信器未初始化。

设备 API 能力

这些属性反映底层 NCCL 的 ncclCommProperties_t 结构。

Communicator.cuda_dev

与此通信器关联的 CUDA 设备 ID。

抛出:

NcclInvalid ------ 如果通信器未初始化。

Communicator.nvml_dev

与此通信器关联的 GPU 的 NVML 设备 ID。

使用 NVML 索引空间,该空间可能与 CUDA 索引不同。

抛出:

NcclInvalid ------ 如果通信器未初始化。

Communicator.device_api_support

此平台是否支持设备端 NCCL 操作。

如果为 False,则无法创建设备通信器。

抛出:

NcclInvalid ------ 如果通信器未初始化。

Communicator.multimem_support

同一 LSA team 中的 rank 是否可以使用 multimem 进行通信。

如果为 False,则无法创建带有 multimem 资源的设备通信器。

抛出:

NcclInvalid ------ 如果通信器未初始化。

Communicator.gin_type

GPU 发起网络(GPU-Initiated Networking,GIN)类型。

如果等于 NcclGinType.NONE,则无法创建带有 GIN 资源的设备通信器。

抛出:

NcclInvalid ------ 如果通信器未初始化。

Communicator.n_lsa_teams

此通信器的加载/存储可访问(Load/Store Accessible,LSA)team 数量。

抛出:

NcclInvalid ------ 如果通信器未初始化。

Communicator.host_rma_support

此通信器是否支持主机 RMA。

抛出:

NcclInvalid ------ 如果通信器未初始化。

Communicator.railed_gin_type

此通信器支持的 railed GIN 类型。

如果等于 NcclGinType.NONE,则无法创建 GIN 连接类型为 NcclGinConnectionType.RAIL 的设备通信器。

抛出:

NcclInvalid ------ 如果通信器未初始化。

Teams

NCCLTeam 表示通信器 rank 的一个带步长的子集。这些属性返回预定义的 team;可将其中一个传递给 TeamRequirement,或传递给下面的 rank 转换器。team 的语义请参见 Teams 一节。

Communicator.team_world

此通信器的 world team。

抛出:

NcclInvalid ------ 如果通信器未初始化。

Communicator.team_lsa

此通信器的 LSA team。

抛出:

NcclInvalid ------ 如果通信器未初始化。

Communicator.team_rail

此通信器的 rail team。

抛出:

NcclInvalid ------ 如果通信器未初始化。

Communicator.team_rank_to_world(team: NCCLTeam , team_rank: int ) → int

team 内的 rank 映射到其在此通信器中的 rank。

team 以当前 rank 为锚点,因此 team.rank 映射回本 rank,相邻 rank 按 team.stride 偏移。

参数:

  • team ------ team_rank 所属的 team,由 team_world、team_lsa 或 team_rail 返回。

  • team_rank ------ team 内的 rank。

返回:

此通信器中对应的 rank。

抛出:

NcclInvalid ------ 如果通信器未初始化。

Communicator.team_rank_to_lsa(team: NCCLTeam , team_rank: int ) → int

team 内的 rank 映射到其在 LSA team 中的 rank。

这是 team_rank_to_world() 的 LSA 相对版本:team.rank 映射回本 rank 在 team_lsa 中的索引。仅当 team_rank 指代与本 rank 共享同一 LSA team 的对端时才有意义。

参数:

  • team ------ team_rank 所属的 team,由 team_world、team_lsa 或 team_rail 返回。

  • team_rank ------ team 内的 rank。

返回:

LSA team 中对应的 rank;如果设备资源状态无法初始化,则返回 -1

抛出:

NcclInvalid ------ 如果通信器未初始化。

05.01.01.02. 创建与生命周期方法

Communicator 上用于创建、拆分、扩展和销毁的方法。

构造

classmethod Communicator.init(nranks: int , rank: int , unique_id: UniqueId | SequenceUniqueId , config: NCCLConfig | None = None ) → Communicator

初始化一个新的 NCCL 通信器。

创建一个连接多个 rank 的通信器。这是一个集合操作:所有 rank 都必须以相同的 nranksunique_id 调用此方法,但各自的 rank 值不同。

参数:

  • nranks ------ 通信器中的 rank 总数。

  • rank ------ 本 rank(必须在 0 到 nranks - 1 之间)。

  • unique_id ------ 所有 rank 共享的唯一标识符。可以传入一个序列以使用 ncclCommInitRankScalable()。

  • config ------ NCCL 配置选项。默认为 None

返回:

一个新的 Communicator 实例。

抛出:

NcclInvalid ------ 如果 unique_id 的类型无效。

classmethod Communicator.init_all(devices: int | Sequenceint | None = None ) → listCommunicator

为单进程多 GPU 操作初始化多个 NCCL 通信器。

在单个进程内创建一组 NCCL 通信器,每个设备一个。此方法针对单机场景进行了优化,此时所有 GPU 都由同一进程控制。与需要多进程协调(例如通过 MPI)的 init() 不同,init_all() 在内部处理所有协调工作。

每个通信器都绑定到其对应的设备,其 rank 等于它在返回列表中的索引。底层 NCCL API 会保留当前设备上下文。所有通信器都必须通过对每个通信器调用 destroy() 来手动销毁。

参数:

devices ------ 指定要初始化哪些设备。None(默认值)初始化所有可见的 CUDA 设备。传入 int 则为设备 [0, 1, ..., devices - 1] 创建通信器。传入 int 序列则使用显式指定的设备 ID。如果最终设备列表为空(devices=0、空序列或没有可见设备),则返回空列表而不调用 NCCL。

返回:

已初始化的通信器列表,每个设备一个。Rank i 使用 devices[i](当 devices 为 int 时使用设备 i)。

抛出:

TypeError ------ 如果 devices 不是 int、int 序列或 None

Communicator.initialize(nranks: int , rank: int , unique_id: UniqueId | SequenceUniqueId , config: NCCLConfig | None = None ) → None

就地初始化此通信器。

init() 类方法的实例方法版本。允许先创建一个空通信器(通过 Communicator()),稍后再初始化。这是一个集合操作;所有 rank 都必须调用此方法。

参数:

  • nranks ------ 通信器中的 rank 总数。

  • rank ------ 本 rank(必须在 0 到 nranks - 1 之间)。

  • unique_id ------ 所有 rank 共享的唯一标识符。

  • config ------ NCCL 配置选项。默认为 None

抛出:

NcclInvalid ------ 如果 unique_id 类型无效,或此通信器已初始化。

引导标识符

UniqueId 由一个 rank(通常是 rank 0)生成,并广播给所有参与的 rank;随后所有 rank 将它传给 Communicator.init()。

class nccl.core.UniqueId(_internal: _nccl_bindings.UniqueId | None = None )

基类:object

用于通信器初始化的 NCCL 唯一标识符。

UniqueId 用于协调多个 rank 之间的通信器初始化。所有 rank 必须使用相同的 UniqueId 才能组成一个通信器。通常由一个 rank 通过 get_unique_id() 生成 UniqueId,再广播给所有其他 rank。支持三种序列化途径:

  • 字节(Bytes) :生产端使用 bytes(uid)(或 as_bytes),接收端使用 from_bytes()。唯一 ID 的字节可以通过任何面向字节的通道传输------TCP socket、共享文件系统等。

  • NumPy :as_ndarray 返回底层缓冲区的就地视图,适用于支持 NumPy 缓冲区的传输方式,例如 mpi4py.MPI.Comm.Bcast(大写 B)。

  • Pickle :实例可以直接被 pickle,因此 mpi4py.MPI.Comm.bcast(小写 b)这类高层对象广播辅助函数可以开箱即用。

property as_bytes*: bytes*

唯一 ID 的字节表示,适合序列化或广播。

property as_ndarray*: numpy.ndarray*

唯一 ID 数据的 NumPy 数组视图。

static from_bytes(b: bytes | bytearray | memoryview ) → UniqueId

从类字节缓冲区重建 UniqueId。

参数:

b ------ UniqueId 的字节表示,通常通过生产 rank 上的 as_bytes 属性获得。

返回:

重建的 UniqueId。

nccl.core.get_unique_id(empty: bool = False ) → UniqueId

生成一个用于通信器初始化的新 NCCL 唯一标识符。

应由一个 rank(通常是 rank 0)调用;生成的 UniqueId 随后必须广播(例如通过 MPI)给所有其他 rank。

参数:

empty ------ 如果为 True,返回一个空的 UniqueId 而不调用 NCCL。当字节将在稍后通过 UniqueId.from_bytes() 填充时很有用。默认为 False。

返回:

一个将在各 rank 间共享的新 UniqueId。

拆分与扩展

Communicator.split(color: int | None = None , key: int = 0 , config: NCCLConfig | None = None ) → Communicator

根据 color 值将此通信器拆分为多个子通信器。

传入相同 color 值的 rank 将属于同一组。如果 colorNone,该 rank 不属于任何组,并收到一个空通信器(ptr=0 的 Communicator 实例)。key 值决定 rank 排序;key 越小,在新通信器中的 rank 越小。如果 key 相等,则由原通信器中的 rank 决定顺序。

这是一个集合操作:通信器中的所有 rank 都必须调用此方法,包括传入 color=None 的 rank。为避免死锁,通信器上不得有未完成的 NCCL 操作。

参数:

  • color ------ 用于对 rank 分组的非负颜色值。传入 None 可将本 rank 排除在所有组之外。默认为 None

  • key ------ 颜色组内的排序键。默认为 0。

  • config ------ 新通信器的配置。如果为 None,则继承父通信器的配置。默认为 None

返回:

新的子通信器;如果 colorNone,则返回空通信器。

抛出:

NcclInvalid ------ 如果通信器未初始化。

另请参阅

ncclCommSplit()

Communicator.shrink(exclude_ranks: Sequenceint | None = None , config: NCCLConfig | None = None , flag: CommShrinkFlag = CommShrinkFlag.DEFAULT ) → Communicator

通过从此通信器中移除指定 rank 来创建新通信器。

exclude_ranks 中列出的 rank 将被排除在新通信器之外;其余 rank 会重新编号为连续的 [0, n) 范围。

这是一个集合操作。所有未被排除的 rank 都必须调用此方法;被排除的 rank 不得调用。使用 DEFAULT 时,为避免死锁,不得有未完成的 NCCL 操作;可结合 config.shrink_share=True 以复用父通信器的资源。使用 ABORT 时,未完成的操作会被自动中止,且不与父通信器共享任何资源。

参数:

  • exclude_ranks ------ 要从新通信器中排除的 rank。默认为 None(不排除任何 rank)。

  • config ------ 新通信器的配置。如果为 None,则继承父通信器的配置。默认为 None

  • flag ------ shrink 行为。正常操作用 DEFAULT,出错后用 ABORT。默认为 DEFAULT。

返回:

不包含被排除 rank 的新通信器。

抛出:

NcclInvalid ------ 如果通信器未初始化。

另请参阅

ncclCommShrink()

Communicator.get_unique_id() → UniqueId

返回一个通信器级别的唯一 ID,供 grow() 使用。

生成一个绑定到此通信器的唯一标识符,可分享给通过 grow() 加入的新 rank。这与用于初始通信器创建的全局 get_unique_id() 不同。只应由一个现有 rank(grow 根)调用此方法。

在前一个 UID 尚未被消费时无法生成新的 UID;每个 UID 只能使用一次,用户必须等待对应的 grow 操作完成后才能再次调用。

返回:

用于 grow 操作的 UniqueId。

抛出:

NcclInvalid ------ 如果通信器未初始化。

Communicator.grow(nranks: int , unique_id: UniqueId | None = None , rank: int | None = None , config: NCCLConfig | None = None ) → Communicator

通过添加新 rank 来扩展通信器。

创建一个新通信器,其中既包含此通信器的现有 rank,也包含加入该组的新 rank。共有三种角色:

  • 现有根(Existing root):调用了 get_unique_id() 的那个现有 rank。

  • 现有非根(Existing non-root):所有其他现有 rank。

  • 新 rank:通过空通信器(Communicator())加入的 rank。

这是一个集合操作。所有 rank(现有的和新的)都必须调用此方法。按角色的用法:

  • 现有根:new_comm = existing_comm.grow(nranks, uid)

  • 现有非根:new_comm = existing_comm.grow(nranks)

  • 新 rank:new_comm = Communicator().grow(nranks, uid, rank=assigned_rank)

UID 在 grow 成功后被消费,不能重复使用。

参数:

  • nranks ------ 新通信器中的 rank 总数(现有加新增)。所有角色必须传入相同的值。

  • unique_id ------ 来自 get_unique_id() 的唯一标识符。现有根和新 rank 必须传入 UniqueId;现有非根必须传入 None。默认为 None

  • rank ------ 本 rank 在新通信器中的 ID。新 rank 必须传入分配给它的 rank,该值必须 >= 父通信器的大小。现有 rank 必须传入 None。默认为 None

  • config ------ 新通信器的配置。默认为 None

返回:

包含所有 rank 的新 Communicator。

抛出:

NcclInvalid ------ 如果新 rank 拿到的是已初始化的通信器,或现有 rank 拿到的是空通信器。

销毁

Communicator.destroy() → None

销毁通信器并释放本地资源。

如果 finalize() 未被显式调用,destroy() 会在内部调用它。如果显式调用了 finalize(),用户必须确保通信器状态变为 ncclSuccess 后才能调用 destroy()。destroy() 返回后不应再访问该通信器。

此通信器拥有的所有资源(已注册的缓冲区、窗口、自定义操作符)都会在销毁前自动关闭。这是一个节点内集合调用:同一节点上的所有 rank 都必须调用它,否则会挂起。推荐的模式是先 finalize() 再 destroy()。

出于安全考虑,清理过程中的错误会被抑制。

另请参阅

ncclCommDestroy()

Communicator.abort() → None

中止通信器并释放资源,终止进行中的操作。

应在发生不可恢复的错误时调用。与 destroy() 不同,此方法会立即中止未完成的操作。所有活跃 rank 都必须调用此函数,才能成功中止 NCCL 通信器。

此通信器拥有的所有资源(已注册的缓冲区、窗口、自定义操作符)都会在中止前自动关闭。出于安全考虑,清理过程中的错误会被抑制。更多细节请参见 NCCL 文档中的容错(Fault Tolerance)一节。

另请参阅

ncclCommAbort()

Communicator.finalize() → None

终结通信器,冲刷未完成的操作和网络资源。

通常在 destroy() 之前调用,以确保所有操作完成。这是一个必须由所有 rank 调用的集合操作。当一个线程管理多个主机本地 rank(每线程多个 GPU)时,对 finalize() 的调用必须放在 group() 内发起,以便所有本地 rank 能一起进入终结流程;否则它们会挂起。

对于非阻塞通信器,此方法本身也是非阻塞的:成功时会将通信器状态设为 ncclInProgress,表示终结正在进行。一旦所有 NCCL 操作完成,通信器状态会转变为 ncclSuccess。用户可以使用 get_async_error() 查询状态。

另请参阅

ncclCommFinalize()

暂停与恢复

Communicator.revoke(flags: int = 0 ) → None

撤销通信器。

停止所有进行中的操作,并将通信器状态标记为 ncclInProgress。当通信器进入静止状态后,状态转变为 ncclSuccess,此后管理操作(destroy()、split()、shrink())可以安全进行。

在 revoke() 之后调用 finalize() 是无效的。撤销期间,通过 split-share / shrink-share 进行的资源共享会被禁用。

参数:

flags ------ 保留供将来使用。目前必须为 0。

抛出:

NcclInvalid ------ 如果通信器未初始化。

Communicator.suspend(flags: CommSuspendFlag = CommSuspendFlag.MEM ) → None

挂起通信器操作以释放资源。

挂起期间,通信器不能用于通信。调用 resume() 可将其恢复。

参数:

flags ------ 控制释放哪些资源的挂起标志。MEM 释放动态 GPU 内存分配。

抛出:

NcclInvalid ------ 如果通信器未初始化。

Communicator.resume() → None

恢复此前挂起的所有通信器资源。

将通过 suspend() 挂起的通信器恢复,使其可以再次用于通信。

抛出:

NcclInvalid ------ 如果通信器未初始化。

标志枚举

CommShrinkFlag

class nccl.core.CommShrinkFlag(*values )

基类:IntEnum

Communicator.shrink() 的行为标志。

DEFAULT = 0

ABORT = 1

CommSuspendFlag

class nccl.core.CommSuspendFlag(*values )

基类:IntFlag

Communicator.suspend() 的行为标志。

MEM = 1

05.01.01.03. 集合通信方法

Communicator 上用于集合通信的方法。对应的 C API 请参见"集合通信函数"。

allreduce

Communicator.allreduce(sendbuf: \[Buffer(https://nvidia.github.io/cuda-python/cuda-core/latest/generated/cuda.core.Buffer.html#cuda.core.Buffer)] | SupportsDLPack | SupportsCAI , recvbuf: \[Buffer(https://nvidia.github.io/cuda-python/cuda-core/latest/generated/cuda.core.Buffer.html#cuda.core.Buffer)] | SupportsDLPack | SupportsCAI , op: NcclRedOp | CustomRedOp , * , stream: \[Stream(https://nvidia.github.io/cuda-python/cuda-core/latest/generated/cuda.core.Stream.html#cuda.core.Stream)] | cuda.core.typing.IsStreamType | int | None = None ) → None

reduce() 的全归约(all-reduce)变体。

等价于 reduce(sendbuf, recvbuf, op, root=None, stream=stream):对所有 rank 的数据进行归约,并将相同的副本存入每个 rank 的 recvbuf。参数语义请参见 reduce()。

另请参阅

reduce()、ncclAllReduce()

broadcast

Communicator.broadcast(sendbuf: \[Buffer(https://nvidia.github.io/cuda-python/cuda-core/latest/generated/cuda.core.Buffer.html#cuda.core.Buffer)] | SupportsDLPack | SupportsCAI | Any , recvbuf: \[Buffer(https://nvidia.github.io/cuda-python/cuda-core/latest/generated/cuda.core.Buffer.html#cuda.core.Buffer)] | SupportsDLPack | SupportsCAI , root: int , * , stream: \[Stream(https://nvidia.github.io/cuda-python/cuda-core/latest/generated/cuda.core.Stream.html#cuda.core.Stream)] | cuda.core.typing.IsStreamType | int | None = None ) → None

将根 rank 上 sendbuf 中的数据复制到所有 rank 的 recvbuf

sendbuf 仅在根 rank 上使用,在其他 rank 上被忽略。

在根 rank 上,两个缓冲区必须具有匹配的数据类型,且 sendcount == recvcount。元素数量从 recvbuf 推断:count = recvcount。当 sendbufrecvbuf 解析到相同的设备内存地址时,执行就地(in-place)操作。

参数:

  • sendbuf ------ 源缓冲区(仅在根 rank 上使用)。

  • recvbuf ------ 接收广播数据的目标缓冲区。

  • root ------ 广播数据的根 rank(0 到 nranks - 1)。

  • stream ------ 用于该操作的 CUDA 流。默认为 None(默认流)。

抛出:

NcclInvalid ------ 如果发送和接收缓冲区的 dtype 不匹配、数量不匹配、位于错误的设备、是无效的规格,或通信器未初始化。

另请参阅

ncclBroadcast()

reduce

Communicator.reduce(sendbuf: \[Buffer(https://nvidia.github.io/cuda-python/cuda-core/latest/generated/cuda.core.Buffer.html#cuda.core.Buffer)] | SupportsDLPack | SupportsCAI , recvbuf: \[Buffer(https://nvidia.github.io/cuda-python/cuda-core/latest/generated/cuda.core.Buffer.html#cuda.core.Buffer)] | SupportsDLPack | SupportsCAI | Any , op: NcclRedOp | CustomRedOp , root: int | None = None , * , stream: \[Stream(https://nvidia.github.io/cuda-python/cuda-core/latest/generated/cuda.core.Stream.html#cuda.core.Stream)] | cuda.core.typing.IsStreamType | int | None = None ) → None

使用指定操作对所有 rank 的数据进行归约。

支持两种模式。在 AllReduce 模式下(rootNone),所有 rank 都在 recvbuf 中收到归约结果。在 Reduce 模式下(指定 root),只有根 rank 收到归约结果;其他 rank 上的 recvbuf 被忽略。

在使用处,两个缓冲区必须具有匹配的数据类型。元素数量从 sendbuf 推断:count = sendcount。在 AllReduce 模式下,所有 rank 必须满足 recvcount >= sendcount;在 Reduce 模式下,只有根 rank 需要满足 recvcount >= sendcount。当 sendbufrecvbuf 解析到相同的设备内存地址时,执行就地操作。

参数:

  • sendbuf ------ 包含待归约数据的源缓冲区。

  • recvbuf ------ 存放归约结果的目标缓冲区。在 Reduce 模式下仅在根 rank 上使用。

  • op ------ 归约操作符(例如 NcclRedOp.SUM、NcclRedOp.MAX、NcclRedOp.MIN、NcclRedOp.AVG、NcclRedOp.PROD,或一个 CustomRedOp)。

  • root ------ 接收归约结果的根 rank(0 到 nranks - 1)。如果为 None,则执行全归约。默认为 None

  • stream ------ 用于该操作的 CUDA 流。默认为 None(默认流)。

抛出:

NcclInvalid ------ 如果发送和接收缓冲区的 dtype 不匹配、数量不匹配、位于错误的设备、是无效的规格,或通信器未初始化。

另请参阅

ncclAllReduce()、ncclReduce()

allgather

Communicator.allgather(sendbuf: \[Buffer(https://nvidia.github.io/cuda-python/cuda-core/latest/generated/cuda.core.Buffer.html#cuda.core.Buffer)] | SupportsDLPack | SupportsCAI , recvbuf: \[Buffer(https://nvidia.github.io/cuda-python/cuda-core/latest/generated/cuda.core.Buffer.html#cuda.core.Buffer)] | SupportsDLPack | SupportsCAI , * , stream: \[Stream(https://nvidia.github.io/cuda-python/cuda-core/latest/generated/cuda.core.Stream.html#cuda.core.Stream)] | cuda.core.typing.IsStreamType | int | None = None ) → None

gather() 的全收集(all-gather)变体。

等价于 gather(sendbuf, recvbuf, root=None, stream=stream):从每个 rank 收集 sendcount 个值,并将拼接结果的相同副本放入每个 rank 的 recvbuf。参数语义请参见 gather()。

另请参阅

gather()、ncclAllGather()

reduce_scatter

Communicator.reduce_scatter(sendbuf: \[Buffer(https://nvidia.github.io/cuda-python/cuda-core/latest/generated/cuda.core.Buffer.html#cuda.core.Buffer)] | SupportsDLPack | SupportsCAI , recvbuf: \[Buffer(https://nvidia.github.io/cuda-python/cuda-core/latest/generated/cuda.core.Buffer.html#cuda.core.Buffer)] | SupportsDLPack | SupportsCAI , op: NcclRedOp | CustomRedOp , * , stream: \[Stream(https://nvidia.github.io/cuda-python/cuda-core/latest/generated/cuda.core.Stream.html#cuda.core.Stream)] | cuda.core.typing.IsStreamType | int | None = None ) → None

对所有 rank 的数据进行归约,并将结果分散到各 rank。

每个 rank 收到归约结果的不同部分:rank i 在其 recvbuf 中收到第 i 块。

两个缓冲区必须具有匹配的数据类型。元素数量从 sendbuf 推断:count = sendcount / nrankssendcount 必须 >= nranks,且 recvcount 必须 >= count。当 recvbuf 解析到 sendbuf_address + rank * count 时,执行就地操作。

参数:

  • sendbuf ------ 源缓冲区(大小 >= nranks * recvcount 个元素)。

  • recvbuf ------ 含有 recvcount 个元素的目标缓冲区。

  • op ------ 归约操作符(例如 NcclRedOp.SUM、NcclRedOp.MAX、NcclRedOp.MIN、NcclRedOp.AVG、NcclRedOp.PROD,或一个 CustomRedOp)。

  • stream ------ 用于该操作的 CUDA 流。默认为 None(默认流)。

抛出:

NcclInvalid ------ 如果发送和接收缓冲区的 dtype 不匹配、sendbuf 太小、位于错误的设备、是无效的规格,或通信器未初始化。

另请参阅

ncclReduceScatter()

alltoall

Communicator.alltoall(sendbuf: \[Buffer(https://nvidia.github.io/cuda-python/cuda-core/latest/generated/cuda.core.Buffer.html#cuda.core.Buffer)] | SupportsDLPack | SupportsCAI , recvbuf: \[Buffer(https://nvidia.github.io/cuda-python/cuda-core/latest/generated/cuda.core.Buffer.html#cuda.core.Buffer)] | SupportsDLPack | SupportsCAI , * , stream: \[Stream(https://nvidia.github.io/cuda-python/cuda-core/latest/generated/cuda.core.Stream.html#cuda.core.Stream)] | cuda.core.typing.IsStreamType | int | None = None ) → None

每个 rank 向其他每个 rank 发送并从其接收 count 个值。

发往目标 rank j 的数据取自 sendbuf + j * count,从源 rank i 接收的数据放在 recvbuf + i * count

两个缓冲区必须具有匹配的数据类型。元素数量从 sendbuf 推断:count = sendcount / nrankssendcount 必须 >= nranks,且 recvcount 必须 >= sendcount

参数:

  • sendbuf ------ 源缓冲区(大小 >= nranks * count 个元素)。

  • recvbuf ------ 目标缓冲区(大小 >= nranks * count 个元素)。

  • stream ------ 用于该操作的 CUDA 流。默认为 None(默认流)。

抛出:

NcclInvalid ------ 如果发送和接收缓冲区的 dtype 不匹配、缓冲区大小与 nranks 不兼容、位于错误的设备、是无效的规格,或通信器未初始化。

另请参阅

ncclAlltoAll()

gather

Communicator.gather(sendbuf: \[Buffer(https://nvidia.github.io/cuda-python/cuda-core/latest/generated/cuda.core.Buffer.html#cuda.core.Buffer)] | SupportsDLPack | SupportsCAI , recvbuf: \[Buffer(https://nvidia.github.io/cuda-python/cuda-core/latest/generated/cuda.core.Buffer.html#cuda.core.Buffer)] | SupportsDLPack | SupportsCAI | Any , root: int | None = None , * , stream: \[Stream(https://nvidia.github.io/cuda-python/cuda-core/latest/generated/cuda.core.Stream.html#cuda.core.Stream)] | cuda.core.typing.IsStreamType | int | None = None ) → None

从所有 rank 收集 sendcount 个值。

支持两种模式。在 AllGather 模式下(rootNone),从所有 rank 收集值,并将结果的相同副本放入每个 recvbuf。在 Gather 模式下(指定 root),值仅被收集到指定的根 rank;其他 rank 上的 recvbuf 被忽略。

在使用处,两个缓冲区必须具有匹配的数据类型。元素数量从 sendbuf 推断:count = sendcount。来自 rank i 的数据放在 recvbuf + i * sendcount。AllGather 模式要求每个 rank 满足 recvcount >= nranks * sendcount;Gather 模式仅要求根 rank 满足该条件。

sendbuf 在 AllGather 模式下解析到 recvbuf_address + rank * sendcount,或在 Gather 模式下解析到 recvbuf_address + root * sendcount 时,执行就地操作。

参数:

  • sendbuf ------ 含有 sendcount 个元素的源缓冲区。

  • recvbuf ------ 目标缓冲区(大小 >= nranks * sendcount 个元素)。在 Gather 模式下仅在根 rank 上使用。

  • root ------ 接收所收集数据的根 rank(0 到 nranks - 1)。如果为 None,则执行全收集。默认为 None

  • stream ------ 用于该操作的 CUDA 流。默认为 None(默认流)。

抛出:

NcclInvalid ------ 如果发送和接收缓冲区的 dtype 不匹配、recvbuf 太小、位于错误的设备、是无效的规格,或通信器未初始化。

另请参阅

ncclAllGather()、ncclGather()

scatter

Communicator.scatter(sendbuf: \[Buffer(https://nvidia.github.io/cuda-python/cuda-core/latest/generated/cuda.core.Buffer.html#cuda.core.Buffer)] | SupportsDLPack | SupportsCAI | Any , recvbuf: \[Buffer(https://nvidia.github.io/cuda-python/cuda-core/latest/generated/cuda.core.Buffer.html#cuda.core.Buffer)] | SupportsDLPack | SupportsCAI , root: int , * , stream: \[Stream(https://nvidia.github.io/cuda-python/cuda-core/latest/generated/cuda.core.Stream.html#cuda.core.Stream)] | cuda.core.typing.IsStreamType | int | None = None ) → None

将数据从根 rank 分散到所有 rank。

每个 rank 从根 rank 接收 count 个元素。在根 rank 上,从 sendbuf + i * count 取出的 count 个元素被发送给 rank isendbuf 在非根 rank 上不使用。

在根 rank 上,两个缓冲区必须具有匹配的数据类型。元素数量从 recvbuf 推断:count = recvcount。根 rank 要求 sendcount >= nrankssendcount / nranks == recvcount。当 recvbuf 解析到 sendbuf_address + root * count 时,执行就地操作。

参数:

  • sendbuf ------ 源缓冲区(仅在根 rank 上使用,大小 >= nranks * count 个元素)。

  • recvbuf ------ 含有 count 个元素的目标缓冲区。

  • root ------ 分散数据的根 rank(0 到 nranks - 1)。

  • stream ------ 用于该操作的 CUDA 流。默认为 None(默认流)。

抛出:

NcclInvalid ------ 如果发送和接收缓冲区的 dtype 不匹配、根 rank 上的 sendbuf 太小、位于错误的设备、是无效的规格,或通信器未初始化。

另请参阅

ncclScatter()

create_pre_mul_sum

Communicator.create_pre_mul_sum(scalar: int | float | numpy.ndarray | \[Buffer(https://nvidia.github.io/cuda-python/cuda-core/latest/generated/cuda.core.Buffer.html#cuda.core.Buffer)] | SupportsDLPack | SupportsCAI , datatype: NcclDataType | None = None ) → CustomRedOp

创建一个 PreMulSum 自定义归约操作符。

执行 output = scalar * sum(inputs),适用于求平均(scalar = 1/N)或加权归约。返回的 CustomRedOp 由通信器跟踪,可以通过其 close() 方法显式释放,或在通信器被销毁或中止时自动释放。

参数:

  • scalar ------ 标量乘数值。Python int 或 float 会使用主机内存转换为 NumPy 数组。NumPy 数组必须恰好包含 1 个元素,并使用主机内存。NcclSupportedBuffer 被视为恰好包含 1 个元素的设备缓冲区。

  • datatype ------ 标量和归约的 NCCL 数据类型。如果为 None,则从 scalar 推断:Python int 变为 int64,Python float 变为 float64(NumPy 的自然 dtype);NumPy 数组使用数组的 dtype;设备缓冲区使用缓冲区的 dtype。

返回:

PreMulSum 操作符的 CustomRedOp。

抛出:

NcclInvalid ------ 如果通信器未初始化;标量类型不受支持;NumPy 数组或设备缓冲区不恰好包含 1 个元素;或请求的 datatype 与设备缓冲区的 dtype 不匹配。

另请参阅

ncclRedOpCreatePreMulSum()

05.01.01.04. 点对点与信号方法

Communicator 上用于点对点及 signal/wait 操作的方法。对应的 C API 请参见"点对点通信函数"。

send

Communicator.send(sendbuf: \[Buffer(https://nvidia.github.io/cuda-python/cuda-core/latest/generated/cuda.core.Buffer.html#cuda.core.Buffer)] | SupportsDLPack | SupportsCAI , peer: int , * , stream: \[Stream(https://nvidia.github.io/cuda-python/cuda-core/latest/generated/cuda.core.Stream.html#cuda.core.Stream)] | cuda.core.typing.IsStreamType | int | None = None ) → None

向对端 rank 发送一个缓冲区。

参数:

  • sendbuf ------ 要发送的源缓冲区。

  • peer ------ 目标 rank ID。

  • stream ------ 用于该操作的 CUDA 流。默认为 None(默认流)。

抛出:

NcclInvalid ------ 如果缓冲区规格无效、缓冲区位于错误的设备,或通信器未初始化。

另请参阅

ncclSend()

recv

Communicator.recv(recvbuf: \[Buffer(https://nvidia.github.io/cuda-python/cuda-core/latest/generated/cuda.core.Buffer.html#cuda.core.Buffer)] | SupportsDLPack | SupportsCAI , peer: int , * , stream: \[Stream(https://nvidia.github.io/cuda-python/cuda-core/latest/generated/cuda.core.Stream.html#cuda.core.Stream)] | cuda.core.typing.IsStreamType | int | None = None ) → None

从对端 rank 接收数据到缓冲区。

参数:

  • recvbuf ------ 用于接收的目标缓冲区。

  • peer ------ 源 rank ID。

  • stream ------ 用于该操作的 CUDA 流。默认为 None(默认流)。

抛出:

NcclInvalid ------ 如果缓冲区规格无效、缓冲区位于错误的设备,或通信器未初始化。

另请参阅

ncclRecv()

signal

Communicator.signal(peer: int , signal_index: int = 0 , context: int = 0 , flags: int = 0 , * , stream: \[Stream(https://nvidia.github.io/cuda-python/cuda-core/latest/generated/cuda.core.Stream.html#cuda.core.Stream)] | cuda.core.typing.IsStreamType | int | None = None ) → None

向对端 rank 发送一个信号。

在指定的 CUDA 流上入队一个信号操作,用于通知目标对端 rank。对端可以使用 wait_signal() 等待该信号。

参数:

  • peer ------ 要发送信号的目标 rank。

  • signal_index ------ 信号索引标识符。目前必须为 0。

  • context ------ 上下文标识符。目前必须为 0。

  • flags ------ 保留供将来使用。目前必须为 0。

  • stream ------ 要在其上入队信号操作的 CUDA 流。默认为 None(默认流)。

抛出:

NcclInvalid ------ 如果通信器未初始化。

另请参阅

ncclSignal()

wait_signal

Communicator.wait_signal(descs: WaitSignalDesc | SequenceWaitSignalDesc , * , stream: \[Stream(https://nvidia.github.io/cuda-python/cuda-core/latest/generated/cuda.core.Stream.html#cuda.core.Stream)] | cuda.core.typing.IsStreamType | int | None = None ) → None

按照信号描述符的描述等待信号。

在指定的 CUDA 流上入队一个等待操作,该操作会阻塞,直到收到来自对端 rank 的所需信号。每个描述符指定一个对端 rank,以及要等待来自该对端的多少个信号操作。

参数:

  • descs ------ 一个或多个 WaitSignalDesc 描述符,指定要等待哪些对端,以及分别从每个对端期待多少个信号。

  • stream ------ 要在其上入队等待操作的 CUDA 流。默认为 None(默认流)。

抛出:

NcclInvalid ------ 如果通信器未初始化。

另请参阅

ncclWaitSignal()

put_signal

Communicator.put_signal(local_buffer: \[Buffer(https://nvidia.github.io/cuda-python/cuda-core/latest/generated/cuda.core.Buffer.html#cuda.core.Buffer)] | SupportsDLPack | SupportsCAI , peer: int , peer_window: RegisteredWindowHandle , peer_window_offset: int = 0 , signal_index: int = 0 , context: int = 0 , flags: int = 0 , * , stream: \[Stream(https://nvidia.github.io/cuda-python/cuda-core/latest/generated/cuda.core.Stream.html#cuda.core.Stream)] | cuda.core.typing.IsStreamType | int | None = None ) → None

将本地缓冲区中的数据放入对端的窗口,并发送一个信号。

在指定的 CUDA 流上入队一个带信号的 put 操作,将本地缓冲区内容传输到目标对端已注册的窗口,并通知该对端。对端可以使用 wait_signal() 等待该信号(也即等待 put 完成)。对端的内存和 local_buffer 都必须通过 register_window() 注册;将对端的窗口句柄作为 peer_window 传入(例如通过对窗口句柄做 allgather 获得)。

参数:

  • local_buffer ------ 其内容将被放入对端的源缓冲区。

  • peer ------ 要放入数据并发送信号的目标 rank。

  • peer_window ------ 对端的 RegisteredWindowHandle(来自 register_window())。

  • peer_window_offset ------ 对端窗口内的偏移量,以元素为单位。默认为 0。

  • signal_index ------ 信号索引标识符。目前必须为 0。

  • context ------ 上下文标识符。目前必须为 0。

  • flags ------ 保留供将来使用。目前必须为 0。

  • stream ------ 要在其上入队 put_signal 操作的 CUDA 流。默认为 None(默认流)。

抛出:

NcclInvalid ------ 如果通信器未初始化,或缓冲区规格无效、缓冲区与通信器位于不同的设备。

另请参阅

ncclPutSignal()

WaitSignalDesc

class nccl.core.WaitSignalDesc(peer: int , op_count: int = 1 , signal_index: int = 0 , context: int = 0 )

基类:LowppSpec

wait-signal 操作的描述符。

描述一个单独的信号等待操作,供 Communicator.wait_signal() 使用。每个描述符指定要等待哪个对端、要等待多少个信号操作,以及该等待操作的附加上下文。

context*: int* = 0

op_count*: int* = 1

signal_index*: int* = 0

peer*: int*

05.01.01.05. 内存注册方法

Communicator 上用于为零拷贝和 RMA 操作注册缓冲区与窗口的方法。返回的句柄类的文档见"通信器资源"。

register_buffer

Communicator.register_buffer(buffer: \[Buffer(https://nvidia.github.io/cuda-python/cuda-core/latest/generated/cuda.core.Buffer.html#cuda.core.Buffer)] | SupportsDLPack | SupportsCAI ) → RegisteredBufferHandle

向此通信器注册一个缓冲区,以用于零拷贝通信。

已注册的缓冲区可以在 NCCL 操作中启用性能优化。缓冲区大小会从缓冲区元素数量和 dtype 自动推导。返回的 RegisteredBufferHandle 由通信器跟踪,可以通过其 close() 方法显式释放,或在通信器被销毁或中止时自动释放。

参数:

buffer ------ 要注册的缓冲区(数组、Buffer 或类缓冲区对象)。

返回:

已注册缓冲区的 RegisteredBufferHandle。

抛出:

NcclInvalid ------ 如果缓冲区位于错误的设备,或通信器未初始化。

另请参阅

ncclCommRegister()

register_window

Communicator.register_window(buffer: \[Buffer(https://nvidia.github.io/cuda-python/cuda-core/latest/generated/cuda.core.Buffer.html#cuda.core.Buffer)] | SupportsDLPack | SupportsCAI , flags: WindowFlag | None = None ) → RegisteredWindowHandle | None

将本地缓冲区集合式地注册到一个 NCCL 窗口中。

这是一个集合调用:通信器中的每个 rank 都必须参与,且默认情况下各 rank 的缓冲区大小必须相等。缓冲区大小会从缓冲区元素数量和 dtype 自动推导。如果在 group 内调用,句柄值可能要到 ncclGroupEnd 完成后才会被填充。对于非阻塞通信器,句柄可能一直保持为 0,直到 get_async_error() 报告成功。

返回的 RegisteredWindowHandle 由通信器跟踪,可以通过其 close() 方法显式释放,或在通信器被销毁或中止时自动释放。

参数:

  • buffer ------ 要注册为窗口的本地缓冲区。

  • flags ------ 窗口注册标志。默认为 None(DEFAULT)。

返回:

已注册窗口的 RegisteredWindowHandle;如果 NCCL 返回 NULL 句柄(例如此平台不支持窗口),则返回 None

抛出:

NcclInvalid ------ 如果缓冲区位于错误的设备,或通信器未初始化。

另请参阅

ncclCommWindowRegister()

WindowFlag

class nccl.core.WindowFlag(*values )

基类:IntFlag

Communicator.register_window() 的窗口注册行为标志。

DEFAULT = 0

COLL_SYMMETRIC = 1

STRICT_ORDERING = 2

05.01.01.06. 设备通信器设置

用于创建 NCCL 设备通信器的主机端方法和资源。设备端通信原语本身只能在 CUDA kernel 中使用,其文档见 C 设备 API(Device API);本页介绍 Python(主机)端为引导它们所暴露的内容。传给 Communicator.create_dev_comm() 的配置对象的文档见"配置"。

create_dev_comm

Communicator.create_dev_comm(requirements: NCCLDevCommRequirements | None = None ) → DevCommResource

创建一个用于设备端 NCCL 操作的设备通信器。

这是一个集合调用:通信器中的每个 rank 都必须参与。当在 group 内调用时,结果可能要到 group 完成后才会被填充。

设备通信器使 GPU kernel 能够直接访问 NCCL 通信原语。可以从一个主机通信器创建多个设备通信器。返回的 DevCommResource 由通信器跟踪,可以通过其 close() 方法显式释放,或在通信器被销毁或中止时自动释放。通过 DevCommResource.ptr 或 resource.dev_comm.ptr 访问设备通信器指针。

参数:

requirements ------ 设备通信器资源分配的配置。如果为 None,则使用默认的 NCCLDevCommRequirements。默认为 None

返回:

设备通信器的 DevCommResource。

抛出:

NcclInvalid ------ 如果通信器未初始化。

另请参阅

ncclDevCommCreate()

GIN 类型枚举

GPU 发起网络(GPU-Initiated Networking,GIN)枚举,描述通信器上可用的设备端网络传输方式,以及用户所需的连接拓扑。

NcclGinType

class nccl.core.NcclGinType(*values )

基类:IntEnum

GIN 传输类型,与 ncclGinType_t 对应。

由 Communicator.gin_type 和 Communicator.railed_gin_type 报告,用于指示通信器上可用的设备端网络传输方式(如果有)。

NONE = 0

PROXY = 2

GDAKI = 3

GPI = 4

NcclGinConnectionType

class nccl.core.NcclGinConnectionType(*values )

基类:IntEnum

GIN 连接拓扑,与 ncclGinConnectionType_t 对应。

在调用 Communicator.create_dev_comm() 之前,设置 NCCLDevCommRequirements 的 gin_connection_type 字段,以声明设备代码必须通过 GIN 可达哪些对端。

NONE = 0

FULL = 1

RAIL = 2

05.01.01.07. 状态与工具方法

Communicator 上用于资源清理和错误/状态查询的方法。

close_all_resources

Communicator.close_all_resources() → None

关闭此通信器拥有的所有资源。

在 destroy() 和 abort() 期间会自动调用,但也可以手动调用。执行尽力而为的清理,忽略资源释放过程中发生的任何错误。该方法是幂等的:可以安全地多次调用。

get_last_error

Communicator.get_last_error() → str

返回此通信器的最近一次错误字符串。

抛出:

NcclInvalid ------ 如果通信器未初始化。

get_async_error

Communicator.get_async_error() → nccl.bindings.nccl.Result

查询异步 NCCL 操作的进度和潜在错误。

不带 stream 参数的操作(例如 finalize())在返回 ncclSuccess 时即完成。带 stream 参数的操作(例如 reduce())在提交时返回 ncclSuccess,但在完成之前可能通过此方法报告错误。如果任何 NCCL 函数返回 ncclInProgress,用户必须查询通信器状态,直到其变为 ncclSuccess,然后才能调用另一个 NCCL 函数。

在状态变为 ncclSuccess 之前,不要在 NCCL 使用的流上发起 CUDA kernel。如果发生错误,用 abort() 销毁通信器;发生错误后,无法对已入队操作的完成情况或正确性做任何假设。

返回:

通信器的当前状态(ncclSuccessncclInProgress 或错误码)。

抛出:

NcclInvalid ------ 如果通信器未初始化。

另请参阅

ncclCommGetAsyncError()

get_mem_stat

Communicator.get_mem_stat(stat: NcclCommMemStat ) → int

查询通信器内存统计信息。

参数:

stat ------ 要查询的内存统计项。

返回:

内存统计值(字节数;GPU_MEM_SUSPENDED 为 0/1)。

抛出:

NcclInvalid ------ 如果通信器未初始化。

NcclCommMemStat

class nccl.core.NcclCommMemStat(*values )

基类:IntEnum

内存统计选择器,与 ncclCommMemStat_t 对应。

用作 Communicator.get_mem_stat() 的 stat 参数,用于标识要查询哪项内存统计。除 GPU_MEM_SUSPENDED 是 0/1 标志外,所有值均以字节为单位返回。

GPU_MEM_SUSPEND = 0

GPU_MEM_SUSPENDED = 1

GPU_MEM_PERSIST = 2

GPU_MEM_TOTAL = 3

get_error_string

模块级辅助函数,用于将 NCCL 结果码渲染为人类可读的字符串。

nccl.core.get_error_string(nccl_result: _nccl_bindings.Result | int ) → str

返回 NCCL 结果码对应的人类可读错误字符串。

参数:

nccl_result ------ NCCL 结果码。

返回:

与结果码对应的人类可读错误消息。

05.01.02. 配置

传给通信器创建方法的配置对象,以及它们所使用的标志枚举。

NCCLConfig

由 Communicator.init()、Communicator.initialize()、Communicator.split()、Communicator.shrink() 和 Communicator.grow() 使用。未设置的字段(None)保持 NCCL 的内部默认值;配置被消费时,各值会由 C 库进行校验。

class nccl.core.NCCLConfig(* , blocking: bool | None = None , cga_cluster_size: int | None = None , min_ctas: int | None = None , max_ctas: int | None = None , net_name: str | None = None , split_share: bool | None = None , traffic_class: int | None = None , comm_name: str | None = None , collnet_enable: bool | None = None , cta_policy: CTAPolicy | None = None , shrink_share: bool | None = None , nvls_ctas: int | None = None , n_channels_per_net_peer: int | None = None , nvlink_centric_sched: bool | None = None , graph_usage_mode: int | None = None , num_rma_ctx: int | None = None , max_p2p_peers: int | None = None , graph_stream_ordering: int | None = None )

基类:LowppSpec

用于通信器初始化的 NCCL 配置。

为 NCCL 通信器提供配置选项,允许对性能和行为特性进行微调。构造函数中未设置的字段保持 NCCL 的内部默认值;配置被消费时,各值会由 C 库进行校验。

另请参阅

ncclConfig_t 中对各字段的描述。

blocking*: bool | None* = None

cga_cluster_size*: int | None* = None

collnet_enable*: bool | None* = None

comm_name*: str | None* = None

cta_policy*: CTAPolicy | None* = None

graph_stream_ordering*: int | None* = None

graph_usage_mode*: int | None* = None

max_ctas*: int | None* = None

max_p2p_peers*: int | None* = None

min_ctas*: int | None* = None

n_channels_per_net_peer*: int | None* = None

net_name*: str | None* = None

num_rma_ctx*: int | None* = None

nvlink_centric_sched*: bool | None* = None

nvls_ctas*: int | None* = None

shrink_share*: bool | None* = None

split_share*: bool | None* = None

traffic_class*: int | None* = None

CTAPolicy

class nccl.core.CTAPolicy(*values )

基类:IntFlag

用于 CTA 调度的 NCCL 性能策略,由 NCCLConfig.cta_policy 使用。

DEFAULT = 0

EFFICIENCY = 1

ZERO = 2

NCCLDevCommRequirements

由 Communicator.create_dev_comm() 使用。未设置的字段(None)保持 NCCL 的内部默认值。

class nccl.core.NCCLDevCommRequirements(* , lsa_multimem: bool | None = None , barrier_count: int | None = None , lsa_barrier_count: int | None = None , rail_gin_barrier_count: int | None = None , lsa_ll_a2a_block_count: int | None = None , lsa_ll_a2a_slot_count: int | None = None , gin_force_enable: bool | None = None , gin_context_count: int | None = None , gin_signal_count: int | None = None , gin_counter_count: int | None = None , gin_connection_type: NcclGinConnectionType | None = None , gin_exclusive_contexts: bool | None = None , gin_queue_depth: int | None = None , gin_traffic_class: int | None = None , world_gin_barrier_count: int | None = None , gin_strong_signals_required: bool | None = None , gin_va_signals_required: bool | None = None , teams: tupleTeamRequirement, ... = () , resources: tupleLsaBarrierRequirement \| GinBarrierRequirement \| LLA2ARequirement, ... = () )

基类:LowppSpec

NCCL 设备通信器需求配置。

这是一个由 Communicator.create_dev_comm() 消费的可复用的高层 Python 请求。按 team 的需求通过 teams 元组声明。每次调用都会把请求快照到独立的底层 ncclDevCommRequirements_t 及链接的 ncclTeamRequirements_t 存储中,其中包括独立的 multimem 输出句柄。NCCL 会在调用返回前复制这些需求和链表节点;生成的 DevCommResource 会保留每个 outMultimemHandle 所引用的存储。因此,该对象可以在两次调用之间修改,而不会影响已创建的设备通信器。不要在与 Communicator.create_dev_comm() 并发的情况下修改它。

另请参阅

ncclDevCommRequirements 中对各字段的描述。

barrier_count*: int | None* = None

gin_connection_type*: NcclGinConnectionType | None* = None

gin_context_count*: int | None* = None

gin_counter_count*: int | None* = None

gin_exclusive_contexts*: bool | None* = None

gin_force_enable*: bool | None* = None

gin_queue_depth*: int | None* = None

gin_signal_count*: int | None* = None

gin_strong_signals_required*: bool | None* = None

gin_traffic_class*: int | None* = None

gin_va_signals_required*: bool | None* = None

lsa_barrier_count*: int | None* = None

lsa_ll_a2a_block_count*: int | None* = None

lsa_ll_a2a_slot_count*: int | None* = None

lsa_multimem*: bool | None* = None

rail_gin_barrier_count*: int | None* = None

resources*: tupleLsaBarrierRequirement \| GinBarrierRequirement \| LLA2ARequirement, ...* = ()

teams*: tupleTeamRequirement, ...* = ()

world_gin_barrier_count*: int | None* = None

需求项

NCCLDevCommRequirements.teams 和 NCCLDevCommRequirements.resources 的元素类型。

TeamRequirement

class nccl.core.TeamRequirement(team: NCCLTeam , multimem: bool = False )

基类:object

设备通信器创建的按 team 需求项。

将这些对象的元组作为 NCCLDevCommRequirements.teams 传入。当 multimem 为 True 时,NCCL 会为该 team 分配一个多播句柄,之后可通过 multimem_handle() 获取。

multimem*: bool* = False

team*: NCCLTeam*

LsaBarrierRequirement

class nccl.core.LsaBarrierRequirement(team: NCCLTeam , n_barriers: int )

基类:object

team 上请求一个带有 n_barriers 个 barrier 的 LSA barrier 资源。

加入 NCCLDevCommRequirements.resources;最终生成的 LsaBarrierHandle 会在 resource_handles 中返回。

team*: NCCLTeam*

n_barriers*: int*

GinBarrierRequirement

class nccl.core.GinBarrierRequirement(team: NCCLTeam , n_barriers: int )

基类:object

team 上请求一个带有 n_barriers 个 barrier 的 GIN barrier 资源。

加入 NCCLDevCommRequirements.resources;最终生成的 GinBarrierHandle 会在 resource_handles 中返回。

team*: NCCLTeam*

n_barriers*: int*

LLA2ARequirement

class nccl.core.LLA2ARequirement(n_blocks: int , max_elements: int , max_element_size: int )

基类:object

请求一个低延迟 all-to-all 资源,包含 n_blocks 个块,大小可容纳最多 max_elements 个元素,每个元素至多 max_element_size 字节。

加入 NCCLDevCommRequirements.resources;最终生成的 LLA2AHandle 会在 resource_handles 中返回。

n_blocks*: int*

max_elements*: int*

max_element_size*: int*

05.01.03. 组操作

用于将 NCCL 操作批量编入组的自由函数和辅助工具。用法细节请参见"组调用"。

group

nccl.core.group() → GeneratorNone, None, None

NCCL 组操作的上下文管理器。

进入时自动调用 group_start(),退出时自动调用 group_end(),即使发生异常也能确保正确清理。

此处不支持模拟模式。如需模拟,请直接调用 group_start() 和 group_end(),并向 group_end() 传入 simulate=True

group_start

nccl.core.group_start() → None

开始一组 NCCL 操作。

此后调用的所有 NCCL 操作将被批量组合,并在调用 group_end() 时执行。这可以让 NCCL 对操作序列进行优化,从而提升性能。

group_end

nccl.core.group_end(* , simulate: bool = False ) → GroupSimInfo | None

结束一组 NCCL 操作。

默认情况下,执行自上次 group_start() 以来排队的所有操作。当 simulate=True 时,排队的操作会被模拟而非真正执行,估计的执行时间通过 GroupSimInfo 返回。

参数:

simulate ------ 为 True 时,对该组进行模拟而非执行,并返回携带估计时间的 GroupSimInfo。默认为 False。

返回:

simulate=False 时返回 Nonesimulate=True 时返回带有模拟结果的 GroupSimInfo。

GroupSimInfo

class nccl.core.GroupSimInfo(estimated_time: float )

基类:object

NCCL 组模拟的结果。

当以 simulate=True 调用 group_end() 时返回。

estimated_time*: float*

05.01.04. 内存管理

由 NCCL 支持的设备内存分配;用法细节请参见"内存分配器"。如需对已有缓冲区进行零拷贝注册,请参见 Communicator.register_buffer() 和 Communicator.register_window()。

mem_alloc

nccl.core.mem_alloc(size: int , device: \[Device(https://nvidia.github.io/cuda-python/cuda-core/latest/generated/cuda.core.Device.html#cuda.core.Device)] | int | None = None ) → \[Buffer(https://nvidia.github.io/cuda-python/cuda-core/latest/generated/cuda.core.Buffer.html#cuda.core.Buffer)]

使用 NCCL 的内存分配器分配 GPU 缓冲区内存。

由于 NCCL 优化对缓冲区粒度的要求,实际分配的大小可能大于请求的大小。返回的缓冲区可以用 mem_free() 显式释放,或在被垃圾回收时自动释放。

参数:

  • size ------ 要分配的字节数。

  • device ------ 目标 CUDA 设备。默认为当前设备。

返回:

由 NCCL 管理的内存支持的 CUDA 缓冲区对象。缓冲区在指定设备上分配;分配完成后会恢复当前设备。

mem_free

nccl.core.mem_free(buf: \[Buffer(https://nvidia.github.io/cuda-python/cuda-core/latest/generated/cuda.core.Buffer.html#cuda.core.Buffer)] ) → None

释放由 mem_alloc() 分配的内存。

显式释放是可选的。当 Buffer 对象被垃圾回收时,内存会自动释放。

参数:

buf ------ 要释放的缓冲区。

05.01.05. 通信器资源

由 Communicator 拥有的资源。下面的 CommResource 子类由其所属通信器跟踪,可通过 close() 显式释放,或在通信器被销毁或中止时自动释放。

CommResource

通信器拥有的资源的抽象基类;它定义了这些资源的 close()is_valid 契约。

class nccl.core.resources.CommResource(comm_ptr: int )

基类:ABC

NCCL 通信器拥有的资源的抽象基类。

资源绑定到特定的通信器。它们可以通过 close() 显式释放,并在所属通信器被销毁或中止时自动释放。

close() → None

显式释放该资源。

幂等:可以安全地多次调用。

property is_valid*: bool*

资源是否已初始化且仍然有效(未关闭)。

RegisteredBufferHandle

class nccl.core.RegisteredBufferHandle(comm_ptr: int , buffer_ptr: int , size: int )

基类:CommResource

用于零拷贝优化通信的 NCCL 已注册缓冲区句柄。

将用户缓冲区注册到通信器,以在 NCCL 操作中启用性能优化。由 Communicator.register_buffer() 创建。注册句柄可以通过 close() 显式释放,或在所属通信器被销毁或中止时自动释放。

close() → None

显式释放该资源。

幂等:可以安全地多次调用。

property handle*: int*

用于 NCCL 操作的注册句柄。

抛出:

RuntimeError ------ 如果缓冲区已被注销或句柄无效。

property is_valid*: bool*

资源是否已初始化且仍然有效(未关闭)。

property size*: int*

已注册缓冲区的大小(字节)。

RegisteredWindowHandle

class nccl.core.RegisteredWindowHandle(comm_ptr: int , buffer_ptr: int , size: int , flags: WindowFlag | None = None )

基类:CommResource

用于远程内存访问(RMA)操作的 NCCL 已注册窗口句柄。

将内存窗口注册到通信器,以支持单边通信模式。由 Communicator.register_window() 创建。注册是集合式的:默认情况下所有 rank 必须以相等的缓冲区大小调用 Communicator.register_window()。注销是本地的。窗口句柄可以通过 close() 显式释放,或在所属通信器被销毁或中止时自动释放。

close() → None

显式释放该资源。

幂等:可以安全地多次调用。

get_lsa_device_pointer(lsa_rank: int , offset: int = 0 ) → int

返回 LSA team 内某个对端的 LSA 设备指针。

返回一个指向对端窗口缓冲区的设备指针,本地 GPU 可通过 LSA(加载/存储可访问)映射对其寻址。

参数:

  • lsa_rank ------ LSA team 内的 rank(0 到 lsa_size - 1)。

  • offset ------ 窗口缓冲区内的字节偏移量。默认为 0。

返回:

以 int 表示的设备指针。

抛出:

RuntimeError ------ 如果窗口已关闭。

get_lsa_multimem_device_pointer(offset: int = 0 ) → int | None

返回此窗口的 LSA 多播设备指针。

返回一个适用于在 LSA(加载/存储可访问)team 上进行多播操作的设备指针。只要窗口和通信器保持存活,该指针就有效。

参数:

offset ------ 窗口缓冲区内的字节偏移量。默认为 0。

返回:

以 int 表示的设备指针;如果不支持 multimem,则返回 None

抛出:

RuntimeError ------ 如果窗口已关闭。

get_multimem_device_pointer(multimem: MultimemHandle , offset: int = 0 ) → int | None

返回此窗口和指定 multimem 的多播设备指针。

与 get_lsa_multimem_device_pointer()(使用 LSA team 的 multimem)不同,此方法针对设备通信器创建期间产生的显式 multimem 句柄解析指针。

参数:

  • multimem ------ 由 DevCommResource.multimem_handle() 返回的 MultimemHandle。

  • offset ------ 窗口缓冲区内的字节偏移量。默认为 0。

返回:

以 int 表示的设备指针;如果不支持 multimem,则返回 None

抛出:

RuntimeError ------ 如果窗口已关闭。

get_peer_device_pointer(peer: int , offset: int = 0 ) → int | None

按 world rank 返回指向某个对端窗口缓冲区的设备指针。

如果该对端无法通过 LSA 到达,则返回 None

参数:

  • peer ------ 对端的 world rank(0 到 nranks - 1)。

  • offset ------ 窗口缓冲区内的字节偏移量。默认为 0。

返回:

以 int 表示的设备指针;如果对端无法通过 LSA 到达,则返回 None

抛出:

RuntimeError ------ 如果窗口已关闭。

property handle*: int*

用于 NCCL 操作的窗口句柄;注销后为 0

property is_valid*: bool*

窗口是否仍然已注册(未关闭,句柄非空)。

property size*: int*

已注册窗口的大小(字节)。

property user_ptr*: int*

随此窗口注册的原始用户缓冲区指针。

抛出:

RuntimeError ------ 如果窗口已被注销。

CustomRedOp

class nccl.core.CustomRedOp(comm_ptr: int , scalar_ptr: int , datatype: NcclDataType , residence: nccl.bindings.nccl.ScalarResidence )

基类:CommResource

NCCL 用户定义的自定义归约操作符。

由 Communicator.create_pre_mul_sum() 创建。PreMulSum 操作符执行 output = scalar * sum(inputs),适用于求平均或加权归约。该操作符可以通过 close() 显式释放,或在所属通信器被销毁或中止时自动释放。

close() → None

显式释放该资源。

幂等:可以安全地多次调用。

property is_valid*: bool*

资源是否已初始化且仍然有效(未关闭)。

property op*: int*

用于归约操作的操作符句柄。

抛出:

RuntimeError ------ 如果操作符已被销毁或无效。

DevCommResource

class nccl.core.DevCommResource(comm_ptr: int , reqs_lowpp: _nccl_bindings.DevCommRequirements , team_multimem_lowpp: dictNCCLTeam, _nccl_bindings.MultimemHandle | None = None , resource_handle_lowpps: tuple_nccl_bindings.LsaBarrierHandle \| _nccl_bindings.GinBarrierHandle \| _nccl_bindings.LLA2AHandle, ... | None = None )

基类:CommResource

用于设备端操作的 NCCL 设备通信器资源。

封装 ncclDevComm_t 并管理其生命周期。由 Communicator.create_dev_comm() 创建。当父通信器被销毁或中止时,设备通信器会被自动销毁。

close() → None

显式释放该资源。

幂等:可以安全地多次调用。

property is_valid*: bool*

资源是否已初始化且仍然有效(未关闭)。

multimem_handle(team: NCCLTeam ) → MultimemHandle

返回为 team 请求的 multimem 句柄。

返回的外观对象(facade)封装了此设备通信器每次创建的输出存储;每次查找都会在同一存储上创建新的外观对象。在设备通信器关闭之前,它始终由该资源支持。

参数:

team ------ 句柄所请求的 team,即创建此设备通信器时所用 teams 需求中的一个条目。

返回:

NCCL 为 team 填充的 MultimemHandle。

抛出:

  • RuntimeError ------ 如果设备通信器已关闭。

  • KeyError ------ 如果在创建此设备通信器所用的需求中,team 未以 multimem=True 请求。

property ptr*: int*

指向底层 ncclDevComm_t 结构的原始指针。

抛出:

RuntimeError ------ 如果设备通信器已被销毁。

property resource_handles*: tupleLsaBarrierHandle \| GinBarrierHandle \| LLA2AHandle, ...*

最终生成的资源句柄,顺序与 resources 一致。

resource_handles[i] 对应 requirements.resources[i],根据需求的不同,分别是 LsaBarrierHandle、GinBarrierHandle 或 LLA2AHandle。在设备通信器关闭之前,每个句柄都由此资源支持。

设备资源句柄

由 DevCommResource.resource_handles 和 DevCommResource.multimem_handle() 返回的句柄。它们是由所属 DevCommResource 支持的视图,而不是可独立关闭的资源。它们仅在该资源保持打开期间有效。将它们传给设备端 API。

MultimemHandle

class nccl.core.MultimemHandle(*args: Any , **kwargs: Any )

multimem 句柄,由 multimem_handle() 为以 multimem=True 请求的 team 返回。将它传给设备端 multimem 操作。

LsaBarrierHandle

class nccl.core.LsaBarrierHandle(*args: Any , **kwargs: Any )

LSA barrier 句柄,由 resource_handles 为每个 LsaBarrierRequirement 返回。将它传给设备端 barrier 会话。

GinBarrierHandle

class nccl.core.GinBarrierHandle(*args: Any , **kwargs: Any )

GIN barrier 句柄,由 resource_handles 为每个 GinBarrierRequirement 返回。将它传给设备端 barrier 会话。

LLA2AHandle

class nccl.core.LLA2AHandle(*args: Any , **kwargs: Any )

低延迟 all-to-all 句柄,由 resource_handles 为每个 LLA2ARequirement 返回。将它传给设备端 all-to-all 会话。

05.01.06. 类型与常量

出现在公开方法签名中的类型封装、预定义值常量和类型别名。

数据类型

NcclDataType

class nccl.core.NcclDataType(*values )

基类:IntEnum

NCCL 数据类型,与 ncclDataType_t 对应。

用作缓冲区规格的 dtype,以及 NCCL 集合操作的 datatype 参数。支持通过 from_numpy_dtype() 和 numpy_dtype 与 NumPy dtype 相互转换。

INT8 = 0

CHAR = 0

UINT8 = 1

INT32 = 2

INT = 2

UINT32 = 3

INT64 = 4

UINT64 = 5

FLOAT16 = 6

HALF = 6

FLOAT32 = 7

FLOAT = 7

FLOAT64 = 8

DOUBLE = 8

BFLOAT16 = 9

FLOAT8E4M3 = 10

FLOAT8E5M2 = 11

classmethod from_numpy_dtype(dtype: numpy.dtype ) → NcclDataType

将 NumPy dtype 映射到其 NCCL 等价类型。

参数:

dtype ------ 一个 NumPy dtype。先按名称映射(用于 ml-dtypes 中的 bfloat16float8_e4m3fnfloat8_e5m2),然后对标准类型按 (kind, itemsize) 映射。

返回:

对应的 NcclDataType 成员。

抛出:

NcclInvalid ------ 如果该 dtype 没有 NCCL 等价类型。

property itemsize*: int*

此数据类型单个元素的字节大小。

property numpy_dtype*: numpy.dtype*

等价的 NumPy dtype。

返回:

与此 NCCL 数据类型对应的 NumPy dtype。对于 BFLOAT16 和 float8 变体,必须安装 ml-dtypes

抛出:

NcclInvalid ------ 如果需要 ml-dtypes 但未安装。

预定义数据类型常量

模块级 NcclDataType 实例,用作缓冲区规格的 dtype 参数。

常量 映射到
nccl.core.INT8 NcclDataType.INT8
nccl.core.CHAR NcclDataType.CHAR
nccl.core.UINT8 NcclDataType.UINT8
nccl.core.INT32 NcclDataType.INT32
nccl.core.INT NcclDataType.INT
nccl.core.UINT32 NcclDataType.UINT32
nccl.core.INT64 NcclDataType.INT64
nccl.core.UINT64 NcclDataType.UINT64
nccl.core.FLOAT16 NcclDataType.FLOAT16
nccl.core.HALF NcclDataType.HALF
nccl.core.FLOAT32 NcclDataType.FLOAT32
nccl.core.FLOAT NcclDataType.FLOAT
nccl.core.FLOAT64 NcclDataType.FLOAT64
nccl.core.DOUBLE NcclDataType.DOUBLE
nccl.core.BFLOAT16 NcclDataType.BFLOAT16
nccl.core.FLOAT8E4M3 NcclDataType.FLOAT8E4M3
nccl.core.FLOAT8E5M2 NcclDataType.FLOAT8E5M2

归约操作符

NcclRedOp

class nccl.core.NcclRedOp(*values )

基类:IntEnum

NCCL 归约操作符,与 ncclRedOp_t 对应。

用作归约集合操作(Communicator.allreduce()、Communicator.reduce()、Communicator.reduce_scatter())的 op 参数。

SUM = 0

PROD = 1

MAX = 2

MIN = 3

AVG = 4

预定义归约操作符

模块级 NcclRedOp 实例,用作归约集合操作的 op 参数。用户自定义操作符通过 Communicator.create_pre_mul_sum() 创建。

常量 映射到
nccl.core.SUM NcclRedOp.SUM
nccl.core.PROD NcclRedOp.PROD
nccl.core.MAX NcclRedOp.MAX
nccl.core.MIN NcclRedOp.MIN
nccl.core.AVG NcclRedOp.AVG

Team

NCCLTeam

class nccl.core.NCCLTeam(n_ranks: int , rank: int , stride: int )

基类:LowppSpec

一个 NCCL team:对通信器的 (n_ranks, rank, stride) 视图。

由 team_world、team_lsa 和 team_rail 产生。

n_ranks*: int*

rank*: int*

stride*: int*

类型别名

命名公开方法所接受参数类型的别名。每个别名展开为一个具体类型的联合,因此任何成员类型的值都可被接受。

nccl.core.NcclBufferSpec

表示一个联合类型

例如 int | str

nccl.core.NcclScalarSpec

表示一个联合类型

例如 int | str

nccl.core.NcclDeviceSpec

表示一个联合类型

例如 int | str

nccl.core.NcclStreamSpec

表示一个联合类型

例如 int | str

异常

NcclInvalid

Python 侧的校验异常,当公开 API 在参数到达 NCCL 本身之前收到畸形参数时抛出。

exception nccl.core.NcclInvalid(msg )

基类:Exception

当传给 NCCL4Py API 的参数无效时抛出。

用于 Python 层在将调用转发给 NCCL 之前检测到的参数校验错误(例如不支持的 dtype、缓冲区数量不匹配、设备错误)。由 NCCL 本身抛出的错误在绑定层以 NCCLError 报告。

05.01.07. 参数

对 NCCL 可调参数的读取访问。每个参数的含义及设置方法请参见"环境变量"。

复制代码
import nccl.core

nccl.core.params["NCCL_DEBUG"]     # 单个参数的值
list(nccl.core.params)             # 所有参数名
nccl.core.dump_params()            # 将所有参数打印到 stdout

params

nccl.core.params

NCCL 参数名到其当前值的只读 Mapping,由 ncclParamGetParameter() 和 ncclParamGetAllParameterKeys() 支持。值以字符串形式返回。查找是实时的:每次访问都会查询 NCCL,而不是读取缓存快照。

dump_params

nccl.core.dump_params() → None

将 NCCL 参数打印到 stdout。设置 NCCL_PARAM_DUMP_ALL=1 可包含内部参数。

05.01.08. 版本

NCCL4Py 提供了检查已安装 NCCL 栈的辅助工具:nccl4py 本身、其绑定所生成的 NCCL 头文件版本,以及已加载的 libnccl.so

复制代码
import nccl.core

nccl.core.show_versions()      # 向 stdout 输出人类可读的版本块
v = nccl.core.get_version()    # 编程式的版本快照

show_versions

nccl.core.show_versions() → None

打印 nccl4py、绑定和已加载 libnccl 的版本信息。

get_version

nccl.core.get_version() → VersionInfo

返回结构化的 nccl4py、绑定和已加载 libnccl 版本。

VersionInfo

class nccl.core.VersionInfo(nccl4py: \[Version(https://packaging.pypa.io/en/stable/version.html#packaging.version.Version)] , nccl_bindings: \[Version(https://packaging.pypa.io/en/stable/version.html#packaging.version.Version)] , libnccl: LibraryInfo | None )

基类:object

nccl4py、其绑定和已加载 libnccl 的版本快照。

nccl4py*: \[Version(https://packaging.pypa.io/en/stable/version.html#packaging.version.Version)]*

nccl_bindings*: \[Version(https://packaging.pypa.io/en/stable/version.html#packaging.version.Version)]*

libnccl*: LibraryInfo | None*

LibraryInfo

class nccl.core.LibraryInfo(version: \[Version(https://packaging.pypa.io/en/stable/version.html#packaging.version.Version)] , cuda_variant: \[Version(https://packaging.pypa.io/en/stable/version.html#packaging.version.Version)] | None , path: Path | None )

基类:object

libnccl 的版本、CUDA 构建变体和加载路径。

version*: \[Version(https://packaging.pypa.io/en/stable/version.html#packaging.version.Version)]*

cuda_variant*: \[Version(https://packaging.pypa.io/en/stable/version.html#packaging.version.Version)] | None*

path*: Path | None*

05.01.09. 框架互操作

惰性加载的辅助工具,用于分配由 NCCL 管理的内存所支持的 CuPy 数组和 PyTorch 张量,以及将框架对象解析为 NCCL 所期望的 (ptr, count, dtype, device_id) 元组的解析器。这些子模块在首次访问属性时通过 nccl.core.cupynccl.core.torch 导入。

CuPy

nccl.core.interop.cupy.empty(shape: int | tupleint, ..., dtype: str | np.dtype | cupy.dtype | type = <class 'float'>, order: Literal'C', 'F' = 'C' ) → cupy.ndarray

创建一个由 NCCL 分配的内存支持的未初始化 CuPy 数组。

使用 NCCL 的内存分配器返回一个填充了未初始化数据的数组。这提供了一个与 CuPy 兼容的接口,同时使用 NCCL 的内存分配器在分布式场景中进行高效的 GPU 内存管理。与 cupy.empty 不同,其底层内存是通过 NCCL 分配的。

内存在数组被垃圾回收时自动释放;无需显式调用释放。如需零拷贝优化,请使用 register_buffer() 或 register_window() 注册该数组。

参数:

  • shape ------ 数组的形状。

  • dtype ------ 数据类型指定符。默认为 float

  • order ------ 内存布局。'C' 表示行主序(C 风格),'F' 表示列主序(Fortran 风格)。默认为 'C'。

返回:

由 NCCL 分配的内存支持的未初始化 CuPy 数组。

抛出:

  • NcclInvalid ------ 如果 order 不是 'C' 或 'F'。

  • ModuleNotFoundError ------ 如果未安装 CuPy。

nccl.core.interop.cupy.resolve_array(array: cupy.ndarray ) → tupleint, int, NcclDataType, int

将 CuPy 数组解析为其 NCCL 缓冲区描述符。

参数:

array ------ 要解析的 CuPy 数组。

返回:

(ptr, count, dtype, device_id) 元组 ------ 设备指针、元素数量、NCCL 数据类型和 CUDA 设备 ID。

抛出:

  • ModuleNotFoundError ------ 如果未安装 CuPy。

  • NcclInvalid ------ 如果 array 不是 CuPy ndarray,或其 dtype 没有 NCCL 等价类型。

PyTorch

nccl.core.interop.torch.empty(*size , dtype: torch.dtype | None = None , device: torch.device | int | str | None = None , morder: Literal'C', 'F' = 'C' ) → torch.Tensor

创建一个由 NCCL 分配的内存支持的未初始化 PyTorch 张量。

使用 NCCL 的内存分配器返回一个填充了未初始化数据的张量。这提供了一个与 PyTorch 兼容的接口,同时使用 NCCL 的内存分配器在分布式场景中进行高效的 GPU 内存管理。与 torch.empty 不同,其底层内存是通过 NCCL 分配的。

内存在张量被垃圾回收时自动释放;无需显式调用释放。如需零拷贝优化,请使用 register_buffer() 或 register_window() 注册该张量。

参数:

  • *size ------ 定义输出张量形状的整数序列。可以是可变数量的参数,也可以是单个列表/元组。

  • dtype ------ 张量所需的数据类型。如果为 None,使用 torch.get_default_dtype()。默认为 None

  • device ------ 张量的设备。如果为 None,使用当前 CUDA 设备。默认为 None

  • morder ------ 内存布局。'C' 表示行主序(C 风格),'F' 表示列主序(Fortran 风格)。默认为 'C'。

返回:

由 NCCL 分配的内存支持的未初始化 PyTorch 张量。

抛出:

  • NcclInvalid ------ 如果 morder 不是 'C' 或 'F',或 device 不是 CUDA 设备。

  • ModuleNotFoundError ------ 如果未安装 PyTorch。

nccl.core.interop.torch.resolve_tensor(tensor: torch.Tensor ) → tupleint, int, NcclDataType, int

将 PyTorch 张量解析为其 NCCL 缓冲区描述符。

参数:

tensor ------ 要解析的 PyTorch 张量。

返回:

(ptr, count, dtype, device_id) 元组 ------ 设备指针、元素数量、NCCL 数据类型和 CUDA 设备 ID。

抛出:

  • ModuleNotFoundError ------ 如果未安装 PyTorch。

  • NcclInvalid ------ 如果 tensor 不是 PyTorch 张量,或其 dtype 没有 NCCL 等价类型。

相关推荐
探索云原生19 小时前
Kueue + HAMi vGPU 实战:显存与算力配额管理
docker·ai·云原生·kubernetes·gpu
cubestudio19 小时前
海光 DCU 怎么接入 Kubernetes 和 AI 平台?CubeStudio 海光 DCU 适配实操(整卡 / 共享 / 两种 vDCU 虚拟化 + DeepSeek 部署)
人工智能·机器学习·gpu
阿里云大数据AI技术2 天前
EMR Serverless Spark:CPU + GPU 异构计算使用指南
人工智能·spark·gpu
Eloudy4 天前
NVTx 主旨介绍
gpu
Eloudy4 天前
NVLS 简介
gpu
Felven4 天前
Intel Core Ultra X9 388H全面性能对比分析报告
cpu·gpu·intel·性能对比·amd
Eloudy5 天前
全文 - 03 part - NVIDIA 集合通信库(NCCL)文档
gpu
每日出拳老爷子5 天前
【AI】Ollama 更新后 skipping CUDA 掉回 CPU
gpu·nvidia·cuda·ollama·本地大模型
晨欣5 天前
NVIDIA GPU 架构演进学习笔记(GPT-5.6 Terra 生成)
笔记·学习·gpu·显卡·nvidia·英伟达