SageMaker管道模式下如何使用Horovod实现多GPU分布式训练

当前，咱们可使用多种技术经过少许数据训练出深度学习模型，具体包括针对图像分类任务的迁移学习、少样本学习甚至是一次性学习等，也能够基于预训练的BERT或GPT2模型对语言模型进行微调。可是，在部分应用用例中咱们仍然须要引入大量训练数据。例如，若是当前图像与ImageNet数据集内的图像彻底不一样，或者当前语言语料库只针对特定领域、而非通用类型，那么单凭迁移学习将很难带来理想的模型性能。

做为深度学习研究人员，咱们可能须要从零开始尝试新的思路或方法。在这种状况下，咱们必须使用大型数据集训练出大型深度学习模型；在找不到最佳训练方法的状况下，整个过程可能须要几天、几周甚至是几个月。

在本文中，咱们将一同了解如何在Amazon SageMaker的单一实例之上运行多GPU训练，并讨论如何在Amazon SageMaker上实现高效多GPU与多节点分布式训练。

Horovod基础知识

在使用大量数据进行模型训练时，最好是将训练做业分配给多个GPU（单一实例或者多个实例）。深度学习框架提供内置方法以支持多GPU训练或分布式训练。但除此以外，还有另一种实现方法，即直接使用分布式深度学习框架（例如Horovod）。

Horovod是Uber公司打造的分布式深度学习开源框架，可以与TensorFlow、Keras、PyTorch以及Apache MXNet等一线热门深度学习工具包协同使用。Horovod使用all-reduce算法取代以往的参数服务器方法进行快速分布式训练，其中还提供多种优化方法以进一步加快分布式训练的执行速度。关于更多详细信息，请参阅碰见Horovod：面向TensorFlow的Uber开源分布式深度学习框架。

为Horovod准备数据

在使用Horovod执行训练做业时，Horovod会为其集群当中的每一个GPU上的工做节点启动独立的进程（每一个GPU对应一个工做节点）。例如，若是使用一个包含4 GPU的训练实例（一台Amazon SageMaker ml.p3.8xlarge 或Amazon Elastic Compute Cloud (Amazon EC2) p3.8xlarge实例）运行Horovod训练做业，则将对应启动4个工做进程。数据集本体已经出于数据并行性的需求而被拆分为多个分片，全部这4个工做节点都将分别读取本身的数据集分片。若是有40000个训练样本，则每一个工做节点将得到10000个互不重复的训练样本。

若是使用Horovod进行分布式训练甚至是多GPU训练，则应事先作好数据分片准备，并指引工做节点从文件系统中读取各个分片。（某些深度学习框架能够自动执行此操做，如PyTorch的DataParallel与DistributedDataParallel）。

下图所示，为进行分片存储的两种可行架构。

咱们能够经过多种不一样方式为Amazon SageMaker训练做业提供数据集。一种最典型的方法就是将全部数据集存储在Amazon Simple Storage Service (Amazon S3)存储桶内，并在须要时进行访问。你们固然可使用共享文件系统（例如Amazon FSx for Lustre或Amazon Elastic File System，简称Amazon EFS)实现数据存储，但经过Amazon SageMaker内置的两种输入模式（文件模式与管道模式）直接从Amazon S3中检索数据可以避免系统产生额外的服务成本。

在文件模式下，当Amazon SageMaker启动训练做业后，数据集将从指定的S3存储桶被传送至训练实例当中，并将其放置在某个特定目录以内。但若是使用的数据集极为庞大，那么将对象从存储桶复制至训练实例的存储上每每须要耗费很长时间，并且直到数据传输完成，训练做业才会真正开始。这会在某些状况下拖慢机器学习（ML）的执行流程，甚至影响到创新或研究工做的项目进度。

另外，你们也能够经过管道模式直接访问存储在Amazon S3中的数据集。管道模式在训练实例与S3存储桶之间建立直接输入管道，并容许训练进程直接访问对象，这就消除了在训练开始以前将全部对象复制至训练实例中的工做。要对给定Amazon S3 URI中的数据集以管道模式加以访问，请在建立Amazon SageMaker Estimator时将输入模式设置为Pipe，具体参见如下代码：

from sagemaker.tensorflow import TensorFlow

tf_estimator = TensorFlow(entry_point='train.py',
 role='SageMakerRole',
 train_instance_type='ml.p3.2xlarge',
 train_instance_count=2,
 framework_version='2.1.0',
 py_version='py3',
 input_mode='Pipe')

在管道模式下，训练数据将做为FIFO流的形式进行交付。TensorFlow扩展的dataset类极大下降了访问流数据集的难度。关于管道模式与TensorFlow的更多详细信息，请参阅在Amazon SageMaker上使用高速管道模式加快模型训练，以及Amazon SageMaker TensorFlow扩展GitHub repo。

配合Horovod使用管道模式

当配合Horovod使用管道模式执行单机多卡或者多机多卡的分布式训练时，有一点须要特别注意。下图所示为这类场景的基本架构。

管道模式将数据从Amazon S3流式的传送到训练实例当中的Unix命名管道/FIFOs当中。一个FIFO文件仅支持一对写入/读取程序，且每轮训练周期内咱们只能为一条通道建立一个FIFO文件。一般，人们会为训练数据集定义一条通道，并为验证或测试数据集定义另外一条单独的通道，然后将这些输入通道做为Amazon SageMaker Estimator 中fit ()函数的参数传递至训练做业。详见如下代码：

from sagemaker.session import s3_input

input_channel = {'train': s3_input('s3://your-bucket-name/train-dataset/')}

tf_estimator.fit(inputs=input_channel)

这种方式在Horovod多GPU训练场景下又会形成怎样的影响？简而言之，使用Horovod在多GPU训练做业中启动的各个进程，会相互争用单一FIFO，致使多个进程没法同时访问这些FIFO。并且因为同一时间内只有单一工做进程可以访问FIFO，且在完成训练做业以前不会释放句柄，就致使全部其余工做进程没法从该FIFO中读取数据，最终令训练做业陷入死锁式的无限循环。若是看到相似于如下形式的重复提示消息，则代表遇到了这样的问题：

[1,0]<stderr>:Stalled ranks:
[1,0]<stderr>:0: [training/Adam/DistributedAdam_Allreduce/HorovodAllreduce_training_Adam_gradients_AddN_11_0, training/Adam/DistributedAdam_Allreduce/HorovodAllreduce_training_Adam_gradients_AddN_12_0, training/Adam/DistributedAdam_Allreduce/HorovodAllreduce_training_Adam_gradients_AddN_14_0, training/Adam/DistributedAdam_Allreduce/HorovodAllreduce_training_Adam_gradients_AddN_15_0, training/Adam/DistributedAdam_Allreduce/HorovodAllreduce_training_Adam_gradients_AddN_18_0, training/Adam/DistributedAdam_Allreduce/HorovodAllreduce_training_Adam_gradients_AddN_19_0 ...]
[1,0]<stderr>:2: [training/Adam/DistributedAdam_Allreduce/HorovodAllreduce_training_Adam_gradients_AddN_11_0, training/Adam/DistributedAdam_Allreduce/HorovodAllreduce_training_Adam_gradients_AddN_12_0, training/Adam/DistributedAdam_Allreduce/HorovodAllreduce_training_Adam_gradients_AddN_14_0, training/Adam/DistributedAdam_Allreduce/HorovodAllreduce_training_Adam_gradients_AddN_15_0, training/Adam/DistributedAdam_Allreduce/HorovodAllreduce_training_Adam_gradients_AddN_18_0, training/Adam/DistributedAdam_Allreduce/HorovodAllreduce_training_Adam_gradients_AddN_19_0 ...]
[1,0]<stderr>:3: [training/Adam/DistributedAdam_Allreduce/HorovodAllreduce_training_Adam_gradients_AddN_11_0, training/Adam/DistributedAdam_Allreduce/HorovodAllreduce_training_Adam_gradients_AddN_12_0, training/Adam/DistributedAdam_Allreduce/HorovodAllreduce_training_Adam_gradients_AddN_14_0, training/Adam/DistributedAdam_Allreduce/HorovodAllreduce_training_Adam_gradients_AddN_15_0, training/Adam/DistributedAdam_Allreduce/HorovodAllreduce_training_Adam_gradients_AddN_18_0, training/Adam/DistributedAdam_Allreduce/HorovodAllreduce_training_Adam_gradients_AddN_19_0 ...]

咱们能够对S3存储桶中的数据集进行分片，且数量与用于训练做业的GPU数量相对应。若是拥有4000个TensorFlow记录文件，且使用一台带有4 GPU的ml.p3.8xlarge实例进行模型训练，则能够为互不重复的1000个TensorFLow记录文件设定不一样的前缀，如如下代码所示：

s3://your-bucket-name/train/0/
s3://your-bucket-name/train/1/
s3://your-bucket-name/train/2/
s3://your-bucket-name/train/3/

使用SharedByS3Key做为Amazon S3数据类型分配方式进行的数据集分片方法，并不彻底适用于Horovod。这是由于在使用SharedByS3Key时，分片只会以实例为单位、而非以工做进程为单位进行，且实例中的工做进程与GPU的数量保持一致。一样的，各个实例仍然只拥有一条输入通道。所以，你们须要将数据集的分片数量，设定为与Horovod集群内GPU数相同。

接下来，咱们须要为Amazon SageMaker训练定义四条输入通道，具体参见如下代码：

from sagemaker.session import s3_input

shuffle_config = sagemaker.session.ShuffleConfig(234)

train_s3_uri_prefix = 's3://your-bucket-name/train'
input_channels = {}

for idx in range(4):
 train_s3_uri = f'{train_s3_uri_prefix}/train/{idx}/'
 train_s3_input = s3_input(train_s3_uri, shuffle_config=shuffle_config)
 input_channels[f'train_{idx}'] = train_s3_input

ShuffleConfig将确保根据每一个训练轮次，对Amazon S3前缀下各文件的使用顺序进行随机分配。关于更多详细信息，请参阅ShuffleConfig。

在Amazon SageMaker Estimator上调用fit方法时，请使用如下通道定义：

tf_estimator.fit(input_channels)

对于验证及测试类任务，咱们只需在单一工做进程上运行（一般使用主工做进程或Rank 0工做进程）。在这里，咱们不须要设置多条验证或测试通道。但若是使用tf.keras.model.fit()函数进行训练，则训练会在只有一个Horovod工做进程进行验证时中止（关于更多详细信息，请参阅Horovod GitHub repo上的issue #600）。若是须要使用tf.keras.model.fit()进行验证，你们还应为各验证数据集提供对应的输入通道（相似于训练输入通道）。请注意，截至2020年7月，管道模式下训练做业的输入通道总数上限为20个。具体请参见如下代码：

validation_s3_uri = 's3://your-bucket-name/validation/'

for idx in range(4):
 validation_s3_input = s3_input(validation_s3_uri)
 input_channels[f'validation_{idx}'] = validation_s3_input

eval_s3_uri = 's3://your-bucket-name/eval/'
eval_s3_input = s3_input(eval_s3_uri)
input_channels['eval'] = eval_s3_input

相较于直接使用S3存储桶前缀，咱们在这里可使用包含有对象键列表的普通ManifestFile。关于更多详细信息，请参阅输入数据。

在训练代码中使用数据通道

在训练脚本中，咱们须要强制要求各个Horovod工做进程只访问属于它本身的数据集分片，确保两个工做进程不会访问同一输入通道。在本文的用例中，咱们将使用从0开始的索引定义各输入通道名称。为此可使用hvd.rank()函数，由其为当前工做进程在集群范围以内提供惟一的排名索引，且排名一样从0开始（请参考如下代码中的第13行）。在本文示例中，咱们使用Amazon SageMaker TensorFlow扩展PipeModeDataset。对于其余深度学习框架，请在每一个训练轮次中从名为/opt/ml/input/data/[channel_name]_${epoch}的FIFO文件中读取数据。关于更多示例，请参见GitHub repo。

1: from sagemaker_tensorflow import PipeModeDataset
 2:
 3: features = {'data': tf.FixedLenFeature([], tf.string),
 4: 'labels': tf.FixedLenFeature([], tf.int64)}
 5:
 6: def parse(record):
 7: parsed = tf.parse_single_example(record, features)
 8: return ({
 9: 'data': tf.decode_raw(parsed['data'], tf.float64)
10: }, parsed['labels'])
11:
12: # For Horovod and Pipe mode, use the input channel allocated to this worker using rank information
13: channel_name = 'train_{}'.format(hvd.rank())
14:
15: ds = PipeModeDataset(channel=channel_name, record_format='TFRecord')
16: ds = ds.map(parse)
17: ds = ds.batch(64)
18: ds = ds.prefetch(10)

在包含一个或多个实例的Horovod集群中，排名分配方式为从0开始，至GPU数量 - 1结束。只要正肯定义了输入通道的名称并从0开始使用索引，咱们就没必要分神管理各实例或名位的排列顺序。

使用Tensorboard进行监控

在对训练进程加以灵活监控方面，咱们能够在每一个训练轮次结束时首先将日志上传至S3存储桶，再经过任意远程计算实例调用Tensorboard。为此，咱们须要建立一项回调以将本地日志推送至S3存储桶路径，此路径仅限于运行在Horovod上的主（Rank 0）计算节点。具体请参见如下代码：

class Sync2S3(tf.keras.callbacks.Callback):
 def __init__(self, logdir, s3logdir):
 super(Sync2S3, self).__init__()
 self.logdir = logdir
 self.s3logdir = s3logdir
 def on_epoch_end(self, batch, logs={}):
 os.system('aws s3 sync '+self.logdir+' '+self.s3logdir)
...
if hvd.rank() == 0:
 logdir = args.output_data_dir + '/' + datetime.now().strftime("%Y%m%d-%H%M%S")
 callbacks.append(TensorBoard(log_dir=logdir))
 callbacks.append(Sync2S3(logdir=logdir, s3logdir=tensorboard_logs))

经过将训练日志存储在S3存储桶内，你们能够在任意服务器上运行Tensorboard，包括EC2实例、Amazon SgaeMaker notebook实例甚至是本地计算机，并为该Tensorboard托管服务器提供访问Amazon S3日志对象的权限。为了支持从Amazon S3源处直接提取日志数据，咱们的Tensorboard必须为1.14.0或者更高版本。如下命令行使用的是位于us-east-1区域内S3存储桶上的日志记录：

S3_REGION=us-east-1
tensorboard --logdir s3://{bucket_name}/tensorboard_logs/

若是是在Amazon SageMaker notebook实例上运行以上命令，则可经过https://<SageMaker-notebook-instance-name>.notebook.<notebook-region>.sagemaker.aws/proxy/6006/完成访问。

资源清理

在完成本文中分布式训练做业以后，请清理相应资源以免其后续产生额外费用，包括S3存储桶、FSx for Lustre以及各个Amazon SageMaker实例。

总结

在Amazon SageMaker上以管道模式使用Horovod的多GPU或分布式训练方法，可以为数据集的各个分片建立独立的训练通道并在数据通道内访问对应分片，借此实现大规模模型训练。这种方式可以缩短在实际训练开始以前将数据集传输至训练实例所占用的时间，所以特别适用于具备大规模训练数据集的Amazon SageMaker训练场景。

关于在Amazon SageMaker上运行的完整训练示例（管道模式加Horovod），请参阅GitHub repo。