不止量化及 LORA - 原生 PyTorch 性能及記憶體優化精解

張愛玲

本科畢業於清華大學電子工程系，後在美國伊利諾伊大學香檳分校獲得計算機科學碩士學位，是AI系統工程領域的資深技術專家。作為深度學習框架PyTorch核心團隊成員，參與了多個關鍵組件的研發與優化工作。同時，她作為技術負責人主導了PyTorch/XLA和Taichi編譯器等多個具有廣泛影響力的開源項目。目前在工業界專註於大規模語言模型訓練的基礎設施開發與性能優化工作。

 

楊占略

清華大學本科，哥倫比亞大學碩士，資深技術專家。曾先後任職於蘋果、百度、太極圖形公司，分別從事深度學習推理框架、訓練框架PaddlePaddle以及Taichi編譯器的開發。現就職於大疆公司，聚焦於端側模型優化、推理部署相關工作。

第 1 章 歡迎來到這場大模型競賽

1.1 模型規模帶來的挑戰

1.2 資料規模帶來的挑戰

1.3 模型規模與資料增長的應對方法

第 2 章 深度學習必備的硬體知識

2.1 CPU 與記憶體

2.1.1 記憶體

2.1.2 CPU

2.2 硬碟

2.3 GPU

2.3.1 CPU 的局限性

2.3.2 GPU 的硬體結構

2.3.3 GPU 程式設計模型及其硬體對應

2.3.4 GPU 的關鍵性能指標

2.3.5 顯示記憶體與記憶體間的資料傳輸

2.4 分散式系統

2.4.1 單機多卡的通訊

2.4.2 多機多卡的通訊

2.4.3 分散式系統的資料儲存

第 3 章 深度學習必備的 PyTorch 知識

3.1 PyTorch 的張量資料結構

3.1.1 張量的基本屬性及建立

3.1.2 存取張量的資料

3.1.3 張量的儲存方式

3.1.4 張量的視圖

3.2 PyTorch 中的運算元

3.2.1 PyTorch 的運算元函數庫

3.2.2 PyTorch 運算元的記憶體分配

3.2.3 運算元的呼叫過程

3.3 PyTorch 的動態圖機制

3.4 PyTorch 的自動微分系統

3.4.1 什麼是自動微分

3.4.2 自動微分的實現

3.4.3 Autograd 擴充自訂運算元

3.5 PyTorch 的非同步執行機制

第 4 章 定位性能瓶頸的工具和方法

4.1 配置性能分析所需的軟硬體環境

4.1.1 減少無關程式的干擾

4.1.2 提升PyTorch 程式的可重複性

4.1.3 控制GPU 頻率

4.1.4 控制CPU 的性能狀態和工作頻率

4.2 精確測量程式執行時間

4.2.1 計量CPU 程式的執行時間

4.2.2 程式預熱和多次執行取平均

4.2.3 計量GPU 程式的執行時間

4.2.4 精確計量GPU 的執行時間

4.3 PyTorch 性能分析器

4.3.1 性能分析

4.3.2 顯示記憶體分析

4.3.3 視覺化性能圖譜

4.3.4 如何定位性能瓶頸

4.4 GPU 專業分析工具

4.4.1 Nsight Systems

4.4.2 Nsight Compute

4.5 CPU 性能分析工具

4.5.1 Py-Spy

4.5.2 strace

4.6 本章小結

第 5 章 資料載入和前置處理專題

5.1 資料連線的準備階段

5.2 資料集的獲取和前置處理

5.2.1 獲取原始資料

5.2.2 原始資料的清洗

5.2.3 資料的離線前置處理

5.2.4 資料的儲存

5.2.5 PyTorch 與第三方函數庫的互動

5.3 資料集的載入和使用

5.3.1 PyTorch 的 Dataset 封裝

5.3.2 PyTorch 的 DataLoader 封裝

5.4 資料載入性能分析

5.4.1 充分利用CPU 的多核心資源

5.4.2 最佳化CPU 上的計算負載

5.4.3 減少不必要的CPU 執行緒

5.4.4 提升磁碟效率

5.5 本章小結

第 6 章 單卡性能最佳化專題

6.1 提高資料任務的平行度

6.1.1 增加資料前置處理的平行度

6.1.2 使用非同步介面提交資料傳輸任務

6.1.3 資料傳輸與GPU 計算任務平行

6.2 提高GPU 計算任務的效率

6.2.1 增大BatchSize

6.2.2 使用融合運算元

6.3 減少CPU 和GPU 間的同步

6.4 降低程式中的額外銷耗

6.4.1 避免張量的建立銷耗

6.4.2 關閉不必要的梯度計算

6.5 有代價的性能最佳化

6.5.1 使用低精度資料進行裝置間拷貝

6.5.2 使用性能特化的最佳化器實現

6.6 本章小結

第 7 章 單卡顯示記憶體最佳化專題

7.1 PyTorch 的顯示記憶體管理機制

7.2 顯示記憶體的分析方法

7.2.1 使用PyTorch API 查詢當前顯示記憶體狀態

7.2.2 使用PyTorch 的顯示記憶體分析器

7.3 訓練過程中的顯示記憶體佔用

7.4 通用顯示記憶體重複使用方法

7.4.1 使用原位操作運算元

7.4.2 使用共用儲存的操作

7.5 有代價的顯示記憶體最佳化技巧

7.5.1 跨批次梯度累加

7.5.2 即時重算前向張量

7.5.3 將GPU 顯示記憶體下放至CPU 記憶體

7.5.4 降低最佳化器的顯示記憶體佔用

7.6 最佳化Python 程式以減少顯示記憶體佔用

7.6.1 Python 垃圾回收機制

7.6.2 避免出現迴圈依賴

7.6.3 謹慎使用全域作用域

7.7 本章小結

第 8 章 分散式訓練專題

8.1 分散式策略概述

8.2 集合通訊基本操作

8.3 應對資料增長的平行策略

8.3.1 資料平行策略

8.3.2 手動實現資料平行算法

8.3.3 PyTorch 的DDP 封裝

8.3.4 資料平行的C/P 值

8.3.5 其他資料維度的切分

8.4 應對模型增長的平行策略

8.4.1 靜態顯示記憶體切分

8.4.2 動態顯示記憶體切分

8.5 本章小結

第 9 章 高級最佳化方法專題

9.1 自動混合精度訓練

9.1.1 浮點數的表示方法

9.1.2 使用低精度資料型態的優缺點

9.1.3 PyTorch 自動混合精度訓練

9.2 自訂高性能運算元

9.2.1 自訂運算元的封裝流程

9.2.2 自訂運算元的後端程式實現

9.2.3 自訂運算元匯入Python

9.2.4 自訂運算元匯入PyTorch

9.2.5 在Python 中使用自訂運算元

9.3 基於計算圖的性能最佳化

9.3.1 torch.compile 的使用方法

9.3.2 計算圖的提取

9.3.3 圖的最佳化和後端程式生成

9.4 本章小結

第 10 章 GPT-2 最佳化全流程

10.1 GPT 模型結構簡介

10.2 實驗環境與機器配置

10.3 顯示記憶體最佳化

10.3.1 基準模型

10.3.2 使用跨批次梯度累加

10.3.3 開啟即時重算前向張量

10.3.4 使用顯示記憶體友善的最佳化器模式

10.3.5 使用分散式方法降低顯示記憶體佔用—FSDP

10.3.6 顯示記憶體最佳化小結

10.4 性能最佳化

10.4.1 基準模型

10.4.2 增加 BatchSize

10.4.3 增加資料前置處理的平行度

10.4.4 使用非同步介面完成資料傳輸

10.4.5 使用計算圖最佳化

10.4.6 使用float16 混合精度訓練

10.4.7 （可選）使用自訂運算元

10.4.8 使用單機多卡加速訓練

10.4.9 使用多機多卡加速訓練

10.4.10 性能最佳化小結

結語