用 Dask XGBoost 解锁多 GPU 模型训练

这篇文章最初发表在 NVIDIA 技术博客上。

作为数据科学家，我们经常面临在大型数据集上训练模型的挑战。一种常用的工具是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 版本，可以通过以下命令进行安装：

r 复制代码

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	夜间急流

表 1。安装正确的 XGBoost，其通道应为 rapidsai 或 rapidsai 夜间

避免手动更新 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 。

javascript 复制代码

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 的技术。

scss 复制代码

GPUs = ','.join([str(i) for i in range(0,8)])
os.environ['CUDA_VISIBLE_DEVICES'] = GPUs

接下来，定义一个助手函数，为多个 GPU 单个节点创建一个本地 GPU cluster。

ini 复制代码

def get_cluster():

    cluster = LocalCUDACluster()

    client = Client(cluster)

    return client

然后，为您的计算创建一个 Dask 客户端。

ini 复制代码

client = get_cluster()

正在加载数据

现在，让我们加载 Otto 数据集。我们将使用 dask_cuDF 的 read_parquet 函数，该函数利用多个 GPU 将 parquet 文件读取到 dask_cuDF.DataFrame 中。

ini 复制代码

users = dask_cudf.read_parquet('/raid/otto/Otto-Comp/pqs/train_v152_*.pq').persist()

该数据集由 152 列组成，这些列代表工程特征，提供了关于用户查看或购买特定产品的频率信息。目标是根据用户的浏览历史来预测用户下一步会点击哪个产品。您可以在writeup中查看此数据集的详细信息。

即使在这个早期阶段，也可能出现内存不足的错误。这个问题通常是由于镶木地板锉刀。为了解决这个问题，我们建议使用较小的行组来重写镶木地板文件。如果您想了解更深入的解释，请参阅 Parquet Large Row Group Demo 笔记本。

加载数据后，我们可以检查其形状和内存使用情况。

scss 复制代码

users.shape[0].compute()

users.memory_usage().sum().compute()/2**30

"点击"列是我们的目标，这意味着如果用户点击了推荐的项目。我们忽略 ID 列，并使用其余列作为功能。

ini 复制代码

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。

bash 复制代码

dtrain = xgb.dask.DaskQuantileDMatrix(client, users[FEATURES], users['clicks'])

XGBoost 模型训练

然后，我们设置 XGBoost 模型参数并开始训练过程。给定目标列"点击"是二进制的，我们使用二进制分类目标。

arduino 复制代码

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 模型了。

输出：

yaml 复制代码

[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：

ini 复制代码

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：

lua 复制代码

os.environ["RAPIDS_NO_INITIALIZE"] = "1"

它阻止 cuDF 在导入时运行各种诊断，这需要创建 NVIDIA CUDA 上下文。当运行分布式并使用 UCX 时，我们必须在创建 CUDA 上下文之前打开网络堆栈（由于各种原因）。通过设置该环境变量，导入 cuDF 的任何子进程都不会在 UCX 有机会创建 CUDA 上下文之前创建。

重新配置群集：

ini 复制代码

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 上。进行这种简单的更改可以显著减少临时目录操作所需的时间，从而优化您的整体工作流程。

最终的集群配置如下：

ini 复制代码

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

表 2。四种优化组合概述

在每个单元中，这些数字表示端到端执行时间、所需的最小 GPU 数量以及可用的总 GPU memory。所有四个演示都完成了加载和训练 110 GB Otto 数据的相同任务。

总结

总之，利用 Dask 和 XGBoost 的多个 GPU 可能是一次激动人心的冒险，尽管偶尔会出现内存不足等问题。

您可以通过以下方式缓解这些记忆挑战，并挖掘多 GPU 模型训练的潜力：

在输入镶木地板文件中仔细配置行组大小等参数
确保 RAPIDS 和 XGBoost 的正确安装
利用 Dask Quantile DMatrix
启用溢出

此外，通过应用 UCX-Py 等高级功能，您可以显著加快训练时间。

阅读原文