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。