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

    使用 NVIDIA cuDNN 9 加速多個 Transformer

    我們提供的 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:FpropBprop

    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

    標簽

    人人超碰97caoporen国产