tvm.runtime.disco

TVM 分布式运行时 API。

class tvm.runtime.disco.DModule(dref: DRef, session: Session)

Disco 会话中的模块。

class tvm.runtime.disco.DPackedFunc(dref: DRef, session: Session)

Disco 会话中的 PackedFunc。

class tvm.runtime.disco.DRef

存在于所有工作进程上的对象。控制器进程为每个对象分配唯一的“注册 ID”，工作进程使用此 ID 来引用驻留在其自身上的对象。

debug_get_from_remote(worker_id: int) → Any

从远程工作进程获取 DRef 的值。它仅用于调试目的。

参数:

worker_id (int) – 要从中获取的工作进程的 ID。

返回值:

value – 寄存器的值。

返回类型:

object

debug_copy_from(worker_id: int, value: ndarray | NDArray) → None

复制 NDArray 值到远程以进行调试。

参数:

  • worker_id (int) – 要复制到的工作进程的 ID。

  • value (Union[numpy.ndarray, NDArray]) – 要复制的值。

class tvm.runtime.disco.ProcessSession(num_workers: int, num_groups: int = 1, entrypoint: str = 'tvm.exec.disco_worker')

由基于管道的多进程支持的 Disco 会话。

class tvm.runtime.disco.Session

Disco 交互式会话。它允许用户使用各种 PackedFunc 调用约定与 Disco 命令队列进行交互。

empty(shape: Sequence[int], dtype: str, device: Device | None = None, worker0_only: bool = False, in_group: bool = True) → DRef

在所有工作进程上创建空的 NDArray 并将其附加到 DRef。

参数:

  • shape (tuple of int) – NDArray 的形状。

  • dtype (str) – NDArray 的数据类型。

  • device (Optional[Device] = None) – NDArray 的设备。

  • worker0_only (bool) – 如果为 False（默认），则在每个工作进程上分配一个数组。如果为 True，则仅在 worker0 上分配一个数组。

  • in_group (bool) – 当 worker0_only 为 True 时生效。如果为 True（默认），则在每个组的第一个工作进程上分配一个数组。如果为 False，则仅在全局 worker0 上分配一个数组。

返回值:

array – 创建的 NDArray。

返回类型:

DRef

shutdown()

关闭 Disco 会话

property num_workers: int

返回会话中的工作进程数

get_global_func(name: str) → DRef

在工作进程上获取全局函数。

参数:

name (str) – 全局函数的名称。

返回值:

func – 全局 PackedFunc 函数

返回类型:

DRef

import_python_module(module_name: str) → None

在每个工作进程中导入一个 Python 模块

这可能在调用之前是必需的

参数:

module_name (str) – Python 模块名称，就像在 python `import` 语句中使用它一样。

call_packed(func: DRef, *args) → DRef

在工作进程上调用 PackedFunc，提供可变参数。

参数:

  • func (PackedFunc) – 要调用的函数。

  • *args (various types) – 在可变参数中，支持的类型包括：- 整数和浮点数；- DLDataType；- DLDevice；- str (C++ 中的 std::string)；- DRef。

返回值:

return_value – 函数调用的返回值。

返回类型:

various types

注意

不支持的类型示例：- NDArray、DLTensor；- TVM 对象，包括 PackedFunc、Module 和 String。

sync_worker_0() → None

将控制器与 worker-0 同步，它将等待直到 worker-0 完成执行所有现有指令。

copy_from_worker_0(host_array: NDArray, remote_array: DRef) → None

将 NDArray 从 worker-0 复制到控制器端的 NDArray。

参数:

  • host_array (numpy.ndarray) – 要复制到 worker-0 的数组。

  • remote_array (NDArray) – worker-0 上的 NDArray。

copy_to_worker_0(host_array: NDArray, remote_array: DRef | None = None) → DRef

将控制器端的 NDArray 复制到 worker-0。

参数:

  • host_array (NDArray) – 要复制到 worker-0 的数组。

  • remote_array (Optiona[DRef]) – worker-0 上的目标 NDArray。

返回值:

output_array – 包含 worker0 上复制数据的 DRef，以及所有其他工作进程上的 NullOpt。如果提供了 remote_array，则此返回值与 remote_array 相同。否则，它是新分配的空间。

返回类型:

DRef

load_vm_module(path: str, device: Device | None = None) → DModule

从文件加载 VM 模块。

参数:

  • path (str) – VM 模块文件的路径。

  • device (Optional[Device] = None) – 加载 VM 模块的设备。默认为每个工作进程的默认设备。

返回值:

module – 加载的 VM 模块。

返回类型:

DModule

init_ccl(ccl: str, *device_ids)

初始化底层通信集合库。

参数:

  • ccl (str) – 通信集合库的名称。当前支持的库有：- nccl - rccl - mpi

  • *device_ids (int) – 底层通信库要使用的设备 ID。

broadcast(src: ndarray | NDArray, dst: DRef | None = None, in_group: bool = True) → DRef

将数组广播到所有工作进程

参数:

  • src (Union[np.ndarray, NDArray]) – 要广播的数组。

  • dst (Optional[DRef]) – 输出数组。如果为 None，则将在每个工作进程上分配一个与 src 的形状和 dtype 匹配的数组。

  • in_group (bool) – 广播操作是默认全局执行还是在组内执行。

返回值:

output_array – 包含所有工作进程上广播数据的 DRef。如果提供了 dst，则此返回值与 dst 相同。否则，它是新分配的空间。

返回类型:

DRef

broadcast_from_worker0(src: DRef, dst: DRef, in_group: bool = True) → DRef

将数组从 worker-0 广播到所有其他工作进程。

参数:

  • src (Union[np.ndarray, NDArray]) – 要广播的数组。

  • dst (Optional[DRef]) – 输出数组。如果为 None，则将在每个工作进程上分配一个与 src 的形状和 dtype 匹配的数组。

  • in_group (bool) – 广播操作是默认全局执行还是在组内执行。

scatter(src: ndarray | NDArray, dst: DRef | None = None, in_group: bool = True) → DRef

将数组分散到所有工作进程。

参数:

  • src (Union[np.ndarray, NDArray]) – 要分散的数组。此数组的第一个维度 src.shape[0] 必须等于工作进程数。

  • dst (Optional[DRef]) – 输出数组。如果为 None，则将在每个工作进程上分配一个具有兼容形状且与 src 相同 dtype 的数组。

  • in_group (bool) – 分散操作是默认全局执行还是在组内执行。

返回值:

output_array – 包含所有工作进程上分散数据的 DRef。如果提供了 dst，则此返回值与 dst 相同。否则，它是新分配的空间。

返回类型:

DRef

scatter_from_worker0(from_array: DRef, to_array: DRef, in_group: bool = True) → None

将数组从 worker-0 分散到所有其他工作进程。

参数:

  • src (Union[np.ndarray, NDArray]) – 要分散的数组。此数组的第一个维度 src.shape[0] 必须等于工作进程数。

  • dst (Optional[DRef]) – 输出数组。如果为 None，则将在每个工作进程上分配一个具有兼容形状且与 src 相同 dtype 的数组。

  • in_group (bool) – 分散操作是默认全局执行还是在组内执行。

gather_to_worker0(from_array: DRef, to_array: DRef, in_group: bool = True) → None

将数组从所有其他工作进程收集到 worker-0。

参数:

  • from_array (DRef) – 要从中收集的数组。

  • to_array (DRef) – 要收集到的数组。

  • in_group (bool) – Gather 操作是全局执行还是在组内执行（默认为组内）。

allreduce(src: DRef, dst: DRef, op: str = 'sum', in_group: bool = True) → DRef

对数组执行 allreduce 操作。

参数:

  • array (DRef) – 要进行规约的数组。

  • op (str = "sum") – 要执行的规约操作。可用选项包括：- “sum” - “prod” - “min” - “max” - “avg”

  • in_group (bool) – Reduce 操作是全局执行还是在组内执行（默认为组内）。

allgather(src: DRef, dst: DRef, in_group: bool = True) → DRef

对数组执行 allgather 操作。

参数:

  • src (DRef) – 要从中收集的数组。

  • dst (DRef) – 要收集到的数组。

  • in_group (bool) – Reduce 操作是全局执行还是在组内执行（默认为组内）。

class tvm.runtime.disco.ThreadedSession(num_workers: int, num_groups: int = 1)

一个由多线程支持的 Disco 会话。

class tvm.runtime.disco.SocketSession(num_nodes: int, num_workers_per_node: int, num_groups: int, host: str, port: int)

一个由基于套接字的多节点通信支持的 Disco 会话。