DeepRec学習・推論フレームワークを基盤とするDeepRec Extensionは、学習タスクの観点から、自動エラスティック学習や分散フォールトトレランスといった機能を提案し、特に大規模スパースモデルの分散学習に重点を置いています。これにより、スパースモデルの学習効率が全体的に向上し、DeepRecエンジンがスパースシナリオにおいてその優位性を発揮できるようになります。

序文

DeepRecは、大規模スパースモデルのトレーニングおよび推論エンジンとして、アリババグループ内やクラウドベースの検索、レコメンデーション、広告の分野で高い評価を得ています。スパースモデルにEmbeddingVariable機能を提供することで、ユーザーは大規模なスパースパラメータを持つモデル構造を容易に構築できます。さらに、分散最適化、ランタイム最適化、グラフおよび演算子最適化など、EmbeddingVariableの進化型機能も提供しており、モデルの有効性とトレーニングスループットの両方のパフォーマンスが向上しています。DeepRecは非同期トレーニング向けにgRPC++やStarServerなどの一連の通信最適化を提案していますが、ユーザーの観点から見ると、分散トレーニングにはフォールトトレランスやリソース選択などの課題が依然として残っています。DeepRec Extensionは、より低コストで高効率な分散モデルトレーニングを実現することを目指しています。

モチベーション

大規模スパースモデルの学習では、ビジネスの拡大に伴いモデルサイズが数百GB/TBレベルにまで増大するため、分散学習が不可欠な選択肢となります。Alibabaグループ内のお客様とAlibaba Cloudのお客様は、TensorFlowのPS（パラメータサーバー）分散学習において、常に多くの課題に直面してきました。

① 複雑な分散モデリングインターフェース。

TensorFlowはグラフ/セッションを作成するための基盤機能を提供しているため、ユーザーは計算グラフを構築する際にパラメータの分割と配置戦略を自由に設定できます。不適切な設定は、PS全体にわたるパラメータ負荷の不均一化につながり、モデルの学習スループットの低下につながる可能性があります。これは、ユーザーが単一マシンのモデルコードを分散バージョンに変更する際のコストを増加させます。

② リソース予測が難しく、柔軟なリソース割り当てができない。

オープンソースのKubeflowエコシステムであれ、業界標準のTensorFlow分散トレーニング製品のインターフェースであれ、各ロールのリソース要求はユーザーに公開されます。タスクを送信する際、ユーザーが要求されるリソース量を判断するのは困難です。多くの場合、リソース要求が少なすぎるとトレーニング中にOutOfMemoryError（OOM）が発生し、リソース要求が多すぎるとリソースが無駄になり、コストが増加します。タスクの実際のリソース使用量は、特にスパースモデルのパラメータが動的に変化するシナリオでは予測が困難です。静的なリソース要求では、動的なリソース変更のニーズを満たすことができません。

③ 分散フォールトトレランスの仕組みが単純すぎる。

現在、TensorFlow のフォールト トレランス メカニズムでは、定期的にチェックポイントを永続化し、フェイルオーバー時にモデルを最新のチェックポイントにロールバックします。DeepRec は、スパース パラメータ シナリオに対して増分チェックポイントを提案しています。これにより、パラメータ ノードが予期せずクラッシュした場合のトレーニング損失を削減できますが、永続ストレージから完全なチェックポイントと増分チェックポイントを読み取ることは避けられません。そのため、特に外部ストレージの読み取り効率が低いシナリオでは、フォールト トレランス時のモデル回復にコストがかかります。一方、PS ノードのメモリ、ディスク、およびネットワーク リソースは比較的アイドル状態であるため、パラメータのバックアップに使用して、完全なモデルを読み取るオーバーヘッドを削減できます。TensorFlow の設計では、チェックポイント中にトレーニング サンプルは保存および復元されません。フェイルオーバー中、モデルは正常に回復しますが、サンプル データの状態情報が失われ、サンプルの損失や再計算などの望ましくない動作につながります。

④ 分散環境が複雑である。

単一のワーカー（PS）を用いた分散非同期トレーニングのシナリオでは、PSは複数のワーカーからサブグラフを実行することがよくあります。単一のPSでマシン障害（分散環境における低速ノード）が発生すると、すべてのワーカーのトレーニング効率に影響を及ぼし、トレーニング全体のスループットが低下します。システムは、低速ノードを迅速に特定して交換する必要があります。単一のPS障害がタスク全体に与える影響は、単一のワーカー障害よりもはるかに大きいため、リソーススケジューリングの観点から、PSはワーカーよりも高い優先度を持つ必要があります。

ディープラーニングフレームワークの本来の目的に立ち返り、私たちはアルゴリズムエンジニアにとって使いやすく効率的なフレームワークを提供することに常に努めてきました。ユーザーはモデル構築に集中でき、フレームワークは基盤となるハードウェアやリソースの情報からユーザーを遮断します。このフレームワークは、分散環境でモデルを容易に実行する方法、より効率的に実行する方法、フォールトトレランスを自動的に処理する方法など、一連のランタイム課題に対処します。ハイパーパラメータ設定が少ないほど、ユーザーにとって使いやすく、フレームワークの最適化の余地が広がります。ハイパーパラメータ設定が多いほど、上級ユーザーにとって使いやすくなり、きめ細かなモデル管理が容易になります。DeepRec Extensionは、上記の問題を解決し、ユーザーが低コストかつ高効率で分散学習を行うことを可能にします。システムの観点から分散学習の効率を向上させる点において、DeepRec Extensionは以下の機能を提供します。

• 分散トレーニングリソースの推定
• 自動弾性トレーニング
• リソース/計算グラフの監視
• 自動バックアップとフォールトトレランス

デザインコンセプトと全体的なアーキテクチャ

DeepRec Extensionの機能には、リソースのスケジューリング、グラフの最適化、グラフの実行が含まれます。これらの機能の設計と実装をオープンソースフレームワークkubeflow/DeepRec/TensorFlowから分離し、特定のフレームワークバージョンへの束縛を回避し、これらのフレームワークの拡張可能なインターフェースを機能設計と実装に最大限活用することを目指しています。DeepRec/TensorFlowの設計思想では、すべての操作をOperationとして抽象化できます。DeepRec Extensionも関連ロジックの実装をOperationとして抽象化し、GraphOptimizationとHooksを使用して元の計算グラフを変更し、計算グラフに機能を適用します。Kubernetesエコシステムは、CRD拡張インターフェースを提供します。kubeflow + TensorFlowのコミュニティバージョンと比較して、DeepRec Extensionの全体的な設計は、CRDをカプセル化し、ネイティブTFJobを拡張し、分散タスクの柔軟なスケーリングとフォールトトレランスをサポートするためにAIMasterタイプのノードを追加します。さらに、当社が提供する機能はシステムのオーバーヘッドを最小限に抑え、IO を削減し、フォールト トレランスや弾性スケーリング時の追加コストを削減します。

これはユーザーが投稿したアプリケーションです。kubeflow と比較すると、このタイプでは AIMaster ノードが追加され、aimaster-operator によって管理されます。まず AIMaster Pod と Service が起動され、次に kubeflow::TFJob が作成されます。AIMaster は、分散トレーニングリソースの推定、柔軟なトレーニング、自動フォールトトレランス機能を担当します。

• 弾性トレーニング

AIMaster は最初にチーフノードを起動します。チーフはタスクのチェックポイントを推定します。チェックポイントが空の場合、ユーザーまたはシステムの初期 PS カウントに基づいて、適切な数の PS/ワーカーノードまたはリソースを初期化します。トレーニング中、各ロール (チーフ、PS、ワーカー) は実行ステータスを AIMaster Gazer に報告します。AIMaster ElasticTrainingController は、実行情報とユーザーが設定したポリシーに基づいて、弾性トレーニングの決定を行います。スケールアップまたはスケールダウンが発生すると、AIMaster はタスクステータスを維持し、各ロールはトレーニングを一時停止し、AIMaster はスケジューラから追加の PS/ワーカーリソースを要求し、すべての PS のサーバー情報を収集します。PS ノードはクラスターサーバー情報を動的に更新し、チーフはパラメーター再配布プロセスを駆動します。最後に、タスクステータスが実行中に設定され、スケールアップまたはスケールダウンプロセスが完了します。PS/ワーカープロセスは、プロセス全体を通して再起動しません。

• 自動フォールトトレランス

AIMasterでは、データマネージャーはチーフによって初期化され、各ワーカーのデータセットで使用されるグローバルサンプル状態情報を保持します。モデルがチェックポイント保存アクションをトリガーすると、データマネージャーもチェックポイント保存を実行し、モデルとの対応を記録します。ワーカーがフェイルオーバーすると、ワーカーには状態情報がないため、スケジューラはワーカーを再起動してトレーニングを続行します。PSがフェイルオーバーすると、障害が発生したPSが再起動しますが、正常なPSは変更されません。定期的なチェックポイント保存中にPS間でパラメータのバックアップが実行されるため、障害が発生したPSは他のPSから最新のチェックポイントデータを取得し、正常なPSはチェックポイントソースから読み取りおよびロード操作を開始することなく、最新のチェックポイントデータをローカルで取得します。

特徴

1. ゲイザー

Gazerは、分散学習タスクにおける各ノードとグラフ情報の監視モジュールとして位置付けられます。その結果は、TensorBoardを介してユーザーに可視化したり、RPCインターフェースを介してAIMasterに報告したりすることで、動的なリソーススケーリングや低速ノードの排除を可能にします。メトリクスは2つのレベルに分かれています。

• CPU やメモリの使用量などのマシンレベルの情報。
• PS サブグラフの実行時間や PS サブグラフのリソース使用量などのグラフ レベルの情報。

実装スキームはOperation + Grapplerです。上記の統計情報を収集するためのカスタムOpが組み込まれています。ユーザーがグラフを構築した後、グラフの書き換えによってOpがサマリーサブグラフに追加されます。実行時には、SummaryHookを利用して定期的にグラフを実行し、AIMasterに情報を報告したり、Summaryに情報を書き込んだりします。ユーザーはTensorBoardを介してこの情報を照会できます。下の画像は、あるタスクのGazer情報を示しており、PS:0ノードが遅いことを示しています。これにより、ユーザーはトレーニングタスクにおける各ノードのステータス情報をリアルタイムで確認できます。ユーザーはGazerモジュールで監視メトリクスをカスタマイズして監視できます。

2. リソースの見積もり

ほとんどの本番環境タスクには定期的なモデルトレーニングが含まれるため、履歴情報からリソース消費データを取得することは直感的なアプローチです。既存のモデルチェックポイントを利用してモデルサイズを分析し、PS数とリソースの初期化に必要な電力貯蔵（PS）リソース数を算出しました。ゼロからトレーニングしたモデルについては、ユーザー定義またはシステムによって割り当てられたデフォルト値に基づいてPSリソースを初期化しました。このリソース推定機能により、ユーザーはトレーニング開始時に適切なリソースを要求できるため、トレーニングコストが削減され、実際のオンライン本番環境タスクにおけるリソースを42%節約できます。

3. 自動弾性トレーニング

ワーカーノードはステートレスであるため、自動エラスティックトレーニングはワーカーノードの弾力性という点では比較的シンプルです。このセクションでは、主にPSエラスティックスケーリングのシステム動作について説明します。原則として、各ロールのタスクとPythonプロセスの再起動は避けます。

Elastic トレーニングは、Elastic GrpcServer、動的グラフ変換とパラメータ移行、AIMaster コントローラ、動的データ配布の各モジュールで構成されています。

• 弾性グループサーバー

DeepRecはElasticGRPCServerを登録します。これは、従来のGRPCServerとは異なり、ノードの追加と削除を含むクラスタ内のノードの動的な更新をサポートし、分散セッション情報を更新することで、新しいクラスタ環境でのグラフ実行の正確性を確保します。クラスタ全体のノード情報はAIMasterによって維持され、ノードの弾性スケーリング後に最新のクラスタ情報がPS/Workersにブロードキャストされます。

• 動的グラフ変換とパラメータ転送

計算グラフは、パラメータの種類に基づいて、密パラメータと疎パラメータに分類できます。疎パラメータはさらに、静的形状変数と動的形状埋め込み変数に分類できます。

パラメータサーバーの負荷分散を考慮すると、PSノード数が変化した場合、VariableとEmbeddingVariableの両方のパラメータを再分配する必要があります。Variableは、TensorFlowネイティブのPartitioner戦略に従って新しいパーティションに変更されます。EmbeddingVariableはキーバリュー属性であるため、疎パラメータを持つIDパーティションは、計算グラフ内のOperationsを通じて変更されます。また、データの一貫性を確保するために、すべてのパラメータを再配置する必要があります。上記はすべて、TensorFlowのOperation + Grapplerを使用して実装され、DeepRec/TensorFlowから分離されています。

• AIマスターコントローラー

AIMasterは、スケールアップおよびスケールダウンのプロセス中、ステートマシンを維持します。すべてのWorker/PS/Chiefは、AIMasterコントローラーからスケールアップおよびスケールダウンのステータスを取得し、Workerが登録したHookメカニズムを組み合わせることで、それぞれのロジックの実行を完了し、AIMasterのエラスティックプロセスの状態を切り替えます。

• 動的データ配信

動的なデータ分散は、ワーカーの弾力性を支えるための必須条件です。AIMasterは、グローバルサンプルの読み取りステータス情報を記録します。初期化はチーフによって完了します。トレーニング中、各チーフ/ワーカーはAIMasterからデータシャードを取得してトレーニングを行い、処理速度の遅いワーカーがトレーニングタスク全体を遅くする状況を効果的に回避します。

自動弾性トレーニング機能は、主に2種類の問題を解決します。まず、継続的に成長するスパースモデルをトレーニングする場合、ユーザーが初期リソースを小さく設定しすぎると、トレーニング中にPS OOMが発生し、トレーニングタスクが失敗します。その場合、ユーザーは手動でリソースを追加申請し、タスクを再開する必要があります。次に、リソースの実際の使用状況が動的に変化するシナリオでは、システムが自動的にリソースを適応させ、モデルのトレーニングコストを削減します。

次の図は、さまざまなモデルで自動弾性トレーニング機能を有効にした結果のタスク リソース節約率と実際のリソース使用率を示しています。

4. 自動フォールトトレランス

フォールト トレランス機能は、データ配布のフォールト トレランスと PS パラメータのフォールト トレランスの 2 つの部分に分かれています。

• データ分散におけるフォールトトレランス

データとモデルの消費状態が組み合わさって、トレーニングタスク全体の状態を構成します。モデルのトレーニング中は、データ状態の永続化とモデルパラメータの永続化が同時に発生します。ノードに障害が発生した場合、モデルとデータの状態を同期的に復元する必要があります。そうしないと、サンプルデータのトレーニングが繰り返されたり、サンプルが失われたりして、モデルのトレーニング精度に影響を及ぼします。

• PS（エラー許容度）

PSパラメータを効率的に復元するために、DeepRec Extensionはノードの残りのメモリまたはディスクリソースをバックアップに利用し、リモート分散ファイルシステムからのチェックポイントのネットワーク間読み取りを回避して、モデルの復元効率を向上させます。通常のトレーニング中は、チェックポイントの保存と同時に冗長PS保存がトリガーされ、グラフベースでローカルおよびリモートパラメータのバックアップが実行されます。このバックアッププロセスはバックグラウンドで実行され、通常のチェックポイントの保存時間には影響しません。パラメータのバックアップ媒体はユーザーが指定し、メモリまたはディスクのいずれかになります。PSノードがフェイルオーバーすると、正常なPSノードは最新のチェックポイントをローカルで素早く復元し、障害が発生したPSノードはスケジュール設定および再起動後、他のPSノードからチェックポイントを素早く復元するため、モデルの復元効率が大幅に向上します。モデルファイルは、メモリ管理を容易にするためにスライステンソルを使用してデバイス間で転送されます。上の図は、保存と復元のグラフを示しています。

左の図は、自動フォールトトレランスを有効にした後の分散トレーニングのE2Eスループット性能の向上を示しています。単一PSノードのモデルサイズは10GB、フェイルオーバー試行回数は0～4回です。右の図は、異なるモデルサイズにおけるPSノードモデルの読み込み時間のパフォーマンス向上を示しています。

要約と展望

DeepRec Extensionは、ユーザーエクスペリエンス、リソース利用率、トレーニング効率、信頼性の観点から、分散トレーニングの効率性に関する課題を体系的に解決します。今後の最適化では、以下の領域に重点を置き、トレーニングパフォーマンスをさらに向上させます。

1. 弾力性戦略の観点から、自動弾力性トレーニングは現在、PSロールとWorkerロールの弾力性スケーリングのみをサポートしています。PSとWorkerの共同弾力性トレーニングにより、リソースをより効率的に活用し、トレーニングスループットを最大化し、ユーザーにとってタスク送信インターフェースをよりシンプルにすることができます。今後、オープンソースコミュニティにおいて議論と設計実装が行われる予定です。
2. PSノードが弾性スケーリングを実行した後、一定量のパラメータを移動する必要がありますが、これは一部のシナリオでは無視できません。スパースパラメータのホット特性とコールド特性を考慮すると、オンデマンドでパラメータを移動する戦略は、ネットワークオーバーヘッドを大幅に削減できます。
3. フォールト トレランスと増分チェックポイントの組み合わせにより、バックアップ パラメータのサイズが削減され、システムのオーバーヘッドがさらに削減され、現在のフォールト トレランスに基づいてフォールト トレランスの効率が向上します。

DeepRec Extensionプロジェクトは、GitHub（ https://github.com/DeepRec-AI/extension ）でオープンソース化されました。開発者はこれを迅速にデプロイし、二次開発を行うことで、監視情報や柔軟なトレーニング戦略をカスタマイズできます。DeepRecコミュニティの発展に貢献していただける開発者の皆様を歓迎いたします。