├── .gitignore ├── README.md ├── data.py ├── data_parallel_train.py ├── ddp_train.py ├── model.py └── single_gpu_train.py /.gitignore: -------------------------------------------------------------------------------- 1 | .vscode/ -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # PyTorch 单机多GPU 训练方法与原理整理 2 | 3 | 这里整理一些PyTorch单机多核训练的方法和简单原理，目的是既能在写代码时知道怎么用，又能从原理上知道大致是怎么回事儿。如果只是炼丹，有时候确实没时间和精力深挖太多实现原理，但又希望能理解简单逻辑。 4 | 5 | PyTorch单机多核训练方案有两种：一种是利用`nn.DataParallel`实现，实现简单，不涉及多进程；另一种是用`torch.nn.parallel.DistributedDataParallel`和`torch.utils.data.distributed.DistributedSampler`结合多进程实现。第二种方式效率更高，但是实现起来稍难，第二种方式同时支持多节点分布式实现。方案二的效率要比方案一高，即使是在单运算节点上。 6 | 7 | 为方便理解，这里用一个简单的CNN模型训练MNIST手写数据集，相关代码： 8 | 9 | - [model.py](./model.py)：定义一个简单的CNN网络 10 | - [data.py](./data.py)：MNIST训练集和数据集准备 11 | - [single_gpu_train.py](./single_gpu_train.py)：单GPU训练代码 12 | 13 | ### 方案一 14 | 15 | 核心在于使用`nn.DataParallel`将模型wrap一下，代码其他地方不需要做任何更改: 16 | 17 | ```python 18 | model = nn.DataParallel(model) 19 | ``` 20 | 21 | 为方便说明，我们假设模型输入为(32, input_dim)，这里的 32 表示batch_size，模型输出为(32, output_dim)，使用 4 个GPU训练。`nn.DataParallel`起到的作用是将这 32 个样本拆成 4 份，发送给 4 个GPU 分别做 forward，然后生成 4 个大小为(8, output_dim)的输出，然后再将这 4 个输出都收集到`cuda:0`上并合并成(32, output_dim)。 22 | 23 | 可以看出，`nn.DataParallel`没有改变模型的输入输出，因此其他部分的代码不需要做任何更改，非常方便。但弊端是，后续的loss计算只会在`cuda:0`上进行，没法并行，因此会导致负载不均衡的问题。 24 | 25 | 如果把`loss`放在模型里计算的话，则可以缓解上述负载不均衡的问题，示意代码如下： 26 | 27 | ```python 28 | 29 | class Net: 30 | def __init__(self,...): 31 | # code 32 | 33 | def forward(self, inputs, labels=None) 34 | # outputs = fct(inputs) 35 | # loss_fct = ... 36 | if labels is not None: 37 | loss = loss_fct(outputs, labels) # 在训练模型时直接将labels传入模型，在forward过程中计算loss 38 | return loss 39 | else: 40 | return outputs 41 | ``` 42 | 43 | 按照我们上面提到的模型并行逻辑，在每个GPU上会计算出一个loss，这些loss会被收集到`cuda:0`上并合并成长度为 4 的张量。这个时候在做backward的之前，必须对将这个loss张量合并成一个标量，一般直接取mean就可以。这在Pytorch官方文档[nn.DataParallel函数]()中有提到： 44 | 45 | > When `module` returns a scalar (i.e., 0-dimensional tensor) in forward(), this wrapper will return a vector of length equal to number of devices used in data parallelism, containing the result from each device. 46 | 47 | 这部分的例子可以参考：[data_parallel_train.py](./data_parallel.py) 48 | 49 | ### 方案二 50 | 51 | 方案二被成为分布式数据并行(distributed data parallel)，是通过多进程实现的，相比与方案一要复杂很多。可以从以下几个方面理解： 52 | 53 | 1. 从一开始就会启动多个进程(进程数等于GPU数)，每个进程独享一个GPU，每个进程都会独立地执行代码。这意味着每个进程都独立地初始化模型、训练，当然，在每次迭代过程中会通过进程间通信共享梯度，整合梯度，然后独立地更新参数。 54 | 55 | 2. 每个进程都会初始化一份训练数据集，当然它们会使用数据集中的不同记录做训练，这相当于同样的模型喂进去不同的数据做训练，也就是所谓的数据并行。这是通过`torch.utils.data.distributed.DistributedSampler`函数实现的，不过逻辑上也不难想到，只要做一下数据partition，不同进程拿到不同的parition就可以了，官方有一个简单的demo，感兴趣的可以看一下代码实现：[Distributed Training](https://pytorch.org/tutorials/intermediate/dist_tuto.html#distributed-training) 56 | 57 | 3. 进程通过`local_rank`变量来标识自己，`local_rank`为0的为master，其他是slave。这个变量是`torch.distributed`包帮我们创建的，使用方法如下： 58 | 59 | ```python 60 | import argparse # 必须引入 argparse 包 61 | parser = argparse.ArgumentParser() 62 | parser.add_argument("--local_rank", type=int, default=-1) 63 | args = parser.parse_args() 64 | ``` 65 | 66 | 必须以如下方式运行代码： 67 | 68 | ```bash 69 | python -m torch.distributed.launch --nproc_per_node=2 --nnodes=1 train.py 70 | ``` 71 | 72 | 这样的话，`torch.distributed.launch`就以命令行参数的方式将`args.local_rank`变量注入到每个进程中，每个进程得到的变量值都不相同。比如使用 4 个GPU的话，则 4 个进程获得的`args.local_rank`值分别为0、1、2、3。 73 | 74 | 上述命令行参数`nproc_per_node`表示每个节点需要创建多少个进程(使用几个GPU就创建几个)；`nnodes`表示使用几个节点，因为我们是做单机多核训练，所以设为1。 75 | 76 | 4. 因为每个进程都会初始化一份模型，为保证模型初始化过程中生成的随机权重相同，需要设置随机种子。方法如下： 77 | 78 | ```python 79 | def set_seed(seed): 80 | random.seed(seed) 81 | np.random.seed(seed) 82 | torch.manual_seed(seed) 83 | torch.cuda.manual_seed_all(seed) 84 | ``` 85 | 86 | 87 | 使用方法通过如下示意代码展示： 88 | 89 | ```python 90 | from torch.utils.data.distributed import DistributedSampler # 负责分布式dataloader创建，也就是实现上面提到的partition。 91 | 92 | # 负责创建 args.local_rank 变量，并接受 torch.distributed.launch 注入的值 93 | parser = argparse.ArgumentParser() 94 | parser.add_argument("--local_rank", type=int, default=-1) 95 | args = parser.parse_args() 96 | 97 | # 每个进程根据自己的local_rank设置应该使用的GPU 98 | torch.cuda.set_device(args.local_rank) 99 | device = torch.device('cuda', args.local_rank) 100 | 101 | # 初始化分布式环境，主要用来帮助进程间通信 102 | torch.distributed.init_process_group(backend='nccl') 103 | 104 | # 固定随机种子 105 | seed = 42 106 | random.seed(seed) 107 | np.random.seed(seed) 108 | torch.manual_seed(seed) 109 | torch.cuda.manual_seed_all(seed) 110 | 111 | # 初始化模型 112 | model = Net() 113 | model.to(device) 114 | criterion = nn.CrossEntropyLoss() 115 | optimizer = optim.SGD(model.parameters(), lr=0.1) 116 | 117 | # 只 master 进程做 logging，否则输出会很乱 118 | if args.local_rank == 0: 119 | tb_writer = SummaryWriter(comment='ddp-training') 120 | 121 | # 分布式数据集 122 | train_sampler = DistributedSampler(train_dataset) 123 | train_loader = torch.utils.data.DataLoader(train_dataset, sampler=train_sampler, batch_size=batch_size) # 注意这里的batch_size是每个GPU上的batch_size 124 | 125 | # 分布式模型 126 | model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.local_rank], output_device=args.local_rank, find_unused_parameters=True) 127 | ``` 128 | 129 | 详细代码参考：[ddp_train.py](./ddp_train.py) 130 | 131 | 132 | ### ddp有用的技巧 133 | 134 | 官方推荐使用方案二(ddp)，所以这里收集ddp使用过程中的一些技巧。 135 | 136 | #### torch.distributed.barrier 137 | 138 | 在读[huggingface/transformers](https://github.com/huggingface/transformers)中的源码，比如`examples/run_ner.py`会看到一下代码： 139 | 140 | ```python 141 | # Load pretrained model and tokenizer 142 | if args.local_rank not in [-1, 0]: 143 | torch.distributed.barrier() # Make sure only the first process in distributed training will download model & vocab 144 | 145 | args.model_type = args.model_type.lower() 146 | config_class, model_class, tokenizer_class = MODEL_CLASSES[args.model_type] 147 | config = config_class.from_pretrained(args.config_name if args.config_name else args.model_name_or_path, 148 | num_labels=num_labels, 149 | cache_dir=args.cache_dir if args.cache_dir else None) 150 | tokenizer = tokenizer_class.from_pretrained(args.tokenizer_name if args.tokenizer_name else args.model_name_or_path, 151 | do_lower_case=args.do_lower_case, 152 | cache_dir=args.cache_dir if args.cache_dir else None) 153 | model = model_class.from_pretrained(args.model_name_or_path, 154 | from_tf=bool(".ckpt" in args.model_name_or_path), 155 | config=config, 156 | cache_dir=args.cache_dir if args.cache_dir else None) 157 | 158 | if args.local_rank == 0: 159 | torch.distributed.barrier() # Make sure only the first process in distributed training will download model & vocab 160 | ``` 161 | 162 | 上述代码要实现预训练模型的下载和读入内存，如果4个进程都分别下载一遍显然是不合理的，那如何才能实现只让一个进程下载呢？这个时候就可以使用`barrier`函数。当slave进程(local_rank!=0)运行到第一个`if`时就被barrier住了，只能等着，但master进程可以往下运行完成模型的下载和读入内存，但在第二个`if`语句时遇到barrier，那会不会被barrier住呢？答案是不会，因为master进程和slave进程集合在一起了(barrier)，barrier会被解除，这样大家都往下执行。当然这时大家执行的进度不同，master进程已经执行过模型读入，所以从第二个`if`往下执行，而slave进程尚未执行模型读入，只会从第一个`if`往下执行。 163 | 164 | 可以看到`barrier`类似一个路障，进程会被拦住，直到所有进程都集合齐了才放行。适合这样的场景：只一个进程下载，其他进程可以使用下载好的文件；只一个进程预处理数据，其他进程使用预处理且cache好的数据等。 165 | 166 | #### 模型保存 167 | 168 | 模型的保存与加载，与单GPU的方式有所不同。这里通通将参数以cpu的方式save进存储, 因为如果是保存的GPU上参数，pth文件中会记录参数属于的GPU号，则加载时会加载到相应的GPU上，这样就会导致如果你GPU数目不够时会在加载模型时报错，像下面这样： 169 | >RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device. 170 | 171 | 模型保存都是一致的，不过时刻记住方案二中你有多个进程在同时跑，所以会保存多个模型到存储上，如果使用共享存储就要注意文件名的问题，当然一般只在rank0进程上保存参数即可，因为所有进程的模型参数是同步的。 172 | 173 | ```python 174 | torch.save(model.module.cpu().state_dict(), "model.pth") 175 | ``` 176 | 177 | 模型的加载： 178 | 179 | ```python 180 | param=torch.load("model.pth") 181 | ``` 182 | 183 | 以下是[huggingface/transformers]()代码中用到的模型保存代码 184 | 185 | ```python 186 | if torch.distributed.get_rank() == 0: 187 | model_to_save = model.module if hasattr(model, "module") else model # Take care of distributed/parallel training 188 | model_to_save.save_pretrained(args.output_dir) 189 | tokenizer.save_pretrained(args.output_dir) 190 | ``` 191 | 192 | #### 同一台机器上跑多个 ddp task 193 | 194 | 假设想在一台有4核GPU的电脑上跑两个ddp task，每个task使用两个核，很可能会需要如下错误： 195 | 196 | ``` 197 | RuntimeError: Address already in use 198 | RuntimeError: NCCL error in: /opt/conda/conda-bld/pytorch_1544081127912/work/torch/lib/c10d/ProcessGroupNCCL.cpp:260, unhandled system error 199 | ``` 200 | 201 | 原因是两个ddp task通讯地址冲突，这时候需要显示地设置每个task的地址 202 | 203 | > specifying a different master_addr and master_port in torch.distributed.launch 204 | 205 | ```bash 206 | # 第一个task 207 | export CUDA_VISIBLE_DEVICES="0,1" 208 | python -m torch.distributed.launch --nproc_per_node=2 --master_addr=127.0.0.1 --master_port=29501 train.py 209 | 210 | # 第二个task 211 | export CUDA_VISIBLE_DEVICES="2,3" 212 | python -m torch.distributed.launch --nproc_per_node=2 --master_addr=127.0.0.2 --master_port=29502 train.py 213 | ``` 214 | 215 | 216 | ### 参考 217 | 218 | [Pytorch 多GPU训练-单运算节点-All you need](https://www.cnblogs.com/walter-xh/p/11586507.html) 219 | 220 | [WRITING DISTRIBUTED APPLICATIONS WITH PYTORCH](https://pytorch.org/tutorials/intermediate/dist_tuto.html) 221 | 222 | [pytorch 多GPU训练总结（DataParallel的使用）](https://blog.csdn.net/weixin_40087578/article/details/87186613) 223 | -------------------------------------------------------------------------------- /data.py: -------------------------------------------------------------------------------- 1 | from torchvision import datasets, transforms 2 | 3 | train_dataset = datasets.MNIST(root='./data/', train=True, transform=transforms.ToTensor()) 4 | test_dataset = datasets.MNIST(root='./data/', train=False, transform=transforms.ToTensor()) 5 | -------------------------------------------------------------------------------- /data_parallel_train.py: -------------------------------------------------------------------------------- 1 | import os 2 | os.environ["CUDA_VISIBLE_DEVICES"]="2,3" # 必须在`import torch`语句之前设置才能生效 3 | import torch 4 | import torch.optim as optim 5 | import torch.nn as nn 6 | from torch.utils.data import DataLoader 7 | from torch.utils.tensorboard import SummaryWriter 8 | 9 | from model import Net 10 | from data import train_dataset 11 | 12 | device = torch.device('cuda') 13 | batch_size = 64 14 | 15 | train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True) 16 | 17 | model = Net() 18 | model = model.to(device) 19 | optimizer = optim.SGD(model.parameters(), lr=0.1) 20 | model = nn.DataParallel(model) # 就在这里wrap一下，模型就会使用所有的GPU 21 | 22 | # training! 23 | tb_writer = SummaryWriter(comment='data-parallel-training') 24 | for i, (inputs, labels) in enumerate(train_loader): 25 | # forward 26 | inputs = inputs.to(device) 27 | labels = labels.to(device) 28 | outputs = model(inputs, labels=labels) 29 | loss = outputs[0] # 对应模型定义中，模型返回始终是tuple 30 | loss = loss.mean() # 将多个GPU返回的loss取平均 31 | # backward 32 | optimizer.zero_grad() 33 | loss.backward() 34 | optimizer.step() 35 | # log 36 | if i % 10 == 0: 37 | tb_writer.add_scalar('loss', loss.item(), i) 38 | tb_writer.close() -------------------------------------------------------------------------------- /ddp_train.py: -------------------------------------------------------------------------------- 1 | import os 2 | os.environ["CUDA_VISIBLE_DEVICES"]="0,1,2,3" 3 | import argparse 4 | import random 5 | import numpy as np 6 | import torch 7 | import torch.optim as optim 8 | import torch.nn as nn 9 | from torch.utils.data.distributed import DistributedSampler 10 | from torch.utils.tensorboard import SummaryWriter 11 | 12 | from model import Net 13 | from data import train_dataset, test_dataset 14 | 15 | parser = argparse.ArgumentParser() 16 | parser.add_argument("--local_rank", type=int, default=-1) 17 | args = parser.parse_args() 18 | 19 | torch.cuda.set_device(args.local_rank) 20 | device = torch.device('cuda', args.local_rank) 21 | torch.distributed.init_process_group(backend='nccl') 22 | 23 | # 固定随机种子 24 | seed = 42 25 | random.seed(seed) 26 | np.random.seed(seed) 27 | torch.manual_seed(seed) 28 | torch.cuda.manual_seed_all(seed) 29 | 30 | batch_size = 64 31 | 32 | model = Net() 33 | model.to(device) 34 | criterion = nn.CrossEntropyLoss() 35 | optimizer = optim.SGD(model.parameters(), lr=0.1) 36 | 37 | # training! 38 | if args.local_rank == 0: 39 | tb_writer = SummaryWriter(comment='ddp-3') 40 | 41 | train_sampler = DistributedSampler(train_dataset) 42 | train_loader = torch.utils.data.DataLoader(train_dataset, sampler=train_sampler, batch_size=batch_size) 43 | 44 | model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.local_rank], output_device=args.local_rank, find_unused_parameters=True) 45 | 46 | for i, (inputs, labels) in enumerate(train_loader): 47 | # forward 48 | inputs = inputs.to(device) 49 | labels = labels.to(device) 50 | outputs = model(inputs) 51 | loss = criterion(outputs, labels) 52 | # backward 53 | optimizer.zero_grad() 54 | loss.backward() 55 | optimizer.step() 56 | # log 57 | if args.local_rank == 0 and i % 5 == 0: 58 | tb_writer.add_scalar('loss', loss.item(), i) 59 | 60 | if args.local_rank == 0: 61 | tb_writer.close() 62 | -------------------------------------------------------------------------------- /model.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | 5 | 6 | class Net(nn.Module): 7 | def __init__(self): 8 | super(Net, self).__init__() 9 | self.conv1 = nn.Conv2d(in_channels=1, out_channels=6, kernel_size=5) 10 | self.pool1 = nn.MaxPool2d(kernel_size=2) 11 | self.conv2 = nn.Conv2d(in_channels=6, out_channels=16, kernel_size=3) 12 | self.pool2 = nn.MaxPool2d(kernel_size=2) 13 | self.fc1 = nn.Linear(16 * 5 * 5, 120) 14 | self.fc2 = nn.Linear(120, 60) 15 | self.fc3 = nn.Linear(60, 10) 16 | 17 | def forward(self, x, labels=None): 18 | x = F.relu(self.conv1(x)) 19 | x = self.pool1(x) 20 | x = F.relu(self.conv2(x)) 21 | x = self.pool2(x) 22 | x = x.view(-1, 16 * 5 * 5) 23 | x = F.relu(self.fc1(x)) 24 | x = F.relu(self.fc2(x)) 25 | x = F.relu(self.fc3(x)) 26 | 27 | x = (x,) 28 | 29 | if labels is not None: 30 | loss_fct = nn.CrossEntropyLoss() 31 | loss = loss_fct(x[0], labels) 32 | x = (loss,) + x 33 | 34 | return x # 模型的输出始终是一个tuple，如果labels不为None，则tuple第一个元素为loss 35 | -------------------------------------------------------------------------------- /single_gpu_train.py: -------------------------------------------------------------------------------- 1 | import os 2 | os.environ["CUDA_VISIBLE_DEVICES"]="2,3" # 必须在`import torch`语句之前设置才能生效 3 | import torch 4 | import torch.optim as optim 5 | import torch.nn as nn 6 | from torch.utils.data import DataLoader 7 | from torch.utils.tensorboard import SummaryWriter 8 | 9 | from model import Net 10 | from data import train_dataset 11 | 12 | device = torch.device('cuda') 13 | batch_size = 64 14 | 15 | train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True) 16 | 17 | model = Net() 18 | model = model.to(device) # 默认会使用第一个GPU 19 | optimizer = optim.SGD(model.parameters(), lr=0.1) 20 | 21 | # training! 22 | tb_writer = SummaryWriter(comment='single-gpu-training') 23 | for i, (inputs, labels) in enumerate(train_loader): 24 | # forward 25 | inputs = inputs.to(device) 26 | labels = labels.to(device) 27 | outputs = model(inputs, labels=labels) 28 | loss = outputs[0] # 对应模型定义中，模型返回始终是tuple 29 | # backward 30 | optimizer.zero_grad() 31 | loss.backward() 32 | optimizer.step() 33 | # log 34 | if i % 10 == 0: 35 | tb_writer.add_scalar('loss', loss.item(), i) 36 | tb_writer.close() --------------------------------------------------------------------------------