GKE 上的 TPU 内存泄漏导致 OOM/“不可用：套接字已关闭”错误

Question

我正在使用抢占式 v2.8 Google Cloud TPU 来执行大规模超参数优化。我使用带有 tensorflow 2.3（Cloud TPU 的最新可用版本）的 GKE 创建了节点。不幸的是，在搜索过程中，我一直在 TPU 节点上遇到内存泄漏。这种内存泄漏似乎最终会导致“Unavailable: Socket Closed”错误（或有时是 OOM 错误），即使在重新部署代码后，TPU 也无法执行任何额外的训练或评估。当我在 CPU 或 GPU 上测试我的代码时，不会出现此问题。

此问题仅出现在 TPU 工作节点上，而不会出现在控制器 CPU 上。 （有一次，由于旧模型的堆积和计算图上不必要的操作，我在 CPU 上遇到了另一次内存泄漏。）tf.backend.clear_session() 和 del model 等方法解决了 CPU 的内存泄漏，但是它仍然存在于 TPU 上。这是 TPU 运行时内存使用情况的图表（最后内存的减少似乎发生在 TPU 断开连接后，因为 GKE 会自动删除它）：

最终，随着 TPU 上使用的内存增加，我收到以下错误：

raise_from tensorflow.python.framework.errors_impl.ResourceExhaustedError: 9 root error(s) found.

Error
2021-08-02T16:36:47.652282141ZHint: If you want to see a list of allocated tensors when OOM happens, add report_tensor_allocations_upon_oom to RunOptions for current allocation info.
Error
2021-08-02T16:36:47.652288611Z
Error
2021-08-02T16:36:47.652296423Z (4) Resource exhausted: {{function_node __inference_train_function_37854}} Attempting to reserve 3.27G at the bottom of memory. That was not possible. There are 3.48G free, 0B reserved, and 1.67G reservable.
Error
2021-08-02T16:36:47.652313550Z [[{{node cluster_train_function/_execute_4_0}}]]
2021-08-02T16:36:47.652921274Z0 successful operations.
Error
2021-08-02T16:36:47.654639274Z0 derived errors ignored.

有时，我会收到“不可用：套接字已关闭”错误或“无法销毁远程张量句柄”错误。

此错误通常仅在训练多个网络后发生。我尝试了其他帖子建议的多种方法来修复错误，例如将我的数据类型转换为 float32，不将我的数据集缓存到内存中，使用更小的 mini batch size 来减少内存消耗，以及在我的成本函数中使用“from_logits=True” .我什至尝试使用多处理来执行网络训练，以便在每次网络评估后清除内存，但由于某种原因，Cloud TPU 无法在我的代码或训练代码中执行任何 for 循环（这个问题我没有使用 GPU 或 CPU、云或其他方式。）较大的网络似乎比较小的网络更快地导致问题发生，这表明旧的、未使用的模型仍保留在 TPU 的内存中。有什么方法可以清除 TPU 上的内存或重置其状态以阻止这种内存泄漏？

这是我写的一个 MVE 来复制这个问题：

import os
import gc
import sys
import random
import numpy as np
import tensorflow as tf
from sklearn import metrics
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import InputLayer, Conv2D, Flatten, Dense
from tensorflow.keras.optimizers import Adam

h = 128
w = 128
channels = 1

mini_batch_size = 256
epochs = 15


using_tpu = True
if using_tpu:
    ## Get tpu name from arguments
    tpu_name = sys.argv[1]
    tpu_name = tpu_name.replace('--tpu=', '')

    ## Initialize TPU
    tpu = tf.distribute.cluster_resolver.TPUClusterResolver(tpu_name)  # TPU detection
    print('Running on TPU ', tpu.cluster_spec().as_dict()['worker'])

    tf.config.experimental_connect_to_cluster(tpu)
    tf.tpu.experimental.initialize_tpu_system(tpu)
    tpu_strategy = tf.distribute.TPUStrategy(tpu)


def create_network():
    strategy = tf.distribute.TPUStrategy(tpu)
    with strategy.scope():

        ## Create random data
        x_train = np.random.randn(1024, 128, 128, 1).astype('float32')  # astype necessary to help prevent Connect to Socket Error
        y_train = np.random.randn(1024, 50).astype('float32')

        x_test = np.random.randn(256, 128, 128, 1).astype('float32')
        y_test = np.random.randn(256, 50).astype('float32')

        model = Sequential()
        model.add(InputLayer((h, w, channels)))


        layers = 5
        ks = [np.random.choice([3, 5, 7]) for l in range(layers)]
        filters = [np.random.choice([64, 128, 256]) for l in range(layers)]

        for l in range(layers):
            model.add(
                    Conv2D(kernel_size=(ks[l], ks[l]), padding='same',
                           filters=filters[l], name='conv' + str(l), activation='relu'))

        model.add(Flatten())
        # Softmax output layer
        model.add(Dense(50))  # Don't need softmax activation because from_logits performs that operation automatically

        lr = 0.001
        opt = Adam(learning_rate=lr, decay=1e-6)

        model.compile(optimizer=opt, loss = tf.keras.losses.CategoricalCrossentropy(from_logits=True), metrics=['accuracy'])

    model.fit(x_train, y_train, epochs=epochs, batch_size=mini_batch_size, shuffle=True, verbose=1)

    ##### memory leak also occurs with dataset API:
    '''
    train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train)).batch(mini_batch_size,
                                                                                 drop_remainder=True)  

    model.fit(train_dataset, epochs=epochs, verbose=1, shuffle=shuffle,
              steps_per_epoch=len(x_train) // mini_batch_size)
    '''
    #######

    y_pred = model(x_test)

    ## Attempt to clear memory
    print(gc.collect())
    del model
    tf.keras.backend.clear_session()


while True:
    create_network()

非常感谢！如果我应该包含任何其他信息，请告诉我。

Answer 1

一些事情：

• 您的错误信息：

(4) Resource exhausted: {{function_node __inference_train_function_37854}} Attempting to reserve 3.27G at the bottom of memory. That was not possible. There are 3.48G free, 0B reserved, and 1.67G reservable.

表示 HBM OOM 而不是内存 OOM。基本上，TPU 在芯片上有自己的一组内存——在这种情况下，你已经耗尽了内存。如果它是 OOM（如 RAM OOM），那么您可能会看到 SocketClosed 错误，您也看到了。

• 话虽如此，您有什么选择？我建议您使用 tf.data 方法，但要进行一些修改：

def get_dataset(is_training: bool):
  def generate_data(_):
    return tf.random.normal([128, 128, 1], dtype=tf.bfloat16)

  dataset = tf.data.Dataset.range(1)
  dataset = dataset.repeat()
  dataset = dataset.map(generate_data, num_parallel_calls=tf.data.experimental.AUTOTUNE)
  dataset = dataset.repeat().batch(mini_batch_size, drop_remainder=is_training)

train_dataset = get_dataset(is_training=True)
eval_dataset = get_dataset(is_training=False)

在这个例子中，我们可以使用 bfloat16 来减少 HBM 上的内存占用，但您可能需要进一步将 minibatch 大小从 1024 减少到 512。或者，您可以从 v2-8 升级到 v3-8，它具有 2 倍的HBM。我不确定基于numpy 的方法是否会导致您看到的OOMs/SocketClosed 错误，但我认为这种方法不应该遇到这种情况。当然，您最终将使用真实数据，在这种情况下，您应该使用tf.data 以获得最佳性能。更多信息here。

• IIUC tf.backend.clear_session() 和 gc.collect() 只会清除主机 VM 上的内存，不会清除 TPU 服务器上的内存。
• PS：您还可以使用steps_per_execution 标志来进一步提高性能。请参阅here 了解更多信息。基本上，这可以防止执行在每一步中不断地从 CPU 切换到 TPU。如果您将其设置为等于一个 epoch 中的训练步数，这将为您提供最佳吞吐量。