调试

注意，本部分的信息可能在未来的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

对于夜间

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运行，我们可以得出结论，TorchXLA已经正确安装。

性能调试

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

性能报告在诊断问题时非常有帮助。如果您有，请尝试将其包含在发送给我们的bug报告中。

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. 退出 分析器 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

如果你看到除了 aten:: 之外的其他 nonzero 和 _local_scalar_dense 操作，那通常意味着 PyTorch/XLA 中缺少了某种下移。请在 GitHub issues 上提交一个功能请求。

清除指标报告

如果你想清除步骤/周期之间的指标，你可以使用

import torch_xla.debug.metrics as met

met.clear_all()

PyTorch/XLA + Dynamo 调试工具

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

性能分析

为了深入剖析您的工作负载以了解瓶颈，请参阅以下资源：

• 官方教程

• Colab笔记本

• 带有性能分析的MNIST训练脚本

• 用于捕获性能概况的实用脚本

简单基准测试

查看 ``examples/train_resnet_benchmark.py` <https://github.com/pytorch/xla/blob/master/examples/train_resnet_benchmark.py>`_，了解如何基准一个 PyTorch/XLA 模型。

已知性能陷阱

PyTorch/XLA的行为在语义上类似于常规的PyTorch，XLA张量与CPU和GPU张量共享完整的张量接口。然而，XLA/硬件的约束以及懒惰评估模型的建议表明某些模式可能会导致性能不佳。

如果你的模型表现不佳，请记住以下几点注意事项：

1. XLA/TPU 在进行过多次重新编译后，性能会下降。

   XLA编译成本高昂。PyTorch/XLA会自动重新编译图，每当遇到新的形状时。 通常模型会在几步内稳定下来，你可以看到训练剩余部分的巨大加速。

   为了避免重新编译，不仅形状必须是常量，而且在所有主机上的XLA设备之间的计算也必须是常量。

   可能的来源:

   • 直接或间接使用 nonzero 引入动态形状；例如，掩码索引 base[index]，其中 index 是一个掩码张量。

   • 具有不同迭代次数的循环在步骤之间可能导致不同的执行图，因此需要重新编译。

   解决方案:

   • 在迭代之间，张量形状应该相同，或者使用少量的形状变化。

   • 在可能的情况下，填充张量以固定大小。

2. 某些操作没有XLA的原生翻译。

   对于这些操作，PyTorch/XLA 自动将数据转移到 CPU 内存中，在 CPU 上进行评估，并将结果返回到 XLA 设备上。在训练过程中执行过多此类操作会导致显著的性能下降。

   可能的来源:

   • The item() operation explicitly asks to evaluate the result. Don’t use it unless it’s necessary.

   解决方案:

   • 对于大多数操作，我们可以将它们降低到XLA来修复它。请查看指标报告部分以找出缺失的操作，并在GitHub上打开一个功能请求。

   • 即使一个PyTorch张量被称为标量，也避免使用tensor.item()。将其保持为张量，并在其中使用张量操作。

   • 使用 torch.where 替代适用的控制流。 例如，clip_grad*norm* 中使用的 item() 控制流存在问题并影响性能，因此我们通过调用 torch.where 来修补 clip_grad_norm_，这给我们带来了显著的性能提升。 .. code-block:: python

     … 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)))

3. ``torch_xla.distributed.data_parallel``` 中的迭代器可能会丢弃输入迭代器中的最后几批数据。

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

   解决方案:

   • 当数据集较小，且步数过少时，这可能会导致无操作的 epoch。因此，在这种情况下，最好使用较小的批量大小。

XLA张量怪癖

1. XLA张量内部是透明的。 XLA张量总是看起来是 连续且没有存储。网络不应尝试检查XLA张量的步长。

2. 在保存XLA张量之前，应将其移动到CPU。 直接保存XLA张量会导致它们被加载回其保存时所在的设备上。如果在加载时某个设备不可用，则加载将失败。在保存XLA张量之前将其移动到CPU可以让您决定加载张量时将其放置在哪些设备上。这在您希望在没有XLA设备的机器上加载张量时是必要的。然而，在保存XLA张量之前将其移动到CPU需要小心，因为不同设备类型之间的张量移动不会保留视图关系。因此，在加载张量后，视图关系需要按需重建。

3. 使用Python的copy.copy复制XLA张量会返回深拷贝，而不是浅拷贝。 使用XLA张量的视图以获取其浅拷贝。

4. 处理共享权重。 可以通过设置一个模块的参数为另一个模块来实现模块间的权重共享。这种“绑定”共享权重的操作应在将模块移动到XLA设备之后进行，否则XLA设备上将会创建两个独立的共享张量副本。

更多调试工具

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

• print(torch_xla._XLAC._get_xla_tensors_text([res])) where res 是结果张量打印出的IR。

• print(torch_xla._XLAC._get_xla_tensors_hlo([res])) where 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文件中的图形格式。可以是text（默认）、dot（_Graphviz_ 格式）或hlo。

• XLA_FLAGS=--xla_dump_to: 如果设置为=/tmp/dir_name，XLA编译器将在每次编译时.dump未优化和优化后的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 VLOGs，并且具有TF_CPP_VMODULE=name=value,...的形式。请注意，对于VLOGs，您必须设置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级别，所以要查看VLOGs请设置TF_CPP_MIN_LOG_LEVEL=0.

• XLA_DUMP_HLO_GRAPH: 如果设置为=1，在编译或执行错误的情况下，违规的HLO图将在由xla_util.cc引发的运行时错误中被dump出来。

常见调试环境变量组合

• 在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）中看到一些测试失败，例如：

To execute this test, run the following from the base repo dir:
    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 应该可以使用。