重要提示

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

激活重计算

网络层的输入激活存储在设备内存中，并在反向传播期间用于计算梯度。当使用长序列长度或大微批大小训练 LLM 时，这些输入激活会迅速耗尽设备内存。检查点保存少量激活并重新计算其余激活是减少设备内存使用的常用技术。

Transformer 层重计算

NeMo 支持 Transformer 层重计算，该技术检查点保存每个 Transformer 层的输入，并重新计算剩余层的激活。此技术显著减少了激活内存的使用。但是，由于重新执行整个层的前向计算，它使每个 Transformer 层的计算成本增加了 30%。NeMo 还支持部分 Transformer 层重计算，当重新计算少量 Transformer 层有助于减少足够的 GPU 内存以适应模型时，这非常有用。这种方法避免了重新计算其余层的需要。

可以通过 Transformer 配置 TransformerConfig 启用重计算配置。

通过设置 recompute_method=full 启用 Transformer 层重计算。可以使用 recompute_num_layers 以及 recompute_method=block 设置要重计算的 Transformer 层数。如果将 recompute_num_layers 设置为总层数，则会检查点保存并重新计算所有 Transformer 层的输入。当使用流水线并行性进行训练时，recompute_num_layers 指示每个流水线阶段的层数。当使用虚拟流水线时，recompute_num_layers 指定每个虚拟流水线阶段的层数。

NeMo 还支持检查点保存多个连续 Transformer 层块的输入，这意味着 Transformer 层块成为重计算粒度。这种方法可以节省激活内存，但会增加重计算缓冲区内存。因此，仅当模型具有许多 Transformer 层或 Transformer 层的中间层保持相对较小的激活存储时，它才有利于节省内存。可以通过设置 recompute_method=uniform 启用此重计算模式，并使用 recompute_num_layers 设置每个重计算块的 Transformer 层数。

from nemo.collections import llm
from functools import partial

# Load train recipe
recipe = partial(llm.llama3_8b.pretrain_recipe)()

recipe.model.config.recompute_method = "block"  # Enable 'block'-wise recomputation
recipe.model.config.recompute_num_layers = 4

自注意力重计算

NeMo 支持自注意力重计算，该技术检查点保存每个自注意力块的输入，并重新计算中间输入激活。这种经济高效的方法以最小的重计算成本实现了高内存节省。自注意力块的中间层占激活内存的大部分。这是因为 softmax、dropout 和 qkv 点积注意力层的输入大小具有序列长度平方的内存复杂度。但是，它们的重计算成本相对小于其他线性投影层，后者与隐藏层大小的平方成线性关系。

当使用 FlashAttention 时，自注意力重计算是硬性启用的，Transformer Engine 中支持 FlashAttention。此外，您可以通过设置 recompute_method=selective 在不使用 FlashAttention 的情况下使用自注意力重计算。

from nemo.collections import llm
from functools import partial

# Load train recipe
recipe = partial(llm.llama3_8b.pretrain_recipe)()

recipe.model.config.recompute_method = "selective"  # Enable selective recomputation

完整和选择性检查点粒度方案

均匀和块检查点方法方案（完整检查点粒度）