故障排除¶

请注意，本节信息可能在未来的 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. 并行加载器 (Parallel loader) 每 x (可配置) 个批次调用 mark_step。3. 退出分析器 StepTrace 区域。

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

由 1-4 引起的执行是预期行为，我们应该避免 5，方法是减少访问张量值的频率或在访问之前手动添加 mark_step。

用户应该预期在前几个步骤中看到 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

如果您看到除了 nonzero 和 _local_scalar_dense 之外的 aten:: 操作，这通常意味着 PyTorch/XLA 中缺少相应的 Lowering。请随时在 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 会引入动态形状；例如，当 index 是掩码张量时的掩码索引 base[index]。
- 步骤之间迭代次数不同的循环可能导致执行图不同，从而需要重新编译。
解决方案:
- 张量形状在迭代之间应保持一致，或者应使用少量形状变体。
- 可能时，将张量填充到固定大小。
某些操作没有针对 XLA 的原生转换。

对于这些操作，PyTorch/XLA 会自动将数据传输到 CPU 内存，在 CPU 上执行，然后将结果传回 XLA 设备。在训练步骤中执行过多此类操作可能导致显著的性能下降。

可能的原因:
- item() 操作会明确要求对结果进行求值。除非必要，请勿使用。
解决方案:
- 对于大多数操作，我们可以将其 Lowering 到 XLA 来解决问题。请查阅指标报告部分，了解缺失的操作并在GitHub 上提交功能请求。
- 即使 PyTorch 张量已知为标量，也要避免使用 tensor.item()。将其保留为张量并对其进行张量操作。
- 在适用时使用 torch.where 替换控制流。例如，clip_grad_norm 中使用 item() 的控制流存在问题并影响性能，因此我们通过调用 torch.where 而不是 item() 来修复了 clip_grad_norm_，从而带来了显著的性能提升。
```
...
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 设备上执行相同的工作量。

解决方案:
- 当数据集很小且步骤太少时，这可能导致一个空操作 (no-op) 的 Epoch。因此，在这些情况下最好使用较小的批次大小。

XLA 张量特性¶

**XLA 张量的内部是不透明的。** XLA 张量始终显示为连续且没有存储。网络不应尝试检查 XLA 张量的步长 (strides)。
**在保存 XLA 张量之前应将其移动到 CPU。** 直接保存 XLA 张量会导致它们在保存时的设备上重新加载。如果在加载时设备不可用，则加载将失败。在保存之前将 XLA 张量移动到 CPU 可以让您决定将加载的张量放在哪个设备上。如果您想在没有 XLA 设备的机器上加载张量，这是必需的。然而，在保存之前将 XLA 张量移动到 CPU 时需要注意，因为跨设备类型移动张量不会保留视图关系。相反，应在加载张量后根据需要重新构建视图。
**使用 Python 的 copy.copy 复制 XLA 张量会返回深拷贝，而不是浅拷贝。** 使用 XLA 张量的视图来获取其浅拷贝。
**处理共享权重。** 模块可以通过将一个模块的 Parameter 设置为另一个模块的 Parameter 来共享权重。这种模块权重的“绑定”应该在将模块移动到 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：启用在创建 IR 节点时捕获 Python 堆栈跟踪，从而帮助理解是哪个 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 张量以 Eager 模式执行，这意味着逐个编译和执行 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

应该能够正常工作。