BERT基座未加要求的gpu ram

Question

我正在处理一个NLP任务，使用BERT，我对GPU内存有一点疑问。

我已经建立了一个模型(使用DistilBERT)，因为我在RTX3090 (24 of的图形处理器的内存，但~20.5gb可用)上出现内存不足的问题。

为了使它发挥作用，我将我的数据限制为训练集中的1.1兆个句子(截短128个单词的句子)，以及大约300 K的验证，但使用的是较大的批处理大小(256)。

现在我有可能在Nvidia A100上重新训练模型(用40 on的gpu的ram)，所以是时候使用伯特基础了，而不是蒸馏版。

我的问题是，如果我将批处理大小(例如从256个减少到64个)，我会有一些可能性来增加我的训练数据的大小(例如，从1.1到2-3毫安)，句子的长度(例如从128到256，或198)，并在A100的40 get上使用bert基(它有很多可训练的参数比蒸馏版本更多)，还是我可能会得到一个OOM错误？

我之所以这样问，是因为我在这个集群上没有无限的尝试，因为我并不是唯一使用它的人(另外，在每一种情况下，我必须以不同的方式准备数据，而且它的大小相当大)，所以我会对可能发生的事情进行评估。

Answer 1

正如您在注释中指出的，您预先标记了数据，并在GPU内存中保持在张量中。

只有当前的批处理应该加载到GPU中，因此您不需要减少培训数据大小(假设您的数据加载和培训例程得到了正确的实现)。为了让您在CPU中训练数据张量，可以使用with tf.device(...):。

但是，考虑到对于CPU内存的大小，训练数据的大小也是巨大的。这方面的一种典型方法是将令牌ID保存在磁盘上，然后从那里加载它们。