Tensorflow 中可扩展、高效的分层 Softmax？

Posted 2023-02-16

【中文标题】Tensorflow 中可扩展、高效的分层 Softmax？

【英文标题】：Scalable, Efficient Hierarchical Softmax in Tensorflow?

【发布时间】：2017-10-23 17:48:06

【问题描述】：

我对实现一个可以处理大型词汇表的分层 softmax 模型感兴趣，比如大约 10M 类。要做到这一点，既可扩展到大类数量又高效的最佳方法是什么？例如，至少one paper 表明，当使用每个节点 sqrt(N) 分类的 2 级树时，HS 可以为大型词汇实现约 25 倍的加速。我也对具有任意分支因子的任意深度树的更通用版本感兴趣。

我在这里看到了几个选项：

1) 为每个批次运行 tf.gather，我们在其中收集索引和拆分。这会产生大批量和胖树的问题，现在系数被大量重复，导致 OOM 错误。

2) 与 #1 类似，我们可以使用 tf.embedding_lookup，它可以帮助解决 OOM 错误，但现在将所有内容都保留在 CPU 上，并大大减慢速度。

3) 使用tf.map_fn 和parallel_iterations=1 分别处理每个样本并返回使用gather。这更具可扩展性，但由于序列化，并没有真正接近 25 倍的加速。

有没有更好的方法来实现 HS？深而窄的树与短而宽的树有不同的方法吗？

【问题讨论】：

它们因任务而异。语言模型的批量较大，大约 400 个，隐藏大小大约 300 个；其他任务可能具有较小的批量大小和较大的隐藏大小，例如 imagenet 分类。相对于问题而言，VRAM 和 RAM 相当大（尽管 GPU RAM 不是）。

我可以看看你在 Tensorflow 中的 HS 实现吗？我现在也需要它。

有点乱，但请看这里：github.com/tansey/sdp/blob/… -- 回想起来，我建议使用 pytorch 或其他动态图框架。

【参考方案1】：

你提到你想要 GPU 级的性能：

但现在将所有内容都保留在 CPU 上并减慢速度

并希望使用 300 个单位的隐藏大小和 1000 万字的字典。

这意味着（假设 float32），您需要 4 * 300 * 10M * 2 字节 = 24 GB 来存储输出层的参数和梯度。

分层 Softmax (HSM) 不会降低内存需求 - 它只是加快了训练速度。

实际上，您将需要更多的 GPU 内存，因为您还需要存储：

其他参数及其梯度

优化器数据，例如动量训练中的速度

激活和反向传播的临时数据

特定于框架的开销

因此，如果您想在 GPU 上进行所有计算，您别无选择，只能将此层分布在多个高内存 GPU 上。

但是，你现在有另一个问题：

为了具体说明，假设您有一个包含 3K 类的 2 级 HSM，每个类有 3K 字（总共 9M 字）。您将 3K 类分布在 8 个 GPU 上，因此每个 GPU 托管 384 个类。

如果一个批次中的所有目标词都来自相同的 384 个类，即它们属于同一个 GPU，该怎么办？一个 GPU 将完成所有工作，而其他 7 个等待它。

问题在于，即使一个批次中的目标词属于不同的 GPU，如果您想在 TensorFlow 中进行此计算（这是因为 TensorFlow是一个“指定并运行”的框架——计算图对于最好的情况和最坏的情况是一样的）

什么是做到这一点的最佳方法，既可扩展到大类数量又高效？

上述模型并行性的低效率（每个 GPU 必须处理整个批次）表明应该尝试将所有内容放在一个位置。

让我们假设您要么在主机上实现所有东西，要么在 1 个巨大的 GPU 上实现。

如果您不是对序列进行建模，或者如果您是，但整个序列只有一个输出，那么与上述内存要求相比，复制您提到的参数的内存开销可以忽略不计：

400 == 批量大小

在这种情况下，您可以简单地使用gather 或embedding_lookup（虽然复制效率低下）

但是，如果您对长度为 100 的模型序列进行建模，并且在每个时间步都有输出，那么参数复制就会成为一个大问题。

在这种情况下，我认为您需要使用 C++ / CUDA C 并将整个层及其渐变实现为自定义操作。

【讨论】：

所以您是说实现此功能的唯一有效方法是使用我在#2 中建议的标准embedding_lookup？这似乎是合理的，但我想知道你真的会在多大程度上看到你在真实世界数据集上描述的 GPU 停滞，这正是我正在寻找的。此外，在我链接的论文中对采样的 softmax 进行了比较，并在许多其他论文中进行了彻底的比较。

另外，如果可以在一个 GPU 上处理所有事情会怎样？例如，假设将来我有一个 32GB 的 GPU。

@WesleyTansey “与采样的 softmax 进行比较” - 我现在看到了。查看此更新和其他更新。

谢谢。所以这似乎只是同意我的观点。我在这里寻找的是一些硬数字，表明有（或没有）比gather 更好的方法来做到这一点。如何防止通过 CUDA 进行复制？性能提升究竟是什么？

@WesleyTansey “所以这似乎只是同意我的观点。”我指出你的记忆问题甚至在你训练网络之前就开始了（这不在你的 Q 中）。我还指出了尝试在 TF 中的多个 GPU 上执行此操作的固有低效率（同上）。虽然我知道 C++ 和 CUDA C，而且对我来说很明显这是可以做到的，但是为你实现这个工作量太大了，抱歉。