（原型）用于调试卡住任务的飞行记录器¶
Created On: Sep 09, 2024 | Last Updated: Dec 17, 2024 | Last Verified: Sep 09, 2024
作者: Chirag Pandya, Junjie Wang
您将学习的内容¶
了解一种新的工具，可用于调试分布式训练过程中卡住的任务。
了解如何启用该工具并使用收集到的数据分析卡住的任务。
概述¶
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 开销。在收集的数据中，*持续时间*字段表示每个集合的执行时间。
附加设置¶
TORCH_SYMBOLIZE_MODE = (dladdr, addr2line, fast)：此设置确定要用来从运行程序中检索 C++ 堆栈踪迹的程序。默认设置为
addr2line。fast是一种新的实验模式，被证明比传统的addr2line快得多。请将此设置与TORCH_NCCL_TRACE_CPP_STACK结合使用以在飞行记录器数据中收集 C++ 堆栈踪迹。
如果您希望将飞行记录器数据存储到您的自定义存储而不是本地磁盘，您可以定义自己的写入类。此类应继承自类
::c10d::DebugInfoWriter(代码)，然后在启动 PyTorch 分布式之前使用::c10d::DebugInfoWriter::registerWriter(代码) 注册新写入器。
通过 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 目录中提供了分析捕获数据的便捷脚本。
要运行便捷脚本，请按以下步骤操作：
将每个 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>]
目前，我们为分析脚本支持两种模式。第一种模式允许脚本对解析的飞行记录器转储文件应用一些启发式方法，以生成识别超时潜在原因的报告。第二种模式则仅输出原始转储文件。默认情况下，脚本打印所有 rank 和所有 ``ProcessGroups``（PG）的飞行记录器转储。这可以通过 –selected-ranks 参数（指定 rank）和 –pg-filters 参数（指定 PG）缩小范围。例如命令如下：
注意：tabulate 模块是必须的，因此您可能需要先安装它。
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诊断工具，名为Flight Recorder。我们讨论了如何启用Flight Recorder以从机器收集诊断数据。此外，我们还探索了如何使用位于PyTorch代码库`tools/flight_recorder <https://github.com/pytorch/pytorch/tree/main/tools/flight_recorder>`__目录中的便捷脚本来分析从Flight Recorder捕获的数据。
