使用 XGBoost 外部存储器版本

处理大型数据集时，训练 XGBoost 模型可能具有挑战性，因为整个数据集需要加载到内存中。这可能成本高昂且有时不可行。从 1.5 版本开始，用户可以定义自定义迭代器来分块加载数据以运行 XGBoost 算法。外部存储器可用于训练和预测，但训练是主要用例，也将是我们本教程的重点。对于预测和评估，用户可以自行遍历数据，而训练需要将整个数据集加载到内存中。在 3.0 版本中，GPU 实现取得了重大进展。我们将在以下部分介绍 CPU 和 GPU 之间的差异。

注意

外部存储器中的数据不支持 exact 树方法训练。

注意

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

外部存储器支持经历了多次开发迭代。与使用 DataIter 的 QuantileDMatrix 不同，XGBoost 使用用户提供的自定义迭代器批量加载数据。然而，与 QuantileDMatrix 不同，外部存储器不会连接批次（除非由 extmem_single_page 指定）。相反，它将所有批次缓存在外部存储器中并按需获取。转到文档末尾查看 QuantileDMatrix 与 ExtMemQuantileDMatrix 的外部存储器版本的比较。

目录

数据迭代器

从 XGBoost 1.5 版本开始，用户可以使用 Python 或 C 接口定义自己的数据加载器。 demo 目录中有一些示例可供快速入门。要启用外部存储器训练，用户需要定义一个带有 2 个类方法的 数据迭代器：next 和 reset，然后将其传递到 DMatrix 或 ExtMemQuantileDMatrix 构造函数。

import os
from typing import List, Callable
import xgboost
from sklearn.datasets import load_svmlight_file

class Iterator(xgboost.DataIter):
  def __init__(self, svm_file_paths: List[str]) -> None:
    self._file_paths = svm_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 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 the iteration
      return False

    # input_data is a function passed in by XGBoost and has the exact same signature of
    # ``DMatrix``
    X, y = load_svmlight_file(self._file_paths[self._it])
    # Keyword-only arguments, see the ``DMatrix`` class for accepted arguments.
    input_data(data=X, label=y)
    self._it += 1
    # Return True to let XGBoost know we haven't seen all the files yet.
    return True

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

it = Iterator(["file_0.svm", "file_1.svm", "file_2.svm"])

# Use the ``ExtMemQuantileDMatrix`` for the hist tree method.
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 参数 max_num_device_pages、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, 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 依赖于异步内存池来减少数据获取的开销。此外，支持 异构内存管理 (HMM) 需要开源 NVIDIA Linux 驱动程序。通常，用户无需更改 ExtMemQuantileDMatrix 参数 max_num_device_pages 和 min_cache_page_bytes，它们会根据设备自动配置，并且不会改变模型精度。但是，如果在构建过程中 ExtMemQuantileDMatrix 设备内存不足，max_quantile_batches 会很有用，有关更多信息，请参阅 QuantileDMatrix 和以下部分。

除了基于批处理的数据获取外，GPU 版本还支持将批处理数据连接成单个块，以提高训练数据的性能。对于通过 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',
}

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

分布式训练

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

最佳实践

在前面的部分中，我们演示了如何使用驻留在外部存储器上的数据训练基于树的模型，并对批量大小提出了一些建议。这里有一些我们认为有用的其他配置。外部存储器功能涉及在树构建期间遍历存储在缓存中的数据批次。为了获得最佳性能，我们建议使用 grow_policy=depthwise 设置，这使得 XGBoost 只需几次批量迭代即可构建完整的树节点层。相反，使用 lossguide 策略要求 XGBoost 对每个树节点遍历数据集，从而导致性能显着降低。

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

由于外部存储器实现侧重于需要 XGBoost 访问整个数据集的训练，因此只有 X 被划分为批次，而其他所有内容都被连接起来。因此，建议用户定义自己的管理代码来遍历数据进行推理，特别是对于 SHAP 值计算。SHAP 结果的大小可能大于 X，从而降低了 XGBoost 中外部存储器的效果。一些框架（如 dask）可以帮助进行数据分块并遍历数据进行推理，同时进行内存溢出。

使用外部存储器时，CPU 训练的性能受限于磁盘 IO（输入/输出）速度。这意味着磁盘 IO 速度主要决定了训练速度。类似地，假设 CPU 内存用作缓存且地址转换服务 (ATS) 不可用，则 PCIe 带宽限制了 GPU 性能。在开发过程中，我们观察到 XGBoost 在 PCIe4x16 上的典型数据传输带宽约为 24GB/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 对内存池使用情况进行分析），请考虑调整 RMM 内存资源，例如结合使用 CudaAsyncMemoryResource 和 BinningMemoryResource(mr, 21, 25)，而不是使用 PoolMemoryResource。另外，ArenaMemoryResource 也是一个很好的选择。

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

注意事项

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

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

然而，有时我们可以减轻这种限制。还应考虑操作系统（主要指 Linux 内核）通常可以将数据缓存在主机内存中。它只有在新数据进入且没有剩余空间时才会逐出页面。实际上，至少一部分数据可以在整个训练过程中持续驻留在主机内存中。我们在优化外部存储器获取器时意识到了这个缓存。压缩缓存通常小于原始输入数据，特别是当输入是密集且没有缺失值时。如果主机内存可以容纳此压缩缓存的很大一部分，则初始化后性能应该不错。我们目前在外部存储器优化方面主要集中在以下几个方面：

  • 在适当的情况下避免遍历数据。

  • 如果操作系统可以缓存数据，性能应该接近内存中训练。

  • 对于 GPU，实际计算应尽可能与内存复制重叠。

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

另一个需要记住的重要一点是，为 XGBoost 创建初始缓存可能需要一些时间。外部存储器的接口是通过自定义迭代器，我们不能假定它是线程安全的。因此，初始化是按顺序执行的。如果您不介意额外的输出，使用带有 verbosity=2 的 config_context() 可以让您了解 XGBoost 在等待期间正在做什么。

与 QuantileDMatrix 比较

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

简要历史

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

  • 基于梯度的采样在 1.1 版本中引入到 GPU hist。

  • 迭代器接口在 1.5 版本中引入，同时对内部框架进行了重大重写。

  • 2.0 引入了 mmap 的使用，并优化了 XBGoost 以实现零拷贝数据获取。

  • 3.0 重做了 GPU 实现，以支持在主机和磁盘上缓存数据，引入了 ExtMemQuantileDMatrix 类，添加了基于分位数的目标支持。

  • 此外，我们在 3.0 版本中开始支持分布式训练。

文本文件输入

警告

这是 1.5 版本之前的外部存储器支持的原始形式，现已弃用，建议用户使用自定义数据迭代器代替。

使用外部存储器版本的文本输入与内存中版本的文本输入没有显着区别。唯一的区别是文件名格式。

外部存储器版本采用以下 URI 格式：

filename?format=libsvm#cacheprefix

filename 是您想要加载的 LIBSVM 格式文件的典型路径，而 cacheprefix 是 XGBoost 用于缓存预处理二进制数据的缓存文件路径。

要从 csv 文件加载，请使用以下语法：

filename.csv?format=csv&label_column=0#cacheprefix

其中 label_column 应指向作为标签的 csv 列。

如果您有一个采用 LIBSVM 格式、存储在类似 demo/data/agaricus.txt.train 文件中的数据集，可以通过以下方式启用外部存储器支持：

dtrain = DMatrix('../data/agaricus.txt.train?format=libsvm#dtrain.cache')

XGBoost 将首先加载 agaricus.txt.train，对其进行预处理，然后将其写入名为 dtrain.cache 的新文件，作为磁盘缓存，用于存储内部二进制格式的预处理数据。有关文本输入格式的更多注意事项，请参阅DMatrix 的文本输入格式。

对于 CLI 版本，只需添加缓存后缀，例如 "../data/agaricus.txt.train?format=libsvm#dtrain.cache"。