跳到主要内容
版本：0.21.0

tvm 运行时的一个模块

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​

一个在所有工作进程上都存在的对象。控制进程为每个对象分配一个唯一的「register ID」，工作进程使用这个 ID 来引用驻留在自身进程上的对象。

debug_get_from_remote(worker_id:int) → Any​

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

  • 参数：worker_id (int) ：需要获取的 worker 的 ID。
  • 返回：value：寄存器的值。
  • 返回类型：object。

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

将 NDArray 值复制到远程以用于调试目的。

  • 参数：
    • worker_id (int) ：要复制到的 worker 的 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​

在所有 worker 上创建一个空的 NDArray 并将它们附加到 DRef。

  • 参数：
    • shape (tupleofint) ：NDArray 的形状。
    • dtype (str) ：NDArray 的数据类型。
    • device (Optional*[***Device]= None) ：NDArray 的设备。
    • worker0_only (bool) ：如果为 False（默认值），则为每个 worker 分配一个数组。如果为 True，则仅在 worker0 上分配一个数组。
    • in_group (bool) ：当worker0_only为 True时生效。如果为 True（默认），则在每个组的第一个 Worker 上分配一个数组。如果为 False，则仅在全局范围内为 worker0 分配一个数组。
  • 返回：array ：创建的 NDArray。
  • 返回类型： DRef。

shutdown()​

关闭 Disco 会话。

property num_workers*:*int​

返回会话中的 worker 数量。

get_global_func(name:str) → DRef​

获取 worker 的全局函数。

  • 参数：name (str) ：全局函数的名称。
  • 返回：func ：全局打包函数。
  • 返回类型：DRef。

import_python_module(module_name:str) → None​

在每个 worker 中导入一个 python 模块。

通话前可能需要执行此操作。

  • 参数：module_name (str) ：Python 模块名称，它将用于 Python 导入语句中。

call_packed(func:DRef, *args) → DRef​

在提供可变参数的 worker 上调用 PackedFunc。

  • 参数：
    • func (PackedFunc) ：要调用的函数。
    • *args (various types) ：DRef。
  • 返回：return_value：函数调用的返回值。
  • 返回类型： 各种类型。

不支持的类型的示例： - 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) ：0的数组。
    • remote_array (NDArray) ：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]) ：0 上的目标 NDArray。
  • 返回：output_array： 包含 worker0 上复制数据的 DRef，以及所有其他 Worker 上的 std::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 模块的设备。默认为每个 Worker 的默认设备。
  • 返回：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​

向所有 worker 广播一个数组。

  • 参数：
    • src (Union*[**np.ndarray,NDArray]*) ：要广播的数组。
    • dst (Optional[DRef]) ：输出数组。如果为 None，则将在每个 Worker 上分配一个与src的形状和数据类型匹配的数组。
    • in_group (bool) ：广播操作默认是全局执行还是按组执行。
  • 返回：output_array： 包含所有 Worker 上广播数据的 DRef。如果提供了 dst ，则此返回值与 dst 相同 。否则，返回新分配的空间。
  • 返回类型： DRef。

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

将数组从 worker-0 广播到所有其他 worker。

  • 参数：
    • src (Union*[**np.ndarray,NDArray]*) ：要广播的数组。
    • dst (Optional[DRef]) ：输出数组。如果为 None，则将在每个 Worker 上分配一个与src的形状和数据类型匹配的数组。
    • in_group (bool) 广播操作默认是全局执行还是按组执行。

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

将数组分散到所有 worker 中。

  • 参数：
    • src (Union*[**np.ndarray,NDArray]*) ：待分散的数组。该数组的第一个维度src.shape[0]必须等于工作线程的数量。
    • dst (Optional[DRef]) ：输出数组。如果为 None，则将在每个 Worker 上分配一个形状兼容且数据类型与 src 相同的数组。
    • in_group (bool) ：分散操作是全局执行还是默认在组中执行。
  • 返回：output_array：包含所有 Worker 的分散数据的 DRef。如果提供了dst ，则此返回值与dst相同 。否则，返回新分配的空间。
  • 返回类型：DRef。

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

将数组从 worker-0 分散到所有其他 workers。

  • 参数：
    • src (Union*[**np.ndarray,NDArray]*) ：待分散的数组。该数组的第一个维度src.shape[0]必须等于工作线程的数量。
    • dst (Optional[DRef]) ：输出数组。如果为 None，则将在每个 Worker 上分配一个形状兼容且数据类型与src相同的数组。
    • in_group (bool) ：分散操作是全局执行还是默认在组中执行。

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

将所有其他 worker 的数组收集到 worker-0。

  • 参数：
    • from_array (DRef) ：要从中收集的数组。
    • to_array (DRef) ：要收集的数组。
    • in_group (bool) ：收集操作是全局执行还是默认以组执行。

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) ：减少操作是否全局执行或默认在组中执行。

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

对一个数组执行 allgather 操作。

  • 参数：
    • src (DRef) ：要收集的数组。
    • dst (DRef) ：要收集的数组。
    • in_group (bool) ：减少操作是否全局执行或默认在组中执行。

classtvm.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 会话。

编辑此页
最后由 wenjing-user 于 更新