torchrun 是 PyTorch 用于分布式训练的命令行工具，旨在简化启动和管理分布式训练任务的过程。下面我将详细讲解 torchrun 的使用方法，并讨论它与分布式数据并行（Distributed Data Parallel, DDP）的区别。
一、torchrun的使用方法
1. 安装
首先，确保你已经安装了 PyTorch 和 torch.distributed 模块。可以使用以下命令安装 PyTorch：
pip install torch
2. 配置环境
在进行分布式训练之前，你需要配置环境变量。通常需要设置以下环境变量：
MASTER_ADDR：主节点的 IP 地址。MASTER_PORT：主节点的端口号。WORLD_SIZE：参与训练的进程总数。RANK：当前进程的排名（从0开始）。
例如，可以在终端中设置这些变量：
export MASTER_ADDR="localhost"
export MASTER_PORT=12355
export WORLD_SIZE=2
export RANK=0
3. 编写训练脚本
编写一个分布式训练脚本。以下是一个简单的示例：
import torch
import torch.distributed as dist
import torch.multiprocessing as mp
from torch.nn.parallel import DistributedDataParallel as DDPdef setup(rank, world_size):dist.init_process_group("gloo", rank=rank, world_size=world_size)def cleanup():dist.destroy_process_group()def demo_basic(rank, world_size):print(f"Running basic DDP example on rank {rank}.")setup(rank, world_size)# Create model and move it to GPU with id rankmodel = torch.nn.Linear(10, 10).to(rank)ddp_model = DDP(model, device_ids=[rank])# Dummy input and targetinput = torch.randn(20, 10).to(rank)target = torch.randn(20, 10).to(rank)# Loss functionloss_fn = torch.nn.MSELoss()optimizer = torch.optim.SGD(ddp_model.parameters(), lr=0.001)optimizer.zero_grad()outputs = ddp_model(input)loss_fn(outputs, target).backward()optimizer.step()cleanup()def main():world_size = 2mp.spawn(demo_basic, args=(world_size,), nprocs=world_size, join=True)if __name__ == "__main__":main()
4. 使用torchrun启动训练
可以使用 torchrun 命令启动分布式训练任务。假设你的脚本名为 train.py，可以使用以下命令启动训练：
torchrun --nproc_per_node=2 train.py
二、torchrun与DDP的区别
torchrun 是一个用于启动和管理分布式训练任务的工具，而 DDP 是 PyTorch 提供的用于实现分布式数据并行的模块。两者的关系可以概括如下：
- torchrun：负责启动和管理分布式训练进程，简化了环境变量的配置和进程的启动。
- DDP：负责在分布式训练中进行模型的同步和梯度的聚合。
使用 torchrun 可以简化分布式训练任务的启动，而 DDP 则是在分布式训练中实现具体的并行计算。通常情况下，二者是配合使用的：torchrun 启动训练任务，DDP 进行分布式训练。
示例总结
结合上述内容，我们可以看到，通过使用 torchrun，你可以轻松启动多个分布式训练进程，而无需手动配置复杂的环境变量。这极大地简化了分布式训练的管理和调试过程。而 DDP 则确保了在多个进程之间同步模型参数和梯度，使得训练过程能够高效地并行化。
希望这些信息能帮助你更好地理解和使用 torchrun 以及 PyTorch 的分布式数据并行。若有任何问题或进一步的需求，请随时提问。
