tensorflow - XLA rng-bit-generator 占用太多内存

Question

XLA 为这个张量分配了 4G 的内存。其大小似乎与批量大小成比例。这对我来说没有意义，它似乎不是要存储在 HBM 中的模型图的一部分。我使用 TPUv3。

除了模型的初始化，我不使用任何随机操作。此外，我为所有权重声明了 bfloat16，但这是一个 u32 张量。

  Largest program allocations in hbm:

  1. Size: 4.00G
     Shape: u32[128,8,1024,1024]{3,2,1,0:T(8,128)}
     Unpadded size: 4.00G
     XLA label: %rng-bit-generator = (u32[2,128]{1,0:T(2,128)}, u32[128,8,1024,1024]{3,2,1,0:T(8,128)}) rng-bit-generator(u32[2,128]{1,0:T(2,128)} %fusion.2446), algorithm=rng_default
     Allocation type: HLO temp
     ==========================

上述分配的原因可能是什么？我使用来自：https ://github.com/kamenbliznashki/pixel_models 的像素钉

问题：

• u32当我所有的权重/模型定义（包括全局环境标志）都使用 BF16 时，为什么这个张量有一个类型？
• 为什么要使用 rng-bit-generator？

标签： tensorflowpytorchtpuxla