facebookresearch/HolisticTraceAnalysis

Holistic Trace Analysis (HTA)

何であるか – HTA は、分散型 PyTorch 学習ジョブの遅さの原因を理解するための Python ライブラリです。PyTorch プロファイラ（Kineto）が生成するトレースファイルを読み取り、GPU が時間をどの部分に費やしているか（計算、通信、メモリ操作、アイドル）を明確に示すデータフレームと可視化を提供します。

主な機能

機能	得られるもの
時間的分解	GPU ごとの時間の内訳（計算、通信、メモリ操作、アイドル）
カーネル分解と分布	各ランクごとの実行時間の長い CUDA カーネルの一覧と、ランク間での実行時間のばらつきの可視化
アイドル時間分析	GPU がアイドル状態になる理由（ホスト待ち、他のカーネル待ち、不明）
通信-計算の重なり	通信と有用な計算が重なっている時間の割合
頻出カーネルパターン	特定の PyTorch 演算子（例：aten::linear）に対して最も繰り返される CUDA カーネルの検出
カーネル起動統計	非常に短い、非常に長い、異常に遅延したカーネル起動のヒストグラム
拡張カウンター	元のプロファイルにキュー長とメモリ帯域幅のトレースを追加し、ルーフライン分析に活用
トレース比較	2つの実行結果を並べて比較し、変更点を確認
CUPTIカウンターAPI（実験的）	GPUの低レベルパフォーマンスカウンターを取得し、より深いルーフライン分析を可能に

対象ユーザー – 複数の GPU/ホストで大規模モデルを学習している人で、手動で原始的なプロファイリング出力を掘り下げる必要なく、体系的かつスクリプト可能な方法でボトルネックを特定したい人。

インストール

# PyPIから（推奨）
pip install HolisticTraceAnalysis

# ソースから
git clone https://github.com/facebookresearch/HolisticTraceAnalysis.git
cd HolisticTraceAnalysis
git submodule update --init
pip install -r requirements.txt
pip install -e .

Linux および macOS で Python ≥ 3.10 で動作します。

一般的なワークフロー

1. トレースの収集 – PyTorch プロファイラを使用して（torch.profiler.profile(..., record_shapes=True, with_stack=True, ...)）
2. すべてのトレースファイルを1つのフォルダにまとめる – HTA は1回の実行ごとに1つのディレクトリを想定しています。
3. ノートブック（またはPythonスクリプト）を開き、TraceAnalysisオブジェクトを作成：

   from hta.trace_analysis import TraceAnalysis
   analyzer = TraceAnalysis(trace_dir="/path/to/traces")
   

4. 必要なメソッドを呼び出す。例：

   temporal = analyzer.get_temporal_breakdown()
   kernels  = analyzer.get_gpu_kernel_breakdown()
   idle     = analyzer.get_idle_time_breakdown()
   overlap  = analyzer.get_comm_comp_overlap()
   

   各メソッドは pandas DataFrame を返し、プロットや調査に利用可能。
5. より深い分析には、examples/ ディレクトリ内のデモノートブック（trace_analysis_demo.ipynb, trace_diff_demo.ipynb）を使用。

設定 – ロギングは hta/configs/logging.config で制御可能。ログレベルの変更やカスタム設定ファイルの指定が可能。

ドキュメントとサポート – 完全なAPIドキュメントは https://hta.readthedocs.io で公開されています。MITライセンスでオープンソースであり、プルリクエストによる貢献を歓迎しています。

---

結論 – HTA は、原始的な PyTorch プロファイラのトレースを、実行可能なパフォーマンスレポートに変換し、分散学習パイプラインの最適化を容易にします。