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

调试卡住作业的飞行记录器#

作者： Chirag Pandya, Junjie Wang

您将学到什么#

  • 了解用于调试分布式训练过程中卡住作业的新工具。

  • 了解如何启用该工具以及使用收集到的数据分析卡住的作业。

先决条件#

  • PyTorch 版本 2.5 或更高版本。

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

概述#

AI 分布式训练作业是指使用连接在网络中的多个设备（如 GPU 或 CPU）训练机器学习模型的过程。这种方法可以更快、更高效地训练需要大量计算资源的大型模型。工程师的目标是尽快完成 AI 训练作业，并不断改进，以便后续训练能够更快地完成。训练好且可用的模型是最终期望的成果。完成训练的最大障碍之一是“卡住的作业”这一概念。

当分布式 AI 训练作业在一段较长的时间内停止取得有意义的进展时，就被认为已“卡住”。

作业可能由于各种原因而卡住

  • 数据饥饿： 当训练作业接收数据的速度未达到预期时，这可能是由于数据管道或数据源的问题所致。

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

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

  • 软件错误或故障： 训练代码或底层库和框架中的错误也可能导致作业卡住。

  • 同步问题： 在分布式训练中，计算的不同部分通常是并行运行的，并在某些点需要同步。如果此同步失败，作业可能会卡住。例如，如果一个或多个 rank 未能加入集体通信，而其余 rank 已加入，则可能发生死锁。这会导致作业无限期等待进展。

飞行记录器顾名思义，会在集体通信运行时捕获诊断信息。捕获的诊断信息用于帮助识别作业卡住时问题的根本原因。飞行记录器由两个核心部分组成：

  • 收集部分：启用后，集体通信的信息会被记录在一个内存中的循环缓冲区中。在作业超时时，或按需，内存缓冲区可以被检索或转储到文件中。

  • 分析器脚本可在 tools/flight_recorder 目录中找到（详细信息如下）。

    分析器脚本使用收集到的数据运行已知的启发式方法，并尝试自动识别导致作业停滞的根本原因。

启用飞行记录器#

有三个必需的环境变量可以使飞行记录器的初始版本正常工作。

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

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

  • TORCH_NCCL_DEBUG_INFO_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 事件，并记录每个集体通信的 _持续时间_。这可能会产生一些 CPU 开销。在收集到的数据中，_duration_ 字段指示每个集体通信的执行时间。

附加设置#

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

    默认设置为 addr2line。

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

  • 如果您不想将飞行记录器数据转储到本地磁盘，而是转储到您自己的存储中，您可以定义自己的 writer 类。该类应继承自 ::c10d::DebugInfoWriter 类 (代码)，然后使用 ::c10d::DebugInfoWriter::registerWriter (代码) 注册新 writer，然后我们初始化 PyTorch distributed。

通过 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/tools/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_ 参数来指定 PGs。一个命令示例是

请注意：需要 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/

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

结论#

在本教程中，我们了解了 PyTorch 的一项新的诊断工具，称为飞行记录器。我们讨论了如何启用飞行记录器来从机器收集诊断数据。此外，我们还探讨了如何使用 PyTorch 存储库的 _tools/flight_recorder_ 目录中的便利脚本来分析从飞行记录器捕获的数据。

在此页面上
PyTorch 库