Introduction to Distributed TensorFlow
Reedbush Hackathon @ The University of Tokyo
2019-01-24

1

目次

データ並列
Parameter ServerAll Reduce
モデル並列
TensorFlow での実装
TensorFlow の標準機能
Horovod2

データ並列データ並列の基本的な考え方

正則化項 + 損失の最小化問題:
min λR(w) + L(w) where L(w) =x

損失項の勾配はデータごとに並列で計算できる[wk+1 = wk − α ∇w R(wk ) +∑ℓi (w)i∑]∇w ℓi (wk )i3

Parameter Server 方式での役割分担

Parameter Server▶ 最新の weight w を保持して必要なときに master や worker に共有

Worker▶ それぞれ異なるデータを使って勾配を計算

Master▶ Worker とだいたい同じ▶ モデルの保存や初期化など実装上必要な一部の処理4

データ並列 + 同期型 + Parameter Server

Updatesum

Parameter
Servers

grad.

grad.

Computegrad.

Masterand
WorkersSharemini-batch samples

5データ並列 + 非同期型 + Parameter Server

UpdateParameter
Servers

grad.

grad.

Computegrad.

Masterand
WorkersSharemini-batch samples

6

それぞれの長所同期型▶ 収束性が良く学習が上手くいきやすい▶

Stale gradients の問題が発生しない

非同期型▶ 対故障性▶

Worker が落ちても影響が小さく学習が継続可能▶ 待ち合わせが無いので高スループット▶

遅い worker に足を引っ張られない

7

All Reduce での役割分担

Worker▶ それぞれ異なるデータを使って勾配を計算▶ Parameter server は無いので各 worker が weight w のコピーを保持

Master▶ Worker とだいたい同じ▶ モデルの保存や初期化など実装上必要な一部の処理

8

データ並列 + 同期型 + All Reduce

gradient with all samples▶ Weight w は各 worker がコピーを持つ▶ Worker 同士で通信して勾配の情報を交換▶ それぞれ w を更新する

mini-batch samples

9

モデル並列モデル並列の考え方▶ データ単位ではなくモデルそのもの (ニューラルネットで言えばネットワー

クの部分) で並列化できる部分を分割して worker に割り当てる▶ まだ発展途上で定石が定まっている状態ではないので創意工夫を凝らしてやっていく形になる▶ Mesh TensorFlow という専用のライブラリも開発されているが事例などはまだほとんど無い段階▶

https://github.com/tensorﬂow/mesh10TensorFlow での実装

TensorFlow で分散学習を実装する時の選択肢
TensorFlow の標準の機能▶ TensorFlow の operation を任意のデバイスに配置する形で実装▶ 自由度が高く頑張れば大抵の構成は作れる▶ データ並列のよくある構成は高レベル API で分散処理を意識せずにやってもらうこともできる
Horovod▶ Uber が開発した OSS▶ データ並列 + All Reduce の形式が前提▶ Reedbush ではこちらの方が実績もあるのでおそらく良い
11TensorFlow の標準機能で実装するときのお作法

ClusterPSPSWorkerWorker

Mastertask.pytask.pytask.pytask.pytask.py▶ クラスタの全ノードは同じスクリプトを実行▶ スクリプト内で master, worker, ps の分岐を記述▶

自分がどれなのかという情報は環境変数に入れておく
12環境変数に入れておく情報の例
1 {
234
5678
9
10
11
12

"cluster": {"ps": ["ps-408715a125-0:2222", "ps-408715a125-1:2222"],"worker": ["worker-408715a125-0:2222", "worker-408715a125-1:2222"],"master": ["master-408715a125-0:2222"]},"task": {"index": 0,"type": "master",...
},
...

13 }

13

TensorFlow 側の処理

1

tf_conf = json.loads(os.environ.get("TF_CONFIG", "{}"))

2 # TF_CONFIG か ら ク ラ ス タ 構 成 の 情 報 を 取 り 出 し て 使 う
3 server = tf.train.Server(45

tf_conf.get("cluster", None), job_name=tf_conf["task"]["type"],task_index=tf_conf["task"]["index"]

6 )7 # TF_CONFIG か ら 自 分 の 役 割 が Master 、Worker 、PS の ど れ か 調 べ て 分 岐 さ せ る

if tf_conf["task"]["type"] == "ps":server.join() # Parameter Server は こ の 処 理 を す る だ け
10 else:11# Master とWorker の 処 理 を 実 行 す る
12...
8

914

高レベル APIs (1/2)

自力で書き切ることもできるが、高レベル API を使うことで分散処理部分を含め面倒な部分を自動でやってもらえる仕組みが存在する
Estimator▶ 分散学習の機能は最も充実している▶ Horovod との併用も可能

Keras▶ Keras 自体に分散学習の機能は無いので Horovod を使う

15

高レベル APIs (2/2)

tf.keras▶ TensorFlow に特化して独自機能が追加された Keras▶ まだ標準機能で multi-node の分散学習には対応していない▶
▶

https://github.com/tensorﬂow/tensorﬂow/issues/23664
tf.keras のモデルを Estimator に変換する機能があるので現状はそれを使用する形になる▶ Horovod とも最近併用可能に▶

https://github.com/uber/horovod/pull/513▶ 将来的には Estimator より tf.keras が主流になる可能性もある

16

Estimator + データ並列▶ Estimator に DistributeStrategy を渡す▶ クラスタ構成は tf.estimator.RunConﬁg のコンストラクタが裏で勝手に環境変数から読み込む
12
3
4
5
6
7
8
9

distribution = tf.contrib.distribute.MirroredStrategy(num_gpus=2)
run_config = tf.estimator.RunConfig(distribute=distribution)
clf = tf.estimator.DNNClassifier(feature_columns=[tf.feature_column.numeric_column("x", shape=[4]),hidden_units=[10, 20, 10], n_classes=3,
model_dir="/tmp/iris_model",config=run_config)
clf.train(...)
17

DistributeStrategy (1/2)

tf.distribute.MirroredStrategy▶ https://www.tensorﬂow.org/guide/distribute_strateg▶ single-node で multi-GPUs の環境で使う All Reduce▶ tf.keras にも対応している▶ multi-node でも動くらしいが推奨はされていない▶

https://github.com/tensorﬂow/tensorﬂow/issues/23664

18

DistributeStrategy (2/2)

tf.contrib.distribute.CollectiveAllReduceStrategy▶ multi-node で multi-GPUs の環境で使う All Reducetf.contrib.distribute.ParameterServerStrategy▶ multi-node で multi-GPUs の環境で使う prameter server 方式

19

補足: DistributeStrategy 以前の書き方▶ tf.estimator.train_and_evaluate を使うと parameter server 方式で学習が走る
1
2
3
4
5
6
7
8
9

run_config = tf.estimator.RunConfig(distribute=distribution)
clf = tf.estimator.DNNClassifier(feature_columns=[tf.feature_column.numeric_column("x", shape=[4]),hidden_units=[10, 20, 10], n_classes=3,model_dir="/tmp/iris_model",config=run_config)...
tf.estimator.train_and_evaluate(clf, ...)

20

TensorFlow と Horovod でデータ並列 + All Reduce (1/2)

12import tensorflow as tf
import horovod.tensorflow as hvd

3
4

hvd.init()

5
6 # 1 プ ロ セ ス に 1 GPU を 割 り 当 て て 他 は 見 え な い よ う に し て お く
7 config = tf.ConfigProto()

config.gpu_options.visible_device_list = str(hvd.local_rank())...
10 # Optimizer を 分 散 学 習 用 の ク ラ ス で wrap す る
11 optimizer = hvd.DistributedOptimizer(optimizer)
12 ...
8921

TensorFlow と Horovod でデータ並列 + All Reduce (2/2)
1 # 初 期 値 を master (rank 0) か ら 他 の プ ロ セ ス に broadcast す る hook
2hooks = [hvd.BroadcastGlobalVariablesHook(0)]3
4 # モ デ ル の 保 存 は master だ け が や れ ば 良 い
5model_dir = "/your/model/dir" if hvd.rank() == 0 else None

▶ あとは Estimator か tf.train.MonitoredTrainingSession を使っていつも通りコードを書けば良い▶ 実行するときは mpirun で実行▶ Epoch 数、学習率などをプロセス数に応じて調整することが多い▶

Epoch 数はプロセス数に反比例、学習率は比例させるなど
22

Keras と Horovod でデータ並列 + All Reduce (2/2)

import horovod.keras as hvd
import keras
3 import tensorflow as tf
1

2

45hvd.init()

67 # 1 プ ロ セ ス に 1 GPU を 割 り 当 て て 他 で は 見 え な い よ う に し て お く

config = tf.ConfigProto()
config.gpu_options.visible_device_list = str(hvd.local_rank())
10 K.set_session(tf.Session(config=config))
8923

Keras と Horovod でデータ並列 + All Reduce (1/2)
1 # 初 期 値 を master か ら 他 の プ ロ セ ス に broadcast す る callback
2callbacks = [hvd.callbacks.BroadcastGlobalVariablesCallback(0)]3
4 # master に だ け モ デ ル 保 存 の callback を 追 加

if hvd.rank() == 0:
callbacks.append(7
keras.callbacks.ModelCheckpoint('./checkpoint-{epoch}.h5')8)56▶ あとは Keras でいつも通りのコードを書いて mpirun で実行▶ Epoch 数、学習率などをプロセス数に応じて調整することが多い▶

Epoch 数はプロセス数に反比例、学習率は比例させるなど
24補足: TensorFlow Eager Mode + Horovod▶ つい最近 eager mode への対応 PR がマージされた▶
▶

https://github.com/uber/horovod/pull/670
TensorFlow v2.0 からは eager mode がデフォルトになるので重要な機能になるかも?

25

おわりに

▶ 今はちょうど分散学習周りが変わっている最中なので方法が色々ある▶ Reedbush 上でやるなら MPI が前提になっているのでおそらく Horovod がやりやすい▶ モデル並列など凝ったことをやる場合は TensorFlow 標準の機能を上手く使う必要がある

26