pytorch 多机多卡训练方法

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

原文地址：https://blog.csdn.net/u012374012/article/details/145284637

免责声明：本站文章内容转载自网络资源，如侵犯了原著者的合法权益，可联系本站删除。更多内容请关注自学内容网（zxcms.com）！