重要提示

您正在查看 NeMo 2.0 文档。此版本引入了 API 的重大更改和一个新的库 NeMo Run。我们目前正在将 NeMo 1.0 的所有功能移植到 2.0。有关先前版本或 2.0 中尚不可用的功能的文档，请参阅 NeMo 24.07 文档。

CPU 卸载#

概述#

NeMo 中的 CPU 卸载是一项通过将激活和非活动权重卸载到 CPU 存储来减少 GPU 峰值内存使用量的功能。NeMo 支持在 Transformer 层级别进行卸载，允许用户指定其语言模型中需要 CPU 卸载的 Transformer 层数。在前向传递期间，NeMo 会在最佳时间卸载激活，并在后向传递期间根据需要重新加载它们。

功能#

  • 通过有效管理激活内存，支持训练具有长序列长度的模型。

  • 通过卸载激活内存，实现每个 GPU 的高批量大小。

  • 在卸载和重新加载期间，将计算与数据传输（Host2Device 和 Device2Host）重叠。

用法#

  • 将 cpu_offloading 设置为 True 以启用 CPU 卸载。

  • 将 cpu_offloading_num_layers 设置为 0 到模型中总层数减 1 之间的值。

  • 根据您的需要设置 cpu_offloading_activations 和 cpu_offloading_weights，以仅卸载激活、仅卸载权重或两者都卸载。

在本页中