在Amazon SageMaker上进行XGBoost分布式训练

背景介绍

XGBoost是一个非常受欢迎的梯度提升决策树算法（Gradient Boosting Decision Tree）开源库，这个开源库的算法Python实现是可以在多台服务器的实例上进行分布式训练的。这一特性对于我们有大量数据要进行训练的场景非常重要，尤其是目前越来越多的组织机构开始在内部采用多种多样的推荐系统，很多算法工程师在起步阶段都会尝试使用梯度提升决策树的算法，而往往XGBoost是他们的首选开源算法框架。数据量大主要表现为数据条目或者通常讲的记录数量多，和数据维度大或者通常讲的列数多。这就导致我们很难在一台服务器上完成训练任务，因为每台物理机或者虚拟机都是有内存限制的，不可能会将所有的数据加载到内存中。同时很多训练任务对训练时间的要求也是很高的，如果一个训练任务无法在规定的时间内完成就有可能造成我们整个机器学习项目的失败。

Amazon SageMaker是亚马逊云科技推出面向开发者和数据科学家快速准备构建、训练和部署机器学习（ML）模型的完全托管的云服务。而在SageMaker中就已经内置了XGBoost算法，这使我们的客户非常容易的在云服务上来进行XGBoost的模型训练，而对于XGBoost的分布式训练也是原生支持的，可以让我们的训练任务扩展到多个计算实例上。

虽然SageMaker可以替我们处理基础架构上的一些工作，实际上XGBoost的分布式训练并不是简单的将训练的机器数量增加就可以了。这其中还是需要我们算法工程师在使用过程中对数据集、超级参数等几个方面进行调整才能真正实现分布式训练，而且能够达到横向资源扩展以及降低训练时间的目标。本文将探讨如何使用SageMaker的内置XGBoost算法来实现分布式训练，以及需要在这个过程中需要如何调整各种数据集和超级参数等配置。

XGBoost分布式训练

在进行XGBoost分布式训练之前让我们简单了解一下它的基本原理是什么，这样能够帮助我们理解在SageMaker上如何运行分布式训练。XGBoost是一种梯度提升模型（GBM）的实现，先进性主要体现在GBM是构建在顺序树上，而XGBoost是构建在并行树上，这就让XGBoost更快。因此XGBoost从诞生之日起就解决了所谓的分布式训练的问题，这些好处包括了：

并行化的树构建方法，可以在训练过程中利用所有的CPU核，这样就能在单个服务器但是多核多CPU的机器上也能得到训练速度的明显提升。
训练十分庞大的模型时使用计算集群进行分布式计算，这就是我们在本文中使用SageMaker要实现的。
对于无法放入单个服务器的内存中的大数据集可以进行横向扩展，这也是要在本文中要实现的目标。
可以实现数据结构的缓存优化，可以令算法能够更好利用硬件的计算能力。
正则化训练目标，训练目标通常包含损失函数与正则化项的合计值，XGBoost恰当的设计了正则化项有效的缓解了模型的复杂性从而避免了过拟合。

基于对XGBoost算法的基本了解，对于分布式训练我们需要注意以下两个问题：第一个是关于数据集的分片的问题，另外一个是XGBoost算法框架中支持的树算法的选择问题。

首先是数据集分片问题。我们通常意义上的数据集切分一般是指训练集、测试集和验证集的切分，而很多算法的分布式训练还需要将训练集按照一定粒度进行再切分才能发挥分布式训练的威力，在这里我们称之为数据集分片。XGBoost的分布式训练就需要对训练数据集按等分的方式进行分片，而分片的份数通常是需要大于等于分布式训练节点且应该是节点数的倍数，这样才能保证在每个节点上的数据量大体相当，防止出现数据倾斜。训练集数据切分是需要在训练的数据处理过程中完成的，后面我们会介绍如何使用SageMaker来方便实现这些分片数据的shard功能。

其次是树算法的选择问题。XGBoost开源算法框架中对于提升树算法的选择有几种，不是所有的树算法都支持分布式训练的，其中对于XGBoost树算法的选择官方有详细的文档加以说明，请参考这个链接。XGBoost使用两个超级参数updater和tree_method来控制对于树算法的选择（依据目前最新版本1.5.1）。主要有下面两种解决方案及其相应的参数设置：

Exact解决方案：可以将tree_method直接设置为exact，就是原版的梯度提升树算法，每次树的分裂都要在整个输入的数据集基础上开始迭代，这种算法虽然更精确但是不能并行，因此这种算法不支持分布式训练。
Approximated解决方案：这个是在原版Exact的提升树算法基础之上演化出来近似训练算法，这种算法会在每个树节点上构筑梯度直方图，并在直方图的基础上进行迭代替换掉原来的整个数据集迭代方式。因此只有这种Approximated方式才能支持分布式训练。要设置为这种方式也有一些具体算法实现差异。
1. 将updater参数设置为grow_local_histmaker，这种方式在实践中不太常用，因此只是在更新方式（updater）的超级参数来设置而不是一个树的算法（tree_method）。在寻找树分裂点时会在当前节点上先使用基于权重的Quantile Sketch算法来寻找候选的点位，并使用hessian矩阵作为权重。由于直方图是在每一个树节点上都要构建一次，因此这种方式在某些情况下会比exact方式快，但是仍然还是计算比较慢，这种方式由于在每个节点上使用局部候选节点较少因此比较适合生长非常深的树。具体可以参看论文原文。
2. 将tree_method参数设置为approx，这种方式是在XGBoost论文中提到的第二种方法，它与grow_local_histmaker方式不同的地方就是不会在每个树节点上去寻找候选节点，而是在构建每个树的时候使用所有的数据行来计算基于权重的Quantile Sketch，其余的做法与grow_local_histmaker方式完全一致。这种使用全局直方图的方式明显不需要多次计算候选节点，但一次要获取比较多的分裂候选节点供后续树的生长使用。
3. 将tree_method参数设置为hist，这种方式与LightGBM算法计算近似树的方式实现非常相似。这种算法只是在训练之前对所有数据进行sketch，并且使用用户提供的权重而不是hessian矩阵，而在接下来的训练过程中每个树节点的直方图都是以这个全局sketch基础上构建的。因为只运行一次sketch因此这种树的算法是最快的。
4. 将tree_method参数设置为gpu_hist，这种方式就是hist方式的GPU实现。

从上面树算法的选择上我们可以得到的结论是，除了Exact这种方式以外我们都可以使用分布式训练的方式来提升训练速度并横向扩展训练资源。

SageMaker分布式训练设置

实验数据集

对于数据集我们将使用Fashion MNIST。这个数据集解决问题的本身是对一些衣服小图片（28×28像素）进行分类，分别是1到10不同的类别（对应t-shirt、trouser、pullover等）。这个数据集的训练集有6万张图片还有1万张验证集。

结合分布式训练的目的是为了将训练扩展到大数据集并降低训练时间，因此为了扩大训练集我们将通过复制训练数据的方式来扩展训练时间以及计算资源。训练数据会按照一定的“重复因子”的倍数进行复制，当然在现实生产环境下我们绝对不会通过这种方式来提升模型的性能或者带来任何其他的好处，除非在这个复制过程中使用一些CV领域内的数据增强的技术，我们在这里这么操作纯粹是为了实验XGBoost分布式训练在SageMaker上是如何实现的。另外这个实验数据集的格式我们采用了SageMaker对XGBoost内置算法支持的CSV格式，所有的SageMaker的内置算法都支持特定的数据集格式，具体情况请参考SageMaker文档。例如我们可以创建如下的训练集：

重复因子1: 127MB，60,000张图片
重复因子2: 254MB，120,000张图片
重复因子4: 508MB，240,000张图片
重复因子8：1016MB，480,000张图片

超级参数及XGBoost版本设置

SageMaker内置算法对于XGBoost的支持非常好，可以直接通过SageMaker的参数来设置所有的XGBoost的超参。而且对XGBoost开源版本支持也会紧跟时代的脚步，到本文撰写的时间点SageMaker支持的最新版本是1.5，而开源最新的版本是1.6，想要获得目前SageMaker支持的最新版本可以参看文档。

为了实验的目的，除非特别说明我们所有的运行训练任务时都是用一套超级参数的设置。另外超级参数调优也不是本文的重点，因此这里给出的超级参数并不一定是最佳的，具体的代码如下：

hyperparameters = {'alpha': 0.0,
       'colsample_bylevel': 0.4083530569296091,
       'colsample_bytree': 0.8040025839325579,
       'eta': 0.11764087266272522,
       'gamma': 0.43319156621549954,
       'lambda': 37.547406128070286,
       'max_delta_step': 10,
       'max_depth': 6,
       'min_child_weight': 5.076838893848415,
       'num_round': 100,  # Not tuned: kept fixed
       'subsample': 0.8915771964367318,
       'num_class': 10,  # Not tuned: defined by Fashion MNIST
       'objective': 'multi:softmax'  # Not tuned: defined by Fashion MNIST
      }

同时，在SageMaker的API中我们将指定使用目前内置算法支持的最新版本1.5：

xgboost_container = sagemaker.image_uris.retrieve("xgboost", region, "1.5-1")

基准训练测试

我们首先使用1到4台的ml.m5.xlarge的计算资源，并使用重复因子为1的训练集进行基准训练。使用前面提供的超级参数以及XGBoost版本，这里需要注意的是对于树算法的选择也是在超级参数内进行调整的，我们在这里没有指定，那么tree_method的默认值就会设置为auto，这是在通常运行XGBoost经常会使用的一个值，根据官方文档描述会根据数据集的大小来选择相应的树算法，小数据集会选择exact大数据集会选择approx，至于数据集大小的划分文档中并没有明确指出，但根据实验结果看一般都会使用exact这个参数，因此我们在后面会指定这个超级参数来看实验的结果。另外一个数据集分片的问题同样在基准测试过程中没有进行，而是无论是单实例还是多实例训练单个实例都会使用全量的数据集。

最终基准测试的结果如下图所示，其中X轴为训练用的实例数量而Y轴为训练所用的时间。

我们看到这个基准测试并不是想要的结果，多台实例训练最终需要的时间竟然比单台还要长。分析其中的原因，能够直接想到的就是因为基准测试过程中没有对数据进行分片，而是每个训练实例都是使用全量的数据集，因此接下来的实验我们将改进这个问题。但这也不是全部的“真相”，后面我们将通过持续的实验来得到更为准确的事实。

数据集分片

我们知道如果在机器学习开源框架中要实现数据集分片读取往往需要自己实现分片的逻辑，虽然并不复杂但仍然对我们原有的训练代码有一定的影响。而使用SageMaker内置的算法我们基本不会自己实现算法代码，而只是通过调用SageMaker的SDK或者API来实现模型训练，因此数据的读取逻辑完全是由内置算法本身来实现的。在SageMaker中进行机器学习训练无论是自己实现算法代码（我们称其为BYOC，Bring Your Own Code）还是内置算法的方式训练我们都可以使用SageMaker封装好的ShardByS3Key数据集分片的方式来进行分布式训练，详细信息请参考API的文档链接。其实这种数据分片的方式实现起来还是非常简单的，首先需要自己将数据集的单个文件进行切片分成几个文件，这里有一些实践的细节需要注意：

第一点是对文件的切分工具的选择，我建议不需要使用特别复杂的工具，直接使用Linux或Mac OS提供的split命令行就可以完美切分文本类的数据集（例如像我们这次实验使用的CSV格式）。根据数据集的行数以及我们要切分文件的个数计算得到每个文件包含的行数再将其作为参数传递给命令行，这里不建议使用-n参数直接将文件切分为想要的份数，这样会导致从一行的中间断掉，比如我们可以用如下命令：

sh-4.2$ split -l 60000 8_factor_fashion-mnist_train.csv split_factor_8

第二点是要切分的数据集的数量选择，通常在这里我们需要考量的是将来需要多少的SageMaker训练的机器实例数量，因为最终这些分片的数据会下载到这些实例中进行分别训练的，因此分片的数量至少要大于或等于将来要运行实例的数量。由于本次实验我们使用的实例数量并不相同，因此在设计这个实验时候我们选择切片的数量是使用最大的训练实例数量即为8，同时其他实验使用的实例数量的选择也是按照2的幂次，数量选择的就是1，2，4，8。这样做是保证我们每次实验在每个训练实例中的数据集的数据量是大体相当的，防止出现数据倾斜的现象。
第三点是我们要使用ShardByS3Key是S3前缀做Shard的方式对数据集进行切片并分配给不同的训练实例使用，因此这里建议将切分好的数据集文件都放入一个S3的目录下，这样在调用SageMaker API时方便指定。

做好上述准备工作后，使用SageMaker ShardByS3Key的方式进行数据集切片实际在API调用时比较简单，只是在以往的data channel的接口中增加一个distribution参数而已。

train_input = TrainingInput(train_url, content_type=content_type, distribution='ShardedByS3Key')

分布式训练测试

在SageMaker中内置算法实现分布式训练非常简单，只需修改一个参数来指定需要运行的实例数量就可以了，具体调用SageMaker的API代码片段如下：

hyperparameters['tree_method'] = 'approx'

# construct a SageMaker estimator that calls the xgboost-container
estimator = sagemaker.estimator.Estimator(image_uri=xgboost_container, 
                                          hyperparameters=hyperparameters,
                                          role=sagemaker.get_execution_role(),
                                          instance_count=2, 
                                          instance_type='ml.m5.xlarge', 
                                          volume_size=5, # 5 GB 
                                          output_path=output_path)

# execute the XGBoost training job
estimator.fit({'train': train_input, 'validation': validation_input})

其中instance_count这个参数就是控制着训练启动的实例数量，我们只要指定需要的数量就可以实现分布式训练。而这段代码的第一句就是我们前面提到的将超级参数中的tree_method设置为approx。

接下来我们将分别指定instance_count参数为1、2、4和8来分别观察训练时间是否会随着实例数量的增加会不会下降。

首先重复因子（Factor）为1的数据集表现还是没有像我们预想的一样随着实例数量的上升训练时间下降，而重复因子为2、4和8确实达到了我们的初期的目的。其实原因也是与数据规模相关的，如果数据规模没有到达一定规模而使用分布式训练是不会降低训练时间的反而因为集群间的协同消耗大量的资源，这也是我们在基准实验中所说的全部“真相”。而且随着数据集的数据量的增加分布式训练的效果也会非常明显，当然随着训练实例的增加带来的边际效果也是递减的而不是线性的，这就需要我们在训练时间和训练资源之间寻找平衡了。

当然我们知道虽然训练的时间下降有很大的提升，对模型的精度的影响也十分重要，因此我们再看一下分布式训练对模型的精度的影响。

我们使用了验证集的平均loss值作为参考指标，同一个复制因子（factor）的不同训练实例数量之间loss值上下有波动但基本持平，这说明没有随着分布式训练增加实例的数量而对我们训练的精度有所损失。而之所以每次训练的loss不太一样也是因为XGBoost算法本身具有很多随机性的，例如在每个训练实例中每次分配到的数据是不同的，以及决策树在节点分裂的时候也是有随机性的，具体信息可以参考官方文档对subsample或者colsample_bylevel的描述。但是随着复制因子的变大loss值在逐渐降低，这也容易解释，主要是因为数据量的变大而且是通过复制得到的数据，虽然loss值在降低并不一定说明我们的模型的性能会更好，而也有可能我们的模型已经过拟合了，因此前面对数据集介绍的时候我们也强调了在实践中是绝对不会使用直接复制这种方式来增加训练集的。

看到了通过分布式训练增加训练实例来提升训练效率的目的达成，我们同时也要关注分布式训练另外一个目的就是突破物理机的内存限制。这里我们可以通过CloudWatch服务监控到SageMaker的内存利用率来观察实验结果。

从上图我们看出对于在每台训练实例的物理机内存占用率上看，分布式训练效果还是很明显的，虽然我们训练的内存使用率没有超过50%的情况，这还是因为我们的数据集不够大，但是随着训练实例数量的增加内存的使用率是在降低的。可以让我们在遇到超过了训练实例物理机内存的情况时，使用分布式训练来训练全量的数据集。

结论

本文通过对SageMaker内置的XGBoost算法对分布式训练的实验，验证了可以通过增加训练实例的数量来达到降低训练时间和内存使用率的目的，同时保证模型的精度不会产生影响。但我们仍然需要关注两个技术细节，一个是要对数据集进行切片并通过SageMaker内置的API对数据按ShardByS3Key的方式进行读取，另外一个是对XGBoost超级参数tree_method的exact值是不支持分布式训练的。

亚马逊AWS官方博客