XGB-13:使用 XGBoost 外部内存版本

在处理大型数据集时，训练 XGBoost 模型可能会面临挑战，因为整个数据集需要加载到内存中。这可能成本高昂，有时也难以实现。从版本 1.5 开始，用户可以定义自定义迭代器以按块加载数据来运行 XGBoost 算法。外部内存可以用于训练和预测，但训练是主要的用例。对于预测和评估，用户可以在训练时自行迭代数据，而训练需要将完整数据集加载到内存中。

在训练过程中，XGBoost 提供了两种不同的外部内存支持模式：

• 一种用于基于 CPU 的算法，如 hist 和 approx
• 另一种用于基于 GPU 的训练算法

注意

不支持使用 exact 树方法对外部内存中的数据进行训练。

外部内存支持经历了多次迭代，并仍在积极开发中。与使用 DataIter 的 QuantileDMatrix 一样，XGBoost 使用用户提供的自定义迭代器逐批加载数据。但与 QuantileDMatrix 不同的是，除非使用 GPU（它采用混合方法），外部内存不会连接这些批次。相反，它会将所有批次缓存到外部内存中，并按需获取它们。

数据迭代器Data Iterator

从 XGBoost 1.5 开始，用户可以使用 Python 或 C 接口定义自己的数据加载器。在demo目录中有一些示例供快速开始。这是文本输入外部内存的一个通用版本，用户不再需要准备 XGBoost 可识别的文本文件。要启用此功能，用户需要定义一个带有 2 个类方法（next 和 reset）的数据迭代器，然后将其传递给 DMatrix构造函数。

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]):
    self._file_paths = svm_file_paths
    self._it = 0
    # XGBoost will generate some cache files under current directory with the prefix
    # "cache"
    super().__init__(cache_prefix=os.path.join(".", "cache"))

  def next(self, input_data: Callable):
    """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 0 to let XGBoost know this is the end of iteration
      return 0

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

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

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

# The ``approx`` also work, but with low performance. GPU implementation is different from CPU.
# as noted in following sections.
booster = xgboost.train({"tree_method": "hist"}, Xy)

上面的片段是 Experimental support for external memory 的简化版本。有关 C 语言的示例，请参阅 demo/c-api/external-memory/。迭代器是在 XGBoost 中使用外部内存的通用接口，可以将生成的 DMatrix 对象传递给训练、预测和评估。

根据可用内存设置批处理大小是很重要的。如果有 64GB 的内存，一个好的起点是将批处理大小设置为每批 10GB。不建议将批处理大小设置为小批次，例如每批 32 个样本，因为这可能会严重影响梯度提升的性能。

CPU 版本

在前一节中，演示了如何使用 CPU 上的 hist 树方法训练基于树的模型。该方法涉及在树构建期间迭代存储在缓存中的数据批次。为了获得最佳性能，建议使用 grow_policy=depthwise 设置，该设置允许 XGBoost 仅通过少量批次迭代就能构建整个树节点层。相反，使用 lossguide 策略需要 XGBoost 为每个树节点迭代数据集，导致性能较慢。

如果使用外部内存，则 CPU 训练的性能受 IO（输入/输出）速度限制。这意味着磁盘 IO 速度主要确定训练速度。在基准测试期间，使用了连接到 PCIe-4 插槽的 NVMe，其他类型的存储对于实际使用可能太慢。此外，系统可能执行缓存以减少文件读取的开销。

GPU 版本（GPU Hist 树方法）

外部内存受 GPU 算法支持（即当设备device设置为 cuda 时）。然而，用于 GPU 的算法与用于 CPU 的算法不同。在 CPU 上训练时，树方法在树构建算法的每一步中迭代来自外部内存的所有批次。另一方面，GPU 算法使用混合方法。在每次迭代的开始时，它会迭代数据并将所有批次连接到 GPU 内存中以提高性能。为了减少总体内存使用，用户可以利用子采样。GPU hist 树方法支持基于梯度的采样，使用户可以设置低采样率而不影响准确性。

param = {
  ...
  'subsample': 0.2,
  'sampling_method': 'gradient_based',
}

注意：

在 GPU 在迭代外部内存时内存不足时，用户可能会收到段错误segfault 而不是内存不足OOM 异常。

备注

在使用 XGBoost 的外部内存时，数据被划分为较小的块，以便在任何给定时间只需要将其存储在内存中的一部分。值得注意的是，此方法仅适用于预测数据（X），而其他数据，如标签和内部运行时结构则是连接的。这意味着在处理 X 明显比其他数据（如 y）大得多的宽数据集时，内存减少效果最显著。

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

然而，有时可以改善这个限制。人们还应考虑到，操作系统（主要是指 Linux 内核）通常可以将数据缓存到主机内存中。只有在新数据进入并且没有剩余空间时，它才会移除页面。实际上，至少部分数据可以在整个训练会话期间持久驻留在主机内存中。在优化外部内存获取器时，意识到了这个缓存。压缩缓存通常比原始输入数据小，特别是当输入是稠密的且没有任何缺失值时。如果主机内存可以容纳这个压缩缓存的相当部分，那么性能在初始化后应该是不错的。到目前为止，开发主要集中在外部内存的两个优化方面：

1. 在适当的时候避免对数据进行迭代
2. 如果操作系统可以缓存数据，则性能应该接近内存中训练的性能

从 XGBoost 2.0 开始，外部内存的实现使用了 mmap。它没有针对诸如断开的网络设备（SIGBUS）等系统错误进行测试。在出现总线错误的情况下，将看到一个严重的崩溃，并且需要清理缓存文件。如果训练会话可能需要很长时间，并且正在使用类似 NVMe-oF 的解决方案，建议定期对模型进行检查点。另外，值得注意的是，大多数测试都是在 Linux 发行版上进行的。

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

与 QuantileDMatrix 相比

将迭代器传递给 QuantileDmatrix 允许直接使用数据块构建 QuantileDmatrix。另一方面，如果将其传递给 DMatrix，则会启用外部内存功能。QuantileDmatrix 在压缩后在内存中连接数据，并且在训练期间不提取数据。另一方面，外部内存 DMatrix 根据需要从外部内存中提取数据批次。当大部分数据都能放入内存中时，尽可能地使用 QuantileDMatrix，训练将比使用外部内存快一个数量级。

文本文件输入

这是外部内存支持的原始形式，建议用户改用自定义数据迭代器。使用外部内存版本的文本输入和内存版本之间没有太大区别。唯一的区别在于文件名的格式。

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

filename?format=libsvm#cacheprefix

filename是要加载的 LIBSVM 格式文件的正常路径，cacheprefix 是 XGBoost 将用于缓存预处理数据的二进制形式的缓存文件的路径。

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

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

其中 label_column 应该指向充当标签的 csv 列。

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

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

XGBoost 首先会加载 agaricus.txt.train，对其进行预处理，然后写入一个名为 dtrain.cache 的新文件，作为一个用于存储内部二进制格式预处理数据的磁盘缓存。

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

参考

• Out-of-Core GPU Gradient Boosting
• https://xgboost.readthedocs.io/en/latest/tutorials/external_memory.html
• https://xgboost.readthedocs.io/en/latest/tutorials/input_format.html