(原型) 用于调试卡死任务的飞行记录器¶

创建于: Sep 09, 2024 | 最后更新于: Dec 17, 2024 | 最后验证于: Sep 09, 2024

你将学到什么¶

了解一种新的工具，用于调试分布式训练期间卡死的任务。
学习如何启用该工具以及如何使用收集的数据分析卡死的任务。

先决条件¶

PyTorch 版本 2.5 或更高。
tabulate。你可以通过运行 pip install tabulate 进行安装。

概述¶

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

当分布式 AI 训练任务长时间没有取得有意义的进展时，就被认为是卡死了。

任务卡死可能有多种原因

数据饥饿： 当训练任务没有以预期的速率接收数据时发生，可能是由于数据管道或数据源的问题。
资源限制： 如果运行任务的系统没有足够的计算资源（例如 CPU、GPU 或内存），任务可能无法继续。
网络问题： 在分布式训练设置中，模型或数据的不同部分可能在不同的设备上处理。如果存在网络问题，这些设备之间的通信可能会中断，导致任务卡死。
软件 Bug 或错误： 训练代码或底层库和框架中的错误也可能导致任务卡死。
同步问题： 在分布式训练中，计算的不同部分通常并行运行，并且需要在某些点进行同步。如果同步失败，任务可能会卡死。例如，如果一个或多个 rank 未能加入集体通信 (collective)，而其余 rank 已加入，则可能发生死锁。这导致任务无限期地等待才能继续。

Flight Recorder，顾名思义，在集体通信运行时捕获诊断信息。捕获的诊断信息用于帮助识别任务卡死时的根本原因。Flight Recorder 由两个核心部分组成

收集部分：启用后，有关集体通信的信息会记录在内存中的循环缓冲区中。在任务超时或按需，可以检索内存中的缓冲区或将其转储到文件。
分析脚本可在 tools/flight_recorder 目录中找到（详细信息如下）。
分析脚本使用收集的数据运行已知的启发式算法，并尝试自动识别导致任务停滞的根本问题。

启用 Flight Recorder¶

要使初始版本的 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 将在 Flight Recorder 中捕获 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 结合使用，可以在 Flight Recorder 数据中收集 C++ 跟踪。
如果你不想将飞行记录器数据转储到本地磁盘，而是转储到你自己的存储，你可以定义自己的写入器类。该类应继承自类 ::c10d::DebugInfoWriter (代码)，然后在我们启动 PyTorch 分布式之前使用 ::c10d::DebugInfoWriter::registerWriter (代码) 注册新的写入器。

通过 API 检索 Flight Recorder 数据¶

你还可以通过 API 调用检索 Flight Recorder 数据。带有默认参数的 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)

Flight Recorder 文件格式¶

Flight Recorder 文件以 pickle 格式转储。文件写入本地磁盘或挂载的共享 NFS 文件夹。

Flight Recorder 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
    },
    ...
    ]
}

分析 Flight Recorder 转储¶

我们在 pytorch/tools/flight_recorder 目录中提供了方便的脚本来分析捕获的数据。

要运行方便脚本，请按照以下步骤操作

将一个 rank 的所有文件复制到一个目录中。
要运行脚本，请使用此命令

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>]

目前，我们支持分析脚本的两种模式。第一种模式允许脚本对解析后的飞行记录器转储应用一些启发式算法，生成一份报告，识别可能导致超时的罪魁祸首。第二种模式仅输出原始转储。默认情况下，脚本打印所有 ranks 和所有 ``ProcessGroups``(PGs) 的飞行记录器转储。可以使用 –selected-ranks 参数指定 ranks，使用 –pg-filters 参数指定 PGs 来缩小范围。示例命令是

注意：需要 tabulate 模块，因此你可能需要先用 pip 安装它。

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]

端到端示例¶

为了演示 Flight Recorder 的使用，我们将使用一个小型程序，其中我们引入不匹配的集体通信。在此示例中，rank0 被编程为执行额外的集体通信。Flight Recorder 转储文件保存到 /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 诊断工具，称为 Flight Recorder。我们讨论了如何启用 Flight Recorder 来从机器收集诊断数据。此外，我们还探讨了如何使用 PyTorch 仓库 tools/flight_recorder 目录中的方便脚本来分析从 Flight Recorder 捕获的数据。