• <xmp id="om0om">
    •
    • <table id="om0om"><noscript id="om0om"></noscript></table>
  • 3 月 19 日下午 2 點，鎖定 NVIDIA AI 網絡中文專場。立即注冊觀看
    Related Resources
    對話式人工智能

    使用 NVIDIA cuDNN 9 加速多個 Transformer

    2024年 5月 24日
    By
    0

    我們提供的 NVIDIA CUDA 深度神經網絡庫(cuDNN) 是一個專門為深度學習應用而設計的 GPU 加速庫，旨在以先進的性能加速深度學習基元。

    cuDNN 與 PyTorch、TensorFlow 和 XLA (加速線性代數) 等熱門深度學習框架集成。這些框架抽象化了直接 GPU 編程的復雜性，使您能夠專注于設計和訓練模型，而無需擔心底層硬件。cuDNN 可作為底層性能引擎，確保以更高的效率執行這些框架上的操作。

    最近，擴展的點積注意力 (SDPA) 已成為大型語言模型 (LLM) 等重要工作負載中的性能關鍵基元。cuDNN 增加了對該基元的支持，并一直在使用閃存注意力和其他優化來提高其性能釋放，同時擴展功能支持面，以支持一系列注意力用例。

    在 NVIDIA H200 Tensor Core GPU 上，cuDNN 可以在 FP8 中實現高達 1.2 PFLOPS 的性能。作為端到端示例，我們的團隊在為 Lama2 70B LoRA 微調啟用 cuDNN FP8 SDPA 后，測得了 1.15 倍的速度提升。該實驗使用了 NVIDIA NeMo 和?NVIDIA Transformer 引擎 在 8-GPU H200 節點上運行。

    軟件堆棧 e2e 加速
    禁用 cuDNN 的 NeMo 1x (基準)
    在 BF16 中使用 cuDNN SDPA 的 NeMo 1.11 倍
    在 FP8 中使用 cuDNN SDPA 的 NeMo 1.15 倍
    表 1.將用于 SDPA 的 cuDNN 用作 8-GPU H200 節點上端到端訓練運行 (Llama2 70B LoRA 微調) 的一部分所產生的影響

    在本文中，我將詳細介紹 cuDNN SDPA 可實現的性能，介紹使用方法，并簡要總結 cuDNN 9 中的其他一些值得注意的新功能。

    經擴展的點積產品注意力性能

    NVIDIA 正式通過在 NVIDIA Omniverse 中開源融合的 Multihead Attention (fMHA) 內核來支持注意力機制。APEX 庫將注意力算法融合到單個內核中。Tri Dao 的創新工作以此內核為起點，實現了閃存注意力的形式，帶來了巨大的性能改進和功能增強。欲了解更多信息，請參閱 Dao-AILab/flash-attention 在 GitHub 上的項目頁面，并查看 FlashAttention-2：通過更好的并行性和作業分區提高注意力 相關論文。

    然后， NVIDIA 通過更快、更靈活的實現，在融合注意力方面取得了進展。此實現現在是 NVIDIA 的最新技術成果，現在已經集成到 NVIDIA Hopper 架構 的 GPU 中。

    XLA 目前提供了 cuDNN SDPA 的路徑。您可以通過 JAX SDPA API 或依靠 XLA 編譯器將 JAX/PyTorch 中的自定義實現降低到 cuDNN SDPA。

    目前，PyTorch 渴望模式 SDPA 不使用 cuDNN，但基于 cuDNN 的實現正在進行中。有關更多信息，請參閱 PyTorch 的相關 Pull Request：Fprop 和 Bprop。

    cuDNN SDPA 實施封裝了以下內容：

    • 深入了解基礎 NVIDIA 硬件架構
    • 實現從非閃存到閃存注意力 v2 的所有先進的 SDPA 算法，以及介于兩者之間的所有算法
    • 啟發式算法，可根據問題大小和目標 GPU 架構自動設置性能旋鈕 (例如圖塊大小)

    這可在 NVIDIA GPU 上實現出色的可用性能。圖 1 和圖 2 顯示，在各種用例中，cuDNN 9 BF16 的速度比現有最佳 GPU 快 2 倍，而 PyTorch 即時實現 中 cuDNN FP8 的實現速度提升高達 3 倍。

    性能提升可實現更長的序列長度以及更短的模型預訓練和微調時間。與其他公開基準（例如 flash-attention/benchmarks）相比，本文僅報告 GPU 時間，不包括主機用度。

    Bar chart shows that cuDNN v9 achieves higher performance for SDPA (forward only) compared to PyTorch eager implementation.
    圖 1.帶因遮罩的 SDPA (僅向前)，頭部尺寸 128
    Bar chart shows that cuDNN v9 achieves higher performance for SDPA (forward plus backward) compared to the PyTorch eager implementation.
    圖 2.帶有因果掩膜的 SDPA (前向和反向)，頭部尺寸 128
    Bar chart shows that cuDNN v9 achieves up to 1.2 PFLOPs SDPA (forward) in FP8.
    圖 3.無因掩膜的 SDPA (僅向前)，頭部尺寸 256

    SDPA 作為 cuDNN 圖形

    cuDNN 中的 SDPA 可指定為張量運算的 cuDNN 圖形。對于給定的圖形，cuDNN 庫擁有一組可以執行該圖形的引擎。

    雖然一些圖形可能沒有合適的引擎，但其目的是為在 GPU 上原子化執行的任何 cuDNN 圖形至少提供一個引擎，通常在一個融合內核中，但有時在一小部分協作內核中。您可以將其視為整個框架圖形的“子圖”。

    SDPA (各種形式的 SDPA) 是此類子圖的理想示例。支持這些模式的引擎旨在盡可能靈活，而不會影響顯著的性能，而是使用現有的最佳算法 (如 Flash Attention)。這意味著您可以更改注意力計算，但仍然可以在 GPU 上高效運行。本節旨在解釋目前提供的靈活性和支持。

    圖 4 顯示了前向傳播 (fprop) 用例的 cuDNN 圖形。這是神經網絡中注意力機制 SDPA 的前向計算所涉及的運算序列。

    A flowchart shows a series of operations including the generation of indices, pointwise operations, and selection operations that collectively represent the causal mask logic.
    圖 4.用于擴展點積注意力的細粒度 fprop 計算 cuDNN 圖形

    因果關系掩膜邏輯降低到單個生成索引、逐點和選擇運算。softmax 邏輯降低到六個 cuDNN 運算 (圖 4)。這種 cuDNN 運算的粒度使您能夠靈活地表達自定義模型。圖 5 顯示了可能的組合。

    A flowchart shows the various configurations supported by cuDNN for attention mechanisms, illustrating the flexibility in customizing attention operations for different deep-learning models.
    圖 5.cuDNN 注意力支持

    圖 5 顯示，在頭部維度更大 (256) 且無因遮罩的情況下，cuDNN FP8 正向閃光注意力可實現高達 1.2 PFLOPS 的運算性能。

    您還可以使用 BMM1 和 Softmax 之間的任意逐點運算構建自定義圖形。這種靈活性支持尚未發現的新變體。

    當創意研究人員調整典型的注意力模式時，他們就不太容易回到優化程度較低的框架實現，從而在性能上達到峰值。圖 2 適用于完整的訓練運行。反向道具注意力圖也具有類似的靈活性。

    SDPA 使用演示

    有多個 API 切入點可用于創建和運行 cuDNN 圖形：

    • 前端 API (包括 C++和 Python 變體)
    • 后端 API (僅支持 C 語言)

    上一節中描述的所有 cuDNN 圖形概念都適用于這兩個級別的 API.但是，cuDNN 團隊建議您使用 Python 或 C++與 cuDNN 前端 API 進行交互，除非您需要 C 接口。前端 API 明顯更簡潔，并增加了一些便利。

    例如，前端 API 擴展了操作節點的概念，使節點能夠封裝多個操作和它們之間的數據流。換言之，它們是方便節點，可以抽象出常見圖形模式 (例如 SDPA) 的細節。節點仍然可以靈活配置眾所周知的變體。此 SDPA 使用演練從最簡單的情況開始，即在 Python 中創建的 SDPA 節點。

    我們在 前端庫中的 SDPA Python 示例 中演示了配置選項和基本使用流程：

    1. 初始化一個具有適當數據類型的 cudnn.pygraph 對象。
    2. 創建一個包含張量維度、布局和數據類型的 Tensor 對象。
    3. 創建縮放的點積閃光注意力節點并提供所需的配置。
    4. 構建圖形并提供設備指針。
    5. 執行圖形。
    # 1. cuDNN graph
    graph = cudnn.pygraph(
              io_data_type=cudnn.data_type.BFLOAT16,
              intermediate_data_type=cudnn.data_type.FLOAT,
              compute_data_type=cudnn.data_type.FLOAT,
            )
     
    # 2. tensor attributes
    q = make_tensor_attr(graph, q_gpu, "q")
    k = make_tensor_attr(graph, k_gpu, "k")
    v = make_tensor_attr(graph, v_gpu, "v")
    attn_scale = make_tensor_attr(graph, attn_scale_cpu, "attn_scale", is_pass_by_value=True)
     
    # 3. sdpa node with the configurations
    o, stats = graph.scaled_dot_product_flash_attention(
                        name="scaled_dot_product_flash_attention",
                        q=q,
                        k=k,
                        v=v,
                        is_inference=false,
                        attn_scale=attn_scal,
                        bias=None,
                        use_alibi_mask=false,
                        use_padding_mask=false,
                        seq_len_q=None,
                        seq_len_kv=None,
                        use_causal_mask=true,
                        dropout=None,
                    )
     
    # 4. build the graph and provide the device pointers
    graph.build()
    variant_pack = {
                     q: q_gpu,
                     k: k_gpu,
                     v: v_gpu,
                     o: o_gpu
                   }
     
    # 5. execute the graph
    graph.execute()

    如果您想獲得比 SDPA 節點默認情況下提供的更高的靈活性，則構建底層 SDPA 圖形的代碼是開源的，因此可以進行自定義。

    例如，以下代碼示例可自定義 SDPA 節點內的擴展節點。有關更多信息，請參閱 scaled_dot_product_flash_attention.h 文件。

    // Optional scale
    if (options.inputs.Attn_scale) {
          // Lower options to scale options
          auto attn_scale_output = std::make_shared();
          attn_scale_output->set_is_virtual(true);
     
          Pointwise_attributes scale_attributes;
          scale_attributes.set_name("attn_scale");
          scale_attributes.set_mode(PointwiseMode_t::MUL);
          scale_attributes.inputs.IN_0 = last_output;
          scale_attributes.inputs.IN_1 = options.inputs.Attn_scale;
          last_output = scale_attributes.outputs.OUT_0 = attn_scale_output;
          auto scale_node = std::make_unique(std::move(scale_attributes), context);
           
          sub_nodes.emplace_back(std::move(scale_node));
    }

    如果您想詳細了解如何加速您的自定義 Transformer，請參閱 NVIDIA cuDNN 的官方文檔。

    cuDNN 9 的其他顯著特性

    除了 SDPA 改進之外，cuDNN 9 還引入了其他幾項重要改進：

    • 混合輸入精度支持矩陣和卷積
    • 改進錯誤報告
    • 硬件前向兼容性
    • 簡化安裝

    混合輸入精度支持矩陣和卷積

    要求輸入操作數的數據類型為同一類型 (FP16、FP32) 的矩陣乘法 (matmul) 和卷積 API 不適合 AWQ (激活感知權重量化) 等情況，其中激活在 FP16 中，權重可能在 INT8 中。假設您想以更高的精度計算，則必須投射數據。如果不是在線的，這將需要額外的內存和轉換成本。

    cuDNN 現在支持混合輸入精度矩陣和卷積，其中 A 和 B 操作數可以是不同的數據類型，并提供用于性能和內存優化的在線融合類型轉換。您可以選擇將操作數轉換為其他類型。cuDNN 在優化的內核中處理所需的轉換。

    圖 6 顯示了 cuDNN 混合輸入精度矩陣與未融合工作流程之間的加速。灰色條表示輸入 A 和 B 分別采用 FP16 和 INT8 精度的情況，其中 A 轉換為 INT8，然后使用 INT32 累加進行 INT8xINT8 矩陣乘法運算。綠色條表示 A 從 INT8 上轉換為 FP16，然后使用 FP32 累加進行 FP16xFP16 矩陣乘法運算的情況。

    Chart shows the performance benefit of cuDNN mixed input precision matrix multiplications over traditional unfused workflows.
    圖 6.cuDNN 中對混合輸入精度矩陣和卷積的支持

    改進錯誤報告

    日志記錄在軟件開發中至關重要，對于使用 cuDNN 的深度學習框架等復雜系統尤其如此。過去，cuDNN 的一個常見痛點是難以調試錯誤和警告。我們一直在不斷改進錯誤報告以解決這一問題。

    cuDNN 9 增加了以下功能：

    • 更具體的錯誤代碼
    • 進行分類以幫助組織數量增加的錯誤碼
    • 符合日志記錄規范的嵌入式日志記錄級別
    • cuDNNGetLastErrorString是一個新的函數，用于以編程方式檢索最新的錯誤消息。

    欲了解更多信息，請參閱 錯誤報告和 API 日志記錄 部分內容，詳見開發者指南。

    硬件前向兼容性

    在版本 9.0.0 之前，cuDNN 庫在庫發布之日支持最新的公開可用 GPU 架構。例如，cuDNN 8.9.x 通過 NVIDIA Hopper (即計算能力 9.0) 提供支持。不支持在未來的 GPU 架構上運行 8.9.x。

    但是，cuDNN 9 對 API 的一大子集具有硬件前向兼容性。這意味著，僅將此 API 子集與 cuDNN v9 結合使用的程序將在未來的 GPU 上正常運行，并且這些程序的用戶不會被強制升級其 cuDNN 安裝以使用未來的 GPU。

    在計算能力大于 9.0 的 GPU 上運行時，前向兼容性不是糾錯，而是意味著庫可以找到功能等效的實現，并使用 PTX JIT 針對新架構進行定位。

    欲了解更多關于支持限制和最佳實踐的信息，請參閱 硬件前向兼容性 部分內容，在開發者指南中。

    簡化安裝

    在 Python 環境中，除了庫之外，您現在還可以使用 pip 安裝新的 Python 前端：

    # cuDNN library for CUDA 12
    pip install nvidia-cudnn-cu12
     
    # cuDNN frontend
    pip install git+https://github.com/NVIDIA/cudnn-frontend.git

    cuDNN 9 還簡化了 RPM 和 Debian 元軟件包的安裝流程。例如，以下命令可在 Ubuntu 22 上安裝 cuDNN：

    wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-keyring_1.1-1_all.deb
    sudo dpkg -i cuda-keyring_1.1-1_all.deb
    sudo apt-get update
    sudo apt-get install -y cudnn

    這個新流程簡化了安裝 Debian 軟件包時的密鑰環設置，并使用 cudnn 元軟件包，以提高安裝效率。

    如果您需要了解更多信息和完整說明，請參閱 cuDNN 安裝指南，以獲取詳細的安裝步驟和配置說明。

    后續步驟

    如果您有反饋、問題或意見，可以在 cuDNN NVIDIA 開發者論壇(適用于 cuDNN 庫主題) 或 NVIDIA/cudnn-frontend 在 GitHub 上發布 (適用于前端主題)。另外，您可以 下載 cuDNN 開始使用。

    請持續關注，了解未來更多新的 cuDNN 功能。本文中概述的進展是一個重要的里程碑，但未來還會有更多進展。

    隨著 AI 不斷將行業推向軟硬件集成的極限， NVIDIA 不斷優化性能并改善用戶體驗，以便 cuDNN 能夠在深度學習框架和圖形編譯器中得到更有效、更廣泛的使用。

    致謝

    NVIDIA cuDNN 團隊與公司內的許多其他團隊緊密合作，為這篇博文提供了技術內容。

    ?

    相關資源

    0

    標簽

    對話式人工智能 | 數據中心/云端 | 生成式人工智能/大語言模型 | HPC / Scientific Computing | cuDNN | NeMo | Intermediate Technical | Deep dive | Conversational AI | HPC/科學計算 | Machine Learning & Artificial Intelligence | 數據中心/云 | 深度學習 | 生成型人工智能

    關于作者

    Matthew Nicely 于 2019 年 3 月加入 NVIDIA ，之前曾在美國 AL-Huntsville 的美國陸軍航空和導彈研究開發與工程中心工作。在那里，他專注于 CUDA 算法開發和 Jetson 系列的優化。在 NVIDIA ，他曾在聯邦部門工作，協助 CUDA 的開發和優化，以及在各種 NVIDIA 工具集上為客戶提供教育和概念驗證，最近轉為 math libraries 產品經理。 2019 年，他獲得了博士學位。計算機工程學位，專注于 GPU 的算法優化。
    人人超碰97caoporen国产