最近リリースされたTensorFlow v2.9では、モデル、データ、スペースパラレル(別名、空間タイル)ディープネットワークトレーニングのための新しいAPIが追加された。DTensorの目的は、シャーディング命令をモデルコードから分離することである。そのためにDTensorは、デバイス間でモデルとバッチパラメータを分割するための高レベルのユーティリティを提供する。この取り組みは、大規模なトレーニング処理を構築するためにかかる開発時間を短縮するための最近の取り組み(GPipe、TF Mesh、GShard、DeepSpeed、Fairscale、ColossalAIなど)の一部である。
トレーニングテストロスは、大規模な(言語)モデルに対して、ネットワークパラメータの数、データサイズ、計算時間に応じて対数的に増加する。そのため、近年、タスクレベルの改善はディープネットワークのサイズに依存していた。また、かつてないほど多いアクセラレーションデバイスの使用によって、そのようなトレーニングプラットフォームの分散的な特性のために、大量のエンジニアリング作業が強いられてきた(つまり、GPUとFPGAのデバイス内メモリユニットは限られている)。代わりにDTensorでは、トレーニングユニットを抽象化し(つまり、デバイスに依存しない)、デバイスのメッシュ設定とテンソル配置のレイアウト設定を行う。
例として、次の図は、3つのレイアウトを持つ2つのメッシュ構成でのテンソルの配置を示している。2つ目のメッシュの場合、列または行ごとのシャーディングを選択できる。その際、どのディメンションを「シャーディングしない」設定にして、代わりにデバイス間で複製するかをを指定する。
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.functionとtf.GradientTapeのAPIでも使うことができる。ただし、現在のTensorFlowバージョンはKeras向けの組み込みトレーニングループをサポートしていない。DTensorシャードモデルをトレーニングするためにカスタムループを作成する必要がある。また、シングルクライアントおよびマルチクライアントのトレーニングジョブをサポートし、複数のプロセスでAPIをネイティブに活用できる。
追加情報は、DTensorの概要ドキュメントから入手できる。TensorFlow Webサイトには、低レベルの分散トレーニングとKerasトレーニングの例が提供されている。