【知识】torchrun 与 torch.multiprocessing.spawn 的对比

转载请注明出处：小锋学长生活大爆炸[xfxuezhagn.cn]

如果本文帮助到了你，欢迎[点赞、收藏、关注]哦~

来自ChatGPT、DeepSeek

有点干，可仅做了解。

torchrun 和 torch.multiprocessing.spawn 都是在 PyTorch 中用于并行化和分布式训练的工具，但它们在使用场景和实现方式上有所不同。

1. 用途和功能

torchrun:
- 主要用于分布式训练，特别是在多机或多卡训练时。
- torchrun 是 PyTorch 提供的一个命令行工具，它自动启动分布式训练环境并启动多个进程。通常用于在多个节点（例如，多个GPU或多个机器）上启动并行训练。
- 它是 torch.distributed.launch 的替代品，提供更简洁的配置和更好的支持。
torch.multiprocessing.spawn:
- 是一个 Python API，用于在单个机器（或单个进程）上启动多个子进程。这些子进程通常是用于在每个进程上运行不同的模型副本或进行数据并行。
- spawn 是在单机多卡（multi-GPU）环境下进行训练时常用的工具，特别适用于分布式数据并行（torch.nn.DataParallel 或 torch.nn.parallel.DistributedDataParallel）。
- 它允许你控制每个进程的启动，并且能确保每个进程有独立的 GPU 资源。

2. 实现方式

torchrun:

它基于 torch.distributed，通常通过传递命令行参数来配置分布式环境。你只需指定 GPU 数量、节点数量、主节点等配置。
它会自动配置并启动各个训练进程，并且处理进程间的通信。
命令行调用的示例：

# script.py
import torch
import torch.distributed as dist

def main():
    dist.init_process_group(backend="nccl")
    rank = dist.get_rank()
    # 训练逻辑

if __name__ == "__main__":
    main()

torchrun --nnodes=1 --nproc_per_node=8 --rdzv_id=1234 --rdzv_backend=c10d --master_addr="localhost" --master_port=29500 script.py

torch.multiprocessing.spawn:

通过 Python 代码调用，每个进程都是通过 multiprocessing.spawn API 启动的。每个子进程可以执行不同的任务。
它通常用来启动多个进程，并在每个进程上执行模型训练代码，能够在单机环境下利用多个 GPU。
代码示例：

import torch
import torch.distributed as dist
from torch.multiprocessing import spawn

def train_fn(rank, world_size, args):
    dist.init_process_group(
        backend="nccl",
        init_method="env://",
        world_size=world_size,
        rank=rank
    )
    # 训练逻辑

if __name__ == "__main__":
    world_size = 4
    spawn(train_fn, args=(world_size, {}), nprocs=world_size)

3. 进程间通信

torchrun:
- 自动设置进程间的通信和同步。它是基于 NCCL（NVIDIA Collective Communications Library）或 Gloo 进行通信，适合大规模分布式训练。
torch.multiprocessing.spawn:
- 你需要手动设置通信（如使用 torch.nn.parallel.DistributedDataParallel 或 torch.distributed 来进行多进程间的数据同步和梯度更新）。
- 更加灵活，但也需要开发者更细致的配置。

4. 跨节点支持

torchrun:
- 支持跨节点训练，可以设置多个机器上的进程，适合大规模多机训练。
torch.multiprocessing.spawn:
- 通常用于单机多卡训练，不直接支持跨节点训练，更多的是集中在本地多个 GPU 上。

5. 效率影响

在 PyTorch 分布式训练中，torchrun 和 torch.multiprocessing.spawn 的底层通信机制（如 NCCL、Gloo）是相同的，因此两者的训练效率（如单步迭代速度）在理想配置下通常不会有显著差异。然而，它们的设计差异可能间接影响实际训练效率，尤其是在环境配置、资源管理和容错机制上。

1. 效率核心因素：无本质差异

通信后端相同：无论是 torchrun 还是 spawn，底层均依赖 PyTorch 的分布式通信库（如 NCCL、Gloo），数据传输效率由后端实现决定，与启动工具无关。
计算逻辑一致：模型前向传播、反向传播的计算逻辑完全由用户代码控制，与启动工具无关。

2. 间接影响效率的场景

场景 1：环境初始化效率

torch.multiprocessing.spawn：
- 需要手动初始化分布式环境（如 init_process_group），若配置错误（如端口冲突、IP 错误）可能导致进程启动延迟或失败。
- 单机多卡场景下简单直接，但多机场景需手动同步 MASTER_ADDR 和 MASTER_PORT，易出错且耗时。
torchrun：
- 自动设置环境变量（如 RANK, WORLD_SIZE, MASTER_ADDR 等），减少配置错误风险。
- 在多机训练中，通过参数（如 --nnodes, --node_rank）快速配置，显著降低初始化时间。

结论：torchrun 在复杂环境（多机）下初始化更高效，减少人为错误导致的延迟。

场景 2：资源管理与进程调度

torch.multiprocessing.spawn：
- 父进程直接管理子进程，若某个子进程崩溃，整个训练任务会直接终止（无容错）。
- 资源分配完全由用户代码控制，缺乏动态调整能力。
torchrun：
- 支持弹性训练（需结合 torch.distributed.elastic），进程崩溃后可自动重启并恢复训练（需用户实现检查点逻辑）。
- 提供更精细的进程监控和资源分配策略（如动态调整 WORLD_SIZE），减少资源闲置。

结论：torchrun 在容错和资源利用率上更优，尤其在长时训练或不稳定环境中，能减少因故障导致的总时间浪费。

场景 3：日志与调试效率

torch.multiprocessing.spawn：
- 各进程日志独立输出，需手动聚合分析（如使用 torch.distributed 的日志工具）。
- 错误堆栈可能分散，调试复杂。
`torchrun``：
- 提供统一的日志输出格式，自动聚合错误信息。
- 支持通过 --redirect 参数重定向日志，便于定位问题。