开始使用 CommDebugMode
¶
作者: Anshul Sinha
在本教程中,我们将探讨如何在分布式训练环境中使用 PyTorch 的 DistributedTensor (DTensor) 中的 CommDebugMode
进行调试,方法是跟踪集体操作。
先决条件¶
Python 3.8 - 3.11
PyTorch 2.2 或更高版本
什么是 CommDebugMode
以及它为什么有用¶
随着模型规模的不断增大,用户正在寻求利用各种并行策略组合来扩展分布式训练。然而,现有解决方案之间缺乏互操作性带来了重大挑战,主要原因是缺乏能够桥接这些不同并行策略的统一抽象。为了解决这个问题,PyTorch 提出了DistributedTensor(DTensor),它抽象了分布式训练中张量通信的复杂性,提供了无缝的用户体验。但是,在处理现有并行解决方案以及使用 DTensor 等统一抽象开发并行解决方案时,缺乏对底层何时以及发生了哪些集体通信的透明性,可能会使高级用户难以识别和解决问题。为了应对这一挑战,CommDebugMode
,一个 Python 上下文管理器将作为 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
参数来控制显示信息的详细程度。以下是每个噪声级别显示的内容
在上面的示例中,您可以看到集体操作 all_reduce 在 MLPModule
的前向传递中发生了一次。此外,您可以使用 CommDebugMode
精确定位 all-reduce 操作发生在 MLPModule
的第二个线性层中。
以下是您可以用来上传自己的 JSON 导出文件的交互式模块树可视化工具
结论¶
在本教程中,我们学习了如何使用 CommDebugMode
调试分布式张量和使用通信集体操作的 PyTorch 并行解决方案。您可以在嵌入式可视化浏览器中使用您自己的 JSON 输出。
有关 CommDebugMode
的更多详细信息,请参阅 comm_mode_features_example.py