用 Dask XGBoost 解锁多 GPU 模型训练

作为数据科学家，我们经常面临在大型数据集上训练模型的挑战。一种常用的工具是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

一些用户可能会注意到 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 。

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

总结

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

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

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

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