?
雖然分析有助于識別效率低下的情況，但高級優化策略對于解決硬件限制和有效擴展 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 rmmfrom rmm.allocators.torch import rmm_torch_allocatorrmm.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 上配置大語言模型訓練。
?
