我的模型有一些 LSTM，当使用一个 GPU 在大型序列上运行时，它们会耗尽 Cuda 内存。所以我将模型的一些组件转移到另一个 GPU 上。我用 Apex AMP 尝试了两件事：
在调用
amp.initialize. 在这种情况下，我在第一次反向传播后很快就得到了 NaN。首先调用
amp.initialize，然后将模型组件移动到另一个 GPU。在这种情况下，它就像模型反向传播在单个 GPU 上运行。它耗尽了 Cuda 内存。
模型训练在没有 Apex 的情况下运行良好，所以我想我错过了在两个 GPU 上反向传播损失的一些步骤。我查看了 Apex 的文档，但是，它只谈论 DataParallelism，而不是 ModelParallelism。
有任何想法吗？