评价此页
★ ★ ★ ★ ★
unstable/flight_recorder_tutorial

用于调试卡死任务的飞行记录器#

作者: Chirag Pandya, Junjie Wang

您将学到什么#

  • 了解一种用于在分布式训练期间调试卡死任务的新工具。

  • 学习如何启用该工具并使用收集到的数据来分析卡死任务。

先决条件#

  • PyTorch 2.5 或更高版本。

  • tabulate。可以通过运行 pip install tabulate 进行安装。

概述#

AI 分布式训练任务是指使用通过网络连接的多个设备（如 GPU 或 CPU）来训练机器学习模型的过程。这种方法能够更快、更高效地训练需要大量计算资源的大型模型。工程师的目标是尽可能快地完成 AI 训练任务，并进行持续改进，以便后续训练能更快完成。获得一个经过训练且可用的模型是最终的预期结果。阻碍训练完成的最大障碍之一是“任务卡死”现象。

当分布式 AI 训练任务在很长一段时间内没有取得实质性进展时，会被视为 卡死 (stuck)。

任务卡死的原因多种多样：

  • 数据匮乏 (Data Starvation)： 当训练任务未按预期速度接收到数据时，会出现这种情况，这可能是由数据流水线或数据源的问题引起的。

  • 资源受限： 如果运行任务的系统没有足够的计算资源（如 CPU、GPU 或内存），任务可能无法继续执行。

  • 网络问题： 在分布式训练设置中，模型或数据的不同部分可能在不同的设备上进行处理。如果存在网络问题，这些设备之间的通信可能会中断，导致任务卡死。

  • 软件错误或 Bug： 训练代码或底层库和框架中的错误也可能导致任务卡死。

  • 同步问题： 在分布式训练中，计算的不同部分通常并行运行，并需要在某些点进行同步。如果同步失败，任务可能会卡死。例如，如果一个或多个 rank 在其他 rank 已加入集合通信 (collective) 的情况下未能加入，则可能发生死锁。这导致任务无限期等待而无法推进。

顾名思义，飞行记录器 (Flight Recorder) 会在集合通信运行时捕获诊断信息。捕获的诊断信息用于帮助识别任务卡死时的根本原因。飞行记录器包含两个核心部分：

  • 收集部分：启用后，有关集合通信的信息会记录在内存环形缓冲区中。在任务超时或有需要时，可以检索内存缓冲区或将其转储到文件。

  • 分析脚本：在 torch/distributed/flight_recorder 目录中提供了一个分析器脚本（详情见下文）。

    分析器脚本使用收集到的数据运行已知的启发式算法，并尝试自动识别导致任务停滞的根本问题。

启用飞行记录器#

要使初始版本的飞行记录器正常工作，需要设置三个环境变量。

  • TORCH_NCCL_TRACE_BUFFER_SIZE = (0, N)：将 N 设置为正数即可启用收集功能。N 表示将在环形缓冲区内部保留的条目数量。我们建议将此值设置为 2000。默认值为 2000。

  • TORCH_NCCL_DUMP_ON_TIMEOUT = (true, false)：将其设置为 true 将在任务超时时将诊断文件写入磁盘。如果启用，任务的运行目录中每个 rank 将输出一个文件。默认值为 false。

  • TORCH_FR_DUMP_TEMP_FILE：设置飞行记录器转储文件的路径及文件前缀。每个 rank 一个文件。默认值为 /tmp/nccl_trace_rank_。

可选设置

  • TORCH_NCCL_TRACE_CPP_STACK = (true, false)：将其设置为 true 可在飞行记录器中捕获 C++ 堆栈跟踪。C++ 堆栈跟踪有助于提供从 PyTorch Python 调用到底层 C++ 实现的确切代码路径。另请参阅附加设置中的 TORCH_SYMBOLIZE_MODE。

  • TORCH_NCCL_ENABLE_TIMING = (true, false)：将其设置为 true 将在每次集合通信开始时启用额外的 CUDA 事件，并记录每次集合通信的 持续时间 (duration)。这可能会产生一些 CPU 开销。在收集的数据中，duration 字段表示每次集合通信执行所需的时间。

附加设置#

  • TORCH_SYMBOLIZE_MODE = (dladdr, addr2line, fast)：此设置决定了用于从运行中的程序检索 C++ 跟踪的程序。

    默认设置为 addr2line。

    fast 是一种新的实验性模式，被证明比传统的 addr2line 快得多。将此设置与 TORCH_NCCL_TRACE_CPP_STACK 结合使用，可在飞行记录器数据中收集 C++ 跟踪。

  • 如果您不希望将飞行记录器数据转储到本地磁盘，而是希望存入您自己的存储空间，则可以定义自己的 writer 类。该类应继承自 ::c10d::DebugInfoWriter (代码)，然后在启动 PyTorch 分布式之前，使用 ::c10d::DebugInfoWriter::registerWriter (代码) 注册新的 writer。

通过 API 检索飞行记录器数据#

您还可以通过 API 调用检索飞行记录器数据。带有默认参数的 API 如下所示：

torch._C._distributed_c10d._dump_nccl_trace(includeCollectives=True, includeStackTraces=True, onlyActive=False)

要查看数据，您可以按如下方式进行 unpickle：

t = pickle.loads(torch._C._distributed_c10d._dump_nccl_trace())
print(t)

飞行记录器文件格式#

飞行记录器文件以 pickle 格式转储。文件被写入本地磁盘或挂载的共享 NFS 文件夹。

飞行记录器 unpickled 后的文件内容如下所示：

{
  "version": "2.5",
  "pg_config": {
    "0": {
    "name": "0",
    "desc": "default_pg",
    "ranks": "[0, 1]"
    }
  },
  "pg_status": {
    "0": {
    "last_enqueued_collective": 2,
    "last_started_collective": -1,
    "last_completed_collective": 2
    }
  },
  "entries": [
  {
    "frames": [
    {
    "name": "test_short_pickle",
    "filename": "pytorch/test/distributed/test_c10d_nccl.py",
    "line": 3647
    },
    {
    "name": "spawn_main",
    "filename": ".conda/envs/pytorch-3.10/lib/python3.10/multiprocessing/spawn.py",
    "line": 116
    },
    {
    "name": "<module>",
    "filename": "<string>",
    "line": 1
    }
    ],
    "record_id": 0,
    "pg_id": 0,
    "process_group": ("0", "default_pg"),
    "collective_seq_id": 1,
    "p2p_seq_id": 0,
    "op_id": 1,
    "profiling_name": "nccl:all_reduce",
    "time_created_ns": 1724779239936775119,
    "input_sizes": [[3, 4]],
    "input_dtypes": ["Float"],
    "output_sizes": [[3, 4]],
    "output_dtypes": ["Float"],
    "state": "completed",
    "time_discovered_started_ns": null,
    "time_discovered_completed_ns": 1724779239975811724,
    "retired": true,
    "timeout_ms": 600000,
    "is_p2p": false
    },
    ...
    ]
}

分析飞行记录器转储文件#

我们在 pytorch/torch/distributed/flight_recorder 目录中提供了用于分析捕获数据的便捷脚本。

要运行此便捷脚本，请执行以下步骤：

  1. 将所有 rank 的文件复制到单个目录中。

  2. 运行该脚本，请使用以下命令：

python fr_trace.py <dump dir containing trace files> [-o <output file>]

如果您安装了 PyTorch 夜间构建版，或者通过设置 USE_DISTRIBUTED=1 从源码构建，则可以直接使用以下命令：

torchfrtrace <dump dir containing trace files> [-o <output file>]

目前，我们为分析器脚本支持两种模式。第一种模式允许脚本将启发式算法应用于解析后的飞行记录器转储文件，以生成一份报告，确定导致超时的潜在罪魁祸首。第二种模式仅输出原始转储内容。默认情况下，该脚本会打印所有 rank 和所有 ``ProcessGroups`` (PGs) 的飞行记录器转储信息。可以通过使用 –selected-ranks 参数指定 rank，以及使用 –pg-filters 参数指定 PG 来缩小范围。示例命令如下：

注意：需要 tabulate 模块，因此您可能需要先 pip install 它。

python fr_trace.py <dump dir containing trace files> -j [--selected-ranks i j k ...] [--pg-filters tp dp]
torchfrtrace <dump dir containing trace files> -j [--selected-ranks i j k ...] [--pg-filters 0 2]

端到端示例#

为了演示飞行记录器的使用，我们将使用一个故意诱发集合通信不匹配的小程序。在此示例中，rank0 被编程为执行额外的集合通信。飞行记录器转储文件保存到 /tmp 目录中。出于演示目的，我们将此程序命名为 crash.py。

注意

请注意，这是一个简化的示例。在实际场景中，过程会更复杂。

import torch
import torch.distributed as dist
import os
from datetime import timedelta

local_rank = int(os.environ["LOCAL_RANK"])
world_size = int(os.environ["WORLD_SIZE"])
assert world_size <= 8, "world size must be less than or equal to 8"
os.environ["TORCH_NCCL_DEBUG_INFO_TEMP_FILE"] = "/tmp/trace_"
os.environ["TORCH_NCCL_DUMP_ON_TIMEOUT"] = "1"
os.environ["TORCH_NCCL_TRACE_BUFFER_SIZE"] = "2000"
device = torch.device(f"cuda:{local_rank}")
print(f"{local_rank=} {world_size=} master addr: {os.environ['MASTER_ADDR']} master port: {os.environ['MASTER_PORT']} {device=}")

# Initialize the process group with a small timeout so that jobs fail quickly
dist.init_process_group("nccl", world_size=world_size, rank=local_rank, timeout=timedelta(seconds=1))

a = torch.full((3, 4), float(local_rank), device=device)
# Write some collectives to populate Flight Recorder data
for i in range(2):
  print(f"calling allreduce on {local_rank=}")
  f = dist.all_reduce(a)

# rank0 is doing an additional collective
if local_rank == 0:
  print("rank0 is doing an allreduce on tensor b, but other ranks forgot")
  b = torch.full((4,5), float(local_rank), device=device)
  f = dist.all_reduce(b)

for i in range(2):
  print(f"calling allreduce on {local_rank=}")
  f = dist.all_reduce(a)

torch.cuda.synchronize(device=device)
print(f"{local_rank=} exiting")

要运行此程序，请使用 torchrun：

torchrun --nnodes=1 --nproc_per_node=2 crash.py

您应该会在 /tmp 目录中看到两个文件：

$ls /tmp/trace*
# Expected output
/tmp/trace_0 /tmp/trace_1

最后，为了分析这两个文件，我们使用 torchfrtrace 命令：

torchfrtrace --prefix "trace_" /tmp/

trace 命令的输出旨在供人类阅读。它包含有关导致失败的集合通信集的信息。上述命令的输出如下所示。我们可以清楚地看到 rank 1 没有加入 “all_reduce” 集合通信。

结论#

在本教程中，我们了解了一个名为飞行记录器 (Flight Recorder) 的 PyTorch 新诊断工具。我们讨论了如何启用飞行记录器从机器收集诊断数据。此外，我们还探索了如何使用位于 PyTorch 仓库 torch/distributed/flight_recorder 目录下的便捷脚本来分析飞行记录器捕获的数据。

在此页面上
PyTorch 库