BT

最新技術を追い求めるデベロッパのための情報コミュニティ

寄稿

Topics

地域を選ぶ

InfoQ ホームページ ニュース TensorFlow DTensor:分散ディープネットワークトレーニング向けの統合API

TensorFlow DTensor:分散ディープネットワークトレーニング向けの統合API

原文(投稿日:2022/05/27)へのリンク

最近リリースされたTensorFlow v2.9では、モデル、データ、スペースパラレル(別名、空間タイル)ディープネットワークトレーニングのための新しいAPIが追加された。DTensorの目的は、シャーディング命令をモデルコードから分離することである。そのためにDTensorは、デバイス間でモデルとバッチパラメータを分割するための高レベルのユーティリティを提供する。この取り組みは、大規模なトレーニング処理を構築するためにかかる開発時間を短縮するための最近の取り組み(GPipeTF MeshGShardDeepSpeedFairscaleColossalAIなど)の一部である。

トレーニングテストロスは、大規模な(言語)モデルに対して、ネットワークパラメータの数、データサイズ、計算時間に応じて対数的に増加する。そのため、近年、タスクレベルの改善はディープネットワークのサイズに依存していた。また、かつてないほど多いアクセラレーションデバイスの使用によって、そのようなトレーニングプラットフォームの分散的な特性のために、大量のエンジニアリング作業が強いられてきた(つまり、GPUとFPGAのデバイス内メモリユニットは限られている)。代わりにDTensorでは、トレーニングユニットを抽象化し(つまり、デバイスに依存しない)、デバイスのメッシュ設定とテンソル配置のレイアウト設定を行う。

例として、次の図は、3つのレイアウトを持つ2つのメッシュ構成でのテンソルの配置を示している。2つ目のメッシュの場合、列または行ごとのシャーディングを選択できる。その際、どのディメンションを「シャーディングしない」設定にして、代わりにデバイス間で複製するかをを指定する。

Fig-1: Ilustration showing placement of a tensor on two different mesh configurations.

DTensorでは、個別にメッシュオブジェクトとレイアウトオブジェクトを作成することにより、デバイス設定をハードコーディングせずにさまざまなトレーニングトポロジを採用する柔軟性を得ることができる。例えば、任意の次元に沿ってテンソルの空間分割を簡単に実装する方法が提供されている。その際、コンピュータービジョンアプリケーション専用のAPIを使用することはない(TensorFlow TPUEstimator空間分割とは異なる)。デバイスメッシュAPIは(論理デバイスメカニズムを介して)TF仮想デバイスで使用されるため、DTensor APIを使ってさまざまなシャーディングシナリオを試すことができることを記しておく。

DTensorには実験的なAPIがあるが、現在は直接Keras統合をサポートしている。以下のコードスニペットは、完全に複製された重み付けのレイアウトが高密度レイヤーに渡される例である。

mesh = dtensor.create_mesh([("batch", 8)], devices=devices)
kernel_layout = dtensor.Layout.replicated(mesh, 2)
bias_layout = dtensor.Layout.replicated(mesh, 1)

model = tf.keras.models.Sequential([
  tf.keras.layers.Flatten(input_shape=(28, 28)),
  tf.keras.layers.Dense(128,
                        activation='relu',
                        name='d1',
                        kernel_layout=kernel_layout,
                        bias_layout=bias_layout),
])

DTensorは、ほとんどのテンソル操作に対して完全互換である。そのため、tf.functiontf.GradientTapeのAPIでも使うことができる。ただし、現在のTensorFlowバージョンはKeras向けの組み込みトレーニングループをサポートしていない。DTensorシャードモデルをトレーニングするためにカスタムループを作成する必要がある。また、シングルクライアントおよびマルチクライアントのトレーニングジョブをサポートし、複数のプロセスでAPIをネイティブに活用できる。

追加情報は、DTensorの概要ドキュメントから入手できる。TensorFlow Webサイトには、低レベルの分散トレーニングKerasトレーニングの例が提供されている。

作者について

この記事に星をつける

おすすめ度
スタイル

特集コンテンツ一覧

BT