在深度学习训练中,使用多机多卡(多台机器和多块 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 distdef 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, DistributedSamplertrain_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
'''