XLA rng位生成器占用太多内存

编程问答 2022-06-09

问题描述

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,128)}) rng-bit-generator(u32[2,128)} %fusion.2446),algorithm=rng_default
     Allocation type: HLO temp
     ==========================

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

的pixelnail

问题：

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

解决方法

暂无找到可以解决该程序问题的有效方法，小编努力寻找整理中！

如果你已经找到好的解决方法，欢迎将解决方案带上本链接一起发送给小编。

小编邮箱:dio#foxmail.com (将#修改为@）

pytorchtensorflowtensorflowtensorflowtpuxla