如何在 TPU 上将模型并行性与自定义 Tensorflow 2.0 模型一起使用？答案

【问题标题】：How to use Model Parallelism with a custom Tensorflow 2.0 model on TPUs?如何在 TPU 上将模型并行性与自定义 Tensorflow 2.0 模型一起使用？
【发布时间】：2021-10-19 12:22:11
【问题描述】：

为了复制Multimodal Few-Shot Learning with Frozen Language Models，我尝试在 TPUv3-32 上训练一个 ~7B 参数子类化 TF2 模型。在 7B 参数中，大约有 6B 参数被冻结。

我想使用模型和数据并行性来尽可能高效地训练它。据我所知，MeshTensorflow只能用于TF1编写的模型。

我尝试使用 TPUStrategy 中的 experimental_device_assignment，但它只将所有变量放在 TPU 的第一个（第 0 个）核心上，这很快就会耗尽内存。

在 TPUv3-8 上，我尝试保持 compute_shape = [2, 2, 1, 2] 和 [1, 1, 1, 2] 和 num_replicas = 1 但没有用。

我也愿意使用 GPU 来训练它。

【问题讨论】：

标签： tensorflow machine-learning google-cloud-platform deep-learning tpu

【解决方案1】：

根据云TPU文档，暂无官方支持：

Cloud TPU 是否支持模型并行性？

目前不支持模型并行性（或在单个 Cloud TPU 设备内的多个内核上执行不同的程序）。

https://cloud.google.com/tpu/docs/faq

潜在的问题可能是 TPUStrategy 中没有计算图的自动分片，因此该图都放置在一个设备上，除非（在模型代码中）您手动将设备放置分配给逻辑设备的权重和操作由DeviceAssignment.build 创建并小心处理跨设备的通信。

也就是说，如果您正在构建一个大型 Transformer，您需要对图形分片友好的层，那么还有另一个 TF2 兼容库（也来自 Google）可以提供帮助：Lingvo。在他们的 Github 中有一个example of sharding a model on a TPU v3-512 node。该库有谷歌的开源 GPipe，它可以帮助加速模型并行训练循环。 Lingvo 也应该使用 GPU。

【讨论】：