在深度学习训练中,使用多机多卡(多台机器和多块 GPU)可以显著加速模型训练过程。 PyTorch 提供了多种方法来实现多机多卡训练,以下是一些常用的方法和步骤:
1. 使用 torch.distributed 包
PyTorch 的 torch.distributed 包提供了分布式训练的支持。以下是使用 torch.distributed 进行多机多卡训练的步骤:
1.1 环境设置
首先,确保每台机器上都安装了相同版本的 PyTorch 和 CUDA 。然后,设置环境变量:
export MASTER_ADDR="主节点的 IP 地址"
export MASTER_PORT="主节点的端口号"
export WORLD_SIZE="总的进程数(机器数 * 每台机器的 GPU 数)"
export RANK="当前进程的全局排名(从 0开始)"
1.2 初始化分布式环境
在代码中,使用 torch.distributed.init_process_group 初始化分布式环境:
import torch
import torch.distributed as dist
def init_distributed_mode():
dist.init_process_group(backend='nccl', init_method='env://')
torch.cuda.set_device(local_rank)
1.3 创建模型和优化器
将模型和优化器移动到 GPU 上,并使用 torch.nn.parallel.DistributedDataParallel 包装模型:
model = MyModel().cuda()
model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[local_rank])
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
1.4 使用 DistributedSampler
在数据加载时,使用 torch.utils.data.distributed.DistributedSampler 来确保每个进程加载不同的数据:
from torch.utils.data import DataLoader, DistributedSampler
train_dataset = MyDataset()
train_sampler = DistributedSampler(train_dataset)
train_loader = DataLoader(train_dataset, batch_size=32, sampler=train_sampler)
1.5 训练循环
在训练循环中,确保每个 epoch 开始时调用 train_sampler.set_epoch(epoch):
for epoch in range(num_epochs):
train_sampler.set_epoch(epoch)
for batch in train_loader:
optimizer.zero_grad()
outputs = model(inputs)
loss = criterion(outputs, targets)
loss.backward()
optimizer.step()
2. 使用 torchrun 启动分布式训练
PyTorch 提供了 torchrun 工具来简化分布式训练的启动过程。以下是使用 torchrun 的步骤:
2.1 编写训练脚本
编写一个标准的训练脚本,确保包含分布式训练的初始化代码。
2.2 使用 torchrun 启动训练
在命令行中使用 torchrun 启动训练:
torchrun --nnodes=2 --nproc_per_node=4 --node_rank=0 --master_addr="主节点的 IP 地址" --master_port="主节点的端口号" train.py
3. 使用 torch.distributed.launch 启动分布式训练
torch.distributed.launch 是另一种启动分布式训练的工具,但它已经被 torchrun 所取代。以下是使用 torch.distributed.launch 的步骤:
3.1 编写训练脚本
编写一个标准的训练脚本,确保包含分布式训练的初始化代码。
3.2 使用 torch.distributed.launch 启动训练
在命令行中使用 torch.distributed.launch 启动训练:
python -m torch.distributed.launch --nproc_per_node=4 --nnodes=2 --node_rank=0 --master_addr="主节点的 IP 地址" --master_port="主节点的端口号" train.py
4. 使用 torch.nn.DataParallel(单机多卡)
如果只需要在单台机器上使用多块 GPU,可以使用 torch.nn.DataParallel:
model = MyModel()
model = torch.nn.DataParallel(model)
model = model.cuda()
for batch in train_loader:
outputs = model(inputs)
loss = criterion(outputs, targets)
loss.backward()
optimizer.step()
5.示例
'''多机多卡分布式训练
第一台机器
python -m torch.distributed.run --nproc_per_node=8 --nnodes=2 --node_rank=0 --master_addr="10.21.73.19" --master_port=12355 分布式训练_多机.py --batch-size 6 --epochs 10 --lr 1e-6 --eval-steps 1000 --max-val-item-count 1000 --use-lora
第二台机器
python -m torch.distributed.run --nproc_per_node=8 --nnodes=2 --node_rank=1 --master_addr="10.21.73.19" --master_port=12355 分布式训练_多机.py --batch-size 6 --epochs 10 --lr 1e-6 --eval-steps 1000 --max-val-item-count 1000 --use-lora
batchsize是每台机器上的batchsize
'''


![77,【1】.[CISCN2019 华东南赛区]Web4](https://i-blog.csdnimg.cn/direct/39879d558c374ab8bb4f4ba7a256805e.png)








![80,【4】BUUCTF WEB [SUCTF 2018]MultiSQL](https://i-blog.csdnimg.cn/direct/0d30aa60303a4b3d91190a5918e807cf.png)







