PyTorch Linux上如何进行分布式训练 - 问答

在Linux上进行PyTorch分布式训练通常涉及以下步骤：

环境准备：
- 确保所有参与训练的节点上都安装了相同版本的PyTorch和依赖库。
- 安装torch.distributed包，它是PyTorch提供的用于分布式训练的工具包。
网络配置：
- 所有节点需要在同一个网络中，并且能够通过IP地址相互通信。
- 配置防火墙规则，确保分布式训练所需的端口是开放的。
代码修改：
- 在你的训练脚本中，需要添加分布式训练的支持。这通常涉及到初始化分布式环境，例如使用torch.distributed.init_process_group函数。
- 修改数据加载器，以便它能够处理分布式数据并行。这可能包括使用torch.utils.data.distributed.DistributedSampler。
启动分布式训练：
- 使用torch.distributed.launch工具或者自定义脚本来启动分布式训练。这个工具会启动多个进程，每个进程对应一个训练节点。
- 你需要指定一些参数，如总的世界大小（world size），即参与训练的进程总数；当前进程的排名（rank），它在世界大小中的索引；以及主机文件（master_ip_file 或 master_addr），它列出了所有参与节点的IP地址和端口。

下面是一个简单的例子，展示了如何在PyTorch中使用torch.distributed.launch进行分布式训练：

# train.py
import torch
import torch.nn as nn
import torch.optim as optim
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader, DistributedSampler
from torchvision import datasets, transforms

def main(rank, world_size):
    # 初始化进程组
    torch.distributed.init_process_group(
        backend='nccl',  # 'nccl' is recommended for distributed GPU training
        init_method='tcp://<master_ip>:<master_port>',  # 替换为你的主节点IP和端口
        world_size=world_size,
        rank=rank
    )

    # 创建模型并将其移动到GPU
    model = nn.Sequential(
        nn.Linear(784, 1024),
        nn.ReLU(),
        nn.Linear(1024, 10)
    ).to(rank)

    # 使用DistributedDataParallel包装模型
    model = DDP(model, device_ids=[rank])

    # 加载数据集
    transform = transforms.Compose([transforms.ToTensor()])
    dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
    sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank)
    loader = DataLoader(dataset, batch_size=64, sampler=sampler)

    # 定义损失函数和优化器
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.SGD(model.parameters(), lr=0.01)

    # 训练模型
    for epoch in range(5):
        sampler.set_epoch(epoch)
        running_loss = 0.0
        for data, target in loader:
            data, target = data.to(rank), target.to(rank)
            optimizer.zero_grad()
            output = model(data.view(data.size(0), -1))
            loss = criterion(output, target)
            loss.backward()
            optimizer.step()
            running_loss += loss.item()
        print(f'Rank {rank}, Epoch {epoch}, Loss: {running_loss/len(loader)}')

    # 清理进程组
    torch.distributed.destroy_process_group()

if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('--world_size', type=int, default=4, help='number of distributed processes')
    parser.add_argument('--rank', type=int, default=0, help='rank of the current process')
    args = parser.parse_args()

    main(args.rank, args.world_size)

启动分布式训练的命令可能如下所示：

python -m torch.distributed.launch --nproc_per_node=<num_gpus_per_node> --nnodes=<num_nodes> --node_rank=<node_rank> train.py

其中<num_gpus_per_node>是每个节点上的GPU数量，<num_nodes>是节点总数，<node_rank>是当前节点的排名。

请注意，这只是一个基本的例子，实际的分布式训练可能需要更复杂的设置，包括模型并行化、更高级的数据加载策略、梯度累积等。此外，随着PyTorch版本的更新，分布式训练的API和最佳实践可能会发生变化，因此建议查阅最新的官方文档以获取最准确的信息。

0 赞

0 踩