分发

torch.distribute 工具。

类

NFSWorkspace

使用网络文件存储（NFS）实现的共享工作区。

函数

get_configs_parallel	使用共享内存（shm）或网络文件系统（NFS）在分布式进程中收集层配置。
get_tensors_parallel	使用共享内存（shm）在分布式进程中收集张量。

class NFSWorkspace

基础类：object

使用网络文件存储（NFS）实现的共享工作区。

NOTE: all read/write/modifition to the NFS dir do not involve any collective

通信无障碍。用户有责任同步所有等级（本地和远程进程）。

此实现使用torch.save和torch.load进行序列化。

Parameters:

workspace_path – 用于后处理跨秩通信的NFS目录路径。 如果未提供，则将使用SharedMemory代替。

__init__(workspace_path=None)

创建NFS工作目录并清理现有的状态文件。

Parameters:

workspace_path (Path | str | None) –

property is_initialized

工作区是否已初始化。

read_configs_and_weights_from_rank(target_rank)

所有等级读取target_rank状态文件。

Parameters:

target_rank (int) – 目标排名

Returns:

模型/模块配置和权重

Return type:

Tuple[Dict[str, Any] | None, Dict[str, Any] | None]

write_configs_and_weights(config_json, weights)

所有等级将状态文件写入共享的NFS目录。

Parameters:

• config_json (Dict[str, Any]) – 模型或模块的配置，以json格式

• weights (Dict[str, Any]) – 模块权重，以torch的state_dict格式表示

get_configs_parallel(config, ranks, group, workspace_path=None)

使用共享内存（shm）或网络文件系统（NFS）在分布式进程中收集层配置。

Parameters:

• config – 每个等级想要传递给第一个等级的配置（可为空）。

• ranks (List[int]) – 排名列表

• group – 屏障同步组。

• workspace_path (Path | str | None) – 用于跨秩通信后处理的NFS目录路径。

Yields:

在所有等级中，第一等级拥有对所有等级配置的完全访问权限。其他等级返回一个空列表。

当提供workspace_path时，将创建一个NFSWorkspace对象以跨等级执行通信。否则，SharedMemory用于本地多进程通信。shm将在使用后被销毁。

get_tensors_parallel(tensor, ranks, group=None)

使用共享内存（shm）在分布式进程中收集张量。

Parameters:

• tensor (Tensor) – 每个等级想要传递给第一个等级的张量。 各个等级的张量需要具有相同的大小。

• ranks (List[int]) – 排名的列表

• group – 屏障同步组。

Yields:

排名第一的等级拥有对所有等级的张量的完全访问权限。 其他等级返回一个空列表

共享内存将在消费后被销毁。