Introduction to Distributed TensorFlow @shuhei_fujiwara last update: 2019-08-05 1

$ whoami 藤原秀平 (FUJIWARA Shuhei) Twitter: @shuhei_fujiwara GitHub: @sfujiwara ▶ Google Developer Expert (Machine Learning) ▶ TensorFlow User Group Tokyo Organizer 2

TensorFlow User Group 活動内容 ▶ 勉強会 ▶ 論文読み会 / ハードウェア系 / スパコン講習会 / etc ▶ 公式ドキュメント翻訳 参加方法 ▶ https://tfug.jp/から Slack に参加 3

分散学習とは GPU:0 GPU:0 GPU:1 GPU:2 GPU:3 GPU:0 GPU:0 GPU:0 GPU:1 GPU:2 GPU:3 GPU:0 GPU:1 GPU:2 GPU:3 ▶ 複数のノードや複数の GPU を上手く使って 学習を高速化 ▶ 処理の並列化 ▶ デバイス間の通信 4

なぜ分散なのか? ▶ 1 台のマシンでできる高速化には限界がある ▶ 2000 GPUs ほしい! ▶ コスト面で有利 (な時もある) ▶ 2 倍のスペックのマシンは 2 倍の値段じゃない ▶ 2 台マシンを用意するなら 2 倍の値段 5

分散学習の種類 1 . 2 ├── Model Parallel 3 └── Data Parallel <-- Main topic of this talk 4 ├── All Reduce 5 │ └── Sync 6 └── Parameter Server 7 ├── Sync 8 └── Aync 6

データ並列分散学習

データ並列の基本的な考え方 ▶ 分散させるためには並列化が必要 ▶ 損失関数の勾配の計算をサンプルごとに並列化できる ▶ 勾配は線形性があるので後で足し合わせれば良い L(w) = ∑ i ℓi(w) ∇wL(w) = ∑ i ∇wℓi(wk) 7

Parameter Server 方式 1 . 2 ├── Model Parallel 3 └── Data Parallel 4 ├── All Reduce 5 │ └── Sync 6 └── Parameter Server <-- 7 ├── Sync 8 └── Aync 8

Parameter Server 方式での役割分担 Parameter Server ▶ 最新の weight w を保持して必要なときに master や worker に共有 Worker ▶ それぞれ異なるデータを使って勾配を計算 Master ▶ Worker とだいたい同じ ▶ モデルの保存や初期化など実装上必要な一部の処理 9

データ並列 + 同期型 + Parameter Server Compute grad. grad. grad. sum Share Update Master and Workers Parameter Servers mini-batch samples 10

データ並列 + 非同期型 + Parameter Server Compute grad. grad. grad. Share Update Master and Workers Parameter Servers mini-batch samples 11

それぞれの長所 同期型 ▶ 収束性が良く学習が上手くいきやすい ▶ Stale gradients の問題が発生しない 非同期型 ▶ 対故障性 ▶ Worker が落ちても影響が小さく学習が継続可能 ▶ 待ち合わせが無いので高スループット ▶ 遅い worker に足を引っ張られない 12

All Reduce 1 . 2 ├── Model Parallel 3 └── Data Parallel 4 ├── All Reduce <-- 5 │ └── Sync 6 └── Parameter Server 7 ├── Sync 8 └── Aync 13

All Reduce での役割分担 Worker ▶ それぞれ異なるデータを使って勾配を計算 ▶ Parameter server は無いので各 worker が weight w のコピーを保持 Master ▶ Worker とだいたい同じ ▶ モデルの保存や初期化など実装上必要な一部の処理 14

データ並列 + 同期型 + All Reduce ▶ Weight w は各 worker が コピーを持つ ▶ Worker 同士で通信して 勾配の情報を交換 ▶ それぞれ w を更新する gradient with all samples mini-batch samples 15

TensorFlow and Horovod

TensorFlow ▶ Google が開発している数値計算の OSS ▶ 行列演算あたりの低レイヤーから自分で書くこともできる ▶ 機械学習用の高レベル API も用意されている ▶ Google 社内で使われていた大規模分散フレームワーク DistBelief をもとに開発 ▶ 標準で分散学習の機能を持っている 16

Horovod ▶ Uber が開発している分散学習の OSS ▶ TensorFlow, Keras, PyTorch, MXNet で書いた既存のコードを All Reduce の分散学習に拡張できる 17

FAQ: TensorFlow の標準機能と Horovod の棲み分け TensorFlow の標準機能 ▶ MPI 環境が前提ではない ▶ Parameter Server 方式はこちらのみ ▶ TensorFlow の operation を任意のデバイスに配置する形で実装 ▶ 自由度が高く頑張れば大抵の構成は作れる Horovod ▶ MPI 環境が前提 ▶ スパコンだとこちらが楽なケースが多い ▶ Horovod が開発された頃 TensorFlow には All Reduce が無かった 18

TensorFlow の高レベル API

Estimator 昔からある TensorFlow の安定した高レベル API ▶ 当面はサポートされるが、今後は後述する tf.keras が主流になる予定 ▶ 標準で分散学習の機能に対応 ▶ Horovod とも併用可能 ▶ Optuna とも併用可能 ▶ https://github.com/pfnet/optuna/pull/292 19

Keras TensorFlow 以外に Theano, CNTK をバックエンドに選べるラッパー ▶ 2017 年に TensorFlow v1.0 が出て、そちらに統合される ことが発表された ▶ ただし、元々の Keras は維持したまま TensorFlow バックエンド に特化した tf.keras と分岐して開発が進む形 ▶ TensorFlow 標準の分散機能には対応していないので 分散学習をするには Horovod を使う 20

tf.keras TensorFlow 本体に組み込まれた Keras ▶ TensorFlow バックエンドに特化して開発されている ▶ 標準で分散学習の機能を持っている ▶ バージョンが古いと対応が追いついていないので注意 ▶ https://github.com/tensorflow/tensorflow/issues/23664 ▶ TPU を使用することができる ▶ 元の Keras との互換性はほどほど ▶ 内部実装はごっそり変わっていたりするので別ライブラリと思った方が安全 ▶ 今後 TensorFlow の高レベル API の主流 (になる予定) ▶ Horovod との併用も可能 ▶ https://github.com/uber/horovod/pull/513 21

分散学習実装のお作法

分散学習を実装するときのお作法 Cluster Worker Worker PS PS Master task.py task.py task.py task.py task.py ▶ クラスタの全ノードは同じスクリプトを実行 ▶ スクリプト内で master, worker, ps の分岐を記述 ▶ 自分がどれなのかという情報を持っておく必要がある 22

自分の役割に関する情報 TensorFlow 標準機能 ▶ 環境変数 TF_CONFIG に必要な情報を入れておく ▶ スパコン環境だとこれを用意するのがちょっと面倒 MPI 環境 ▶ プロセスごとに rank と local rank を持っている ▶ Rank: 全体で何番目のプロセスか ▶ Local Rank: ノードごとに何番目のプロセスか 23

TensorFlow 標準機能を使った実装

環境変数 TF_CONFIG に入れておく情報の例 1 { 2 "cluster": { 3 "ps": ["ps-408715a125-0:2222", "ps-408715a125-1:2222"], 4 "worker": ["worker-408715a125-0:2222", "worker-408715a125-1:2222"], 5 "master": ["master-408715a125-0:2222"] 6 }, 7 "task": { 8 "index": 0, 9 "type": "master", 10 ... 11 }, 12 ... 13 } 24

TF_CONFIG の最近の変更 ▶ master ==> chief ▶ 後述する DistributeStrategy を使う場合はこの名称変更が必要 ▶ evaluator ▶ Evaluation 処理だけを行うノードを作れる ▶ 今のところ最大 1 ノードまで (のはず 25

高レベル API を使わない場合の TensorFlow 側の処理 1 tf_conf = json.loads(os.environ.get("TF_CONFIG", "{}")) 2 # TF_CONFIG か ら ク ラ ス タ 構 成 の 情 報 を 取 り 出 し て 使 う 3 server = tf.train.Server( 4 tf_conf.get("cluster", None), job_name=tf_conf["task"]["type"], 5 task_index=tf_conf["task"]["index"] 6 ) 7 # TF_CONFIG か ら 自 分 の 役 割 が Master、Worker、PS の ど れ か 調 べ て 分 岐 さ せ る 8 if tf_conf["task"]["type"] == "ps": 9 server.join() # Parameter Server は こ の 処 理 を す る だ け 10 else: 11 # MasterとWorkerの 処 理 を 実 行 す る 12 ... 26

Estimator + DistributeStrategy ▶ RunConfig に DistributeStrategy を渡すだけ ▶ クラスタ構成は tf.estimator.RunConfig のコンストラクタが裏で勝手に 環境変数 TF_CONFIG から読み込む 1 distribution = tf.distribute.MirroredStrategy(num_gpus=2) 2 run_config = tf.estimator.RunConfig(distribute=distribution) 3 clf = tf.estimator.DNNClassifier( 4 feature_columns=[tf.feature_column.numeric_column("x", shape=[4]), 5 hidden_units=[10, 20, 10], n_classes=3, 6 model_dir="/tmp/iris_model", 7 config=run_config 8 ) 9 clf.train(...) 27

tf.keras + DistributeStrategy ▶ scope でコンテキストを作るだけ 1 mirrored_strategy = tf.distribute.MirroredStrategy() 2 with mirrored_strategy.scope(): 3 model = tf.keras.Sequential([ 4 tf.keras.layers.Dense(1, input_shape=(1,)) 5 ]) 6 model.compile(loss='mse', optimizer='sgd') 28

DistributeStrategy (1/2) 各種 API と DistributeStrategy の対応状況: https://www.tensorflow.org/beta/guide/distribute_strategy#types_of_strategies tf.distribute.MirroredStrategy ▶ https://www.tensorflow.org/guide/distribute_strategy ▶ single-node で multi-GPUs の環境で使う All Reduce ▶ multi-node でも動くらしいが推奨はされていない ▶ https://github.com/tensorflow/tensorflow/issues/23664 29

DistributeStrategy (2/2) tf.distribute.experimental.MultiWorkerMirroredStrategy ▶ multi-node で multi-GPUs の環境で使う All Reduce ▶ CollectiveAllReduceStrategy から名前が変わったので注意 tf.distribute.experimental.ParameterServerStrategy ▶ Prameter server 方式 30

補足: DistributeStrategy 以前の書き方 ▶ tf.estimator.train_and_evaluate を使うと parameter server 方式で学習が走る 1 run_config = tf.estimator.RunConfig(distribute=distribution) 2 clf = tf.estimator.DNNClassifier( 3 feature_columns=[tf.feature_column.numeric_column("x", shape=[4]), 4 hidden_units=[10, 20, 10], n_classes=3, 5 model_dir="/tmp/iris_model", 6 config=run_config 7 ) 8 ... 9 tf.estimator.train_and_evaluate(clf, ...) 31

Horovod を使った実装

TensorFlow と Horovod でデータ並列 + All Reduce (1/2) 1 import tensorflow as tf 2 import horovod.tensorflow as hvd 3 4 hvd.init() 5 6 # 1 プ ロ セ ス に 1 GPUを 割 り 当 て て 他 は 見 え な い よ う に し て お く 7 config = tf.ConfigProto() 8 config.gpu_options.visible_device_list = str(hvd.local_rank()) 9 ... 10 # Optimizer を 分 散 学 習 用 の ク ラ ス で wrap す る 11 optimizer = hvd.DistributedOptimizer(optimizer) 12 ... 32

TensorFlow と Horovod でデータ並列 + All Reduce (2/2) 1 # 初 期 値 を master (rank 0) か ら 他 の プ ロ セ ス に broadcast す る hook 2 hooks = [hvd.BroadcastGlobalVariablesHook(0)] 3 4 # モ デ ル の 保 存 は master だ け が や れ ば 良 い 5 model_dir = "/your/model/dir" if hvd.rank() == 0 else None ▶ あとは Estimator か tf.train.MonitoredTrainingSession を使って いつも通りコードを書けば良い ▶ 実行するときは mpirun で実行 ▶ Epoch 数、学習率などをプロセス数に応じて調整することが多い ▶ Epoch 数はプロセス数に反比例、学習率は比例させるなど 33

Keras と Horovod でデータ並列 + All Reduce (1/2) 1 import horovod.keras as hvd 2 import keras 3 import tensorflow as tf 4 5 hvd.init() 6 7 # 1 プ ロ セ ス に 1 GPU を 割 り 当 て て 他 で は 見 え な い よ う に し て お く 8 config = tf.ConfigProto() 9 config.gpu_options.visible_device_list = str(hvd.local_rank()) 10 K.set_session(tf.Session(config=config)) 34

Keras と Horovod でデータ並列 + All Reduce (2/2) 1 # 初 期 値 を master か ら 他 の プ ロ セ ス に broadcast す る callback 2 callbacks = [hvd.callbacks.BroadcastGlobalVariablesCallback(0)] 3 4 # master に だ け モ デ ル 保 存 の callback を 追 加 5 if hvd.rank() == 0: 6 callbacks.append( 7 keras.callbacks.ModelCheckpoint('./checkpoint-{epoch}.h5') 8 ) ▶ あとは Keras でいつも通りのコードを書いて mpirun で実行 ▶ Epoch 数、学習率などをプロセス数に応じて調整することが多い 35

補足: TensorFlow Eager Mode + Horovod ▶ しばらく前に eager mode への対応 PR がマージされた ▶ https://github.com/uber/horovod/pull/670 ▶ TensorFlow v2.0 からは eager mode がデフォルトになるので重要 36

おわりに ▶ 環境や構成に応じて標準機能と Horovod を使い分けると良い ▶ スパコン環境では Horovod がおすすめ ▶ クラウドの場合などは標準機能の方が楽だったりする ▶ アップデートが速い機能なので最新ドキュメントやリポジトリを チェックしよう 37