BERT如何利用TPU存储器？

Question

README在Google的BERT存储库中表示，即使对于BERT-Large模型，长度为512的单个句子也无法放在12 GB的Titan X中。

但是在BERT论文中，它说使用64个TPU芯片来训练BERT-Large 最大长度为512，批处理大小为256。如何将大于256倍的批处理放入仅多171倍的内存中？

从另一个角度来看，我们可以在每个样本的内存使用情况下比较这两种配置：

• TPU：假设在预训练中使用了TPUv3，则总TPU内存为32 GB /芯片* 64芯片= 2048 GB。根据该论文，批量大小为256且最大长度为512的这种配置在此配置中效果很好，这意味着 8 GB内存可以容纳单个样本。此外，每个样本的内存使用量将减少到仅如果使用GPUv2，则为4 GB。
• GPU：一个12 GB的Titan X甚至无法容纳一个长度为512的样本。

为什么GPU上的内存消耗要大得多？这是否意味着对TPU的内存消耗进行了优化，比对GPU进行了优化？

Answer 1

这可能归因于TPU随附的高级编译器，并且针对tensorflow ops进行了优化。正如BERT中的readme - out-of-memory issues所说，

  

在DNN训练期间，GPU / TPU内存的主要用途是在向前传递中缓存中间激活，这对于向后传递中的有效计算是必需的。

但是，在TPU编译中，有一条特殊的XLA

指令（优化了TensorFlow计算的线性代数的特定域编译器）指令

  

可以将来自不同TensorFlow操作的多个指令合并为一个计算。与融合中的根指令相对应的TensorFlow操作被用作融合操作的命名空间。

另一方面，在具有普通TF的GPU上运行基本上没有（或非常有限）优化。