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 的值。它仅用于调试目的。
debug_copy_from(worker_id:int, value: ndarray |NDArray) → None
将 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 的全局函数。
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。
- 参数:
- 返回: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)
初始化底层通信集合库。
broadcast(src: ndarray |NDArray, dst:DRef|None= None, in_group:bool= True) → DRef
向所有 worker 广播一个数组。
- 参数:
- 返回:output_array: 包含所有 Worker 上广播数据的 DRef。如果提供了 dst ,则此返回值与 dst 相同 。否则,返回新分配的空间。
- 返回类型: DRef。
broadcast_from_worker0(src:DRef, dst:DRef, in_group:bool= True) → DRef
将数组从 worker-0 广播到所有其他 worker。
- 参数:
scatter(src: ndarray |NDArray, dst:DRef|None= None, in_group:bool= True) → DRef
将数组分散到所有 worker 中。
- 参数:
- 返回:output_array:包含所有 Worker 的分散数据的 DRef。如果提供了dst ,则此返回值与dst相同 。否则,返回新分配的空间。
- 返回类型:DRef。
scatter_from_worker0(from_array:DRef, to_array:DRef, in_group:bool= True) → None
将数组从 worker-0 分散到所有其他 workers。
- 参数:
gather_to_worker0(from_array:DRef, to_array:DRef, in_group:bool= True) → None
将所有其他 worker 的数组收集到 worker-0。
allreduce(src:DRef, dst:DRef, op:str= 'sum', in_group:bool= True) → DRef
对数组执行 allreduce 操作。
- 参数:
allgather(src: DRef, dst: DRef, in_group: bool = True)→ DRef
对一个数组执行 allgather 操作。
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 会话。