评价此页
CommDebugMode">

CommDebugMode 入门#

创建于:2024 年 8 月 19 日 | 最后更新:2024 年 10 月 8 日 | 最后验证:2024 年 11 月 5 日

作者Anshul Sinha

在本教程中,我们将探讨如何在 PyTorch 的 DistributedTensor (DTensor) 中使用 CommDebugMode 进行调试,通过在分布式训练环境中跟踪集体操作。

先决条件#

  • Python 3.8 - 3.11

  • PyTorch 2.2 或更高版本

什么是 CommDebugMode 以及它有何用处#

随着模型规模的不断增大,用户正寻求利用各种并行策略的组合来扩展分布式训练。然而,现有解决方案之间缺乏互操作性带来了严峻的挑战,主要原因是缺少能够连接这些不同并行策略的统一抽象。为了解决这个问题,PyTorch 推出了 DistributedTensor(DTensor),它抽象了分布式训练中张量通信的复杂性,提供了无缝的用户体验。然而,当处理现有的并行解决方案和使用 DTensor 等统一抽象开发并行解决方案时,由于缺乏对底层集体通信发生的时间和原因的透明度,可能会使高级用户在识别和解决问题时面临挑战。为了应对这一挑战,Python 上下文管理器 CommDebugMode 将作为 DTensor 的主要调试工具之一,使用户能够查看在使用 DTensor 时集体操作何时以及为何发生,从而有效地解决这个问题。

使用 CommDebugMode#

以下是使用 CommDebugMode 的方法

# The model used in this example is a MLPModule applying Tensor Parallel
comm_mode = CommDebugMode()
    with comm_mode:
        output = model(inp)

# print the operation level collective tracing information
print(comm_mode.generate_comm_debug_tracing_table(noise_level=0))

# log the operation level collective tracing information to a file
comm_mode.log_comm_debug_tracing_table_to_file(
    noise_level=1, file_name="transformer_operation_log.txt"
)

# dump the operation level collective tracing information to json file,
# used in the visual browser below
comm_mode.generate_json_dump(noise_level=2)

这是 MLPModule 在噪声级别 0 下的输出示例

Expected Output:
    Global
      FORWARD PASS
        *c10d_functional.all_reduce: 1
        MLPModule
          FORWARD PASS
            *c10d_functional.all_reduce: 1
            MLPModule.net1
            MLPModule.relu
            MLPModule.net2
              FORWARD PASS
                *c10d_functional.all_reduce: 1

要使用 CommDebugMode,您必须将运行模型的代码包装在 CommDebugMode 中,并调用您想要用来显示数据的 API。您还可以使用 noise_level 参数来控制显示信息的详细程度。以下是每个噪声级别显示的内容:

0. 打印模块级别的集体操作计数
1. 打印 DTensor 操作(不包括微不足道的操作),模块分片信息
2. 打印张量操作(不包括微不足道的操作)
3. 打印所有操作

在上面的示例中,您可以看到集体操作 all_reduce 在 MLPModule 的前向传播中发生了一次。此外,您可以使用 CommDebugMode 来精确地指出 all_reduce 操作发生在 MLPModule 的第二个线性层中。

以下是您可以用于上传自己的 JSON 文件的交互式模块树可视化工具

CommDebugMode 模块树 - PyTorch 文档
将文件拖放到此处

结论#

在本教程中,我们学习了如何使用 CommDebugMode 来调试使用 PyTorch 的通信集体操作的 Distributed Tensors 和并行解决方案。您可以在嵌入式可视化浏览器中使用自己的 JSON 输出。

有关 CommDebugMode 的更详细信息,请参阅 comm_mode_features_example.py