Successfully reported this slideshow.
We use your LinkedIn profile and activity data to personalize ads and to show you more relevant ads. You can change your ad preferences anytime.

Mesh tensorflow

1,028 views

Published on

第一回 Tensorflow内部構造勉強会で話した内容です。
「Mesh Tensorflow」

Published in: Technology
  • Be the first to comment

Mesh tensorflow

  1. 1. Mesh Tensorflow 2019/3/4 kuroko1t
  2. 2. 自己紹介 kuroko(@kuroko1t) ● 仕事 ○ 前は..次世代スパコンの CPUの論理回路設計 ○ 最近は..深層学習向けフレームワークの調査したり( Tensorflow,horovod...) ● その他 ○ GoでDeep Learningを試みたり(link) ○ horovodの記事書いてみたり (Qiita) ○ horovodにTensorflowのeager対応のPRだしたり
  3. 3. グラフから考える分散学習 データ並列 モデル並列 グラフを複製(replicaを作成) ex) horovodだと1プロセス1GPU GPU0 GPU1 GPU0 GPU1 グラフを分割 (複数GPUで1モデルを共有)
  4. 4. Mesh Tensorflowって? ● データ並列,モデル並列両方をカバー ● python only ● DNNモデル ○ Mesh Tensorflowのメインターゲットは Transformer ■ 言語モデルのパラメータが5億近くなってきていて,モデル並列しないと計算できない ○ exampleにはmnist(cnn)のみ ● 独自の演算API ○ 本家Tensorflowとは異なるAPIを用いる ● single Process
  5. 5. Mesh Tensorflowの概念 概念図の用語 ● Mesh Shape ○ coreとMeshの割当の定義 ● Layout Rules ○ TensorとMeshの割当の定義 ● mtf.Dimension ○ Tensorのshapeの名前付け ● Meshとは? ○ プロセッサのn次元配列として定義 され,それぞれのTensorがMeshに 割当られます.通常は Tensorは coreに割り当てられる(tf.deviceを利 用) ● Layout Rules? ○ Tensorの各次元をMeshのどの次 元に配置するか定義します (Layout Rulesで表現される) (“batch”,3) Mesh0 core0 core1 core2 core3 Mesh1 Tensor T1 (“col”,2) (“row”,2) mtf.Dimension Mesh shape (batch, mesh1)→core1,core2, core3でバッチ並列 Layout Rules
  6. 6. サンプルソースを見てみる。。
  7. 7. Mesh Tensorflow実行例 graph = mtf.Graph() mesh = mtf.Mesh(graph, "my_mesh") row_x = mtf.Dimension("row_x", 6) col_x = mtf.Dimension("col_x", 4) row_y = mtf.Dimension("row_y", 4) col_y = mtf.Dimension("col_y", 5) x = mtf.import_tf_tensor( mesh, tf.random.uniform([6,4]), mtf.Shape([row_x, col_x])) y = mtf.import_tf_tensor( mesh, tf.random.uniform([4,5]), mtf.Shape([row_y, col_y])) result = mtf.matmul(x, y, output_shape=[row_x, col_y]) mesh_shape = "b1:2" mesh_layout = "col_x:b1" mesh_shape = mtf.convert_to_shape(mesh_shape) mesh_devices = [""] * mesh_shape.size mesh_impl = mtf.placement_mesh_impl.PlacementMeshImpl( mesh_shape, mesh_layout, mesh_devices) lowering = mtf.Lowering(graph, {mesh: mesh_impl}) restore_hook = mtf.MtfRestoreHook(lowering) tf_result = lowering.export_to_tf_tensor(result) with tf.train.MonitoredTrainingSession(hooks=[restore_hook]) as sess: print(sess.run(tf_result).shape) mtf用のGraph, Mesh の作成 tfのtensor→mtfのtensorへ変換API mtfの演算APIの呼び出し mtf→tfへop,tensorの変換(lowering) Session run()を実行
  8. 8. Meshのshape, layout定義 mesh_shape = [("rows:2"), ("cols:2")] layout_rules = [("batch", "rows"), ("hidden", "cols")] mesh_impl = mtf.placement_mesh_impl.PlacementMeshImpl( mesh_shape, layout_rules, mesh_devices) ● meshのcore割当数を定義. ○ coreの割当数を定義(ex: rowsという名前のMeshに2個coreを割り当て) ○ [("rows:2"), ("cols:2")] ● meshとTensor Shapeの対応 ○ mtf.Dimensionとmeshの対応を定義する ○ [("batch", "rows"), ("hidden", "cols")] ● PlacementMeshImpl(CPU,GPU向け) ○ 上記layout ruleに応じてvariableをデバイスに割当(session起動時) ○ 集合演算の実装がされている ○ TPU向けにSimdMeshImplというのも用意されている Tensor “cols” “rows” mtf.Dimension
  9. 9. 演算のデバイス割り振り def parallel(devices, fn, *args, **kwargs): ... for i, device in enumerate(devices): with tf.device(device): with tf.variable_scope("parallel_%d" % i): my_args = [x[i] for x in args] my_kwargs = {k: v[i] for k, v in six.iteritems(kwargs)} ret.append(fn(*my_args, **my_kwargs)) 内部的にwith tf.deviceで各デバイスに演算を割り当てている
  10. 10. Mesh Tensorflowの演算API
  11. 11. 演算API # tensor importの例 x = mtf.import_tf_tensor( mesh, tf.ones([50, 10]), mtf.Shape([col_x, row_x])) # 演算APIの例 f1 = mtf.relu(mtf.conv2d_with_blocks( x, kernel1, strides=[1, 1, 1, 1], padding="SAME", h_blocks_dim=row_blocks_dim, w_blocks_dim=col_blocks_dim)) ● 演算APIにおける次元定義 ○ 次元の定義はtf.Dimension()で定義した名前で行う。 ○ この段階だとTensorの次元がどのように Mesh分割されるかは定義されない ● mtf.layer系 API ○ mtf.layer.dense, attention系など30個 ○ Attention専用APIが多い(mtf.layer.masked_local_attention_1d()...) ● mtf 系 API ○ mtf.add, mtf.conv2d, mtf.maximum… 演算系APIが120個
  12. 12. 演算APIの例:mtf.matmul() ● 例えば。A(2,4) x B(4,5) = C(2,5) で Aのcolに2core割当したとする.. core1 =xcore01 1 4 4 5 core0 /core1 ● A行列を行方向で分割(Aのrow,Bのcolは同じDimension) 1. core0,core1で並列演算 a. core0:A(1,4) x B(4x5) = C0(1,5) b. core1 : A(1,4) x B(4,5) = C1(1,5) 2. C0,C1をまとめる a. Cをexport_to_tf_tensor(mtf→tfへ変換)するとtf.concat実行してC0,C1をまとめる core0 core1 A B C C0 C1 1 1 5
  13. 13. 演算APIの例:mtf.matmul() core1 =xcore02 2 2 5 core0 2 core1 2 2 5 A B C(core0/core1) = C0(core0) + C1(core1) ● A行列を列方向で分割 (Aのrow,Bのcolは同じDimension) 1. core0,core1で並列演算 a. core0:A(2,2) x B(2,5) = C0 b. core1: A(2,2) x B(2,5) = C1 2. core0, core1での演算を合計する a. C(2,5) = C0(2,5) + C1(2,5) b. allreduce関数を利用して計算 (後ほど説明します) core0 /core1
  14. 14. Mesh Tensorflowの集合通信
  15. 15. デバイス間の集合通信 PlacementMeshImpl内に定義 ● allreduce ○ 各TensorのSumをとる ○ add, optimizer, layer.batch_norm, reduce系のopで利用されている ● allconcat ○ 各Tensorをconcatenateする(like a allgather) ● alltoall ○ MPIのalltoallと同等.
  16. 16. Allreduce(CPU,GPU) GPU0 GPU1 GPU2 GPU3 GPU4 GPU5 left center right center tf.add tf.add tf.add tf.add SUM値 SUM値 tf.identitytf.identity 分配分配 for i in xrange(1, left_center + 1): with tf.device(devices[i]): left_sum = binary_reduction(left_sum, xs[i]) right_sum = xs[n-1] 演算デバイスを 制御 tf.add tf.add 1. 2グループに分割 2. 各グループ内でtf.add 3. 2グループ間でtf.add 4. tf.identityで分配 allreduceの方法
  17. 17. Allreduce(TPU) ● 内部のアルゴリズムは不明 ○ CrossReplicaSumというOpkernelが実行される??... ■ 実装がないため内部のアルゴリズム不明... ■ 結果としては各slice(各デバイス)のSUMを返す. ○ またなぜか実装がTensorflow内部にある.. ■ tensorflow/contrib/tpu/ops/gen_tpu_ops.py
  18. 18. mtf.Lowering, optimizer # optimizerの登録 optimizer = mtf.optimize.SgdOptimizer() update_ops = optimizer.apply_grads(var_grads, graph.trainable_variables) lowering = mtf.Lowering(graph, {mesh: mesh_impl}) # optimizer operationのlowering tf_update_ops = [lowering.lowered_operation(op) for op in update_ops] tf_update_ops.append(tf.assign_add(global_step, 1)) train_op = tf.group(tf_update_ops) ● OptimizerはSGDとAdafactorのみサポート ● Lowering()の機能としてはmtf→tfの橋渡しで,tensorflowのAPIに変換したり ,実行デバイスの割当をおこなう ○ 演算ごとにどのようにloweringするか定義されている
  19. 19. 性能 Transformerの性能 モデル並列の方が性能がスケーリングしている→Mesh Tensorflowを利用し てモデル並列を簡単に実装できるし,性能出ていいよねって感じだと(特に説 明は見当たらないが..) 出展:README.md
  20. 20. おわり

×