此页面由 Cloud Translation API 翻译。

TensorFlow - TPU 问题排查

本指南以及常见问题解答为在 Cloud TPU 上训练 TensorFlow 模型的用户提供了问题排查帮助。如果您要排查 Pytorch 或 JAX 训练问题，则可以参考这些框架的问题排查文档：

PyTorch - TPU 问题排查
JAX - TPU 问题排查

重要提示：Cloud TPU 支持两种 TPU 架构：TPU 虚拟机或 TPU 节点。系统架构中介绍了这两种虚拟机架构。本文档中的问题排查步骤适用于这两种配置，但还有一个单独的 TPU 虚拟机部分，专门用于排查仅在 TPU 虚拟机配置时发生的问题。

如需获取有关如何使用 Cloud TPU 的更多常规指南，请参阅：

概览

Cloud TPU 的常见问题可分为以下类别：

连接到 TPU 的问题
调试常见错误
减少内存使用量
提高训练速度
针对模型准确率降低进行调试

无法连接到 TPU 服务器

本部分介绍如何排查 TensorFlow 在连接到 TPU 时停止响应或输出错误的情况。大型模型的 TPU 图编译步骤可能需要很长时间，因此请先让脚本执行至少 5 分钟，然后再判断它是否已停止响应。

首先，验证问题在于服务器本身，还是在于 TensorFlow 训练流水线。为此，请使用您的 TPU 服务器网址运行 MNIST 教程，并验证它是否可以正常运行。如果使用 MNIST 教程时仍然存在连接问题，则表明问题在于 TPU 服务器。在此示例中：

运行以下命令以列出可用的 TPU：将 zone 和 project-id 替换为您的可用区和项目 ID。

(vm)$ gcloud compute tpus list --zone zone --project project-id

这会显示如下输出：

NAME       ZONE           ACCELERATOR_TYPE  NETWORK_ENDPOINT   NETWORK  RANGE          STATUS
demo-tpu   us-central1-b  v2-8              10.240.1.2:8470    default  10.240.1.0  READY

验证您将正确的值传递给 --tpu（上例中的 demo-tpu），并且此 TPU 被列为 READY。
如果 TPU 未列为 READY，或者您仍然无法建立连接，请使用以下命令手动重启服务器：
```
(vm)$ gcloud compute tpus stop $TPU_SERVER_NAME && gcloud compute tpus start $TPU_SERVER_NAME
```
在上面的示例中，$TPU_SERVER_NAME 为 demo-tpu。这可能需要几分钟才能完成。
重新运行上面的 ... tpus list 命令并等待 TPU 处于 READY 状态。这可能需要几分钟。
尝试再次运行 MNIST 教程。
如果您在运行 MNIST 教程时仍遇到问题，请使用获取支持中介绍的某种机制寻求帮助。

如果 MNIST 示例可正确运行但您的模型仍然停止响应，则问题可能出在您的训练流水线上。要进行调试，请先将代码中的 TPUStrategy 替换为默认策略。使用默认策略时，无论您使用 strategy.scope() 还是 strategy.run()，模型都在 CPU（或 GPU，如果存在）上运行，而不是在 TPU 上运行。如果模型在 CPU 上运行，而不是在 TPU 上运行，则必须存在 TPU 特定问题。如果它仍未运行，最佳做法是在 CPU 上调试问题。

训练期间 `ssh` 连接丢失

在长时间运行的训练期间，您与 Cloud TPU 的 ssh 连接可能会超时（尤其是在您使用 Cloud Shell 时）。此时，TPU 控制台没有输出的内容，就像 TPU 已停止训练一样。为避免这种情况，请使用终端多路复用器或会话管理工具（如 tmux 或 screen）运行训练会话。这将使 ssh 连接保持活跃状态，无论训练的长度如何。

调试常见错误

无法创建 TPU

创建 Cloud TPU 时，您可能会看到以下错误：

googleapiclient.errors.HttpError: < HttpError 403 when requesting https://content-tpu.googleapis.com/v1/projects/{PROJECT}/locations/{ZONE}/nodes/{TPU_NAME}?alt=json returned "Request had insufficient authentication scopes."

这是一个权限问题，可通过运行以下命令来解决：

gcloud auth login --update-adc

此命令会更新您的应用默认凭据 (ADC)，应该可以解决此问题。如需了解详情，请参阅 gcloud auth login。

无法使用本地文件系统

错误消息

InvalidArgumentError: Unimplemented: File system scheme '[local]' not implemented

受影响的框架和配置

使用 TPU 节点架构进行 TensorFlow 训练时，可能会出现此消息。

详情

所有输入文件和模型目录都必须使用 Cloud Storage 存储桶路径 (gs://bucket-name/...)，并且此存储桶必须可通过 TPU 服务器访问。请注意，所有数据处理和模型检查点都在 TPU 服务器上执行，而不是在本地机器上执行。如需了解如何正确配置 Cloud Storage 以便与 TPU 搭配使用，请参阅连接到 Cloud Storage 存储分区指南。

不支持的数据类型

错误消息

TypeError: DataType is not a supported TPU infeed type.

受影响的框架和配置

使用 TPU 节点架构进行 TensorFlow 训练时，可能会出现此消息。

详情

目前，TPU 仅支持 tf.float32、tf.int32、tf.bfloat16 和tf.bool 数据类型。其他常见数据类型（例如 tf.uint8、tf.string 和 tf.int64）必须在数据预处理期间（即在 tf.data.Dataset 流水线）转换为某种受支持的数据类型。

请参阅 MNIST 训练中使用的 decode_image 函数中的转换示例。

不支持动态形状

错误消息

ValueError: shape [Shape] must have a fixed size for dimension
d that is known at graph construction time.

受影响的框架和配置

此消息仅在使用 TensorFlow 进行 XLA 编译期间出现。

详情

如需在 TPU 上执行模型，TensorFlow 使用 XLA 编译器编译该模型。虽然此编译步骤可显著提高训练速度和改进内存用量，但图中所有张量的形状（维度大小）必须在图编译时已知。如果有任何形状在编译时无法确定，则 TPU 编译将失败，并显示类似上面所示的错误。

返回动态形状的一个常用操作是 dataset.batch(batch_size)，因为数据流中剩余的样本数可能小于批量大小。因此，在 TPU 上进行训练时，请为 dataset.batch 设置 drop remainder=True。这可能会丢弃文件中的最后几个样本，以确保每个批量都具有静态形状 (batch_size)。例如：

dataset = tf.data.Dataset.range(8)
dataset = dataset.batch(3, drop_remainder=True)

不可用的 TensorFlow 操作

错误消息

NotFoundError: No registered 'OpName' OpKernel for XLA_TPU_JIT
devices compatible with node

受影响的框架和配置

使用 TensorFlow 进行训练时，可能会出现此消息。

详情

模型使用目前无法在 TPU 上执行的 TensorFlow 操作。

如需了解 TPU 上可用操作的列表，以及未来支持计划和解决方法建议，请参阅可用的 TensorFlow 操作指南。

内存不足错误消息

错误消息

ResourceExhaustedError: Ran out of memory in memory space hbm; used:
YYY; limit: 7.48G.

受影响的框架和配置

使用 TensorFlow、PyTorch 或 JAX 进行训练时，可能会出现此消息。

详情

每个 Cloud TPU 都由 8 个 TPU 核心组成，v2 TPU 具有 8GB，v3 TPU 具有 16GB 的 RAM（或 HBM，即高带宽内存）。这些内存用于存储权重（可变）张量，以及梯度计算所需的中间结果张量。如果模型太大而无法放入 TPU RAM 中，初始化将失败并显示上述错误消息。如需更多帮助，请参阅减少内存使用量。

减少内存使用的提示：

检查是否有过度张量填充
使用 bfloat16 格式
如果输入大小或模型太大，您可以使用 TensorFlow 的实验模型并行性来解决模型大小问题。

停止执行时出现问题

如果 TensorFlow 在 TPU 执行期间遇到错误，则脚本有时会停止响应，而不是退出 shell。如果发生这种情况，请按键盘上的 CTRL+\ 以触发 SIGQUIT，这会使 Python 立即退出。

同样，在 TPU 执行过程中按 CTRL+C 并不会立即关闭 TensorFlow，而是一直等到当前迭代循环结束才彻底退出。

如果您在以这种方式退出后重新连接到 TPU 时遇到任何新的错误，请使用以下命令手动重置 TPU 服务器：

gcloud compute tpus stop tpu-name --zone=zone
gcloud compute tpus start tpu-name --zone=zone

其中，tpu-name 取自 gcloud compute tpus list 命令显示的第一列，zone 是第二列中显示的可用区。

过度张量填充

内存问题的可能原因

TPU 内存中的张量经过了填充，也就是说，TPU 将存储在内存中的张量的大小向上舍入，以便更高效地执行计算。此填充操作在硬件级别以透明方式进行，不会影响结果。但是，在某些情况下，填充可能会导致内存使用量和执行时间显著增加。

如何减少内存使用量

TPU 软件会尝试在内存中存放张量，以便最大限度地提高计算效率并减少填充。但是，此内存存放进程很复杂，为了获得最佳结果，该模型应遵循以下经验法则。为了最大限度地减少内存开销并提高计算效率，必须满足以下条件之一：

总批量大小应为 64 的倍数（每个 TPU 核心 8 个），并且特征维度应为 128 的倍数。

或者
总批量大小应为 1024 的倍数（每个 TPU 核心 128 个），并且特征维度应为 8 的倍数。

使用 1024 作为批量大小以及 128 的倍数作为特征维度可以获得最佳效率，但这并非适用于所有模型。为清楚起见，“特征维度”是指全连接层的隐藏大小或卷积中的输出通道的数量。并非所有层都符合此规则，尤其是网络的第一层和最后一层。这没什么问题，而且大多数模型预期都需要一定量的填充。

减少内存使用量

如果您在 TPU 上执行模型时遇到内存不足错误，则必须采取相应措施来减少该模型的内存使用量。

减少内存用量最有效方法是：

减少过度张量填充
减小批次大小

批量大小或模型过大

内存问题的可能原因

在 CPU、GPU 或 TPU 上训练神经网络时，内存用于两个方面：

内存用量与模型中的权重数量成正比。
存储来自于前向传导、计算反向传导所需的中间激活。内存使用量与批量大小、层大小和层数成正比。

综上所述，模型所需的内存在很大程度上取决于批量大小。

模型所需的内存取决于网络中的层数。

TPU 运行时会尝试优化运算符，以便使模型能够存放在内存中（称为再具体化，类似于梯度检查点），但有时无法做到这一点。

如何减少内存使用量

缓慢减小批量大小，直到它能够存放在内存中为止，同时确保总批量大小为 64 的倍数（每个核心的批量大小应为 8 的倍数）。请注意，较大的批量大小在 TPU 上更为有效。通常，总批量大小为 1024（每个核心 128 个）是不错的起点。

如果即使批量较小（例如 64），模型也无法在 TPU 上运行，请尝试减少层数或层大小。

提高训练速度

如果您的模型能够在 TPU 上成功运行，但训练速度低于预期，可以参考本部分概述的几种有助于提高训练速度的方法。如需了解如何提高训练性能的其他建议，请参阅性能指南。

每个训练循环的每次执行步数太少

性能问题描述

将参数 steps_per_execution 传递给 Model.compile 可控制主机回调之间执行的训练步数。每个主机回调都需要在 TPU 服务器的主机 CPU 和 TPU 设备之间进行大量通信，因此如果 steps_per_execution 太小，则可能会减慢训练速度。

如何确定您的模型是否受到影响

如果 TPU 配置文件显示 TPU 设备步骤之间的频繁主机 CPU 回调，则训练可以从更大的 steps_per_execution 值中受益。

如何缓解

将 steps_per_execution 设置为更大的值。请注意，steps_per_execution 可以设置为较大的值，但请记住，记录消息和保存检查点只能在运行指定步数后才发生。

输入处理瓶颈

性能问题描述

当 TPU 在特定数据块上进行训练时，输入处理函数会在 CPU 上准备下一个数据块。如果输入函数花费的时间长于模型函数，则当输入函数检索数据时，TPU 保持空闲状态。

如何确定您的模型是否受到影响

按照 Cloud TPU 工具：输入流水线分析器中的说明在 TensorBoard 中查看流水线分析：

输入流水线分析页面会显示一份清晰的摘要，其中显示了您的模型是否因输入处理而遭遇瓶颈。此页面还显示每个操作的执行时间，便于您找出有问题的操作。

如何缓解

使用 Dataset 加载数据时，有几种可能的缓解措施：

将您的数据以一组 tf.train.Example 结构的形式存储在 TFRecord 文件中，并使用 TFRecordDataset 加载它们。如需查看示例，请参阅 Dataset API 教程或 ResNet 教程。
使用 dataset.cache() 和/或 dataset.prefetch() 来缓冲输入数据。这可以防止因文件访问中偶尔出现的速度缓慢问题而导致产生瓶颈。
指定 dataset.map() 函数的 num_parallel_calls 参数以启用多线程 map() 操作。对 num_parallel_calls 值的简单启发法是使用可用 CPU 核心的数量。
离线执行数据预处理，而不是在每次训练的每个周期执行预处理。此方法的优点是一劳永逸，缺点是离线执行数据预处理的费用高昂。

所有输入处理都由 TPU 服务器上的 CPU 执行，而不是本地机器，因此本地机器的速度不是一个影响因素。

单步用时缓慢，MXU 利用率低

性能问题描述

Cloud TPU 能够以极快的速度执行矩阵乘法和卷积。大多数其他 TensorFlow 操作也可在 TPU 上高效地实现，但这些并不是 TPU 相对于其他硬件的主要优势。因此，模型应该以矩阵乘法或卷积为主，才能使 TPU 发挥最大作用。

如何确定您的模型是否受到影响

在这种情况下，您将看到的症状是单步用时缓慢，以及分析性能时显示的 MXU 利用率低。

如何缓解

尝试减少不是矩阵乘法的操作数量。减少矩阵乘法的数量后，请重新进行基准测试以了解 TPU 上的性能是否可接受。

过度张量填充

性能问题描述

TPU 会在内存中填充张量，以便高效地使用其计算单元。填充可能会增加内存和内存带宽的使用量。如需理解填充和解决张量填充问题方面的帮助，请参阅张量填充。

吞吐量缓慢和内存用量低

性能问题描述

一般来讲，使用较大的批量大小可提高 TPU 上的训练速度（以样本数/秒表示）。

如何确定您的模型是否受到影响

任何模型的批量大小应至少为 64（每个 TPU 核心 8 个），因为 TPU 总是会将张量填充到此大小。在 TPU 上训练时，理想的批量大小为 1024（每个 TPU 核心 128 字节），因为这样可以消除内存传输和填充导致的低效问题。

如何缓解

最佳做法是使用可存放在内存中的最大批量大小，并且为 64 的倍数。实现此目的的最简单方法是从 1024 开始，如果该值会导致出现内存不足的错误，则尝试减小批量大小，直到模型成功运行为止。更改模型的批量大小可能需要调整其他超参数以实现同样的模型准确率（例如学习速率），但必须根据具体情况对其进行评估，不能一概而论。

层大小太小

性能问题描述

即使模型以矩阵乘法或卷积为主，如果输入张量很小，TPU 也可能无法全速运行。与其他硬件相比，当批量大小和层大小都很大时（例如维度 >= 512），TPU 运行的效率最高。

如何确定您的模型是否受到影响

一般来讲，小于 128 的层大小在 TPU 上的效率低下，因为 128 是 TPU 矩阵乘法单元的原生维度。对于全连接层，建议采用 512 作为隐藏大小下限，以实现高效率。注意，卷积层通常无需与全连接层一样大，即可实现相同的效率水平。

如何缓解

如果模型中较小层的主要动机是训练速度，请在 TPU 上使用较大的层重新对模型进行基准测试。例如，将层的输出大小从 256 增加到 512 可能只会将训练时间增加 20%，即使模型执行的计算量是原来的 2 倍。

操作级模型分析

为了识别性能瓶颈，测量操作级执行时间和内存使用情况通常很有用。如需了解如何执行此操作，
请参阅 Cloud TPU 工具：Trace Viewer 指南。

针对模型准确率降低进行调试

Cloud TPU 生态系统的目标之一是：目前在 CPU 或 GPU 上进行训练的任何模型在 TPU 上训练时都能达到非常接近的准确率，并且可能只需对批量大小和学习速率等超参数进行细微调整。但是，用户有时会发现在 TPU 上训练模型时准确率有所降低。由于神经网络训练的随机性，针对这些问题进行调试可能会让人很头痛。本部分提供一些指导，帮助您确定在将模型移植到 TPU 后模型准确率下降的根本原因。

了解数据分片（并行数据）

TensorFlow 的主要目标之一是每个操作应该产生几乎相同的结果，无论它是在 CPU、GPU 还是 TPU 上执行。但有一些例外情况，例如随机操作。一般来讲，如果您发现 TPU 与 CPU 上的非随机操作输出存在任何明显差异，请将其作为 bug 上报。

但是，对于整个训练流水线，在 CPU/GPU 上与在 TPU 上进行训练存在一项重要的差异。在 TPU 上进行训练时，TensorFlow 会执行数据分片。每个 Cloud TPU 包含 8 个 TPU 核心，这些核心作为独立的处理单元运行。对于训练中的每个步骤，每个 TPU 核心都会接收一批数据，计算权重梯度，与其他 TPU 核心交换梯度，然后计算权重更新。默认情况下，系统会在所有核心上计算损失平均值，但也可以通过更改 CrossShardOptimizer 的参数来计算总损失。

如果计算模型的总损失可以通过独立的每样本损失的平均值（或总和）来计算，则此过程在数学上等同于针对单个大批次进行训练。

最常见的非独立按样本操作是批次归一化，它单独针对每个核心批次运行。例如，如果总批次大小为 128，则每个核心批次大小为 16，并且 8 个核心中的每个核心都对自己的 16 个样本执行批次归一化。我们已发现在某些情况下对小批次（例如小于 32 个）进行批次归一化会导致准确率降低。理想情况下，总批次大小应很大（例如 256 到 1024 个）。如果这种批次大小太大而导致模型无法存放到内存中，则必须根据具体情况评估分片的影响。

确定性训练

针对模型准确率差异进行调试很困难，其中一个原因是在不同的框架（TensorFlow、PyTorch、JAX）之间，训练软件在每次训练模型时都使用不同的权重初始化和数据重排。将训练过程修改为确定性的，可保证多次运行会产生几乎相同的模型。本部分演示如何确定性地运行 MNIST 教程：

通过在 CPU 上运行一步来生成初始检查点文件。这一步用于实现确定性权重初始化。另外，请确保对模型中的任何随机函数使用固定的随机种子。

# Run training for 1 step to create an initial checkpoint.
python mnist_tpu.py \
  --use_tpu=False \
  --data_dir=${STORAGE_BUCKET}/data/ \
  --model_dir=${STORAGE_BUCKET}/init_output \
  --random_seed=12345 \
  --iterations=1
  --train_steps=1

修改输入函数中的任何数据重排函数以使用随机种子。 MNIST 教程中已完成了这一步。这适用于输入数据处理操作，因为它们始终在 CPU 上运行。模型函数中的随机操作在 TPU 与 CPU 之间可能不具有确定性。每当您调用随机操作时，请传递固定的种子，以确保两次运行之间得到相同的结果。例如：

# In the flag definitions
tf.flags.DEFINE_integer("batch_size", None, "Random seed for training")

# In the input_fn
if FLAGS.random_seed is not None:
dataset = dataset.shuffle(seed=FLAGS.random_seed)

在 CPU 上运行同一个模型两次，以验证训练具有确定性。请注意，训练必须运行合理的步数（例如 1000 步），但不需要运行到收敛。

由于 CPU 训练与单核 TPU 训练相当，因此请使用适合单个 TPU 核心的批量大小（通常是用完整批量大小除以 8）。TensorFlow 不保证各次运行之间的逐位确定性，但损失应该非常接近：

复制初始权重

gsutil mkdir ${STORAGE_BUCKET}/cpu_output_1
gsutil cp -f ${STORAGE_BUCKET}/init_output/* ${STORAGE_BUCKET}/cpu_output_1
gsutil mkdir ${STORAGE_BUCKET}/cpu_output_2
gsutil cp -f ${STORAGE_BUCKET}/init_output/* ${STORAGE_BUCKET}/cpu_output_2

第 1 回合

python mnist_tpu.py \
  --use_tpu=False \
  --data_dir=${STORAGE_BUCKET}/data/ \
  --model_dir=${STORAGE_BUCKET}/cpu_output_1 \
  --batch_size=128 \
  --random_seed=12345 \
  --train_steps=2000 \
  --eval_steps=10

输出 1

accuracy = 0.9910644, global_step = 1000, loss = 0.025323588

第 2 回合

python mnist_tpu.py \
  --use_tpu=False \
  --data_dir=${STORAGE_BUCKET}/data/ \
  --model_dir=${STORAGE_BUCKET}/cpu_output_1 \
  --batch_size=128 \
  --random_seed=12345 \
  --train_steps=2000 \
  --eval_steps=10

输出 2

accuracy = 0.9910644, global_step = 1000, loss = 0.025323414

单核 TPU 训练

可以确定性地运行 MNIST 教程以后，下一步是在 TPU 上重现 CPU 训练的结果，使用单个 TPU 核心来确定问题是在于数据分片还是在于 TPU 执行引擎本身。

下面说明了如何针对 MNIST 教程执行单核训练和评估：

使用与 CPU 相同的权重初始化

gsutil cp -f ${STORAGE_BUCKET}/init_output/* ${STORAGE_BUCKET}/tpu_output

运行 1000 步的训练

python mnist.py \
    --use_tpu=True \
    --master=$GRPC_SERVER \
    --train_file=${STORAGE_BUCKET}/data/train.tfrecords \
    --model_dir=${STORAGE_BUCKET}/tpu_output \
    --random_seed=12345 \
    --num_shards=1 \
    --batch_size=128 \
    --train_steps=1000 \
    --eval_steps=10

输出

  accuracy = 0.9910644, global_step = 1000, loss = 0.02514153

损失与 CPU 训练的模型并不完全一致，但应该很接近。如果您的模型的损失不接近，则可能表明您的 TPU 执行引擎中出现了 bug。在提交 bug 报告之前，请核实以下事项：

您为 TPUConfig 传递的是 num_shards=1。
您的模型函数中没有任何随机操作，并且输入函数中的所有随机操作都已使用了正确的种子。
您为 CPU 和 TPU 训练使用的是相同的初始检查点文件。

调试多核 TPU 训练

如果您的模型在 CPU 与单核 TPU 上实现了相同的损失，那么问题可能属于以下某一种：

(a) 降级是由于采用不同初始化训练神经模型时的自然随机方差所导致。

(b) 降级是由于与 TPU 上的数据分片相关的问题所导致。

如需确定问题是否属于 (a)，请使用相同的权重初始化在 CPU/GPU 和多核 TPU 上重新训练完整模型（如上所述）。

如果您确信准确率的降低具备统计显著性，那么最可能与数据分片相关的问题有：

如果您的模型使用批量归一化，则总批量大小小于 256（即每个核心少于 32 个）可能会降低准确率。
批量损失函数受分片的影响。此类损失函数通常非常专业化。例如，2017 年 Karras 等人在训练生成对抗网络时使用了批量判别器。

TPU 虚拟机问题排查

以下问题和解决方案仅适用于 TPU 虚拟机配置。

gcloud 设置问题排查

题目: gcloud components update 会显示以下错误消息

ERROR: (gcloud.components.update)
You cannot perform this action because the Cloud SDK component manager is
disabled for this installation.

解决方案: 如需将 gcloud 与 TPU 虚拟机搭配使用，您需要使用不通过软件包管理器管理的 gcloud 安装。如需从源代码安装 gcloud，请按以下步骤操作

  sudo apt-get remove google-cloud-sdk
  curl -O https://dl.google.com/dl/cloudsdk/channels/rapid/downloads/google-cloud-sdk-311.0.0-linux-x86_64.tar.gz
  tar -xzf google-cloud-sdk-311.0.0-linux-x86_64.tar.gz
  ./google-cloud-sdk/install.sh
  source ~/.bashrc

问题

gcloud compute tpus tpu-vm ssh ${TPU_NAME} --zone ${ZONE} 命令会显示以下错误消息：

Waiting for SSH key to propagate.
ssh: connect to host 34.91.136.59 port 22: Connection timed out
ssh: connect to host 34.91.136.59 port 22: Connection timed out
ssh: connect to host 34.91.136.59 port 22: Connection timed out
ERROR: (gcloud.compute.tpus.tpu-vm.ssh) Could not SSH into the instance.  It is possible that your SSH key has not propagated to the instance yet. Try running this command again.  If you still cannot connect, verify that the firewall and instance are set to accept ssh traffic.

解决方案

SSH 密钥传播可能存在问题。您可以尝试将自动生成的密钥移动到备份位置，以强制 gcloud 重新创建它们：

mv ~/.ssh/google_compute_engine ~/.ssh/old-google_compute_engine
mv ~/.ssh/google_compute_engine.pub ~/.ssh/old-google_compute_engine.pub

调试日志

受支持的 Cloud TPU 框架（JAX、PyTorch 和 TensorFlow）通过每个 TPU 虚拟机上存在的名为 libtpu 的共享库访问 TPU。此库包含用于编译 TPU 程序的 XLA 编译器、用于运行已编译程序的 TPU 运行时，以及该运行时用于对 TPU 的低级层访问的 TPU 驱动程序。

可用于调试的 libtpu 库记录信息。默认情况下，这些日志写入每个 Cloud TPU 虚拟机上的 /tmp/tpu_logs。在开始训练以修改日志记录行为之前，您可以设置以下环境变量：

TPU_LOG_DIR：写入日志的目录: 目录位置默认为 /tmp/tpu_logs。如果该目录尚不存在，则会创建该目录，但不会创建父目录。如果发现或创建指定的目录时出错，则系统会将消息输出到 stderr，但不会停止程序并停用日志记录。将目录名称设置为“已停用”，以完全停用对磁盘的日志记录。
TPU_MIN_LOG_LEVEL：将记录到磁盘的最低严重性: 选项包括 0 (INFO)、1 (WARNING)、2 (ERROR) 和 3 (FATAL)。默认值为 0
TPU_STDERR_LOG_LEVEL：除了磁盘之外，还将记录到 stderr 的最低严重级别（如存在）: 选项与 TPU_MIN_LOG_LEVEL 相同。默认值为 3
TPU_MAX_LOG_SIZE_MB：每个日志文件的大小上限（以 MB 为单位）: 当上一个日志文件达到大致大小时，系统会自动启动新的日志文件。默认值为 1024.

TensorFlow - TPU 问题排查

概览

无法连接到 TPU 服务器

训练期间 ssh 连接丢失

调试常见错误

无法创建 TPU

无法使用本地文件系统

不支持的数据类型

不支持动态形状

不可用的 TensorFlow 操作

内存不足错误消息

停止执行时出现问题

过度张量填充

减少内存使用量

批量大小或模型过大

提高训练速度

每个训练循环的每次执行步数太少

输入处理瓶颈

单步用时缓慢，MXU 利用率低

过度张量填充

吞吐量缓慢和内存用量低

层大小太小

操作级模型分析

针对模型准确率降低进行调试

了解数据分片（并行数据）

确定性训练

复制初始权重

第 1 回合

输出 1

第 2 回合

输出 2

单核 TPU 训练

使用与 CPU 相同的权重初始化

运行 1000 步的训练

输出

调试多核 TPU 训练

TPU 虚拟机问题排查

gcloud 设置问题排查

调试日志

训练期间 `ssh` 连接丢失