Amazon SageMaker 分散学習ライブラリ - 新しいデータ並列・モデル並列の分散学習

データ並列およびモデル並列ができる分散学習について紹介しました。データ並列は以下の論文を参考にしています。

I. Thangakrishnan, et al., Herring: Rethinking the Parameter Server at Scale for the Cloud, SC20
https://assets.amazon.science/ba/69/0a396bd3459294ad940a705ad7f5/herring-rethinking-the-parameter-server-at-scale-for-the-cloud.pdf

Transcript

1. 1 © 2020 Amazon Web Services, Inc. or its affiliates. All rights reserved | 1 © 2020 Amazon Web Services, Inc. or its affiliates. All rights reserved | アマゾン ウェブ サービス ジャパン株式会社 ソリューションアーキテクト 鮫島 正樹 Amazon SageMaker 分散学習ライブラリ 新しいデータ並列・モデル並列の分散学習
2. 2 © 2020 Amazon Web Services, Inc. or its affiliates. All rights reserved | Background • データセットの巨⼤化・機械学習モデルの⼤規模化によって、より多くの学 習時間が必要になってきている 例) 8個の NVIDIA V100 GPU を搭載した ml.p3dn.24xlarge を利⽤した 場合の学習時間 • COCO Dataset による Mask RCNN (物体検出モデル): 6時間以上 • BERT (⾃然⾔語処理モデル) の事前学習: 100時間以上 • 単⼀の GPU メモリに読み込むことが難しいような⼤規模なモデルも開発さ れている
3. 3 © 2020 Amazon Web Services, Inc. or its affiliates. All rights reserved | SageMaker Distributed Training Library データ並列 (SageMaker Data Parallelism; SDP) データセットを分割して複数のGPUで並列して勾配を計算し、 最終的に勾配を⼀つにまとめてモデルを更新する モデル並列 (SageMaker Model Parallelism; SMP) モデルを分割して複数のGPUに展開し、それぞれのモデルの計算 をなるべく並列化することで、必要なメモリ・計算を分散させる • 2020/12 に SageMaker 独⾃の⾼効率な分散学習ライブラリを発表 • PyTorch, TensorFlow を⽤いた GPU での学習に対応
4. 4 © 2020 Amazon Web Services, Inc. or its affiliates. All rights reserved | SageMaker データ並列による分散学習 I. Thangakrishnan, et al., Herring: Rethinking the Parameter Server at Scale for the Cloud, SC20 https://assets.amazon.science/ba/69/0a396bd3459294ad940a705ad7f5/herring-rethinking-the-parameter-server-at-scale-for-the-cloud.pdf • Amazon が SC20 で発表した分散学習アルゴリズム Herring の採⽤ • トレンドとなっていた All-Reduce 形式ではなく、Parameter Serverの形式 • Parameter Server • 複数の Worker にデータを分散して勾配を計算し、それを Parameter Server に集約する • All-Reduce • 複数の Worker にデータを分散して勾配を計算し、Worker 全体が互い に共有する
5. 5 © 2020 Amazon Web Services, Inc. or its affiliates. All rights reserved | なぜ Parameter Servers か https://d2l.ai/chapter_computational-performance/parameterserver.html • Worker がどれだけ増えても2ホップで 勾配を集約できる • 勾配の集約を⼯夫することで、⾮同期 的な学習を⾏うことができる • Worker に必要な帯域が⼩さくてすむ
6. 6 © 2020 Amazon Web Services, Inc. or its affiliates. All rights reserved | Parameter Servers の課題 • Parameter Server が複数ある場合 は、それぞれに異なるパラメータ （重み）を保持させる • パラメータのサイズが異なるので、 Parameter Server が受け取るデー タサイズに不均衡が⽣じてしまう • Parameter Server が管理するパラ メータが計算されるときしか、 Parameter Server は動かない Parameter Server Parameter ある時間帯でのみ Parameter Server は動く Parameterのサイズは ばらばら
7. 7 © 2020 Amazon Web Services, Inc. or its affiliates. All rights reserved | Balanced Fusion Buffers (BFB) • Parameter Server 全体にパラメー タ保持の領域 (BFB) を⽤意し、保 持されたパラメータを等分割して Parameter Server が受け取る • Worker は、BFB 内の決められた 位置にパラメータを保存する（そ のための効率的な⽅法も提案） • Parameter Server の前に Proxy Parameter Server を置くこ とで並列接続性を改善
8. 8 © 2020 Amazon Web Services, Inc. or its affiliates. All rights reserved | Result – BERTモデル の Training
9. 9 © 2020 Amazon Web Services, Inc. or its affiliates. All rights reserved | SageMaker モデル並列による分散学習 • モデル分割は、単⼀のGPU に乗りきらないモデル を学習するための⽅式として知られていたが、ど こで分割するかが難しい問題 • 2つの観点で分割を決める • 分割したモデルの計算量が等しくなるように • モデル間のデータのやり取りを最⼩限に抑える • モデルを例えばAとBに分割し、A→Bと計算が⾏わ れる場合は実⾏時間がかかる。なるべく並列して 計算できるようにしたい。 https://www.amazon.science/latest-news/the-science-of-amazon-sagemakers-distributed-training- engines
10. 10 © 2020 Amazon Web Services, Inc. or its affiliates. All rights reserved | モデルの⾃動分割 https://docs.aws.amazon.com/sagemaker/latest/dg/model-parallel-core-features.html • メモリ容量の⼤きい CPU 上にモデルを読み込み、モ デルの解析を⾏って、モデルを分割する。このため、 最初にオーバーヘッドの時間がある。 • Tensorflow は DAG を構築するので、DAG をベース に分割を⾏う。PyTorch はそのような機能がないの で、ライブラリでTreeを構築して分割を⾏う。 • CPU上でモデルが分割されると、GPU にロードされ て学習が⾏われる。 Tensorflow の DAG
11. 11 © 2020 Amazon Web Services, Inc. or its affiliates. All rights reserved | Pipelining https://docs.aws.amazon.com/sagemaker/latest/dg/model-parallel-core-features.html • 単純にモデルを分割 (X on GPU0 と Y on GPU1) してX→Yと計算すると、 • 分割前と同様にXとYの計算時間もかかる • XとYの間のデータのやり取りで時間がかかる → 結果的にトータルの実⾏時間は増えてしまう。 • ミニバッチをさらに細かいマイクロバッチにわけて、XとYをなるべく効率的 に計算する → Pipelining Interleaved pipeline
12. 12 © 2020 Amazon Web Services, Inc. or its affiliates. All rights reserved | Interleaved Pipeline Forward/Backwardを並⾏して⾏うことで効率化
13. 13 © 2020 Amazon Web Services, Inc. or its affiliates. All rights reserved | SageMaker モデル並列学習のパラメータ https://sagemaker.readthedocs.io/en/stable/api/training/smd_model_parallel_general.html partitions (int) いくつに分けるか microbatches (int) ミニバッチを何分割するか (1のとき分割しない) pipeline (interleaved or simple) Forward と Backward を混ぜて⾼速化する interleaved と、 固めて⾏う simple optimize (memory or speed) memory: 変数の数にもとづいて分割 speed: 計算量にもとづいて分割 placement_strategy (cluster or spread) 分割したモデルを複数 GPU で学習する場合、そのモデルを 近くの GPU (cluster) か、遠くの GPU (spread) に配置 auto_partition ⾃動分割を使うかどうか default_partition ⾃動分割しない場合で、演算をどの Partition に割り当てる か指定しない場合、default_partitionに割り当てられる
14. 14 © 2020 Amazon Web Services, Inc. or its affiliates. All rights reserved | モデル並列のための実装の⼀部 (PyTorch) https://github.com/aws/amazon-sagemaker- examples/tree/master/training/distributed_training/pytorch/model_parallel/bert import smdistributed.modelparallel.torch as smp model = smp.DistributedModel(model) optimizer = smp.DistributedOptimizer(optimizer) @smp.step def smp_step(args, device, input_ids, ... criterion, step): # calculation loss return loss for step, batch in enumerate(train_iter): loss_mbs = smp_step(args, device, input_ids,..., model, optimizer, criterion, step) loss = loss_mbs.reduce_mean() # ライブラリインポート # 分散並列⽤のモデルに変換 # 分散⽤の Optimizer を設定 # ロス計算の定義 # 学習ループで分散学習を⾏う
15. 15 © 2020 Amazon Web Services, Inc. or its affiliates. All rights reserved | データ並列 vs モデル並列 https://docs.aws.amazon.com/sagemaker/latest/dg/distributed-training.html • まずはデータ並列を選択する。データ並列が実⾏可能であれば、性能的にも データ並列のほうが優れる。 • メモリに乗り切らない場合にモデル並列を検討する。ただし、本当にモデル並 列を⾏う前に以下も検討する。 • ハイパーパラメータの変更 • バッチサイズを⼩さくする。バッチサイズ1でもメモリ不⾜が発⽣する場 合はモデル並列を検討する。 • 他にも • Nvidia Tensor Core GPU であれば混合精度学習でメモリを節約可能 • ⾃然⾔語処理なら、利⽤する⽂の⻑さを短くする • 画像認識なら、画像の解像度を下げる
16. 16 © 2020 Amazon Web Services, Inc. or its affiliates. All rights reserved | まとめ • ⼤規模なデータセット・モデルの学習で時間がかかっていたり、GPU メモ リで処理しきれない場合、分散学習での効率化が可能 • 分散学習にはデータ並列とモデル並列があり、基本的にはデータ並列を選 択する。モデルがメモリに乗らない場合にモデル並列を試す。 • SageMaker のライブラリには、分散学習に関する複雑な処理を⾃動化する 機能があるため、少ない実装での分散学習が可能