作為數據科學家,我們經常面臨在大型數據集上訓練模型的挑戰。一種常用的工具是XGBoost,這是一種穩健且高效的梯度提升框架,因其在處理大型表格數據時的速度和性能而被廣泛采用。
理論上,使用多個 GPU 可以顯著提高計算能力,從而加快模型訓練。然而,許多用戶發現,當試圖通過 Dask 和 XGBoost 進行訓練時,Dask 是一個用于并行計算的靈活的開源 Python 庫,而 XGBoost 則提供 Dask API 來訓練 CPU 或 GPU 的 Dask DataFrames。
訓練 Dask XGBoost 的一個常見障礙是處理不同階段的內存不足(OOM)錯誤,包括
- 加載訓練數據
- 將 DataFrame 轉換為 XGBoost 的 DMatrix 格式
- 在實際模型培訓期間
解決這些記憶問題可能很有挑戰性,但非常有益,因為多 GPU 訓練的潛在好處很誘人。
頂級外賣
這篇文章探討了如何在多個 GPU 上優化 Dask XGBoost 并管理內存錯誤。在大型數據集上訓練 XGBoost 帶來了各種挑戰。我使用Otto Group Product Classification Challenge 數據集來演示 OOM 問題以及如何解決它。該數據集有 1.8 億行和 152 列,加載到內存中時總計 110 GB。
我們要解決的關鍵問題包括:
- 使用最新版本的 RAPIDS 以及正確版本的 XGBoost 進行安裝。
- 設置環境變量。
- 處理 OOM 錯誤。
- 利用 UCX-py 實現更多加速。
請務必按照每個章節隨附的筆記本進行操作。
先決條件
利用 RAPIDS 的力量進行多 GPU 訓練的第一步是正確安裝 RAPIDS 庫。需要注意的是,有幾種安裝這些庫的方法 – pip,conda,docker,和從源代碼構建,每種方法都與 Linux 和 Windows Subsystem for Linux 兼容。
每種方法都有其獨特的考慮因素。對于本指南,我建議在遵守 conda 安裝說明的同時使用 Mamba。 Mamba 提供了與 conda 類似的功能,但速度要快得多,尤其是在依賴關系解析方面。具體來說,我選擇了 全新安裝 mamba。
安裝最新的 RAPIDS 版本
作為最佳實踐,我們建議您始終安裝最新的 RAPIDS 庫以使用最新的功能。您可以在 RAPIDS 安裝指南 中找到最新的安裝說明。
這篇文章使用 23.04 版本,可以通過以下命令進行安裝:
mamba create -n rapids-23.04 -c rapidsai -c conda-forge -c nvidia \ rapids=23.04 python=3.10 cudatoolkit=11.8
本說明安裝所有所需的庫,包括 Dask、Dask-cuDF 、XGBoost 等。特別是,您需要檢查使用以下命令安裝的 XGBoost 庫:
曼巴列表 xgboost
輸出如表 1 所示:
名稱 | 版本 | 建筑 | 頻道 |
XGBoost | 1.7.1 月 24 日星期四 | CUDA _11_py310_3 | 夜間急流 |
避免手動更新 XGBoost
一些用戶可能會注意到 XGBoost 的版本不是最新的,它是 1.7.5。使用 pip 或 conda-forge 手動更新或安裝 XGBoost? 與 UCX 一起訓練 XGBoost 時出現問題。
錯誤消息將顯示如下內容:
異常:“XGBoostError(’ 14:14.27 /opt/conda/conda-bld/work/rabit/include/rabit/internal/utils.h:86:Allreduce 失敗’)”
相反,請使用從 RAPIDS 安裝的 XGBoost。驗證 XGBoost 版本正確性的快速方法是曼巴列表 xgboost并檢查 xgboost 的“通道”,該通道應為“rapidsai”或“Rapidsainightly”。
rapidsai 通道中的 XGBoost 是在啟用 RMM 插件的情況下構建的,在多 GPU 訓練方面提供了最佳性能。
多- GPU 訓練演練
首先,我將瀏覽 Otto 數據集的 multi-GPU 訓練筆記本,并介紹使其工作的步驟。稍后,我們將討論一些高級優化,包括 UCX 和溢出。
您還可以在 XGB-186-CLICKS-DASK GitHub 上找到筆記本。或者,我們還提供了一個具有完全命令行可配置性的 python 腳本。
我們要使用的主要庫是 xgboost、dask、dask_ CUDA 和 dask- cuDF 。
import os import dask import dask_cudf import xgboost as xgb from dask.distributed import Client from dask_cuda import LocalCUDACluster
環境設置
首先,讓我們設置我們的環境變量,以確保我們的 GPU 是可見的。此示例使用八個 GPU ,每個 GPU 32 GB 內存,這是在沒有 OOM 復雜性的情況下運行此筆記本的最低要求。在下面的啟用內存溢出一節中,我們將討論將此要求降低到 4 GPU 的技術。
GPUs = ','.join([str(i) for i in range(0,8)]) os.environ['CUDA_VISIBLE_DEVICES'] = GPUs
接下來,定義一個助手函數,為多個 GPU 單個節點創建一個本地 GPU cluster。
def get_cluster(): cluster = LocalCUDACluster() client = Client(cluster) return client
然后,為您的計算創建一個 Dask 客戶端。
client = get_cluster()
正在加載數據
現在,讓我們加載 Otto 數據集。我們將使用 dask_cuDF 的 read_parquet 函數,該函數利用多個 GPU 將 parquet 文件讀取到 dask_cuDF.DataFrame 中。
users = dask_cudf.read_parquet('/raid/otto/Otto-Comp/pqs/train_v152_*.pq').persist()
該數據集由 152 列組成,這些列代表工程特征,提供了關于用戶查看或購買特定產品的頻率信息。目標是根據用戶的瀏覽歷史來預測用戶下一步會點擊哪個產品。您可以在writeup中查看此數據集的詳細信息。
即使在這個早期階段,也可能出現內存不足的錯誤。這個問題通常是由于鑲木地板銼刀。為了解決這個問題,我們建議使用較小的行組來重寫鑲木地板文件。如果您想了解更深入的解釋,請參閱 Parquet Large Row Group Demo 筆記本。
加載數據后,我們可以檢查其形狀和內存使用情況。
users.shape[0].compute() users.memory_usage().sum().compute()/2**30
“點擊”列是我們的目標,這意味著如果用戶點擊了推薦的項目。我們忽略 ID 列,并使用其余列作為功能。
FEATURES = users.columns[2:] TARS = ['clicks'] FEATURES = [f for f in FEATURES if f not in TARS]
接下來,我們創建了一個用于訓練 xgboost 模型的輸入數據格式,即DaskQuantileDMatrix。當使用直方圖樹方法時,DaskQuantileDMatrix是 DaskDMatrix 的一個替代品,它有助于減少總體內存使用量。
這一步驟對于避免 OOM 錯誤至關重要。如果我們使用 DaskDMatrix,即使使用 16 個 GPU 也會發生 OOM 錯誤。相反,DaskQuantileDMatrix 能夠在沒有 OOM 錯誤的情況下以八個或更少的 GPU 訓練 xgboot。
dtrain = xgb.dask.DaskQuantileDMatrix(client, users[FEATURES], users['clicks'])
XGBoost 模型訓練
然后,我們設置 XGBoost 模型參數并開始訓練過程。給定目標列“點擊”是二進制的,我們使用二進制分類目標。
xgb_parms = { 'max_depth':4, 'learning_rate':0.1, 'subsample':0.7, 'colsample_bytree':0.5, 'eval_metric':'map', 'objective':'binary:logistic', 'scale_pos_weight':8, 'tree_method':'gpu_hist', 'random_state':42 }
現在,您已經準備好使用全部八個 GPU 來訓練 XGBoost 模型了。
輸出:
[99] train-map:0.20168 CPU times: user 7.45 s, sys: 1.93 s, total: 9.38 s Wall time: 1min 10s
就是這樣!您已經完成了使用多個 GPU 訓練 XGBoost 模型。
啟用內存溢出
在上一個XGB-186-CLICKS-DASK筆記本電腦中,我們在 Otto 數據集上訓練 XGBoost 模型,至少需要八個 GPU。假設該數據集占用 110GB 的內存,而每個 V100 GPU 提供 32GB,那么數據與 GPU 內存的比率僅為 43%(計算為 110/(32*8))。
最理想的情況是,我們只需要使用四個 GPU 就可以將其減半。然而,在我們之前的設置中直接減少 GPU 總是會導致 OOM 錯誤。此問題源于創建生成DaskQuantileDMatrix從 Dask cuDF 數據幀和在訓練 XGBoost 的其他步驟中。這些變量本身消耗了 GPU 內存的相當大的份額。
優化相同的 GPU 資源以訓練更大的數據集
在XGB-186-CLICKS-DASK-SPILL 筆記本電腦中,我介紹了之前設置的一些小調整。現在,通過啟用溢出,您只需使用四個 GPU 就可以在同一數據集上進行訓練。這項技術允許您使用相同的 GPU 資源來訓練更大的數據。
溢出是一種技術,當 GPU 內存中所需的空間被其他數據幀或序列占用,導致原本會成功的操作耗盡內存時,該技術會自動移動數據。它能夠在不適合內存的數據集上進行核心外計算。 RAPIDS cuDF 和 dask- cuDF 現在支持從 GPU 溢出到 CPU 內存。
實現溢出非常容易,我們只需要用兩個新參數重新配置集群,設備內存限制和jit_unspill:
def get_cluster(): ip = get_ip() cluster = LocalCUDACluster(ip=ip, device_memory_limit='10GB', jit_unspill=True) client = Client(cluster) return client
device_memory_limit=’10GB’設置在觸發溢出之前每個 GPU 可以使用的 GPU memory 的量的限制。我們的配置有意為設備內存限制10GB,基本上小于 GPU 的總 32GB。這是一種深思熟慮的策略,旨在搶占 XGBoost 訓練期間的 OOM 錯誤。
同樣重要的是要理解 XGBoost 的內存使用不是由 Dask-CUDA 或 Dask-cuDF 直接管理的。因此,為了防止內存溢出,Dask- CUDA 和 Dask- cuDF 需要在 XGBoost 操作達到內存限制之前啟動溢出過程。
Jit_unspill啟用實時取消溢出,這意味著當 GPU 內存不足時,集群將自動將數據從 GPU’內存溢出到主內存,并及時取消溢出以進行計算。
就這樣!筆記本的其余部分與上一個筆記本相同。現在,它只需四個 GPU 就可以進行訓練,節省了 50%的計算資源。
想要了解更多詳細信息,請參閱 XGB-186-CLICKS-DASK-SPILL 筆記本。
使用統一通信 X(UCX)實現最佳數據傳輸
UCX-py 是一種高性能通信協議,特別適用于 GPU – GPU 通信,能夠提供優化的數據傳輸能力。
為了有效地使用 UCX,我們需要設置另一個環境變量RAPIDS _NO_INITIALIZE:
os.environ["RAPIDS_NO_INITIALIZE"] = "1"
它阻止 cuDF 在導入時運行各種診斷,這需要創建 NVIDIA CUDA 上下文。當運行分布式并使用 UCX 時,我們必須在創建 CUDA 上下文之前打開網絡堆棧(由于各種原因)。通過設置該環境變量,導入 cuDF 的任何子進程都不會在 UCX 有機會創建 CUDA 上下文之前創建。
重新配置群集:
def get_cluster(): ip = get_ip() cluster = LocalCUDACluster(ip=ip, device_memory_limit='10GB', jit_unspill=True, protocol="ucx", rmm_pool_size="29GB" ) client = Client(cluster) return client
protocol=’cx’參數將 ucx 指定為用于在集群中的工作程序之間傳輸數據的通信協議。
使用 prmm_pool_size=’29GB’參數為每個工作線程設置 RAPIDS 內存管理器(RMM)池的大小。RMM 允許有效地使用 GPU 存儲器。在這種情況下,池大小被設置為 29GB,這小于 32GB 的總 GPU 內存大小。這種調整至關重要,因為它解釋了 XGBoost 創建某些存在于 RMM 池控制之外的中間變量的事實。
通過簡單地啟用 UCX,我們的訓練時間得到了顯著的加速——在溢出時,速度提高了 20%,而在不需要溢出時,速度提高了 40.7%。請參閱XGB-186-CLICKS-DASK-UCX-SPILL筆記本了解詳細信息。
配置本地目錄
有時會出現警告消息,例如“UserWarning:創建臨時目錄花費了驚人的長時間。”?磁盤性能正成為一個瓶頸。
為了規避這個問題,我們可以設置本地目錄屬于dask-CUDA,指定本地計算機上存儲臨時文件的路徑。這些臨時文件在 Dask 的磁盤溢出操作中使用。
建議的做法是設置本地目錄到快速存儲設備上的某個位置。例如,我們可以設置本地目錄到/raid/dask_dir如果它在高速本地 SSD 上。進行這種簡單的更改可以顯著減少臨時目錄操作所需的時間,從而優化您的整體工作流程。
最終的集群配置如下:
def get_cluster(): ip = get_ip() cluster = LocalCUDACluster(ip=ip, local_directory=’/raid/dask_dir’ device_memory_limit='10GB', jit_unspill=True, protocol="ucx", rmm_pool_size="29GB" ) client = Client(cluster) return client
結果
如表 2 所示,兩種主要的優化技術是 UCX 和溢出。我們設法用四個 GPU 和 128GB 的內存來訓練 XGBoost。我們還將很好地展示性能擴展到更多 GPU 。
? | 溢出 | 溢出 |
UCX 關閉 | 135s/8 GPU /256 GB | 270s/4 GPU /128 GB |
UCX 打開 | 80s/8 GPU /256 GB | 217s/4 GPU /128 GB |
在每個單元中,這些數字表示端到端執行時間、所需的最小 GPU 數量以及可用的總 GPU memory。所有四個演示都完成了加載和訓練 110 GB Otto 數據的相同任務。
總結
總之,利用 Dask 和 XGBoost 的多個 GPU 可能是一次激動人心的冒險,盡管偶爾會出現內存不足等問題。
您可以通過以下方式緩解這些記憶挑戰,并挖掘多 GPU 模型訓練的潛力:
- 在輸入鑲木地板文件中仔細配置行組大小等參數
- 確保 RAPIDS 和 XGBoost 的正確安裝
- 利用 Dask Quantile DMatrix
- 啟用溢出
此外,通過應用 UCX-Py 等高級功能,您可以顯著加快訓練時間。