### 使用 `torchrun` 进行分布式训练 为了实现多机多卡的高效并行计算，PyTorch 提供了内置工具 `torchrun` 来简化启动过程。此命令可以用于执行基于 `DistributedDataParallel (DDP)` 的脚本[^1]。 #### 基础配置与环境设置 确保所有节点上的Python环境一致，并安装相同版本的PyTorch和其他依赖库。对于跨机器通信，则需保证网络畅通无阻，通常通过SSH免密登录来达成这一目标。 #### 启动参数详解 当利用 `torchrun` 发起任务时，有几个重要的选项需要注意： - `-nnodes`: 总共参与运算的工作站数目。 - `-nproc_per_node`: 每个工作站上使用的GPU数量。 - `--master_addr`: 主控服务器地址，默认为localhost；如果是集群环境下，请指定实际IP。 - `--node_rank`: 当前工作站编号，在单一作业内唯一标识一台计算机。 - `--rdzv_id`, `--rdzv_backend`, 和 `--rdzv_endpoint`: 配置rendezvous机制以便各进程间同步初始化信息。 下面是一个典型的调用实例，假设有一个名为train.py的文件准备运行于由两台设备组成的环境中，每台配备四张NVIDIA GPU显卡： ```bash torchrun --nnodes=2 \ --nproc_per_node=4 \ --node_rank=$RANK \ --master_addr="main_node_ip" \ train.py ``` 其中 `$RANK` 是当前节点在整个集群中的序号（从0开始），而 `"main_node_ip"` 则代表主节点的公网可达IP地址。 #### 编写兼容 DDP 的 Python 脚本 为了让程序能够适应分布式的架构设计模式，需要做如下调整： - 导入必要的模块：`import torch.distributed as dist` - 初始化进程组：`dist.init_process_group(backend='nccl')` - 将模型封装到 `DistributedDataParallel` 类中：`model = nn.parallel.DistributedDataParallel(model)` - 设置随机种子以保持实验可重复性 - 修改数据加载器以支持分片读取：采用 `DistributedSampler` 以下是经过改造后的简单示例代码片段： ```python import os import torch from torchvision import datasets, transforms, models from torch.nn.parallel import DistributedDataParallel as DDP from torch.utils.data.distributed import DistributedSampler def main(): rank = int(os.environ["LOCAL_RANK"]) transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,)) ]) dataset = datasets.MNIST('.', download=True, transform=transform) sampler = DistributedSampler(dataset) dataloader = torch.utils.data.DataLoader( dataset, batch_size=64, shuffle=False, num_workers=8, pin_memory=True, sampler=sampler ) model = models.resnet18(pretrained=False).cuda(rank) ddp_model = DDP(model, device_ids=[rank]) criterion = torch.nn.CrossEntropyLoss().cuda(rank) optimizer = torch.optim.SGD(ddp_model.parameters(), lr=0.01) for epoch in range(epochs): running_loss = 0.0 for i, data in enumerate(dataloader): inputs, labels = data[0].to(f'cuda:{rank}'), data[1].to(f'cuda:{rank}') outputs = ddp_model(inputs) loss = criterion(outputs, labels) optimizer.zero_grad() loss.backward() optimizer.step() if __name__ == '__main__': from torch.cuda.amp import GradScaler scaler = GradScaler() # 自动混合精度优化 main() ``` 上述代码展示了如何构建一个基本的支持DDP特性的MNIST分类器训练流程。