トラブルシューティングは難しいですか? xpu_timer は大規模モデルのシームレスなトレーニングを保証します!

著者について: Zhang Ji は、システムレベル/ネットワーク最適化に重点を置いて、検索と推奨/LLM のトレーニング最適化に取り組んでいます。

背景

大規模モデルのパラメータ数が数十億から数兆へと急増するにつれ、学習規模の急速な拡大はクラスターコストの大幅な増加につながるだけでなく、システムの安定性、特にマシン障害の頻発という無視できない問題を引き起こしています。大規模な分散学習タスクでは、トラブルシューティングとパフォーマンス最適化のために、可観測性が不可欠となっています。そのため、大規模モデル学習の分野で働く人々は、必然的に以下の課題に直面しています。

トレーニング中は、ネットワークやコンピューティングのボトルネックなどのさまざまな要因によりパフォーマンスが不安定になり、変動や劣化が生じる可能性があります。
分散トレーニングでは、複数のノードが連携して動作します。いずれかのノードに障害が発生した場合（ソフトウェア、ハードウェア、ネットワークカード、GPUの問題など）、トレーニングプロセス全体を一時停止する必要があり、トレーニング効率に重大な影響を与え、貴重なGPUリソースを無駄に消費してしまいます。

ただし、これらの問題は、主に次の理由により、実際の大規模モデルのトレーニング中に検出することが困難です。

トレーニングプロセスは同期処理であるため、全体的なパフォーマンス指標に基づいてどのマシンに問題があるのかを特定することは困難です。遅いマシンは、トレーニング全体の速度を低下させる可能性があります。
トレーニングのパフォーマンスが遅い場合、トレーニングロジックやフレームワークの問題ではなく、環境が原因となることがよくあります。トレーニング関連のモニタリングデータがなければ、タイムラインを印刷しても実際には意味がなく、タイムラインファイルの保存にも大きなストレージ容量が必要になります。
分析ワークフローは複雑です。例えば、トレーニングがハングした場合、トーチがタイムアウトする前にすべてのスタックトレースを出力し、その後分析を開始する必要があります。大規模なタスクでは、これをトーチのタイムアウト時間内に完了するのは困難です。

大規模な分散学習ジョブでは、トラブルシューティングとパフォーマンス向上のために可観測性が不可欠です。Ant Financialは、大規模学習におけるAI学習の可観測性要件を満たすため、xpu_timerライブラリを開発しました。将来的にはDLRoverでxpu_timerをオープンソース化する予定ですので、皆様の開発へのご協力をお待ちしております :) xpu_timerライブラリは、cublas/cudartライブラリを傍受することで、学習中の行列乗算/集合通信操作の時間をcudaEventで計測するプロファイリングツールです。また、タイムライン分析、ハング検出、ハングスタック分析機能も備えており、様々な異種プラットフォームに対応できるように設計されています。このツールには以下の特徴があります。

コードに干渉せず、トレーニングのパフォーマンスを低下させることがなく、トレーニングプロセスに永続的に存在することができます。
ユーザーの影響を受けず、フレームワークに依存しない
低損失/高精度
指標を集約して配信し、さらなるデータの処理と分析を容易にします。
高い情報保存効率
ユーザーフレンドリーなインターフェース: 他のシステムとの統合やユーザーによる直接操作を容易にするユーザーフレンドリーなインターフェースを提供し、洞察と意思決定のプロセスを加速します。

デザインスキーム

まず、トレーニングのハング/パフォーマンス低下の問題に対処するために、常駐カーネルタイマーを設計しました。

多くの場合、トレーニングのハングはnccl操作によって発生します。通常は、行列乗算を記録し、通信を設定するだけで済みます。
単一マシンのパフォーマンス低下（ECC、MCE）については、行列乗算を記録するだけで済みます。行列乗算を解析することで、ユーザーの行列形状が科学的であるかどうかも確認できるため、TensorCoreのパフォーマンスを最大限に引き出すことができます。様々なフレームワークで行列乗算を実装する際に、Cublasを直接使用できます。

そのため、インターセプションはカーネル起動レイヤーで発生するように設計し、実行時にLD_PRELOADを設定することで、対象となる操作のトレースを実行できます。この方法は、現在のトレーニングフレームワークのほとんどが動的リンクを使用しているため、動的リンクのシナリオでのみ使用できます。NVIDIA GPUの場合、次のシンボルに注目できます。

ibcudart.so

cudaLaunchKernel
cudaLaunchKernelExC

libcublas.so

cublasGemmEx
cublasGemmStridedBatchedEx
cublasLtMatmul
cublasSgemm
cublasSgemmStridedBatched

異なるハードウェアに適応する場合、異なるトレース機能を実装するために異なるテンプレートクラスが使用されます。

ワークフロー

PyTorchを例に挙げると、Launch ThreadはTorchのメインスレッドであり、Working Threadはライブラリ内のワーカースレッドです。ここでは、前述の7つのカーネルを抽出します。

使用方法と効果

前提条件

NCCLはlibtorch_cuda.soに静的にコンパイルされます。
torchはlibcudart.soを動的にリンクします

NCCLが動的にリンクされている場合、カスタム関数オフセットを指定できます。これは実行時に動的に解決されます。Pythonパッケージをインストールすると、以下のコマンドラインツールが利用できるようになります。

xpu_timer_gen_syms	nccl 解析を動的に生成するために使用される動的に注入されたライブラリ関数のオフセット
xpu_timer_gen_trace_timeline	Chromeトレースを生成するために使用
xpu_timer_launch	フックパッケージを取り付けるために使用
xpu_timer_stacktrace_viewer	タイムアウト後にビジュアルスタックを生成するために使用されます。
xpu_timer_print_env	libevent.so のアドレスを出力し、コンパイル情報を出力します。
xpu_timer_dump_timeline	タイムラインダンプをトリガーするために使用される

 LD_PRELOAD の使用法: XPU_TIMER_XXX=xxx LD_PRELOAD=`libevent_hook.so へのパス` python xxx

タイムラインのリアルタイムダイナミックキャプチャ

各ランクには独自のポートサービスがあり、コマンドはすべてのランクに同時に送信する必要があります。サービスの開始に使用されるポートはbrpcポートです。各ランクのトレースデータサイズは32バイトで、1000件のレコードが格納され、合計32KBになります。生成されるタイムラインJSONのサイズは150KB * ワールドサイズで、Torchのタイムラインの基本的な使用方法よりも大幅に小さくなります。

使用法: xpu_timer_dump_timeline [-h]
--host HOST ダンプするホスト
--rank RANKはホストのランクに対応します
[--port PORT] ダンプ先のポート。デフォルトは 18888 です。ノードがすべてのカードを使用している場合は、これを変更する必要はありません。 [--dump-path DUMP_PATH] ダンプ先のアドレス。絶対パスを使用し、長さは 1000 文字を超えてはなりません。
[--dump-count DUMP_COUNT] ダンプするトレースの数 [--delay DELAY] このコマンドを開始してからダンプを開始するまでの待機秒数
[--dry-run] パラメータを印刷する

独立した状況

 xpu_timer_dump_timeline \
--ホスト 127.0.0.1 \
--rank "" \
--遅延 3 \
--dump-path /root/lizhi-test \
--dump-count 4000
マルチマシンシナリオ # 下の図に示すように、ジョブにマスター/ワーカーの混合構成がある場合（マスターもトレーニングに参加します）
# --host xxx-master --rank 0 と書くことができます
# まだ不明な場合は、--dry-run を使用してください
xpu_timer_dump_timeline \
--ホストワーカー\
--ランク0-3 \
--delay 3 --dump-path /nas/xxx --dump-count 4000
xpu_timer_dump_timeline \
--ホストワーカー --ランク 1-3 \
--host マスター --rank 0 --dry-run
/root/timeline に 1000 でダンプする
ダンプホスト ['worker-1:18888', 'worker-2:18888', 'worker-3:18888', 'master-0:18888']
その他のデータ {'dump_path': '/root/timeline', 'dump_time': 1715304873, 'dump_count': 1000, 'reset': False}

次のファイルが対応するタイムラインフォルダーに追加されます。

次に、このディレクトリで xpu_timer_gen_trace_timeline を実行します。

 xpu_timer_gen_trace_timeline

次の 3 つのファイルが生成されます。

merged_tracing_kernel_stack 補助ファイル、元のフレームグラフファイル
trace.json からマージされたタイムライン
tracing_kernel_stack.svg、行列乗算/nccl の呼び出しスタック。

ラマレシピの32カードSFT分析事例

タイムラインはおおよそ以下のようになります。各ランクごとにmatmul/ncclデータが2行表示され、すべてのランクが表示されます。ここでは前後の情報は表示されません。前後の時間は順方向の時間の2倍であるため、期間で大まかに判断できます。

タイムラインの前進、約87ミリ秒

タイムラインを約173ミリ秒逆戻り

合計48層あり、合計時間は(173+87)*48 = 12480msです。lmhead、埋め込み、その他の演算を加えると約13秒かかるため、全体の時間は一致します。さらに、タイムラインを見ると、通信時間が計算時間よりもはるかに長く、通信がボトルネックになっていることが分かります。

ハングスタック分析

pip を使用してパッケージをインストールした後、コマンドラインツールを使用して解析できます。デフォルトでは、カーネルは300秒後に詳細なスタック情報を出力します。SVG画像はChromeにドラッグすることで表示できます。`pstack` と `py-spy` を使用して対応するスタックを出力します。結果はトレーニングプロセスの標準エラー出力に保存されます。conda経由でgdbをインストールした場合、gdb Python APIを使用してLWP名を含むスタックを取得します。デフォルトでインストールされたgdb 8.2では、スタックの取得に失敗する場合があります。conda gdbのデフォルトの場所は `/opt/conda/bin/gdb` です。以下は、NCCLタイムアウトをシミュレートする2つのGPUのスタックです。

以下は、8 つの GPU と Illama 7B SFT を搭載した単一のマシンでのトレーニングの例です。

Pythonパッケージが提供するツールを使うと、集約スタックのフレームスタック図を生成できます。ここでは、ランク1のスタックが存在しないことがわかります。これは、8枚のカードを使ったトレーニング中に、-STOP rank1をkillすることでハングをシミュレートしているため、ランク1が停止状態になっているためです。

 xpu_timer_stacktrace_viewer --path /path/to/stack
実行後、パスに cpp_stack.svg と py_stack.svg の 2 つの SVG ファイルが生成されます。

スタックをマージする際、同一の呼び出しパスはマージ可能であると想定しています。つまり、スタックトレースは完全に一貫しています。したがって、メインスレッドがスタックするポイントはほぼ同じになります。ただし、ループやアクティブなスレッドが存在する場合、出力されるスタックトップは異なる場合がありますが、基盤となるスタックは同じです。例えば、Pythonスタックでは、すべてのスレッドが `[email protected]` でスタックします。さらに、フレームグラフのサンプル数は意味を持ちません。ハングが検出されると、すべてのランクで対応するスタックトレースファイルが生成されます（ランク1はサスペンドされているため、生成されません）。各ファイルには、完全なPython/C++スタックが含まれます。

以下に、スタックの種類を区別するために異なる色で表した、統合されたスタックを示します。Pythonでは、スタックの色はシアンと緑のみとなります。

シアンはCPython/Python
赤はC/その他のシステム関連情報を示します。
緑はTorch/NCCLを表します。
黄色は C++ を表します。

Pythonスタックは以下のとおりです。青いボックスが実際のスタックを表しています。命名規則は func@source_path@stuck_rank|leak_rank です。

`func` は現在の関数の名前です。gdb が関数名を取得できない場合は、`??` が表示されます。
`source_path` は、プロセス内の so/source ファイルのアドレスを参照します。
`stuck_rank` は、このスタックにプッシュされたランクを示します。連続するランク番号は、開始-終了のシーケンスに折り畳まれます（例：rank 0,1,2,3 -> 0-3）。
`leak_rank` は、どのスタックがこの場所に入らなかったかを示します。ここでのランク番号も折りたたまれます。

したがって、この図は、ランク0、ランク2～7がすべて同期ブロックの下にスタックしており、ランク1が入らなかったことを示しています。これは、ランク1に問題がある（実際には停止している）ことを示唆しています。この情報はスタックの先頭にのみ追加されます。

同様に、cppスタックから、メインスレッドが同期ブロックでスタックし、最終的にcuda.soの取得時間でスタックしたことがわかります。同様に、rank1のみにこのスタックトレースが欠落しており、これは__libc_start_mainを含むスタックがプロセスのエントリポイントを表していることを示唆しています。

一般的に、スタックには最も深いリンクが1つだけあると想定できます。フォークが発生した場合、異なるランクが異なるリンクにスタックされていることを意味します。

カーネルコールスタック分析

コールスタックを持つTorchのタイムラインとは異なり、タイムラインは作成時にtracing_kernel_stack.svgという対応するスタックファイルを生成します。このファイルをChromeにドラッグして表示できます。

緑は NCCL の動作を示します。
赤いテキストは matmul 操作を表します。
青いのは Python スタックです。

Grafanaダッシュボードの表示

今後の計画

NCCL/eBPF などのよりきめ細かいトレースを追加すると、トレーニング中のハングアップ問題の根本原因をより正確に分析および診断できるようになります。
さまざまな国産グラフィックカードを含む、より多くのハードウェアプラットフォームをサポートします。

DLRoverについて

DLRover（分散型ディープラーニングシステム）は、Ant GroupのAIインフラチームが管理するオープンソースコミュニティです。クラウドネイティブ技術を基盤としたインテリジェントな分散型ディープラーニングシステムです。DLRoverを利用することで、開発者はハードウェアアクセラレーションや分散実行といったエンジニアリングの詳細に煩わされることなく、モデルアーキテクチャ設計に集中できます。また、最適化ツールなど、ディープラーニング学習アルゴリズムを開発することで、学習をより効率的かつインテリジェントなものにすることができます。現在、DLRoverはKubernetesとRayを用いたディープラーニング学習タスクの自動運用・保守をサポートしています。AIインフラ技術に関する詳細は、DLRoverプロジェクトをご覧ください。

DLRover DingTalk技術交流グループに参加してください: 31525020959

DLRoverスター:

https://github.com/intelligent-machine-learning/dlrover

記事のおすすめ

AIトレーニングの計算効率向上：Ant FinancialのDLRover障害自己修復技術の革新的な実践

AIインフラアーキテクトを知る：大型高速模型「レーシングカー」のホイール交換を担当する人

[オンラインリプレイ] NVIDIA GTC 2024 カンファレンス | AIエンジニアリングのコスト削減策とは？トレーニングから推論まで、Ant Groupのフルスタック実践

HUOXIU