VQNet的分布式计算模块¶

分布式计算是指通过多台设备（如GPU/CPU节点）协同完成神经网络的训练或推理任务，利用并行处理加速计算并扩展模型规模。其核心是通过分布式接口（如MPI、NCCL）协调设备间的通信与同步

VQNet的分布式计算模块使用mpi启动多进程并行计算, 使用nccl进行GPU之间通信。该功能仅在linux操作系统下能够使用。

重要

VQNet的分布式计算模块仅支持Linux。

环境部署¶

以下介绍VQNet分别基于CPU、GPU分布式计算所需的Linux系统下环境的部署.该部分必须有MPI的支持, 以下介绍MPI的环境部署。

MPI安装¶

MPI为CPU间通信的常用库, VQNet中CPU的分布式计算功能则基于MPI进行实现,以下将介绍如何在Linux系统中对MPI进行安装(目前基于CPU的分布式计算功能仅在Linux上实现)。

本软件当前编译依赖的是mpicxx==4.1.2,您可以通过 conda 或其他方式安装。

conda install conda-forge::mpich-mpicxx==4.1.2

此外我们还必须安装 mpi4py 库。通过pip install完成mpi4py的安装即可, 若是出现以下类似错误

为mpi4py与 Python 版本之间不兼容的问题, 可以通过以下方法解决

# 通过下列代码暂存当前python环境的编译器
pushd /root/anaconda3/envs/$CONDA_DEFAULT_ENV/compiler_compat && mv ld ld.bak && popd

# 再次安装
pip install mpi4py

# 还原
pushd /root/anaconda3/envs/$CONDA_DEFAULT_ENV/compiler_compat && mv ld.bak ld && popd

NCCL安装¶

NCCL为GPU间通信的常用库, VQNet中GPU的分布式计算功能则基于NCCL进行实现,本软件默认在安装时候同时安装NCCL的动态链接库, 一般不需要安装NCCL。

节点间通信环境部署¶

在多节点上实现分布式计算,首先 需要保证多节点上mpich环境的一致,Python 环境一致 ,其次,需要设置 节点间的免密通信 。

假设需要设置node0(主节点)、node1、node2三个节点的免密通信。

# 在每个节点上执行
ssh-keygen -t rsa

# 之后一直回车,在.ssh文件夹下生成一个公钥(id_rsa.pub)和一个私钥(id_rsa)
# 使用ssh-copy-id将各节点的公钥复制到主节点node0
# 在子节点node1上执行
ssh-copy-id node0

# 在子节点node2上执行
ssh-copy-id node0

# 在主节点node0上,将node0自己的公钥也加入authorized_keys
cat ~/.ssh/id_rsa.pub >> ~/.ssh/authorized_keys

# 再将node0上的authorized_keys文件分发到各子节点
scp ~/.ssh/authorized_keys node1:~/.ssh/authorized_keys
scp ~/.ssh/authorized_keys node2:~/.ssh/authorized_keys

# 保证三个不同节点生成的公钥都在authorized_keys文件中,即可实现节点间的免密通信

可选的, 最好还设置一个共享目录,使得改变共享目录下的文件时,不同节点中文件也会进行更改,预防多节点运行模型时不同节点中的文件不同步的问题。使用nfs-utils和rpcbind实现共享目录。

# 安装软件包
yum -y install nfs* rpcbind

# 编辑主节点上配置文件
vim /etc/exports
/data/mpi *(rw,sync,no_all_squash,no_subtree_check)

# 主节点上启动服务
systemctl start rpcbind
systemctl start nfs

# 在所有子节点node1,node2上mount主节点node0的共享目录
mount node0:/data/mpi /data/mpi

分布式启动¶

使用分布式计算接口,通过 vqnetrun 命令启动, 接下来介绍 vqnetrun 的各个参数.

n, np¶

vqnetrun 接口中可以通过 -n, -np 参数控制启动的进程数,执行样例如下:

Example:

from pyvqnet.distributed import CommController
Comm_OP = CommController("mpi") # init mpi controller

rank = Comm_OP.getRank()
size = Comm_OP.getSize()
print(f"rank: {rank}, size {size}")

# vqnetrun --backend nccl --nproc_per_node 2 python test.py
# vqnetrun -np 2 python test.py

backend¶

vqnetrun 接口中可以通过 --backend 参数选择分布式后端,支持 mpi (默认) 和 nccl 两种模式。 MPI 模式支持多节点跨主机执行, NCCL 模式用于单机多 GPU 场景,采用去中心化启动模型(每个节点独立运行 vqnetrun)。

Example:

from pyvqnet.distributed import CommController
Comm_OP = CommController("nccl")

rank = Comm_OP.getRank()
size = Comm_OP.getSize()
print(f"rank: {rank}, size {size}")

# 单机多GPU:
# vqnetrun --backend nccl --nproc_per_node 2 python test.py

# 多节点情况下,每个节点独立执行:
# 节点0: vqnetrun --backend nccl --nproc_per_node 2 --nnodes 2 --node_rank 0 --master_addr 10.0.0.1 python test.py
# 节点1: vqnetrun --backend nccl --nproc_per_node 2 --nnodes 2 --node_rank 1 --master_addr 10.0.0.1 python test.py

nproc_per_node¶

vqnetrun 接口中可以通过 --nproc_per_node 参数控制每个节点启动的进程数,仅在 NCCL 模式下使用。

Example:

from pyvqnet.distributed import CommController
Comm_OP = CommController("nccl")

rank = Comm_OP.getRank()
size = Comm_OP.getSize()
print(f"rank: {rank}, size {size}")

# vqnetrun --backend nccl --nproc_per_node 4 python test.py

nnodes¶

vqnetrun 接口中可以通过 --nnodes 参数控制总节点数,仅在 NCCL 模式下使用,默认值为 1。

Example:

from pyvqnet.distributed import CommController
Comm_OP = CommController("nccl")

rank = Comm_OP.getRank()
size = Comm_OP.getSize()
print(f"rank: {rank}, size {size}")

# vqnetrun --backend nccl --nproc_per_node 2 --nnodes 2 --node_rank 0 --master_addr 10.0.0.1 python test.py

node_rank¶

vqnetrun 接口中可以通过 --node_rank 参数控制当前节点的 rank,仅在 NCCL 模式下使用,默认值为 0。

Example:

from pyvqnet.distributed import CommController
Comm_OP = CommController("nccl")

rank = Comm_OP.getRank()
size = Comm_OP.getSize()
print(f"rank: {rank}, size {size}")

# 节点0: vqnetrun --backend nccl --nproc_per_node 2 --nnodes 3 --node_rank 0 --master_addr 10.0.0.1 python test.py
# 节点1: vqnetrun --backend nccl --nproc_per_node 2 --nnodes 3 --node_rank 1 --master_addr 10.0.0.1 python test.py
# 节点2: vqnetrun --backend nccl --nproc_per_node 2 --nnodes 3 --node_rank 2 --master_addr 10.0.0.1 python test.py

master_addr¶

vqnetrun 接口中可以通过 --master_addr 参数指定主节点地址,仅在 NCCL 模式下使用,默认值为 127.0.0.1。

Example:

from pyvqnet.distributed import CommController
Comm_OP = CommController("nccl")

rank = Comm_OP.getRank()
size = Comm_OP.getSize()
print(f"rank: {rank}, size {size}")

# vqnetrun --backend nccl --nproc_per_node 2 --nnodes 2 --node_rank 0 --master_addr 10.0.0.1 python test.py

master_port¶

vqnetrun 接口中可以通过 --master_port 参数指定主节点端口,仅在 NCCL 模式下使用,默认值为 29500。

Example:

from pyvqnet.distributed import CommController
Comm_OP = CommController("nccl")

rank = Comm_OP.getRank()
size = Comm_OP.getSize()
print(f"rank: {rank}, size {size}")

# vqnetrun --backend nccl --nproc_per_node 2 --master_port 29600 python test.py

nccl_socket_ifname¶

vqnetrun 接口中可以通过 --nccl_socket_ifname 参数指定 NCCL 使用的网络接口名,仅在 NCCL 模式下使用。该参数对应环境变量 NCCL_SOCKET_IFNAME,用于在多网卡环境下指定 NCCL 通信使用的网卡。

Example:

from pyvqnet.distributed import CommController
Comm_OP = CommController("nccl")

rank = Comm_OP.getRank()
size = Comm_OP.getSize()
print(f"rank: {rank}, size {size}")

# vqnetrun --backend nccl --nproc_per_node 2 --nccl_socket_ifname eth0 python test.py

H, hosts¶

vqnetrun 接口中可以通过 -H, --hosts 指定节点以及进程分配来跨节点执行(在跨节点运行时必须将节点的环境配置成功, 在相同的环境,相同的路径下执行),执行样例如下:

Example:

from pyvqnet.distributed import CommController, get_host_name
Comm_OP = CommController("mpi") # init mpi controller

rank = Comm_OP.getRank()
size = Comm_OP.getSize()
print(f"rank: {rank}, size {size}")
print(f"LocalRank {Comm_OP.getLocalRank()} hosts name {get_host_name()}")

# vqnetrun -np 4 -H node0:1,node2:1 python test.py
# vqnetrun -np 4 --hosts node0:1,node2:1 python test.py

hostfile, f, hostfile¶

vqnetrun 接口中可以通过指定hosts文件来指定节点以及进程分配来跨节点(在跨节点运行时必须将节点的环境配置成功, 在相同的环境,相同的路径下执行), 命令行参数为 -hostfile, -f, --hostfile.

文件内每行的格式必须为:<hostname> slots=<slots> 如；

node0 slots=1

node2 slots=1

执行样例如下

Example:

from pyvqnet.distributed import CommController, get_host_name
Comm_OP = CommController("mpi") # init mpi controller

rank = Comm_OP.getRank()
size = Comm_OP.getSize()
print(f"rank: {rank}, size {size}")
print(f"LocalRank {Comm_OP.getLocalRank()} hosts name {get_host_name()}")

# vqnetrun -np 4 -f hosts python test.py
# vqnetrun -np 4 -hostfile hosts python test.py
# vqnetrun -np 4 --hostfile hosts python test.py

output-filename¶

vqnetrun 接口中可以通过命令行参数 --output-filename 来将输出结果保存到指定文件.

执行样例如下

Example:

from pyvqnet.distributed import CommController, get_host_name
Comm_OP = CommController("mpi") # init mpi controller

rank = Comm_OP.getRank()
size = Comm_OP.getSize()
print(f"rank: {rank}, size {size}")
print(f"LocalRank {Comm_OP.getLocalRank()} hosts name {get_host_name()}")

# vqnetrun -np 4 --hostfile hosts --output-filename output  python test.py

verbose¶

vqnetrun 接口中可以通过命令行参数 --verbose 来对节点间的通信进行检测,并额外输出检测结果。

执行样例如下

Example:

from pyvqnet.distributed import CommController, get_host_name
Comm_OP = CommController("mpi") # init mpi controller

rank = Comm_OP.getRank()
size = Comm_OP.getSize()
print(f"rank: {rank}, size {size}")
print(f"LocalRank {Comm_OP.getLocalRank()} hosts name {get_host_name()}")

# vqnetrun -np 4 --hostfile hosts --verbose python test.py

start-timeout¶

vqnetrun 接口中可以通过命令行参数 --start-timeout 来指定超时前执行所有检查并启动进程。默认值为 30 秒。

执行样例如下

Example:

from pyvqnet.distributed import CommController, get_host_name
Comm_OP = CommController("mpi") # init mpi controller

rank = Comm_OP.getRank()
size = Comm_OP.getSize()
print(f"rank: {rank}, size {size}")
print(f"LocalRank {Comm_OP.getLocalRank()} hosts name {get_host_name()}")

# vqnetrun -np 4 --start-timeout 10 python test.py

disable-cache¶

vqnetrun 接口中可以通过 --disable-cache 标志禁用初始化检查缓存。默认情况下, vqnetrun 会将初始化检查的结果缓存 60 分钟, 如果禁用了缓存, 则每次运行都会重新执行初始化检查。

执行代码如下

Example:

from pyvqnet.distributed import CommController, get_host_name
Comm_OP = CommController("mpi") # init mpi controller

rank = Comm_OP.getRank()
size = Comm_OP.getSize()
print(f"rank: {rank}, size {size}")
print(f"LocalRank {Comm_OP.getLocalRank()} hosts name {get_host_name()}")

# vqnetrun -np 4 --disable-cache python test.py

cb, check-build¶

vqnetrun 接口中可以通过 -cb, --check-build 标志查看当前构建支持的分布式后端和通信库。

执行代码如下

Example:

# vqnetrun -cb
# vqnetrun --check-build

CommController¶

class pyvqnet.distributed.ControlComm.CommController(backend, rank=None, world_size=None)¶

CommController用于控制在cpu、gpu下数据通信的控制器, 通过设置参数 backend 来生成cpu(mpi)、gpu(nccl)的控制器。(目前分布式计算的功能仅支持linux操作系系统下使用)

参数:

backend – 用于生成cpu或者gpu的数据通信控制器。
rank – 该参数仅在非pyvqnet后端下有用, 默认值为: None。
world_size – 该参数仅在非pyvqnet后端下有用, 默认值为: None。

返回:

CommController 实例。

Example:

from pyvqnet.distributed import CommController
Comm_OP = CommController("nccl") # init nccl controller

# Comm_OP = CommController("mpi") # init mpi controller

getRank()¶

用于获得当前进程的进程号。

返回:: 返回当前进程的进程号。

Example:

from pyvqnet.distributed import CommController
Comm_OP = CommController("nccl") # init nccl controller

Comm_OP.getRank()

getSize()¶

用于获得总共启动的进程数。

返回:: 返回总共进程的数量。

Example:

from pyvqnet.distributed import CommController
Comm_OP = CommController("nccl") # init nccl controller

Comm_OP.getSize()
# vqnetrun --backend nccl --nproc_per_node 2 python test.py

getLocalRank()¶

用于获得当前机器上的当前进程号。

返回:: 当前机器上的当前进程号。

Example:

from pyvqnet.distributed import CommController
Comm_OP = CommController("nccl") # init nccl controller

Comm_OP.getLocalRank()
# vqnetrun --backend nccl --nproc_per_node 2 python test.py

barrier()¶

同步。

返回:: 同步操作。

Example:

from pyvqnet.distributed import CommController
Comm_OP = CommController("nccl")

Comm_OP.barrier()

get_device_num()¶

用于获得当前节点上的显卡数量, (仅支持gpu下使用)。

返回:: 返回当前节点上显卡数量。

Example:

from pyvqnet.distributed import CommController
Comm_OP = CommController("nccl")

Comm_OP.get_device_num()
# python test.py

allreduce(tensor, c_op='avg')¶

支持对数据作allreduce通信。

参数:

tensor – 输入数据.
c_op – 计算方式.

Example:

from pyvqnet.distributed import CommController
from pyvqnet.tensor import tensor
import numpy as np
Comm_OP = CommController("mpi")

num = tensor.to_tensor(np.random.rand(1, 5))
print(f"rank {Comm_OP.getRank()}  {num}")

Comm_OP.allreduce(num, "sum")
print(f"rank {Comm_OP.getRank()}  {num}")
# vqnetrun --backend nccl --nproc_per_node 2 python test.py

reduce(tensor, root=0, c_op='avg')¶

支持对数据作reduce通信。

参数:

tensor – 输入数据。
root – 指定数据返回的节点。
c_op – 计算方式。

Example:

from pyvqnet.distributed import CommController
from pyvqnet.tensor import tensor
import numpy as np
Comm_OP = CommController("mpi")

num = tensor.to_tensor(np.random.rand(1, 5))
print(f"rank {Comm_OP.getRank()}  {num}")

Comm_OP.reduce(num, 1)
print(f"rank {Comm_OP.getRank()}  {num}")
# vqnetrun --backend nccl --nproc_per_node 2 python test.py

broadcast(tensor, root=0)¶

将指定进程root上的数据广播到所有进程上。

参数:

tensor – 输入数据。
root – 指定的节点。

Example:

from pyvqnet.distributed import CommController
from pyvqnet.tensor import tensor
import numpy as np
Comm_OP = CommController("mpi")

num = tensor.to_tensor(np.random.rand(1, 5))
print(f"rank {Comm_OP.getRank()}  {num}")

Comm_OP.broadcast(num, 1)
print(f"rank {Comm_OP.getRank()}  {num}")
# vqnetrun --backend nccl --nproc_per_node 2 python test.py

allgather(tensor)¶

将所有进程上数据allgather到一起。

参数:: tensor – 输入数据。

Example:

from pyvqnet.distributed import CommController
from pyvqnet.tensor import tensor
import numpy as np
Comm_OP = CommController("mpi")

num = tensor.to_tensor(np.random.rand(1, 5))
print(f"rank {Comm_OP.getRank()}  {num}")

num = Comm_OP.allgather(num)
print(f"rank {Comm_OP.getRank()}  {num}")
# vqnetrun --backend nccl --nproc_per_node 2 python test.py

send(tensor, dest)¶

p2p通信接口。

参数:

tensor – 输入数据.
dest – 目的进程.

Example:

from pyvqnet.distributed import CommController,get_rank
from pyvqnet.tensor import tensor
import numpy as np
Comm_OP = CommController("mpi")

num = tensor.to_tensor(np.random.rand(1, 5))
recv = tensor.zeros_like(num)

if get_rank() == 0:
    Comm_OP.send(num, 1)
elif get_rank() == 1:
    Comm_OP.recv(recv, 0)
print(f"rank {Comm_OP.getRank()}  {num}")
print(f"rank {Comm_OP.getRank()}  {recv}")

# vqnetrun --backend nccl --nproc_per_node 2 python test.py

recv(tensor, source)¶

p2p通信接口。

参数:

tensor – 输入数据.
source – 接受进程.

Example:

from pyvqnet.distributed import CommController,get_rank
from pyvqnet.tensor import tensor
import numpy as np
Comm_OP = CommController("mpi")

num = tensor.to_tensor(np.random.rand(1, 5))
recv = tensor.zeros_like(num)

if get_rank() == 0:
    Comm_OP.send(num, 1)
elif get_rank() == 1:
    Comm_OP.recv(recv, 0)
print(f"rank {Comm_OP.getRank()}  {num}")
print(f"rank {Comm_OP.getRank()}  {recv}")

# vqnetrun --backend nccl --nproc_per_node 2 python test.py

split_groups(rankL)¶

根据入参设置的进程号列表用于划分多个通信组。

参数:: rankL – 进程组序号列表。
返回:: 当后端为 nccl 返回的是进程组序号元组，当后端为 mpi 返回一个列表，其长度等于分组个数；每个元素是二元组 (comm, rank)，其中 comm 为该分组的 MPI 通信器，rank 为组内序号。

Example:

from pyvqnet.distributed import CommController,get_rank,get_local_rank
from pyvqnet.tensor import tensor
import numpy as np
Comm_OP = CommController("mpi")

groups = Comm_OP.split_groups([[0, 1],[2,3]])
print(groups)
#[[<mpi4py.MPI.Intracomm object at 0x7f53691f3230>, [0, 3]], [<mpi4py.MPI.Intracomm object at 0x7f53691f3010>, [2, 1]]]

# mpirun -n 4 python test.py

allreduce_group(tensor, c_op='avg', group=None)¶

组内allreduce通信接口。

参数:

tensor – 输入数据.
c_op – 计算方法.
group – 当使用mpi后端时候，输入由 init_groups 或 split_groups 生成的组对应通信组，当使用nccl后端时候输入`split_groups` 生成的组序号。

Example:

from pyvqnet.distributed import CommController,get_rank,get_local_rank
from pyvqnet.tensor import tensor
import numpy as np
from pyvqnet import kcomplex64
Comm_OP = CommController("nccl")

groups = Comm_OP.split_groups([[0, 1]])

complex_data = tensor.QTensor([3+1j, 2, 1 + get_rank()],dtype=kcomplex64).reshape((3,1)).toGPU(1000+ get_local_rank())

print(f"allreduce_group before rank {get_rank()}: {complex_data}")

Comm_OP.allreduce_group(complex_data, c_op="sum",group = groups[0])
print(f"allreduce_group after rank {get_rank()}: {complex_data}")
# vqnetrun --backend nccl --nproc_per_node 2 python test.py

reduce_group(tensor, root=0, c_op='avg', group=None)¶

组内reduce通信接口。

参数:

tensor – 输入数据.
root – 指定进程号.
c_op – 计算方法.
group – 当使用mpi后端时候，输入由 init_groups 或 split_groups 生成的组对应通信组，当使用nccl后端时候输入`split_groups` 生成的组序号。

Example:

from pyvqnet.distributed import CommController,get_rank,get_local_rank
from pyvqnet.tensor import tensor
import numpy as np
from pyvqnet import kcomplex64
Comm_OP = CommController("nccl")

groups = Comm_OP.split_groups([[0, 1]])

complex_data = tensor.QTensor([3+1j, 2, 1 + get_rank()],dtype=kcomplex64).reshape((3,1)).toGPU(1000+ get_local_rank())

print(f"reduce_group before rank {get_rank()}: {complex_data}")

Comm_OP.reduce_group(complex_data, c_op="sum",group = groups[0])
print(f"reduce_group after rank {get_rank()}: {complex_data}")
# vqnetrun --backend nccl --nproc_per_node 2 python test.py

broadcast_group(tensor, root=0, group=None)¶

组内broadcast通信接口。

参数:

tensor – 输入数据.
root – 指定从哪个进程号广播，默认为0.
group – 当使用mpi后端时候，输入由 init_groups 或 split_groups 生成的组对应通信组，当使用nccl后端时候输入`split_groups` 生成的组序号。

Example:

from pyvqnet.distributed import CommController,get_rank,get_local_rank
from pyvqnet.tensor import tensor
import numpy as np
from pyvqnet import kcomplex64
Comm_OP = CommController("nccl")

groups = Comm_OP.split_groups([[0, 1]])

complex_data = tensor.QTensor([3+1j, 2, 1 + get_rank()],dtype=kcomplex64).reshape((3,1)).toGPU(1000+ get_local_rank())

print(f"broadcast_group before rank {get_rank()}: {complex_data}")

Comm_OP.broadcast_group(complex_data,group = groups[0])
Comm_OP.barrier()
print(f"broadcast_group after rank {get_rank()}: {complex_data}")
# vqnetrun --backend nccl --nproc_per_node 2 python test.py

allgather_group(tensor, group=None)¶

组内allgather通信接口。

参数:

tensor – 输入数据.
group – 当使用mpi后端时候，输入由 init_groups 或 split_groups 生成的组对应通信组，当使用nccl后端时候输入`split_groups` 生成的组序号。

Example:

from pyvqnet.distributed import CommController,get_rank
from pyvqnet.tensor import tensor
from pyvqnet import kcomplex64
Comm_OP = CommController("mpi")
group = init_groups([[0,1]])
#mpi init group internally
# A list of lists, where each sublist contains a communicator and the corresponding rank list.
complex_data = tensor.QTensor([3+1j, 2, 1 + get_rank()],dtype=kcomplex64).reshape((3,1))
print(f" before rank {get_rank()}: {complex_data}")
for comm_ in group:
    if Comm_OP.getRank() in comm_[1]:
        complex_data = Comm_OP.allgather_group(complex_data, comm_[0])
        print(f"after rank {get_rank()}: {complex_data}")
# mpirun -n 2 python test.py

from pyvqnet.distributed import CommController,get_rank,get_local_rank
from pyvqnet.tensor import tensor
Comm_OP = CommController("nccl")
groups = Comm_OP.split_groups([[0, 1]])
complex_data = tensor.QTensor([3+1j, 2, 1 + get_rank()],dtype=kcomplex64).reshape((3,1)).toGPU(1000+ get_local_rank())
print(f" before rank {get_rank()}: {complex_data}")
complex_data = Comm_OP.allgather_group(complex_data, group = groups[0])
print(f"after rank {get_rank()}: {complex_data}")
# vqnetrun --backend nccl --nproc_per_node 2 python test.py

nccl_async_all_gather( output, input, group=None, async_op=False):

对GPU上数据使用NCCL进行异步或同步all_gather。

参数:

output – QTensor all_gather的目标数据。
input – QTensor 待gather的数据。
group – 通信进程组，group 是一个包含组索引的元组。默认值：None，不使用任何组。
async_op – 此操作是否为异步操作，默认值：False。

返回:

Work，一个异步通信句柄。使用 wait() 等待此操作完成。

Example:

from pyvqnet import tensor
import pyvqnet
from pyvqnet import kcomplex64
from pyvqnet.distributed import CommController,get_rank,get_local_rank
Comm_OP = CommController("nccl")

complex_data = tensor.QTensor([3+1j, 2, 1 + get_rank()],dtype=kcomplex64).reshape((3,1)).toGPU(1000+ get_local_rank())

out_data = tensor.empty([2,3,1],dtype=pyvqnet.kcomplex64).toGPU(pyvqnet.DEV_GPU_0+ get_local_rank())
work = Comm_OP.nccl_async_all_gather(out_data, complex_data, group = None,async_op=True)
work.wait()

nccl_async_all_reduce(tensor, c_op="avg",group=None, async_op=False):

对GPU上数据使用NCCL进行异步或同步allreduce。

参数:

tensor – 需要归约的 QTensor。
c_op – 计算方法，可以是 “sum” 或 “avg”，默认值为 “avg”。
group – 通信进程组，group 是一个包含组索引的元组。默认值：None，不使用任何组。
async_op – 此操作是否为异步操作，默认值：False。

返回:

Work，一个异步通信句柄。使用 wait() 等待此操作完成。

Example:

import pyvqnet
from pyvqnet.tensor import tensor
from pyvqnet.distributed.ControlComm import CommController,get_local_rank

Comm_OP = CommController("nccl")
complex_data = tensor.ones([500,500],dtype=pyvqnet.kcomplex64).toGPU(pyvqnet.DEV_GPU_0 + get_local_rank())
work = Comm_OP.nccl_async_all_reduce(complex_data, "sum",group = None,async_op = True)
work.wait()

nccl_async_reduce( tensor_, dest, c_op="avg", group=None, async_op=False ):

对GPU上数据使用NCCL进行异步或同步reduce。

参数:

tensor – 需要归约的 QTensor。
dest – 要归约的 QTensor 的目标进程号/rank。
c_op – 计算方法，可以是 “sum” 或 “avg”，默认值为 “avg”。
group – 通信进程组，group 是一个包含组索引的元组。默认值：None，不使用任何组。
async_op – 此操作是否为异步操作，默认值：False。

返回:

Work，一个异步通信句柄。使用 wait() 等待此操作完成。

Example:

from pyvqnet import tensor
import pyvqnet
from pyvqnet.distributed import CommController,get_rank,get_local_rank
Comm_OP = CommController("nccl")

complex_data = tensor.ones([500,500],dtype=pyvqnet.kcomplex64).toGPU(pyvqnet.DEV_GPU_0 + get_local_rank())
work = Comm_OP.nccl_async_reduce(complex_data, 0,"sum",group = None,async_op = True)
work.wait()

nccl_async_broadcast(tensor, src, group=None, async_op=False)¶

对GPU上数据使用NCCL进行异步或同步广播。

参数:

tensor – 需要广播的 QTensor。
src – 要广播的 QTensor 的源排名。
group – 通信进程组，group 是一个包含组索引的元组。默认值：None，不使用任何组。
async_op – 此操作是否为异步操作，默认值：False。

返回:

Work，一个异步通信句柄。使用 wait() 等待此操作完成。

Example:

import pyvqnet
from pyvqnet.tensor import tensor
from pyvqnet.distributed.ControlComm import CommController,get_local_rank
Comm_OP = CommController("nccl")
if get_local_rank() == 1:
    complex_data = tensor.ones([5,5],dtype=pyvqnet.kcomplex64).toGPU(pyvqnet.DEV_GPU_0+ get_local_rank())
else:
    complex_data = tensor.zeros([5,5],dtype=pyvqnet.kcomplex64).toGPU(pyvqnet.DEV_GPU_0+ get_local_rank())
work = Comm_OP.nccl_async_broadcast(complex_data, 1, group = None,async_op=True)

work.wait()

nccl_async_send( t, dest, async_op=False ):

对GPU上数据使用NCCL进行异步或同步P2P send。

参数:

t – 需要发送的 QTensor。
dest – 要发送 QTensor 的目标进程号/rank。
async_op – 此操作是否为异步操作，默认值：False。

返回:

Work，一个异步通信句柄。使用 wait() 等待此操作完成。

Example:

import pyvqnet
from pyvqnet.tensor import tensor
from pyvqnet.distributed.ControlComm import CommController,get_local_rank
Comm_OP = CommController("nccl")

if get_local_rank() == 0:
    complex_data = tensor.ones([5000,5000],dtype=pyvqnet.kcomplex64).toGPU(pyvqnet.DEV_GPU_0+ get_local_rank())*10
else:
    complex_data = tensor.ones([5000,5000],dtype=pyvqnet.kcomplex64).toGPU(pyvqnet.DEV_GPU_0+ get_local_rank())

if get_local_rank() == 0:
    work = Comm_OP.nccl_async_send(complex_data, 1 ,True)
else:
    work = Comm_OP.nccl_async_recv(complex_data, 0 ,True)
work.wait()

nccl_async_recv( t, src, async_op=False ):

对GPU上数据使用NCCL进行异步或同步P2P recv。

参数:

t – 需要接受的 QTensor。
src – 要接受 QTensor 的目标进程号/rank。
async_op – 此操作是否为异步操作，默认值：False。

返回:

Work，一个异步通信句柄。使用 wait() 等待此操作完成。

Example:

import pyvqnet
from pyvqnet.tensor import tensor
from pyvqnet.distributed.ControlComm import CommController,get_local_rank
Comm_OP = CommController("nccl")

if get_local_rank() == 0:
    complex_data = tensor.ones([5000,5000],dtype=pyvqnet.kcomplex64).toGPU(pyvqnet.DEV_GPU_0+ get_local_rank())*10
else:
    complex_data = tensor.ones([5000,5000],dtype=pyvqnet.kcomplex64).toGPU(pyvqnet.DEV_GPU_0+ get_local_rank())

if get_local_rank() == 0:
    work = Comm_OP.nccl_async_send(complex_data, 1 ,True)
else:
    work = Comm_OP.nccl_async_recv(complex_data, 0 ,True)
work.wait()

destroy()¶

销毁 NCCL 通信器资源 / Destroy NCCL communicators explicitly.

在进程退出前应调用此方法以正确清理 NCCL 资源。若不调用，NCCL 通信器可能会泄漏内存或在 CUDA 驱动关闭时引发错误。

This method should be called before process exit to properly clean up NCCL resources. If not called, NCCL communicators may leak memory or cause errors during CUDA driver shutdown.

Example:

from pyvqnet.distributed import CommController
Comm_OP = CommController("nccl")
# ... 使用 Comm_OP / use Comm_OP ...
Comm_OP.destroy()  # 退出前调用 / Call before exit

split_data¶

在多进程中,使用 split_data 根据进程数对数据进行切分,返回相应进程上数据。

pyvqnet.distributed.datasplit.split_data(x_train, y_train, shuffle=False)¶

设置分布式计算参数。

参数:

x_train – np.array - 训练数据.
y_train – np.array - 训练数据标签.
shuffle – bool - 是否打乱后再进行切分,默认值是False.

返回:

切分后的训练数据和标签。

Example:

from pyvqnet.distributed import split_data
import numpy as np

x_train = np.random.randint(255, size = (100, 5))
y_train = np.random.randint(2, size = (100, 1))

x_train, y_train= split_data(x_train, y_train)

get_local_rank¶

pyvqnet.distributed.ControlComm.get_local_rank()¶

得到当前节点上进程号。例如你在第2个节点的第3个进程,每个节点5个进程,则返回2。

返回:: 当前机器上的当前进程号。

Example:

from pyvqnet.distributed.ControlComm import get_local_rank

print(get_local_rank())
# vqnetrun --backend nccl --nproc_per_node 2 python test.py

get_rank¶

pyvqnet.distributed.ControlComm.get_rank()¶

用于获得当前进程的全局进程号。例如你在第2个节点的第3个进程,每个节点5个进程,则返回7。

返回:: 当前进程的进程号。

Example:

from pyvqnet.distributed.ControlComm import get_rank

print(get_rank())
# vqnetrun --backend nccl --nproc_per_node 2 python test.py

init_groups¶

pyvqnet.distributed.ControlComm.init_groups(rank_lists)¶

根据给出的进程数列表来对基于 mpi 后端的进程组进行初始化。

警告

该接口只支持分布式后端为 mpi 。

参数:: rank_lists – 通信进程组列表.
返回:: 返回一个列表，其长度等于分组个数；每个元素是二元组 (comm, rank)，其中 comm 为该分组的 MPI 通信器，rank 为组内序号。

Example:

from pyvqnet.distributed import *
import numpy as np
Comm_OP = CommController("mpi")
num = tensor.to_tensor(np.random.rand(1, 5))
print(f"rank {Comm_OP.getRank()}  {num} before allreduce")

group_l = init_groups([[0,2], [1]])

for comm_ in group_l:
    if Comm_OP.getRank() in comm_:
        Comm_OP.allreduce_group(num, "sum", group = comm_)
        print(f"rank {Comm_OP.getRank()}  {num} after allreduce")

# vqnetrun -n 3 python test.py

PipelineParallelTrainingWrapper¶

class pyvqnet.distributed.pp.PipelineParallelTrainingWrapper(args, join_layers, trainset)¶

Pipeline Parallel Training Wrapper 实现了 1F1B训练。仅在 Linux 平台上,且具有 GPU 的情况下可用。更多算法细节可以在(https://www.deepspeed.ai/tutorials/pipeline/)找到。

参数:

args – 参数字典。参见示例。
join_layers – Sequential 模块的列表。
trainset – 数据集。

返回:

PipelineParallelTrainingWrapper 实例。

以下使用 CIFAR10数据集 CIFAR10_Dataset,在2块GPU上训练AlexNet上的分类任务。本例子中分成两个流水线并行进程 pipeline_parallel_size = 2。批处理大小为 train_batch_size = 64, 单GPU 上为 train_micro_batch_size_per_gpu = 32。其他配置参数可见 args。此外,每个进程需要在 __main__ 函数中配置环境变量的 LOCAL_RANK。

os.environ["LOCAL_RANK"] = str(dist.get_local_rank())

调用 train_batch 进行训练。

Example:

import os
import pyvqnet

from pyvqnet.nn import Module,Sequential,CrossEntropyLoss
from pyvqnet.nn import Linear
from pyvqnet.nn import Conv2D
from pyvqnet.nn import activation as F
from pyvqnet.nn import MaxPool2D
from pyvqnet.nn import CrossEntropyLoss

from pyvqnet.tensor import tensor
from pyvqnet.distributed.pp import PipelineParallelTrainingWrapper
from pyvqnet.distributed.configs import comm as dist
from pyvqnet.distributed import *


pipeline_parallel_size = 2

num_steps = 1000

def cifar_trainset_vqnet(local_rank, dl_path='./cifar10-data'):
    transform = pyvqnet.data.TransformCompose([
        pyvqnet.data.TransformResize(256),
        pyvqnet.data.TransformCenterCrop(224),
        pyvqnet.data.TransformToTensor(),
        pyvqnet.data.TransformNormalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ])

    trainset = pyvqnet.data.CIFAR10_Dataset(root=dl_path,
                                            mode="train",
                                            transform=transform,layout="HWC")

    return trainset

class Model(Module):
    def __init__(self):
        super(Model, self).__init__()
        self.features = Sequential(
        Conv2D(input_channels=3, output_channels=8, kernel_size=(3, 3), stride=(1, 1), padding='same'),
        F.ReLu(),
        MaxPool2D([2, 2], [2, 2]),

        Conv2D(input_channels=8, output_channels=16, kernel_size=(3, 3), stride=(1, 1), padding='same'),
        F.ReLu(),
        MaxPool2D([2, 2], [2, 2]),

        Conv2D(input_channels=16, output_channels=32, kernel_size=(3, 3), stride=(1, 1), padding='same'),
        F.ReLu(),

        Conv2D(input_channels=32, output_channels=64, kernel_size=(3, 3), stride=(1, 1), padding='same'),
        F.ReLu(),

        Conv2D(input_channels=64, output_channels=64, kernel_size=(3, 3), stride=(1, 1), padding='same'),
        F.ReLu(),
        MaxPool2D([3, 3], [2, 2]),)

        self.cls = Sequential(
        Linear(64 * 27 * 27, 512),
        F.ReLu(),

        Linear(512, 256),
        F.ReLu(),
        Linear(256, 10) )

    def forward(self, x):
        x = self.features(x)
        x = tensor.flatten(x,1)
        x = self.cls(x)

        return x

def join_layers(vision_model):
    layers = [
        *vision_model.features,
        lambda x: tensor.flatten(x, 1),
        *vision_model.cls,
    ]
    return layers


if __name__ == "__main__":


    args = {
    "backend":'nccl',
    "train_batch_size" : 64,
    "train_micro_batch_size_per_gpu" : 32,
    "epochs":5,
"optimizer": {
    "type": "Adam",
    "params": {
    "lr": 0.001
    }},
    "local_rank":dist.get_local_rank(),
    "pipeline_parallel_size":pipeline_parallel_size, "seed":42, "steps":num_steps,
    "loss":CrossEntropyLoss(),
    }
    os.environ["LOCAL_RANK"] = str(dist.get_local_rank())
    trainset = cifar_trainset_vqnet(args["local_rank"])
    w = PipelineParallelTrainingWrapper(args,join_layers(Model()),trainset)

    all_loss = {}

    for i in range(args["epochs"]):
        w.train_batch()

    all_loss = w.loss_dict

ZeroModelInitial¶

class pyvqnet.distributed.ZeroModelInitial(args, model, optimizer)¶

Zero1 API 接口, 目前仅用于 Linux 平台下基于 GPU 并行计算。

参数:

args – 参数字典。参见示例。
model – 输入模型。
optimizer – 优化器。

返回:

Zero1 Engine.

以下使用 MNIST 数据库, 在2块GPU上训练一个MLP模型上的分类任务。

批处理大小为 train_batch_size = 64, zero_optimization 的阶段 stage 设置为1. 若Optimizer为None, 则采用 args 中 optimizer 的设置. 其他配置参数可见 args。此外,每个进程需要在配置环境变量的 LOCAL_RANK。

os.environ["LOCAL_RANK"] = str(dist.get_local_rank())

Example:

from pyvqnet.distributed import *
from pyvqnet import *
from time import time
import pyvqnet.optim as optim
import pyvqnet.nn as nn
import pyvqnet
import sys
import pyvqnet
import numpy as np
import os
import struct

def load_mnist(dataset="training_data",
            digits=np.arange(2),
            path="./"):
    """
    load mnist data
    """
    from array import array as pyarray
    if dataset == "training_data":
        fname_image = os.path.join(path, "train-images.idx3-ubyte").replace(
            "\\", "/")
        fname_label = os.path.join(path, "train-labels.idx1-ubyte").replace(
            "\\", "/")
    elif dataset == "testing_data":
        fname_image = os.path.join(path, "t10k-images.idx3-ubyte").replace(
            "\\", "/")
        fname_label = os.path.join(path, "t10k-labels.idx1-ubyte").replace(
            "\\", "/")
    else:
        raise ValueError("dataset must be 'training_data' or 'testing_data'")

    flbl = open(fname_label, "rb")
    _, size = struct.unpack(">II", flbl.read(8))

    lbl = pyarray("b", flbl.read())
    flbl.close()

    fimg = open(fname_image, "rb")
    _, size, rows, cols = struct.unpack(">IIII", fimg.read(16))
    img = pyarray("B", fimg.read())
    fimg.close()

    ind = [k for k in range(size) if lbl[k] in digits]
    num = len(ind)
    images = np.zeros((num, rows, cols),dtype=np.float32)

    labels = np.zeros((num, 1), dtype=int)
    for i in range(len(ind)):
        images[i] = np.array(img[ind[i] * rows * cols:(ind[i] + 1) * rows *
                                cols]).reshape((rows, cols))
        labels[i] = lbl[ind[i]]

    return images, labels


train_images_np, train_labels_np = load_mnist(dataset="training_data", digits=np.arange(10),path="../data/MNIST_data/")
train_images_np = train_images_np / 255.

test_images_np, test_labels_np = load_mnist(dataset="testing_data", digits=np.arange(10),path="../data/MNIST_data/")
test_images_np = test_images_np / 255.

local_rank = pyvqnet.distributed.get_rank()

from pyvqnet.distributed import ZeroModelInitial

class MNISTClassifier(nn.Module):

    def __init__(self):
        super(MNISTClassifier, self).__init__()
        self.fc1 = nn.Linear(28*28, 512)
        self.fc2 = nn.Linear(512, 256)
        self.fc3 = nn.Linear(256, 128)
        self.fc4 = nn.Linear(128, 64)
        self.fc5 = nn.Linear(64, 10)
        self.ac = nn.activation.ReLu()

    def forward(self, x:pyvqnet.QTensor):

        x = x.reshape([-1, 28*28])
        x = self.ac(self.fc1(x))
        x = self.fc2(x)
        x = self.fc3(x)
        x = self.fc4(x)
        x = self.fc5(x)
        return x

model = MNISTClassifier()

model.to(local_rank + 1000)

Comm_op = CommController("nccl")
Comm_op.broadcast_model_params(model, 0)

batch_size = 64

criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

args_ = {
        "train_batch_size": batch_size, # 等效的总batch
        "optimizer": {
            "type": "adam",
            "params": {
            "lr": 0.001,
            }
        },
        "zero_optimization": {
            "stage": 1,
        }
    }

os.environ["LOCAL_RANK"] = str(get_local_rank())
model = ZeroModelInitial(args=args_, model=model, optimizer=optimizer)

def compute_acc(outputs, labels, correct, total):
    predicted = pyvqnet.tensor.argmax(outputs, dim=1, keepdim=True)
    total += labels.size
    correct += pyvqnet.tensor.sums(predicted == labels).item()
    return correct, total

train_acc = 0
test_acc = 0
epochs = 5
loss = 0
time1 = time()

for epoch in range(epochs):
    model.train()
    total = 0
    correct = 0
    step = 0

    num_batches = (train_images_np.shape[0] + batch_size - 1) // batch_size

    for i in range(num_batches):

        data_ = tensor.QTensor(train_images_np[i*batch_size: (i+1) * batch_size,:], dtype = kfloat32)
        labels = tensor.QTensor(train_labels_np[i*batch_size: (i+1) * batch_size,:], dtype = kint64)

        data_ = data_.to(local_rank + 1000)
        labels = labels.to(local_rank + 1000)

        outputs = model(data_)
        loss = criterion(labels, outputs)

        model.backward(loss) # 基于返回的model做backward、step
        model.step()

        correct, total = compute_acc(outputs, labels, correct, total)
        step += 1
        if step % 50 == 0:
            print(f"Train : rank {get_rank()} Epoch [{epoch+1}/{epochs}], step {step} Loss: {loss.item():.4f} acc {100 * correct / total}")
            sys.stdout.flush()

    train_acc = 100 * correct / total

time2 = time()
print(f'Accuracy of the model on the 10000 Train images: {train_acc}% time cost {time2 - time1}')

ColumnParallelLinear¶

class pyvqnet.distributed.ColumnParallelLinear(input_size, output_size, weight_initializer, bias_initializer, use_bias, dtype, name, tp_comm)¶

张量并行计算,列并行线性层

线性层定义为 Y = XA + b。其二维并行为 A = [A_1,…,A_p]。

参数:

input_size – 矩阵 A 的第一个维度。
output_size – 矩阵 A 的第二个维度。
weight_initializer – callable 默认为 normal。
bias_initializer – callable 默认为0。
use_bias – bool - 默认为 True。
dtype – 默认 None,使用默认数据类型。
name – 模块名称,默认为“”。
tp_comm – 通讯控制器。

以下使用 MNIST 数据库, 在2块GPU上训练一个MLP模型上的分类任务。

使用时与经典的Linear层的使用相似

多进程使用时基于 # vqnetrun –backend nccl –nproc_per_node 2 python test.py 的方式进行

Example:

import pyvqnet.distributed
import pyvqnet.optim as optim
import pyvqnet.nn as nn
import pyvqnet
import sys
from pyvqnet.distributed.tensor_parallel import ColumnParallelLinear, RowParallelLinear
from pyvqnet.distributed import *
from time import time

import pyvqnet
import numpy as np
import os
from pyvqnet import *
import pytest

Comm_OP = CommController("nccl")

import struct
def load_mnist(dataset="training_data",
            digits=np.arange(2),
            path="./"):
    """
    load mnist data
    """
    from array import array as pyarray
    # download_mnist(path)
    if dataset == "training_data":
        fname_image = os.path.join(path, "train-images-idx3-ubyte").replace(
            "\\", "/")
        fname_label = os.path.join(path, "train-labels-idx1-ubyte").replace(
            "\\", "/")
    elif dataset == "testing_data":
        fname_image = os.path.join(path, "t10k-images-idx3-ubyte").replace(
            "\\", "/")
        fname_label = os.path.join(path, "t10k-labels-idx1-ubyte").replace(
            "\\", "/")
    else:
        raise ValueError("dataset must be 'training_data' or 'testing_data'")

    flbl = open(fname_label, "rb")
    _, size = struct.unpack(">II", flbl.read(8))

    lbl = pyarray("b", flbl.read())
    flbl.close()

    fimg = open(fname_image, "rb")
    _, size, rows, cols = struct.unpack(">IIII", fimg.read(16))
    img = pyarray("B", fimg.read())
    fimg.close()

    ind = [k for k in range(size) if lbl[k] in digits]
    num = len(ind)
    images = np.zeros((num, rows, cols),dtype=np.float32)

    labels = np.zeros((num, 1), dtype=int)
    for i in range(len(ind)):
        images[i] = np.array(img[ind[i] * rows * cols:(ind[i] + 1) * rows *
                                cols]).reshape((rows, cols))
        labels[i] = lbl[ind[i]]

    return images, labels

train_images_np, train_labels_np = load_mnist(dataset="training_data", digits=np.arange(10),path="./data/MNIST/raw/")
train_images_np = train_images_np / 255.

test_images_np, test_labels_np = load_mnist(dataset="testing_data", digits=np.arange(10),path="./data/MNIST/raw/")
test_images_np = test_images_np / 255.

local_rank = pyvqnet.distributed.get_rank()

class MNISTClassifier(nn.Module):
    def __init__(self):
        super(MNISTClassifier, self).__init__()
        self.fc1 = RowParallelLinear(28*28, 512, tp_comm = Comm_OP)
        self.fc2 = ColumnParallelLinear(512, 256, tp_comm = Comm_OP)
        self.fc3 = RowParallelLinear(256, 128, tp_comm = Comm_OP)
        self.fc4 = ColumnParallelLinear(128, 64, tp_comm = Comm_OP)
        self.fc5 = RowParallelLinear(64, 10, tp_comm = Comm_OP)
        self.ac = nn.activation.ReLu()

    def forward(self, x:pyvqnet.QTensor):

        x = x.reshape([-1, 28*28])
        x = self.ac(self.fc1(x))
        x = self.fc2(x)
        x = self.fc3(x)
        x = self.fc4(x)
        x = self.fc5(x)
        return x


model = MNISTClassifier()
total_params = sum(p.numel() for p in model.parameters() if p.requires_grad)

model.to(local_rank + 1000)

Comm_OP.broadcast_model_params(model, 0)

criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

def compute_acc(outputs, labels, correct, total):
    predicted = pyvqnet.tensor.argmax(outputs, dim=1, keepdim=True)
    total += labels.size
    correct += pyvqnet.tensor.sums(predicted == labels).item()
    return correct, total

train_acc = 0
test_acc = 0
epochs = 5
loss = 0

time1 = time()
for epoch in range(epochs):
    model.train()
    total = 0
    correct = 0
    step = 0

    batch_size = 64
    num_batches = (train_images_np.shape[0] + batch_size - 1) // batch_size

    for i in range(num_batches):
        data_ = tensor.QTensor(train_images_np[i*batch_size: (i+1) * batch_size,:], dtype = kfloat32)
        labels = tensor.QTensor(train_labels_np[i*batch_size: (i+1) * batch_size,:], dtype = kint64)

        data_ = data_.to(local_rank + 1000)
        labels = labels.to(local_rank + 1000)

        optimizer.zero_grad()

        outputs = model(data_)
        loss = criterion(labels, outputs)

        loss.backward()
        optimizer.step()

        correct, total = compute_acc(outputs, labels, correct, total)
        step += 1
        if step % 50 == 0:
            print(f"Train : rank {get_rank()} Epoch [{epoch+1}/{epochs}], step {step} Loss: {loss.item():.4f} acc {100 * correct / total}")
            sys.stdout.flush()

    train_acc = 100 * correct / total
time2 = time()

print(f'Accuracy of the model on the 10000 Train images: {train_acc}% time cost {time2 - time1}')

RowParallelLinear¶

class pyvqnet.distributed.RowParallelLinear(input_size, output_size, weight_initializer, bias_initializer, use_bias, dtype, name, tp_comm)¶

张量并行计算,行并行线性层。

线性层的定义为 Y = XA + b。A 沿其一维并行,X 沿其二维并行。 A = transpose([A_1 … A_p]) X = [X_1, …, X_p]。

参数:

input_size – 矩阵 A 的第一个维度。
output_size – 矩阵 A 的第二个维度。
weight_initializer – callable 默认为 normal。
bias_initializer – callable 默认为0。
use_bias – bool - 默认为 True。
dtype – 默认 None,使用默认数据类型。
name – 模块名称。
tp_comm – 通讯控制器。

以下使用 MNIST 数据库, 在2块GPU上训练一个MLP模型上的分类任务。使用时与经典的Linear层的使用相似

多进程使用时基于 # vqnetrun –backend nccl –nproc_per_node 2 python test.py 的方式进行

Example:

import pyvqnet.distributed
import pyvqnet.optim as optim
import pyvqnet.nn as nn
import pyvqnet
import sys
from pyvqnet.distributed.tensor_parallel import ColumnParallelLinear, RowParallelLinear
from pyvqnet.distributed import *
from time import time

import pyvqnet
import numpy as np
import os
from pyvqnet import *
import pytest

Comm_OP = CommController("nccl")

import struct
def load_mnist(dataset="training_data",
            digits=np.arange(2),
            path="./"):
    """
    load mnist data
    """
    from array import array as pyarray
    # download_mnist(path)
    if dataset == "training_data":
        fname_image = os.path.join(path, "train-images-idx3-ubyte").replace(
            "\\", "/")
        fname_label = os.path.join(path, "train-labels-idx1-ubyte").replace(
            "\\", "/")
    elif dataset == "testing_data":
        fname_image = os.path.join(path, "t10k-images-idx3-ubyte").replace(
            "\\", "/")
        fname_label = os.path.join(path, "t10k-labels-idx1-ubyte").replace(
            "\\", "/")
    else:
        raise ValueError("dataset must be 'training_data' or 'testing_data'")

    flbl = open(fname_label, "rb")
    _, size = struct.unpack(">II", flbl.read(8))

    lbl = pyarray("b", flbl.read())
    flbl.close()

    fimg = open(fname_image, "rb")
    _, size, rows, cols = struct.unpack(">IIII", fimg.read(16))
    img = pyarray("B", fimg.read())
    fimg.close()

    ind = [k for k in range(size) if lbl[k] in digits]
    num = len(ind)
    images = np.zeros((num, rows, cols),dtype=np.float32)

    labels = np.zeros((num, 1), dtype=int)
    for i in range(len(ind)):
        images[i] = np.array(img[ind[i] * rows * cols:(ind[i] + 1) * rows *
                                cols]).reshape((rows, cols))
        labels[i] = lbl[ind[i]]

    return images, labels

train_images_np, train_labels_np = load_mnist(dataset="training_data", digits=np.arange(10),path="./data/MNIST/raw/")
train_images_np = train_images_np / 255.

test_images_np, test_labels_np = load_mnist(dataset="testing_data", digits=np.arange(10),path="./data/MNIST/raw/")
test_images_np = test_images_np / 255.

local_rank = pyvqnet.distributed.get_rank()

class MNISTClassifier(nn.Module):
    def __init__(self):
        super(MNISTClassifier, self).__init__()
        self.fc1 = RowParallelLinear(28*28, 512, tp_comm = Comm_OP)
        self.fc2 = ColumnParallelLinear(512, 256, tp_comm = Comm_OP)
        self.fc3 = RowParallelLinear(256, 128, tp_comm = Comm_OP)
        self.fc4 = ColumnParallelLinear(128, 64, tp_comm = Comm_OP)
        self.fc5 = RowParallelLinear(64, 10, tp_comm = Comm_OP)
        self.ac = nn.activation.ReLu()


    def forward(self, x:pyvqnet.QTensor):

        x = x.reshape([-1, 28*28])
        x = self.ac(self.fc1(x))
        x = self.fc2(x)
        x = self.fc3(x)
        x = self.fc4(x)
        x = self.fc5(x)
        return x

model = MNISTClassifier()
total_params = sum(p.numel() for p in model.parameters() if p.requires_grad)

model.to(local_rank + 1000)
Comm_OP.broadcast_model_params(model, 0)

criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

def compute_acc(outputs, labels, correct, total):
    predicted = pyvqnet.tensor.argmax(outputs, dim=1, keepdim=True)
    total += labels.size
    correct += pyvqnet.tensor.sums(predicted == labels).item()
    return correct, total

train_acc = 0
test_acc = 0
epochs = 5
loss = 0

time1 = time()
for epoch in range(epochs):
    model.train()
    total = 0
    correct = 0
    step = 0

    batch_size = 64
    num_batches = (train_images_np.shape[0] + batch_size - 1) // batch_size

    for i in range(num_batches):
        data_ = tensor.QTensor(train_images_np[i*batch_size: (i+1) * batch_size,:], dtype = kfloat32)
        labels = tensor.QTensor(train_labels_np[i*batch_size: (i+1) * batch_size,:], dtype = kint64)

        data_ = data_.to(local_rank + 1000)
        labels = labels.to(local_rank + 1000)

        optimizer.zero_grad()

        outputs = model(data_)
        loss = criterion(labels, outputs)

        loss.backward()
        optimizer.step()

        correct, total = compute_acc(outputs, labels, correct, total)
        step += 1
        if step % 50 == 0:
            print(f"Train : rank {get_rank()} Epoch [{epoch+1}/{epochs}], step {step} Loss: {loss.item():.4f} acc {100 * correct / total}")
            sys.stdout.flush()

    train_acc = 100 * correct / total
time2 = time()

print(f'Accuracy of the model on the 10000 Train images: {train_acc}% time cost {time2 - time1}')

比特重排序¶

量子比特重排序技术是比特并行中的技术，其核心是通过改变比特并行过程中量子逻辑门的排列顺序，减少比特并行中需要执行比特变换的次数，以下是基于比特并行构建大比特量子线路时需要的模块。参照论文 Lazy Qubit Reordering for Accelerating Parallel State-Vector-based Quantum Circuit Simulation 。以下接口需要通过 mpi 启动多个进程进行计算。

DistributedQMachine¶

class pyvqnet.distributed.qubits_reorder.DistributedQMachine(num_wires, dtype=kcomplex64, grad_mode='')¶

用于比特并行中的变分量子计算的模拟类，包含每个节点包含的部分比特上的量子态。通过MPI,每个节点都申请一个该类，进行分布式的量子变分线路模拟，N的值必须等于2的分布式并行的比特个数 global_qubit 的幂次方，可通过 set_qr_config 进行配置。

参数:

num_wires – 完整量子线路的比特的数量。
dtype – 计算数据的数据类型。默认是 pyvqnet.kcomplex64，对应的参数精度是 pyvqnet.kfloat32。
grad_mode – 在 DistQuantumLayerAdjoint 进行反传时设置为 adjoint。

备注

输入的比特数是整个量子线路所需要的比特数量，通过DistributedQMachine会根据全局比特数构建量子模拟器, 其比特数量为 num_wires - global_qubit，反传必须基于 DistQuantumLayerAdjoint。

警告

该接口只支持在Linux下运行；

必须对 DistributedQMachine 中比特并行中参数进行配置，如样例中所示，包括：

qm.set_just_defined(True)
qm.set_save_op_history_flag(True) # open save op
qm.set_qr_config({'qubit': 总比特个数, 'global_qubit': 分布式比特个数})

Example:

from pyvqnet.distributed import get_rank
from pyvqnet import tensor
from pyvqnet.qnn.vqc import rx, ry, cnot, MeasureAll,rz
import pyvqnet
from pyvqnet.distributed.qubits_reorder import DistributedQMachine,DistQuantumLayerAdjoint
pyvqnet.utils.set_random_seed(123)


qubit = 10
batch_size = 5

class QModel(pyvqnet.nn.Module):
    def __init__(self, num_wires, dtype, grad_mode=""):
        super(QModel, self).__init__()

        self._num_wires = num_wires
        self._dtype = dtype
        self.qm = DistributedQMachine(num_wires, dtype=dtype, grad_mode=grad_mode)

        self.qm.set_just_defined(True)
        self.qm.set_save_op_history_flag(True) # open save op
        self.qm.set_qr_config({"qubit": num_wires, # open qubit reordered, set config
                                "global_qubit": 1}) # global_qubit == log2(nproc)

        self.params = pyvqnet.nn.Parameter([qubit])

        self.measure = MeasureAll(obs={
            "X5":1.0
        })

    def forward(self, x, *args, **kwargs):
        self.qm.reset_states(x.shape[0])

        for i in range(qubit):
            rx(q_machine=self.qm, params=self.params[i], wires=i)
            ry(q_machine=self.qm, params=self.params[3], wires=i)
            rz(q_machine=self.qm, params=self.params[4], wires=i)
        cnot(q_machine=self.qm,  wires=[0, 1])
        rlt = self.measure(q_machine=self.qm)
        return rlt


input_x = tensor.QTensor([[0.1, 0.2, 0.3]], requires_grad=True).toGPU(pyvqnet.DEV_GPU_0+get_rank())

input_x = tensor.broadcast_to(input_x, [2, 3])

input_x.requires_grad = True

quantum_model = QModel(num_wires=qubit,
                    dtype=pyvqnet.kcomplex64,
                    grad_mode="adjoint").toGPU(pyvqnet.DEV_GPU_0+get_rank())

adjoint_model = DistQuantumLayerAdjoint(quantum_model)
adjoint_model.train()

batch_y = adjoint_model(input_x)
batch_y.backward()

print(batch_y)
# mpirun -n 2 python test.py

DistQuantumLayerAdjoint¶

class pyvqnet.distributed.qubits_reorder.DistQuantumLayerAdjoint(vqc_module, name)¶

使用伴随矩阵方式对比特并行计算中的参数进行梯度计算的DistQuantumLayer层

参数:

vqc_module – 输入的蕴含 DistributedQMachine 模块。
name – 模块名称。

备注

输入的vqc_module模块必须包含 DistributedQMachine，基于 DistributedQMachine 进行比特并行下的adjoint反传梯度计算。

警告

该接口只支持在Linux下运行；

Example:

from pyvqnet.distributed import get_rank
from pyvqnet import tensor
from pyvqnet.qnn.vqc import rx, ry, cnot, MeasureAll,rz
import pyvqnet
from pyvqnet.distributed.qubits_reorder import DistributedQMachine,DistQuantumLayerAdjoint
pyvqnet.utils.set_random_seed(123)


qubit = 10
batch_size = 5

class QModel(pyvqnet.nn.Module):
    def __init__(self, num_wires, dtype, grad_mode=""):
        super(QModel, self).__init__()

        self._num_wires = num_wires
        self._dtype = dtype
        self.qm = DistributedQMachine(num_wires, dtype=dtype, grad_mode=grad_mode)

        self.qm.set_just_defined(True)
        self.qm.set_save_op_history_flag(True) # open save op
        self.qm.set_qr_config({"qubit": num_wires, # open qubit reordered, set config
                                    "global_qubit": 1}) # global_qubit == log2(nproc)

        self.params = pyvqnet.nn.Parameter([qubit])

        self.measure = MeasureAll(obs={
            "X5":1.0
        })

    def forward(self, x, *args, **kwargs):
        self.qm.reset_states(x.shape[0])

        for i in range(qubit):
            rx(q_machine=self.qm, params=self.params[i], wires=i)
            ry(q_machine=self.qm, params=self.params[3], wires=i)
            rz(q_machine=self.qm, params=self.params[4], wires=i)
        cnot(q_machine=self.qm,  wires=[0, 1])
        rlt = self.measure(q_machine=self.qm)
        return rlt


input_x = tensor.QTensor([[0.1, 0.2, 0.3]], requires_grad=True).toGPU(pyvqnet.DEV_GPU_0+get_rank())

input_x = tensor.broadcast_to(input_x, [2, 3])

input_x.requires_grad = True

quantum_model = QModel(num_wires=qubit,
                    dtype=pyvqnet.kcomplex64,
                    grad_mode="adjoint").toGPU(pyvqnet.DEV_GPU_0+get_rank())

adjoint_model = DistQuantumLayerAdjoint(quantum_model)
adjoint_model.train()

batch_y = adjoint_model(input_x)
batch_y.backward()

print(batch_y)
# mpirun -n 2 python test.py