使用 XGBoost 外部内存版本

概述

当处理大型数据集时，训练 XGBoost 模型可能具有挑战性，因为整个数据集需要加载到主内存中。这可能成本高昂，有时甚至不可行。

外部内存训练有时被称为核外训练。它指的是 XGBoost 能够选择性地将数据缓存到主处理器（无论是 CPU 还是 GPU）外部的位置。XGBoost 本身不支持网络文件系统。因此，对于 CPU，外部内存通常指硬盘。对于 GPU，它指的是主机内存或硬盘。

用户可以定义自定义迭代器以分块加载数据来运行 XGBoost 算法。外部内存可用于训练和预测，但训练是主要用例，我们将在此教程中重点关注它。对于预测和评估，用户可以自行迭代数据，而训练则需要将整个数据集加载到内存中。在模型训练期间，XGBoost 分批获取缓存以构建决策树，从而避免将整个数据集加载到主内存中，并实现更好的垂直扩展（在同一节点内扩展）。

3.0 版本中对 GPU 实现取得了重大进展。我们将在以下章节中介绍 CPU 和 GPU 之间的区别。

注意

exact 树方法不支持从外部内存中训练数据。出于性能考虑，我们建议使用默认的 hist 树方法。

注意

该功能被认为是实验性的，但在 3.0 中已准备好进行公开测试。尚不支持向量叶。

外部内存支持经过了多次开发迭代。请参阅以下章节以了解简史。

数据迭代器

要开始使用外部内存，用户需要定义一个数据迭代器。数据迭代器接口已于 1.5 版本添加到 Python 和 C 接口，并于 3.0.0 版本添加到 R 接口。与带有 DataIter 的 QuantileDMatrix 类似，XGBoost 使用用户提供的自定义迭代器分批加载数据。然而，与 QuantileDMatrix 不同，外部内存不会连接批次（除非通过 GPU 的 extmem_single_page 指定）。相反，它将所有批次缓存到外部内存中并按需获取。请参阅文档末尾以查看 QuantileDMatrix 和外部内存版本的 ExtMemQuantileDMatrix 之间的比较。

一些示例在 demo 目录中，供快速入门。要启用外部内存训练，自定义数据迭代器需要有两个类方法：next 和 reset。

import os
from typing import List, Callable

import numpy as np
import xgboost

class Iterator(xgboost.DataIter):
    """A custom iterator for loading files in batches."""

    def __init__(
        self, device: Literal["cpu", "cuda"], file_paths: List[Tuple[str, str]]
    ) -> None:
        self.device = device

        self._file_paths = file_paths
        self._it = 0
        # XGBoost will generate some cache files under the current directory with the
        # prefix "cache"
        super().__init__(cache_prefix=os.path.join(".", "cache"))

    def load_file(self) -> Tuple[np.ndarray, np.ndarray]:
        """Load a single batch of data."""
        X_path, y_path = self._file_paths[self._it]
        # When the `ExtMemQuantileDMatrix` is used, the device must match. GPU cannot
        # consume CPU input data and vice-versa.
        if self.device == "cpu":
            X = np.load(X_path)
            y = np.load(y_path)
        else:
            import cupy as cp

            X = cp.load(X_path)
            y = cp.load(y_path)

        assert X.shape[0] == y.shape[0]
        return X, y

    def next(self, input_data: Callable) -> bool:
        """Advance the iterator by 1 step and pass the data to XGBoost.  This function
        is called by XGBoost during the construction of ``DMatrix``

        """
        if self._it == len(self._file_paths):
            # return False to let XGBoost know this is the end of iteration
            return False

        # input_data is a keyword-only function passed in by XGBoost and has the similar
        # signature to the ``DMatrix`` constructor.
        X, y = self.load_file()
        input_data(data=X, label=y)
        self._it += 1
        return True

    def reset(self) -> None:
        """Reset the iterator to its beginning"""
        self._it = 0

定义迭代器后，我们可以将其传递给 DMatrix 或 ExtMemQuantileDMatrix 构造函数。

it = Iterator(device="cpu", file_paths=["file_0.npy", "file_1.npy", "file_2.npy"])

# Use the ``ExtMemQuantileDMatrix`` for the hist tree method, recommended.
Xy = xgboost.ExtMemQuantileDMatrix(it)
booster = xgboost.train({"tree_method": "hist"}, Xy)

# The ``approx`` tree method also works, but with lower performance and cannot be used
# with the quantile DMatrix.
Xy = xgboost.DMatrix(it)
booster = xgboost.train({"tree_method": "approx"}, Xy)

上述片段是外部内存的实验性支持的简化版本。有关 C 语言示例，请参阅 demo/c-api/external-memory/。迭代器是使用 XGBoost 外部内存的通用接口，您可以将生成的 DMatrix 对象用于训练、预测和评估。

ExtMemQuantileDMatrix 是 QuantileDMatrix 的外部内存版本。这两个类是专门为 hist 树方法设计的，以减少内存使用和数据加载开销。有关更多信息，请参阅各自的参考文献。

根据可用内存设置批次大小非常重要。对于 CPU，如果您有 64GB 内存，一个好的起点是将批次大小设置为每批 10GB。不建议将批次大小设置得很小，例如每批 32 个样本，因为这会严重损害梯度提升的性能。有关 GPU 版本和其他最佳实践的信息，请参阅以下章节。

GPU 版本（GPU Hist 树方法）

GPU 算法支持外部内存（即，当 device 设置为 cuda 时）。从 3.0 版本开始，默认的 GPU 实现类似于 CPU 版本的实现。它还支持在采用 hist 树方法（默认）时使用 ExtMemQuantileDMatrix。对于 GPU 设备，主内存是设备内存，而外部内存可以是磁盘或 CPU 内存。XGBoost 默认将缓存暂存到 CPU 内存。用户可以通过在 DataIter 中指定 on_host 参数来将后端存储更改为磁盘。但是，不建议使用磁盘，因为它可能会使 GPU 比 CPU 慢。此选项仅用于实验目的。此外，ExtMemQuantileDMatrix 参数 min_cache_page_bytes 和 max_quantile_batches 可以帮助控制数据放置和内存使用。

ExtMemQuantileDMatrix 的输入（通过迭代器）必须在 GPU 上。以下是外部内存的实验性支持中的一个片段。

import cupy as cp
import rmm
from rmm.allocators.cupy import rmm_cupy_allocator

# It's important to use RMM for GPU-based external memory to improve performance.
# If XGBoost is not built with RMM support, a warning will be raised.
# We use the pool memory resource here for simplicity, you can also try the
# `ArenaMemoryResource` for improved memory fragmentation handling.
mr = rmm.mr.PoolMemoryResource(rmm.mr.CudaAsyncMemoryResource())
rmm.mr.set_current_device_resource(mr)
# Set the allocator for cupy as well.
cp.cuda.set_allocator(rmm_cupy_allocator)
# Make sure XGBoost is using RMM for all allocations.
with xgboost.config_context(use_rmm=True):
    # Construct the iterators for ExtMemQuantileDMatrix
    # ...
    # Build the ExtMemQuantileDMatrix and start training
    Xy_train = xgboost.ExtMemQuantileDMatrix(it_train, max_bin=n_bins)
    # Use the training DMatrix as a reference
    Xy_valid = xgboost.ExtMemQuantileDMatrix(it_valid, max_bin=n_bins, ref=Xy_train)
    booster = xgboost.train(
        {
            "tree_method": "hist",
            "max_depth": 6,
            "max_bin": n_bins,
            "device": device,
        },
        Xy_train,
        num_boost_round=n_rounds,
        evals=[(Xy_train, "Train"), (Xy_valid, "Valid")]
    )

在外部内存训练时，所有内存分配都使用 RAPIDS Memory Manager (RMM) 和异步内存资源至关重要。XGBoost 依靠异步内存池来减少数据获取的开销。此外，Heterogeneous memory management (HMM) 支持需要开源的 NVIDIA Linux 驱动程序。通常，用户无需更改 ExtMemQuantileDMatrix 参数，例如 min_cache_page_bytes，它们会根据设备自动配置，并且不会改变模型精度。但是，如果 ExtMemQuantileDMatrix 在构建过程中耗尽设备内存，max_quantile_batches 可能会很有用，请参阅 QuantileDMatrix 和以下章节了解更多信息。目前，我们主要关注支持 NVLink-C2C 的设备，以实现基于 GPU 的外部内存支持。

除了基于批次的数据获取之外，GPU 版本还支持将批次连接成一个单一的 Blob 用于训练数据，以提高性能。对于通过 PCIe 而非 nvlink 连接的 GPU，基于批次的训练性能开销非常大，特别是对于非密集数据。总的来说，它可能比内存内训练慢至少五倍。连接页面可以使性能更接近内存内训练。此选项应与子采样结合使用以减少内存使用。在连接过程中，子采样会移除一部分样本，从而减少训练数据集的大小。GPU hist 树方法支持梯度采样，允许用户设置较低的采样率而不影响精度。在 3.0 之前，带子采样的连接是基于 GPU 的外部内存的唯一选择。3.0 之后，XGBoost 默认使用常规批次获取，而页面连接可以通过以下方式启用：

param = {
  "device": "cuda",
  "extmem_single_page": true,
  'subsample': 0.2,
  'sampling_method': 'gradient_based',
}

有关采样算法及其在外部内存训练中使用的更多信息，请参阅这篇论文。最后，请参阅以下章节以了解最佳实践。

NVLink-C2C

较新的 NVIDIA 平台，如 Grace-Hopper，使用 NVLink-C2C，它促进了 CPU 和 GPU 之间的快速互连。通过主机内存作为数据缓存，XGBoost 可以以显著较低的开销检索数据。当输入数据密集时，训练的性能损失微乎其微甚至没有，除了 ExtMemQuantileDMatrix 的初始构建。初始构建会遍历输入数据两次，因此与内存内训练相比，最大的开销是在数据密集时额外读取一次数据。请注意，该平台有多种变体，它们具有不同的 C2C 带宽。在该功能的初始开发期间，我们使用了 LPDDR5 480G 版本，其主机到设备传输带宽约为 350GB/s。在选择用于训练 XGBoost 模型的变体时，应特别注意 C2C 带宽。

在此，我们提供一个简单的示例作为使用外部内存进行训练的起点。我们将此示例用于其中一个基准测试。在 GH200（通过芯片到芯片链路连接到 Grace CPU 的 H200 GPU）系统上训练一个包含 2 ^ 29 个 32 位浮点样本、512 个特征（总计 1TB）的模型。可以从以下开始： - 将数据均匀分成 128 个批次，每个批次 8GB。 - 按照前面所述定义自定义迭代器。 - 将 ExtMemQuantileDMatrix 的 max_quantile_batches 参数设置为 32（每个子流 256GB 用于量化）。加载数据。 - 使用 device=cuda 开始训练。

要在这些平台上运行实验，需要版本 >=565.47 的开源 NVIDIA Linux 驱动程序，它应该与 CTK 12.7 及更高版本一起提供。最后，Linux 6.11 存在一个已知问题，可能导致 CUDA 主机内存分配失败并出现 invalid argument 错误。

自适应缓存

从 3.1 版本开始，XGBoost 引入了用于基于 GPU 的外部内存训练的自适应缓存。此功能有助于将数据缓存分为主机缓存和设备缓存。通过在 GPU 上保留一部分缓存，当有足够的 GPU 内存时，我们可以减少训练期间的数据传输量。此功能可以通过 xgboost.ExtMemQuantileDMatrix 中的 cache_host_ratio 参数进行控制。当设备具有完整的 C2C 带宽时，此功能被禁用，因为它不需要。在带宽受限或具有 PCIe 连接的设备上，除非明确指定，否则会根据设备内存大小和数据集大小自动估算比率。

然而，此参数增加了内存碎片，因为 XGBoost 需要大小不规则的大内存页面。因此，您可能会在 DMatrix 构建之后但在实际训练开始之前看到内存不足错误。

作为参考，我们使用 NVIDIA A6000 GPU（配备 48GB 设备内存）测试了具有 128GB（512 个特征）密集 32 位浮点数据集的自适应缓存。cache_host_ratio 估计约为 0.3，这意味着约 30% 的量化缓存位于主机上，其余 70% 实际上位于内存中。考虑到这个比率，开销极小。但是，随着数据量的增长，估计的比率也会增加。

非统一内存访问 (NUMA)

在多插槽系统中，NUMA 通过优先访问每个插槽本地的内存来优化数据访问。在这些系统上，设置正确的亲和性以减少跨插槽数据访问的开销至关重要。由于核外训练将数据缓存暂存到主机上并使用 GPU 训练模型，因此训练性能对数据读取带宽特别敏感。举例来说，在 GB200 机器上，从 GPU 访问错误的 NUMA 节点可以将 C2C 带宽降低一半。即使您不使用分布式训练，也应该注意 NUMA 控制，因为无法保证您的进程具有正确的配置。

我们测试了两种 NUMA 配置方法。第一种（也是推荐的）方法是使用 Linux 发行版上可用的 numactl 命令行

numactl --membind=${NODEID} --cpunodebind=${NODEID} ./myapp

要获取节点 ID，您可以通过 nvidia-smi 检查机器拓扑

nvidia-smi topo -m

“NUMA Affinity” 列出了每个 GPU 的 NUMA 节点 ID。在下面显示的示例输出中，GPU0 与 0 节点 ID 相关联。

        GPU0    GPU1    NIC0    NIC1    NIC2    NIC3    CPU Affinity    NUMA Affinity   GPU NUMA ID
GPU0     X      NV18    NODE    NODE    NODE    SYS     0-71            0               2
GPU1    NV18     X      SYS     SYS     SYS     NODE    72-143          1               10
NIC0    NODE    SYS      X      PIX     NODE    SYS
NIC1    NODE    SYS     PIX      X      NODE    SYS
NIC2    NODE    SYS     NODE    NODE     X      SYS
NIC3    SYS     NODE    SYS     SYS     SYS      X

或者，也可以使用 hwloc 命令行界面，请确保使用严格标志。

hwloc-bind --strict --membind node:${NODEID} --cpubind node:${NODEID} ./myapp

另一种方法是使用 CPU 亲和性。dask-cuda 项目通过使用 nvml 库以及 Linux 调度例程为 Dask 接口配置最佳 CPU 亲和性。这有助于指导内存分配策略，但并不强制执行。因此，当内存压力较大时，操作系统可能会在不同的 NUMA 节点上分配内存。另一方面，它更容易使用，因为像 LocalCUDACluster 这样的启动器已经集成了该解决方案。

我们使用第一种方法进行基准测试，因为它具有更好的强制性。

分布式训练

分布式训练与内存内学习类似，但框架集成工作仍在进行中。请参阅分布式外部内存训练的实验性支持，了解使用通信器构建简单管道的示例。由于用户可以定义自己的自定义数据加载器，XGBoost 中现有的分布式框架接口不太可能满足所有用例，该示例可以作为拥有自定义基础设施的用户的起点。

最佳实践

在前面的章节中，我们演示了如何使用驻留在外部内存中的数据训练基于树的模型。此外，我们还就批次大小和 NUMA 提出了一些建议。以下是一些我们认为有用的其他配置。外部内存功能涉及在树构建期间迭代存储在缓存中的数据批次。为了获得最佳性能，我们建议使用 grow_policy=depthwise 设置，它允许 XGBoost 仅通过几次批次迭代就构建整个树节点层。相反，使用 lossguide 策略要求 XGBoost 针对每个树节点迭代数据集，从而导致性能显著降低（树大小与深度呈指数关系）。

此外，应优先选择 hist 树方法而不是 approx 树方法，因为前者不会在每次迭代中重新创建直方图 bin。创建直方图 bin 需要加载原始输入数据，这会非常昂贵。专为 hist 树方法设计的 ExtMemQuantileDMatrix 可以显著加快外部内存的初始数据构建和评估。

由于外部内存的实现侧重于 XGBoost 需要访问整个数据集的训练，因此只有 X 被分成批次，而其他所有数据都连接在一起。因此，建议用户定义自己的管理代码来迭代数据进行推理，特别是对于 SHAP 值计算。SHAP 矩阵的大小可能大于特征矩阵 X，这使得 XGBoost 中的外部内存效率较低。

使用外部内存时，CPU 训练的性能受限于磁盘 I/O（输入/输出）速度。这意味着磁盘 I/O 速度主要决定训练速度。类似地，假设使用 CPU 内存作为缓存且地址转换服务（ATS）不可用，PCIe 带宽限制了 GPU 性能。在开发过程中，我们观察到 XGBoost 中 PCIe4x16 的典型数据传输带宽约为 24GB/s，PCIe5 约为 42GB/s，这显著低于 GPU 处理性能。而对于支持 C2C 的机器，训练中的数据传输和处理性能彼此接近。

运行推理的计算密集度远低于训练，因此速度快得多。因此，推理的性能瓶颈又回到了数据传输。对于 GPU，即使有 C2C 链接可用，从主机读取数据到设备所需的时间也完全决定了运行推理所需的时间。

Xy_train = xgboost.ExtMemQuantileDMatrix(it_train, max_bin=n_bins)
Xy_valid = xgboost.ExtMemQuantileDMatrix(it_valid, max_bin=n_bins, ref=Xy_train)

此外，由于 GPU 实现依赖于异步内存池，即使使用 CudaAsyncMemoryResource，它也容易出现内存碎片。您可能希望从一个全新的内存池开始训练，而不是在 ETL 过程之后立即开始训练。如果您遇到内存不足错误，并且确信内存池尚未满（可以使用 nsight-system 分析内存池使用情况），请考虑使用 ArenaMemoryResource 内存资源。或者，结合 CudaAsyncMemoryResource 和 BinningMemoryResource(mr, 21, 25)，而不是默认的 PoolMemoryResource。

在 CPU 基准测试期间，我们使用了连接到 PCIe-4 插槽的 NVMe。其他类型的存储对于实际使用来说可能太慢了。但是，您的系统可能会执行一些缓存以减少文件读取的开销。请参阅以下部分以了解备注。

备注

当 XGBoost 使用外部内存时，数据被分成较小的块，以便在任何给定时间只需将其中一小部分存储在内存中。重要的是要注意，此方法仅适用于预测器数据 (X)，而其他数据（如标签和内部运行时结构）则会连接起来。这意味着内存减少在处理宽数据集时最有效，其中 X 的大小显著大于其他数据（如 y），而对窄数据集影响甚微。

正如人们可能预期的那样，按需获取数据对存储设备造成了巨大的压力。当今的计算设备处理的数据量远超存储设备在单个时间单位内可以读取的数据量。这个比例是数量级的。GPU 能够在瞬间处理数百 GB 的浮点数据。另一方面，连接到 PCIe-4 插槽的四通道 NVMe 存储通常具有约 6GB/s 的数据传输速率。因此，训练很可能会受到存储设备的严重限制。在采用外部内存解决方案之前，进行一些粗略的计算可能有助于您确定其可行性。例如，如果您的 NVMe 驱动器每秒可以传输 4GB 数据（一个相当实际的数字），并且您的压缩 XGBoost 缓存中有 100GB 数据（对应于一个约 200GB 的密集 float32 numpy 数组）。一个深度为 8 的树在参数最优时需要至少 16 次遍历数据。在不考虑其他开销并假设计算与 I/O 重叠的情况下，您需要大约 14 分钟来训练一棵树。如果您的数据集碰巧达到 TB 级别大小，您可能需要数千棵树才能获得泛化模型。这些计算可以帮助您估算预期的训练时间。

然而，有时我们可以缓解这种限制。还应该考虑操作系统（主要是指 Linux 内核）通常可以将数据缓存到主机内存中。它只在有新数据进入且没有剩余空间时才逐出页面。在实践中，至少一部分数据可以在整个训练会话中保留在主机内存中。我们在优化外部内存获取器时已经意识到这个缓存。压缩缓存通常比原始输入数据小，特别是当输入是密集且没有任何缺失值时。如果主机内存可以容纳这个压缩缓存的很大一部分，那么在初始化之后性能应该不错。我们目前为止的开发主要集中在以下几个方面来优化外部内存：

在适当的时候避免迭代数据。
如果操作系统可以缓存数据，则性能应接近内存内训练。
对于 GPU，实际计算应尽可能与内存复制重叠。

从 XGBoost 2.0 开始，CPU 外部内存实现使用 mmap。它尚未针对断开的网络设备（SIGBUS）等系统错误进行测试。如果发生总线错误，您将看到硬崩溃并需要清理缓存文件。如果训练会话可能需要很长时间并且您使用 NVMe-oF 等解决方案，我们建议您定期检查点模型。此外，值得注意的是，大多数测试都是在 Linux 发行版上进行的。

需要记住的另一点是，为 XGBoost 创建初始缓存可能需要一些时间。外部内存的接口是通过自定义迭代器实现的，我们不能假设它是线程安全的。因此，初始化是按顺序执行的。如果您的输出量没问题，使用 config_context() 并设置 verbosity=2 可以为您提供一些关于 XGBoost 在等待期间正在做什么的信息。

与 QuantileDMatrix 比较

将迭代器传递给 QuantileDMatrix 可以直接通过数据块构建 QuantileDMatrix。另一方面，如果将其传递给 DMatrix 或 ExtMemQuantileDMatrix，它将启用外部内存功能。QuantileDMatrix 在压缩后将数据连接到内存中，并且在训练期间不获取数据。另一方面，外部内存 DMatrix (ExtMemQuantileDMatrix) 根据需要从外部内存中获取数据批次。当您可以将大部分数据放入内存时，请使用 QuantileDMatrix（如有必要，带迭代器）。对于许多平台来说，训练速度可能比外部内存快一个数量级。

简史

长期以来，外部内存支持一直是一个实验性功能，并且经历了多次开发迭代。以下是主要更改的简要总结：

1.1 版本在 GPU hist 中引入了基于梯度的采样。
1.5 版本引入了迭代器接口，并对内部框架进行了重大重写。
2.0 版本引入了 mmap 的使用，并对 XGBoost 进行了优化以实现零拷贝数据获取。
3.0 版本重写了 GPU 实现，以支持在主机和磁盘上缓存数据，引入了 ExtMemQuantileDMatrix 类，并添加了基于分位数的目标支持。
此外，我们在 3.0 版本中开始支持分布式训练。
3.1 版本增加了对分区缓存页面的支持。一个缓存页面可以部分位于 GPU 中，其余部分位于主机内存中。此外，当数据稀疏时，XGBoost 可以与 Grace Blackwell 硬件解压缩引擎协同工作。
文本文件缓存格式已在 3.1.0 中删除。