问题排查¶

请注意，本节中的信息可能会在 PyTorch/XLA 软件的未来版本中删除，因为其中许多信息是特定于给定内部实现的，而内部实现可能会发生变化。

健全性检查¶

在执行任何深入调试之前，我们希望对已安装的 PyTorch/XLA 进行健全性检查。

检查 PyTorch/XLA 版本¶

PyTorch 和 PyTorch/XLA 版本应匹配。请查看我们的 README，了解有关可用版本的更多详细信息。

vm:~$ python
>>> import torch
>>> import torch_xla
>>> print(torch.__version__)
2.1.0+cu121
>>> print(torch_xla.__version__)
2.1.0

执行简单计算¶

vm:~$ export PJRT_DEVICE=TPU
vm:~$ python3
>>> import torch
>>> import torch_xla.core.xla_model as xm
>>> t1 = torch.tensor(100, device=xm.xla_device())
>>> t2 = torch.tensor(200, device=xm.xla_device())
>>> print(t1 + t2)
tensor(300, device='xla:0')

使用虚假数据运行 Resnet¶

对于 nightly 版本

vm:~$ git clone https://github.com/pytorch/xla.git
vm:~$ python xla/test/test_train_mp_imagenet.py --fake_data

对于发布版本 x.y，您需要使用分支 rx.y。例如，如果您安装了 2.1 版本，您应该执行

vm:~$ git clone --branch r2.1 https://github.com/pytorch/xla.git
vm:~$ python xla/test/test_train_mp_imagenet.py --fake_data

如果您可以运行 resnet，我们可以得出结论，torch_xla 已正确安装。

性能调试¶

为了诊断性能问题，我们可以使用 PyTorch/XLA 提供的执行指标和计数器。当模型速度缓慢时，要检查的第一件事是生成指标报告。

指标报告对于诊断问题非常有帮助。如果您有指标报告，请尝试将其包含在您发送给我们的错误报告中。

PyTorch/XLA 调试工具¶

您可以通过设置 PT_XLA_DEBUG_LEVEL=2 来启用 PyTorch/XLA 调试工具，该工具提供了一些有用的调试功能。您也可以将调试级别降低到 1 以跳过执行分析。

执行自动指标分析¶

调试工具将分析指标报告并提供摘要。一些示例输出如下所示

pt-xla-profiler: CompileTime too frequent: 21 counts during 11 steps
pt-xla-profiler: TransferFromDeviceTime too frequent: 11 counts during 11 steps
pt-xla-profiler: Op(s) not lowered: aten::_ctc_loss, aten::_ctc_loss_backward,  Please open a GitHub issue with the above op lowering requests.
pt-xla-profiler: CompileTime too frequent: 23 counts during 12 steps
pt-xla-profiler: TransferFromDeviceTime too frequent: 12 counts during 12 steps

编译和执行分析¶

调试工具将分析您模型的每次编译和执行。一些示例输出如下所示

Compilation Analysis: ================================================================================
Compilation Analysis: Compilation Cause
Compilation Analysis:   mark_step in parallel loader at step end
Compilation Analysis: Graph Info:
Compilation Analysis:   Graph Hash: c74c3b91b855b2b123f833b0d5f86943
Compilation Analysis:   Number of Graph Inputs: 35
Compilation Analysis:   Number of Graph Outputs: 107
Compilation Analysis: Python Frame Triggered Execution:
Compilation Analysis:   mark_step (/workspaces/dk3/pytorch/xla/torch_xla/core/xla_model.py:1055)
Compilation Analysis:   next (/workspaces/dk3/pytorch/xla/torch_xla/distributed/parallel_loader.py:44)
Compilation Analysis:   __next__ (/workspaces/dk3/pytorch/xla/torch_xla/distributed/parallel_loader.py:32)
Compilation Analysis:   train_loop_fn (/workspaces/dk3/pytorch/xla/examples/train_decoder_only_base.py:48)
Compilation Analysis:   start_training (/workspaces/dk3/pytorch/xla/examples/train_decoder_only_base.py:65)
Compilation Analysis:   <module> (/workspaces/dk3/pytorch/xla/examples/train_decoder_only_base.py:73)
Compilation Analysis: --------------------------------------------------------------------------------
Compilation Analysis: ================================================================================

Post Compilation Analysis: ================================================================================
Post Compilation Analysis: Graph input size: 1.548000 GB
Post Compilation Analysis: Graph output size: 7.922460 GB
Post Compilation Analysis: Aliased Input size: 1.547871 GB
Post Compilation Analysis: Intermediate tensor size: 12.124478 GB
Post Compilation Analysis: Compiled program size: 0.028210 GB
Post Compilation Analysis: --------------------------------------------------------------------------------
Post Compilation Analysis: ================================================================================

Execution Analysis: ================================================================================
Execution Analysis: Execution Cause
Execution Analysis:   mark_step in parallel loader at step end
Execution Analysis: Graph Info:
Execution Analysis:   Graph Hash: c74c3b91b855b2b123f833b0d5f86943
Execution Analysis:   Number of Graph Inputs: 35
Execution Analysis:   Number of Graph Outputs: 107
Execution Analysis: Python Frame Triggered Execution:
Execution Analysis:   mark_step (/workspaces/dk3/pytorch/xla/torch_xla/core/xla_model.py:1055)
Execution Analysis:   next (/workspaces/dk3/pytorch/xla/torch_xla/distributed/parallel_loader.py:44)
Execution Analysis:   __next__ (/workspaces/dk3/pytorch/xla/torch_xla/distributed/parallel_loader.py:32)
Execution Analysis:   train_loop_fn (/workspaces/dk3/pytorch/xla/examples/train_decoder_only_base.py:48)
Execution Analysis:   start_training (/workspaces/dk3/pytorch/xla/examples/train_decoder_only_base.py:65)
Execution Analysis:   <module> (/workspaces/dk3/pytorch/xla/examples/train_decoder_only_base.py:73)
Execution Analysis: --------------------------------------------------------------------------------
Execution Analysis: ================================================================================

编译/执行的一些常见原因有：1. 用户手动调用 mark_step。2. 并行加载器为每个 x（可配置）批次调用 mark_step。3. 退出 profiler StepTrace 区域。

Dynamo 决定编译/执行图。5. 用户尝试在 mark_step 之前访问张量的值（通常是由于日志记录）。

由 1-4 引起的执行是预期的，我们希望通过降低访问张量值的频率或在访问之前手动添加 mark_step 来避免 5。

用户应该期望在最初几个步骤中看到此 Compilation Cause + Executation Cause 对。在模型稳定后，用户应该期望只看到 Execution Cause（您可以通过 PT_XLA_DEBUG_LEVEL=1 禁用执行分析）。为了有效地使用 PyTorch/XLA，我们期望每个步骤都运行相同的模型代码，并且每次图只发生一次编译。如果您一直看到 Compilation Cause，您应该尝试按照本节导出 IR/HLO，并比较每个步骤的图，了解差异的来源。

以下部分将解释如何获取和理解更详细的指标报告。

获取指标报告¶

在您的程序中放入以下行以生成报告

import torch_xla.debug.metrics as met

# For short report that only contains a few key metrics.
print(met.short_metrics_report())
# For full report that includes all metrics.
print(met.metrics_report())

理解指标报告¶

该报告包括以下内容：- 我们发出 XLA 编译的次数以及花费在发出编译上的时间。- 我们执行的次数以及花费在执行上的时间 - 我们创建/销毁的设备数据句柄的数量等。

此信息以样本的百分位数报告。一个例子是

Metric: CompileTime
  TotalSamples: 202
  Counter: 06m09s401ms746.001us
  ValueRate: 778ms572.062us / second
  Rate: 0.425201 / second
  Percentiles: 1%=001ms32.778us; 5%=001ms61.283us; 10%=001ms79.236us; 20%=001ms110.973us; 50%=001ms228.773us; 80%=001ms339.183us; 90%=001ms434.305us; 95%=002ms921.063us; 99%=21s102ms853.173us

我们还提供计数器，这些计数器是跟踪内部软件状态的命名整数变量。例如

Counter: CachedSyncTensors
  Value: 395

在此报告中，任何以 aten:: 开头的计数器都表示 XLA 设备和 CPU 之间的上下文切换，这可能是模型代码中的潜在性能优化区域。

计数器有助于了解哪些操作被路由回 PyTorch 的 CPU 引擎。它们使用其 C++ 命名空间完全限定

Counter: aten::nonzero
  Value: 33

如果您看到 aten:: 操作，除了 nonzero 和 _local_scalar_dense 之外，这通常意味着 PyTorch/XLA 中缺少降低。请随时在 GitHub issues 上打开功能请求。

清除指标报告¶

如果您想清除步骤/epoch 之间的指标，可以使用

import torch_xla.debug.metrics as met

met.clear_all()

PyTorch/XLA + Dynamo 调试工具¶

您可以通过设置 XLA_DYNAMO_DEBUG=1 来启用 PyTorch/XLA + Dynamo 调试工具。

性能分析¶

要深入分析您的工作负载以了解瓶颈，请查看以下资源

简单基准测试¶

请查看

examples/debug/train_resnet_benchmark.py，了解如何基准测试 PyTorch/XLA 模型。

已知的性能注意事项¶

PyTorch/XLA 在语义上与常规 PyTorch 类似，XLA 张量与 CPU 和 GPU 张量共享完整的张量接口。但是，XLA/硬件中的约束和延迟求值模型表明，某些模式可能会导致性能不佳。

如果您的模型显示性能不佳，请记住以下注意事项

XLA/TPU 在重新编译次数过多时会降低性能。

XLA 编译开销很大。每次遇到新形状时，PyTorch/XLA 都会自动重新编译图。通常模型应在几个步骤内稳定下来，您可以在训练的其余部分看到巨大的加速。

为了避免重新编译，不仅形状必须恒定，而且所有主机中 XLA 设备上的计算也应恒定。

可能的来源:
- 直接或间接使用 nonzero 会引入动态形状；例如，掩码索引 base[index]，其中 index 是掩码张量。
- 步骤之间迭代次数不同的循环可能导致不同的执行图，因此需要重新编译。
解决方案:
- 张量形状在迭代之间应相同，或者应使用少量形状变化。
- 尽可能将张量填充到固定大小。
某些操作没有到 XLA 的原生转换。

对于这些操作，PyTorch/XLA 会自动传输到 CPU 内存，在 CPU 上求值，并将结果传输回 XLA 设备。在训练步骤中执行过多的此类操作可能会导致明显的减速。

可能的来源:
- item() 操作显式地要求评估结果。除非必要，否则不要使用它。
解决方案:
- 对于大多数操作，我们可以将其降低到 XLA 以进行修复。查看指标报告部分，找出缺少的操作，并在 GitHub 上打开功能请求。
- 即使 PyTorch 张量已知为标量，也要避免使用 tensor.item()`。将其保留为张量，并在其上使用张量操作。
- 在适用时使用 torch.where 替换控制流。例如，在 clip_grad_norm 中使用的带有 item() 的控制流存在问题并影响性能，因此我们修补了 clip_grad_norm_，方法是调用 torch.where，这为我们带来了显着的性能提升。
```
...
else:
  device = parameters[0].device
  total_norm = torch.zeros([], device=device if parameters else None)
  for p in parameters:
    param_norm = p.grad.data.norm(norm_type) ** norm_type
    total_norm.add_(param_norm)
  total_norm = (total_norm ** (1. / norm_type))
clip_coef = torch.tensor(max_norm, device=device) / (total_norm + 1e-6)
for p in parameters:
  p.grad.data.mul_(torch.where(clip_coef < 1, clip_coef, torch.tensor(1., device=device)))
```
``torch_xla.distributed.data_parallel`` 中的迭代器可能会丢弃输入迭代器中的最后几个批次。

这是为了确保我们在所有 XLA 设备上执行相同数量的工作。

解决方案:
- 当数据集较小且步骤太少时，这可能会导致空操作 epoch。因此，在这些情况下最好使用小批量大小。

XLA 张量怪癖¶

XLA 张量内部是不透明的。 XLA 张量始终显示为连续且没有存储。网络不应尝试检查 XLA 张量的步幅。
XLA 张量应在保存之前移动到 CPU。 直接保存 XLA 张量会导致它们加载回保存它们的设备。如果设备在加载时不可用，则加载将失败。在保存之前将 XLA 张量移动到 CPU 可让您决定将加载的张量放在哪些设备上。如果您想在没有 XLA 设备的机器上加载张量，则这是必需的。但是，在保存之前将 XLA 张量移动到 CPU 时应注意，因为跨设备类型移动张量不会保留视图关系。相反，应在加载张量后根据需要重建视图。
使用 Python 的 copy.copy 复制 XLA 张量会返回深层副本，而不是浅层副本。 使用 XLA 张量的视图来获取其浅层副本。
处理共享权重。 模块可以通过将一个模块的参数设置为另一个模块来共享权重。模块权重的这种“绑定”应在模块移动到 XLA 设备之后完成。否则，将在 XLA 设备上创建共享张量的两个独立副本。

更多调试工具¶

我们不希望用户使用本节中的工具来调试他们的模型。但是，当您提交错误报告时，我们可能会要求您提供这些工具，因为它们提供了指标报告没有的额外信息。

print(torch_xla._XLAC._get_xla_tensors_text([res]))，其中 res 是结果张量，会打印出 IR。
print(torch_xla._XLAC._get_xla_tensors_hlo([res]))，其中 res 是结果张量，会打印出生成的 XLA HLO。

请注意，这些函数必须在 mark_step() 之前调用，否则张量将已被物化。

环境变量¶

还有许多环境变量控制 PyTorch/XLA 软件堆栈的行为。

设置此类变量会导致不同程度的性能下降，因此应仅在调试时启用它们。

XLA_IR_DEBUG：启用 Python 堆栈跟踪以捕获创建 IR 节点的代码位置，从而允许了解哪个 PyTorch 操作负责生成 IR。
XLA_HLO_DEBUG：当 XLA_IR_DEBUG 处于活动状态时，启用捕获的 Python 堆栈帧，以传播到 XLA HLO 元数据。
XLA_SAVE_TENSORS_FILE：将用于转储执行期间的 IR 图的文件路径。请注意，如果该选项保持启用状态且 PyTorch 程序运行时间过长，则该文件可能会变得非常大。图会附加到文件中，因此为了在每次运行中获得干净的状态，应显式删除该文件。
XLA_SAVE_TENSORS_FMT：存储在 XLA_SAVE_TENSORS_FILE 文件中的图的格式。可以是 text（默认）、dot（Graphviz 格式）或 hlo。
XLA_FLAGS=--xla_dump_to：如果设置为 =/tmp/dir_name，XLA 编译器将转储每次编译的未优化和优化 HLO。
XLA_METRICS_FILE：如果设置，则为本地文件的路径，内部指标将在每个步骤保存到该文件。指标将附加到文件中（如果已存在）。
XLA_SAVE_HLO_FILE：如果设置，则为本地文件的路径，如果发生编译/执行错误，则有问题的 HLO 图将保存到该文件。
XLA_SYNC_WAIT：强制 XLA 张量同步操作等待其完成，然后再移动到下一步。
XLA_USE_EAGER_DEBUG_MODE：强制 XLA 张量急切执行，这意味着逐个编译和执行 torch 操作。这对于绕过漫长的编译时间很有用，但总体步骤时间会慢得多，并且内存使用量会更高，因为所有编译器优化都将被跳过。
TF_CPP_LOG_THREAD_ID：如果设置为 1，则 TF 日志将显示线程 ID，从而帮助调试多线程进程。
TF_CPP_VMODULE：用于 TF VLOG 的环境变量，格式为 TF_CPP_VMODULE=name=value,...。请注意，对于 VLOG，您必须设置 TF_CPP_MIN_LOG_LEVEL=0。
TF_CPP_MIN_LOG_LEVEL：要打印的消息级别。TF_CPP_MIN_LOG_LEVEL=0 将打开 INFO 日志记录，TF_CPP_MIN_LOG_LEVEL=1 将打开 WARNING 日志记录，依此类推。我们的 PyTorch/XLA TF_VLOG 默认使用 tensorflow::INFO 级别，因此要查看 VLOG，请设置 TF_CPP_MIN_LOG_LEVEL=0。
XLA_DUMP_HLO_GRAPH：如果设置为 =1，则在发生编译或执行错误时，有问题的 HLO 图将作为 xla_util.cc 引发的运行时错误的一部分转储。

常用调试环境变量组合¶

以 IR 格式记录图执行

XLA_IR_DEBUG=1 XLA_HLO_DEBUG=1 XLA_SAVE_TENSORS_FMT="text" XLA_SAVE_TENSORS_FILE="/tmp/save1.ir"

以 HLO 格式记录图执行

XLA_IR_DEBUG=1 XLA_HLO_DEBUG=1 XLA_SAVE_TENSORS_FMT="hlo" XLA_SAVE_TENSORS_FILE="/tmp/save1.hlo"

显示运行时和图编译/执行的调试 VLOG

TF_CPP_MIN_LOG_LEVEL=0 TF_CPP_VMODULE="xla_graph_executor=5,pjrt_computation_client=3"

重现 PyTorch/XLA CI/CD 单元测试失败。¶

您可能会看到 PR 的一些测试失败，例如

要执行此测试，请从基本仓库目录运行以下命令

PYTORCH_TEST_WITH_SLOW=1 python ../test/test_torch.py -k test_put_xla_uint8

直接在命令行中运行此命令不起作用。您需要将环境变量 TORCH_TEST_DEVICES 设置为您的本地 pytorch/xla/test/pytorch_test_base.py。例如

TORCH_TEST_DEVICES=/path/to/pytorch/xla/test/pytorch_test_base.py PYTORCH_TEST_WITH_SLOW=1 python ../test/test_torch.py -k test_put_xla_uint8

应该可以工作。