torch 多进程训练和分布式训练

发布时间 2023-12-12 14:56:43作者: xuyv

通常来说,多进程没有涉及到梯度同步的概念。
分布式则设计梯度同步。分布式中,如果用cpu,则指定gloo后端。用gpu,则指定nccl后端。

多进程训练

只需要mp.spawn即可,每个进程只负责自己的模型训练,不涉及梯度同步。

例子

https://www.geeksforgeeks.org/multiprocessing-in-python-and-pytorch/
(也可以用model.share_memory()来同步参数的更新)

chatgpt

分布式训练

需要init_process_group,以及使用DDP类包装model,使得梯度能够自动同步(DDP负责)。
还有,dataloader需要指定为dist_dataloadr版本,指定后,dis_dataloader自动为训练划分数据,从而做到自动数据并行的分布式训练。
因此,设计三部分修改:

  1. init_process_group 负责进程初始化
  2. DDP(DistributedDataParallel) 负责梯度同步
  3. 切分数据集或者指定dataloader(use_ddp=True), 负责切分数据集

例子

https://docs.dgl.ai/en/0.8.x/tutorials/multi/2_node_classification.html

https://github.com/pyg-team/pytorch_geometric/blob/master/examples/multi_gpu/distributed_sampling.py

chatgpt