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 实例公开了若干用于检查的属性(ptr、nranks、device、rank,以及与设备 API 相关的属性,如 cuda_dev、nvml_dev、device_api_support、multimem_support、gin_type、n_lsa_teams、host_rma_support、railed_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 都必须以相同的 nranks 和 unique_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 将属于同一组。如果 color 为 None,该 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。
返回:
新的子通信器;如果 color 为 None,则返回空通信器。
抛出:
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。当 sendbuf 和 recvbuf 解析到相同的设备内存地址时,执行就地(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 模式下(root 为 None),所有 rank 都在 recvbuf 中收到归约结果。在 Reduce 模式下(指定 root),只有根 rank 收到归约结果;其他 rank 上的 recvbuf 被忽略。
在使用处,两个缓冲区必须具有匹配的数据类型。元素数量从 sendbuf 推断:count = sendcount。在 AllReduce 模式下,所有 rank 必须满足 recvcount >= sendcount;在 Reduce 模式下,只有根 rank 需要满足 recvcount >= sendcount。当 sendbuf 和 recvbuf 解析到相同的设备内存地址时,执行就地操作。
参数:
-
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 / nranks。sendcount 必须 >= 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 / nranks。sendcount 必须 >= 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 模式下(root 为 None),从所有 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 i。sendbuf 在非根 rank 上不使用。
在根 rank 上,两个缓冲区必须具有匹配的数据类型。元素数量从 recvbuf 推断:count = recvcount。根 rank 要求 sendcount >= nranks 且 sendcount / 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推断:Pythonint变为int64,Pythonfloat变为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() 销毁通信器;发生错误后,无法对已入队操作的完成情况或正确性做任何假设。
返回:
通信器的当前状态(ncclSuccess、ncclInProgress 或错误码)。
抛出:
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 时返回 None;simulate=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 中的 bfloat16、float8_e4m3fn、float8_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.cupy 和 nccl.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 等价类型。