虽然分析有助于识别效率低下的情况，但高级优化策略对于解决硬件限制和有效扩展 AI 工作负载至关重要。在本文中，我们将探讨 CPU 卸载、Unified Memory、Automatic Mixed Precision 和 FP8 训练等技术。这些方法不仅能提高性能，还能助力研究人员突破 LLM 的极限。
激活函数的 CPU 卸载
在处理大型模型时，有效管理 GPU 内存至关重要。解决内存限制的一种策略是 CPU 卸载的激活函数。此技术涉及在模型训练或推理期间临时将中间激活张量从 GPU 内存移动到 CPU 内存。本节将探讨这种方法的优缺点。
了解 CPU 卸载
CPU 卸载激活函数可以处理更大的批量大小或训练更大的模型，而不会耗尽 GPU 内存。将激活函数分流到 CPU 可腾出空间用于对训练至关重要的其他运算，例如向前和向后传递。从本质上讲，卸载提供了一种在内存容量有限的 GPU 上训练更大模型的方法。
虽然卸载可能有利于memory-constrained环境，但必须了解其潜在的缺点，包括：
- 同步用度增加：在 CPU 和 GPU 之间移动数据会增加同步用度，因为这些设备必须进行协调才能确保数据一致性。
- 降低 GPU 利用率：CPU 和 GPU 之间频繁传输数据可能会导致 GPU 在等待数据时处于闲置状态。
- CPU 瓶颈：如果 CPU 资源成为瓶颈，或者 CPU 和 GPU 之间达到带宽限制，卸载可能无法很好地扩展。
要为 Llama 2 模型的八个层启用激活卸载，请按以下方式更改启动命令：
nsys profile ...
model.activations_checkpoint_granularity=null \
...
+model.cpu_offloading=True \
+model.cpu_offloading_num_layers=8
CPU 卸载的实际应用
引入 CPU 卸载具有一些明显的效果 (Figures 2 和 3) ，包括：
- CPU 活动增加：CPU 活动显著增加，尤其是在
pt_autograd_0和pt_main_thread等线程中。这种 CPU 活动的增加表明 CPU 积极参与管理卸载激活。
- GPU 显存占用减少：随着 CPU 上临时存储的数据增加，CUDA 硬件中的 GPU 显存占用减少。
同步用度增加：图 2 中的时间轴显示 GPU 内核执行中出现更频繁的缺口 (CUDA 硬件中的蓝条) ，这可能是由于 CPU 和 GPU 之间的同步增加。
一些关键见解包括：
- 激活卸载正在增加 CPU 活动。具体来说，
pt_autograd_0和pt_main_thread显示了高频活动，表明 CPU 积极参与管理卸载活动。
- 图 3 中的时间轴显示了内核执行中出现的显著且更频繁的差距，这可能会导致 CPU 和 GPU 之间的同步增加。
- 显著减少了 CUDA 硬件中的 GPU 显存占用。
这些效果凸显了为 LLM 训练启用 CPU 卸载所涉及的基本权衡。虽然卸载激活有助于缓解 GPU 显存限制，使训练更大的模型或使用更大的批量大小成为可能，但它通过将一些工作负载转移到 CPU 并增加 CPU 和 GPU 之间的数据移动频率来实现这一点。因此，由于同步开销增加和潜在的 CPU 瓶颈，整体训练速度可能会降低。
图 2 和 3 清楚地说明了这些动态，不仅显示了 GPU 显存占用的减少，而且还显示了 CPU 所需的活动和协调性的增加。了解这些交互对于旨在优化资源分配以及在扩展 LLM 工作负载时实现内存效率和计算性能之间的最佳平衡的从业者来说至关重要。最终目标是在 GPU 显存和利用率之间做出权衡。
在 NVIDIA Grace Hopper 上利用统一内存
本节将探讨底层硬件架构及其如何影响性能。作为最后一项实验，我们在 Grace Hopper 平台上利用 Unified Memory (UM) 运行相同的脚本。UM 为扩展深度学习工作负载带来了机遇和挑战。
统一内存概述
UM 提供了一个可供 CPU 和 GPU 访问的统一内存空间，从而简化内存管理并有望提高性能。在传统系统中，CPU 内存和 GPU 内存是不同的，它们之间需要显式数据传输。使用 UM 时，此边界模糊不清，允许 CPU 和 GPU 直接访问相同的内存位置。这提供了一些优势，包括：
- 简化内存管理：开发者无需明确指定应驻留在 CPU 还是 GPU 内存中，即可分配内存。系统会根据需要自动处理两者之间的数据迁移。
- 自动数据迁移：当 CPU 访问当前 GPU 内存中的数据 (反之亦然) 时，系统会自动将该数据迁移到 CPU 内存。如果 GPU 需要访问 CPU 内存中驻留的数据，则适用相同的原则。
借助 UM，您可以更轻松地处理无法加载到 GPU 上的超大型数据集。当 GPU 尝试使用超出其现有容量的显存时，它可能会耗尽显存或导致其他错误。通过使用 UM，数据可以根据需要从系统内存无缝传输到 GPU 并返回。
由于系统内存通常比 GPU 内存大得多，因此使用 UM 可简化内存容量限制的处理，并允许用户处理相对更大的数据集。
import rmm
from rmm.allocators.torch import rmm_torch_allocator
rmm.reinitialize(pool_allocator=True, managed_memory=True)
torch.cuda.memory.change_current_allocator(rmm_torch_allocator)
确保在 Python 脚本中的任何其他导入之前添加此代码段。
通过 CLI 启动脚本时，请设置 model.peft.peft_scheme=none，因为它会禁用 LoRA 微调并专注于监督式微调 (SFT) 。
图 4 比较了两种工作负载的性能：Llama 2 7B 的 LoRA 微调 (左) 和监督式微调 (右) 。这两个实验都利用 UM，但它们的行为和资源利用率差异很大。
LoRA 微调
- 一致的 GPU 利用率：LoRA 微调工作负载在尽可能减少空闲时间的情况下展现出一致的 GPU 利用率。
- 最小的内存迁移：只有 1.1% 的内存操作涉及 UM，这表明 LoRA 的参数高效方法将大多数数据保留在 GPU 上，从而减少内存传输开销。
监督性微调
- 频繁的 GPU 空闲周期：完整的微调工作负载显示频繁的 GPU 空闲周期，这表明数据访问瓶颈。
- 显存迁移显著：9.8% 的显存操作涉及 Unified Memory 活动，主机到设备传输繁重。
利用 UM 在单个 GH200 Grace Hopper Superchip 上拟合更大的模型，这在以前是不可能实现的。图 4 中的时间轴显示，LoRA 具有一致的 GPU 活动和最少的内存传输，这表明大多数数据存储在 GPU 上。相反，监督式微调工作流程强调了 UM 如何根据内存访问模式以不同的效率处理更大的工作负载。
尽管在 Grace Hopper 上利用 UM 使更大的模型适合单个芯片，但使用 UM 并不是扩展深度学习工作流程那么简单。
大规模训练有三种主要方法：在不卸载的情况下使用多个 GPU、在卸载时使用更少的 GPU，或在卸载时利用单个 GPU (如果可行) 。每种方法都需要权衡取舍：
- 在不卸载的情况下使用多个 GPU 可能会缩短解决问题的时间，但效率可能会降低。
- 通过卸载减少 GPU 可以优化资源使用，但可能会增加同步和内存迁移 overhead。
- 具有卸载功能的单个 GPU 仅适用于某些应用程序，并且由于频繁的内存迁移，通常会延长训练时间。
最终，选择取决于具体的应用和限制条件。对于某些应用程序而言，效率可能是满足严格时间预算（例如在几周或几个月内完成训练）的次要因素。
其他优化
虽然目前讨论的技术可以显著提高性能，但 NVIDIA NeMo 框架中随时可以使用更多优化策略，下面将详细介绍其中两个策略。
自动混合精度
NeMo 为 Automatic Mixed Precision (AMP) 提供内置支持，只需更改最少的代码即可实现混合精度训练。通过启用 AMP，您可以利用 NVIDIA GPU 中的 Tensor Cores 来加速计算、减少内存占用并提高吞吐量。
要启用 AMP，请在 NeMo 配置文件或训练脚本中配置精度设置。查找模型或训练器配置中的精度或 AMP 标志。这使得部分代码能够使用较低 (但精度较低) 的数据类型 (即 bfloat16 或 float16) 进行计算。
FP8 训练
NeMo 框架旨在支持 FP8 训练等先进功能，可显著减少内存占用并加速计算。虽然这确实需要修改，但 FP8 通常会显著提升性能。在使用 FP8 时，建议使用新的 Transformer Engine，NVIDIA Hopper 架构通过其内置功能提供了该引擎。
结论
在不断发展的 LLM 环境中，对计算资源的需求不断增长，这使得识别和解决性能瓶颈的能力成为研究人员和工程师等的关键能力。本文和在 NVIDIA Grace Hopper 上分析 LLM 训练工作流所探讨的技术包括战略性地调整超参数、仔细评估 CPU 卸载的权衡，以及在 NVIDIA Grace Hopper 超级芯片等先进硬件上驾驭 Unified Memory 的复杂性。
这些技术是持续优化过程的起点。通过结合使用NVIDIA Nsight Systems等工具和NVIDIA NeMo等框架提供的高级功能进行细致的分析，如果想要达到峰值性能，利用所有资源的能力至关重要。
优化和使用出色的软件将效率和性能提升到新的水平，使您能够训练更大、更复杂的模型，并突破 AI 领域的极限。
如需了解更多信息，请点播观看 GTC 会议：在 Grace Hopper Superchip 上配置大语言模型训练。
