debug-graph-breaks

debug-graph-breaks

説明

スクリプトを実行して出力を読み取り、修正を直接適用して再実行し、修正を検証することで、torch.compile のグラフブレークを診断して解消します。目標が fullgraph=True の達成でも、影響の大きいブレークを減らすことでも対応します。

Claude への指示

概要

このスキルが呼び出されたら、ユーザーが torch.compile でコンパイルしたコードのグラフブレークを見つけ、理解し、修正できるよう支援してください。ワークフローは 検出、診断、修正 の 3 つのフェーズで構成されます。

グラフブレークのドキュメントサイト（https://meta-pytorch.org/compile-graph-break-site/）には、既知のグラフブレークタイプごとの詳細なページがあります。ユーザーが特定のグラフブレーク（例：GB0059）に遭遇した場合は、対応するページ（例：`https://meta-pytorch.org/compile-graph-break-site/gb/gb0059.html`）を取得し、そのブレークに固有の詳細な背景情報、例、修正案を確認してください。ユーザーがグラフブレークのウェブサイトリポジトリのクローンへのローカルパスを提示した場合は、代わりにそのディレクトリからドキュメントファイルを直接読んでください。

フェーズ 1：検出 — グラフブレークを見つける

ユーザーがスクリプトのパスを提示し、既存のログがない場合は、TORCH_LOGS="graph_breaks" を指定してスクリプトを実行してください。

TORCH_LOGS="graph_breaks" python your_script.py 2>&1

次の内容がグラフブレークごとに出力されます。

• ユーザーコード内でブレークを引き起こしたファイルと行
• グラフブレークの理由（GB タイプ、説明、ヒントを含む）
• ユーザーコードのトレースバック
• そのグラフブレークに固有のドキュメントへのリンク

出力が曖昧な場合や、ブレークの発生元が不明な場合は、詳細モードで再実行してください。

TORCH_LOGS="graph_breaks" TORCHDYNAMO_VERBOSE=1 python your_script.py 2>&1

詳細モードでは、内部の Dynamo スタックトレースと直近のバイトコード命令が追加されるため、通常の出力が曖昧な場合に役立ちます。

ユーザーが既存の TORCH_LOGS 出力を提示した場合は、そのまま解析し、フェーズ 2 に進んでください。

ユーザーからの要望がある場合、または関連する成果物が提示された場合は、次の検出方法も利用できます。

オプション B：tlparse レポート

ユーザーが TORCH_TRACE による構造化 HTML レポートである tlparse レポートを持っている場合、次のものを提示できます。

• tlparse の出力があるローカルディレクトリのパス
• レポートの URL

tlparse レポートには次の内容が含まれます。

• index.html：すべてのコンパイルをツリー表示する スタックトライ を含むメインページです。ステータスは色分けされ（緑＝成功、ライム＝グラフブレーク、赤＝エラー）、各コンパイルは [0/0]（フレーム 0、最初のコンパイル）や [1/0]（フレーム 1、最初のコンパイル。多くの場合、グラフブレーク後の再開）などの コンパイル ID で識別されます。
• failures_and_restarts.html：すべての RestartAnalysis イベント（グラフブレーク）とコンパイル失敗の表で、それぞれのグラフブレークの理由の詳細が記載されています。
• コンパイルごとのディレクトリ（例：-_0_0_0/、-_1_0_1/）：ビルド成果物が含まれます。
  • dynamo_graph_break_reason_N.txt — ユーザーコード内の位置、理由、説明、ヒント、ユーザースタックトレース、内部の Dynamo トレースバックなど、グラフブレークの詳細
  • dynamo_output_graph_N.txt — そのコンパイルで生成された FX グラフ
  • compilation_metrics_N.html — コンパイル時間、グラフのメトリクス（演算子、ノード、入力）、再起動の理由、ガード数、キャッシュのメトリクス
  • *_ORIGINAL_BYTECODE_N.txt / *_MODIFIED_BYTECODE_N.txt — Dynamo による変更前後のバイトコード

tlparse レポートを読む際は、次の手順に従ってください。

1. まず failures_and_restarts.html を確認し、すべてのグラフブレークの概要を把握します。
2. 詳細を確認するには、各 dynamo_graph_break_reason_N.txt を読みます。
3. サブグラフごとのコンパイル時間への影響を把握するには、compilation_metrics_N.html を使用します。
4. _1 のサフィックスが付いたコンパイル ID（例：[0/0_1]）は再起動を示します。_1 は、試行 _0 でグラフブレークが見つかった後に Dynamo が解析を再開したことを意味します。

オプション C：fullgraph=True（一度に 1 つずつブレークを修正）

すべてのブレークを一度に修正するのではなく、1 つずつ修正したい場合は、fullgraph=True を使用してください。

torch.compile(fn, fullgraph=True)(*args)

最初に遭遇したグラフブレークでエラーが発生します。それを修正して再実行し、次のブレークを修正してください。

フェーズ 2：診断 — 各グラフブレークを理解する

見つかったグラフブレークごとに、以下を行います。

1. グラフブレークの理由を解析します。 出力から GB タイプの文字列（例：「Failed to trace builtin operator」「Unsupported Tensor.item() call with capture_scalar_outputs=False」）と、GB のドキュメント URL を抽出します。

2. グラフブレークのドキュメントページを取得します。 ログ出力にある URL（例：https://meta-pytorch.github.io/compile-graph-break-site/gb/gb0059.html）を使って、そのグラフブレークタイプの詳細情報、例、推奨される修正方法を確認します。グラフブレークのサイトのローカルクローンが提供されている場合は、代わりにローカルディレクトリ内の該当ファイルを読みます。

3. ログ出力のヒントカテゴリを使って、グラフブレークを分類します。

   • USER_ERROR：ユーザーのコードにバグがあり、イーガーモードでも失敗します。torch.compile を使わずにスクリプトを実行して確認します。
   • FUNDAMENTAL：Dynamo はこのパターンを根本的にトレースできません（例：ジェネレーター、特定の動的制御フロー）。ユーザーはコードを再構成する必要があります。
   • SUPPORTABLE：Dynamo は理論上このパターンをサポートできますが、まだ対応していません。ユーザーは PyTorch の Issue を作成するか、回避策を探すことができます。
   • CAUSED_BY_EARLIER_GRAPH_BREAK：先に以前のグラフブレークを修正します。こちらはその過程で解消される可能性があります。
   • DYNAMO_BUG：Dynamo のバグである可能性が高いです。Issue の作成を提案します。

4. ユーザーコード内の場所を特定します。 ユーザーのスタックトレースを使って、ブレークの原因となっているユーザーコードの正確な行を特定します。そのコードを直接読んで、状況を把握します。

5. ユーザーのスクリプトを読みます。 周辺の状況を把握するため、ユーザーのスクリプト内の関連コードを必ず読みます。グラフブレークのドキュメントページにある例と比較し、ブレークの原因となっている正確なパターンを特定して、最も的確な修正方法を判断します。

6. 優先順位を付けます。 すべてのグラフブレークが同等に重要なわけではありません。どれから修正するかを決めます。

   • 簡単なものから先に解決します：print() 呼び出し、ログ出力、単純な .item() 呼び出しには、簡単な修正方法があります。
   • 影響の大きいものを優先します：ホットループ内やコアのフォワードパス内のグラフブレークは、初期化や頻度の低いコードパス内のものより重要です。
   • CAUSED_BY_EARLIER_GRAPH_BREAK は最後にします：先行するグラフブレークを修正すると、自動的に解消されることがあります。
   • 連鎖するブレーク：ループ内のグラフブレーク（GB7000）は、フレーム全体をイーガーモードへフォールバックさせます。内側のブレークを修正すると、両方が解消されます。
   • ユーザーが fullgraph=True を必要としていない場合は、パフォーマンスへの影響が大きいブレークと、そのままにできるブレークを見極めます。

フェーズ 3：修正 — グラフブレークを解消する

グラフブレークごとに、そのタイプに応じた修正を適用します。変更を加える前に必ずユーザーのコードを読み、そのうえでコードを直接変更します。

サードパーティライブラリ（例：HuggingFace transformers）内のグラフブレークを修正する

グラフブレークがサードパーティライブラリ（ユーザーのコードではない）に起因する場合、ユーザーがライブラリを直接変更することはできません。代わりに、次の方法を使います。

1. ライブラリのソースを調べます。 ユーザーがインストール済みパッケージへのパス（例：transformers の site-packages ディレクトリ）を提供した場合は、該当するソースファイルを読んで、問題のあるコードとその呼び出し元を把握します。

2. サブクラスでオーバーライドします。 グラフブレークを起こしているメソッドをコンパイル可能な実装でオーバーライドするモデルのサブクラスを作成します。これにより、ライブラリを変更せずに済み、さまざまな環境で利用できます。

from transformers import BartModel

class CompilableBartModel(BartModel):
    def _override_problematic_method(self, *args, **kwargs):
        # Compile-compatible reimplementation
        ...

model = CompilableBartModel(config)
compiled_model = torch.compile(model, fullgraph=True)

3. fullgraph=True で検証します。 サブクラスによる修正を適用した後、fullgraph=True で再実行し、すべてのグラフブレークが解消されたことを確認します。

この方法は、自己完結型で再利用可能な修正を実現できるため、モンキーパッチよりも推奨されます。

エスケープハッチ（コンパイルに安全な代替手段）

以下の API を使うと、グラフブレークを起こさずにコンパイル領域内でコードを実行できます。

torch._higher_order_ops.print(format_str, *args, **kwargs)

• print() のコンパイルに安全な代替手段です。Python のフォーマット文字列構文を使用します。
• 例：torch._higher_order_ops.print("Activated shape: {}", h.shape)

torch._dynamo.config.reorderable_logging_functions

• print、logging.info、カスタムロガーなどのロギング関数をこのセットに追加します。Dynamo はロギング呼び出しを維持したまま、グラフブレークを避けるように関数の順序を変更します。
• 例：torch._dynamo.config.reorderable_logging_functions.add(print)

torch._dynamo.config.ignore_logging_functions

• Dynamo がトレース中にロギング関数を完全に無視するよう、このリストに追加します（呼び出しは破棄されます）。
• 例：torch._dynamo.config.ignore_logging_functions = [print]

@torch._dynamo.decorators.leaf_function

• Dynamo と AOT autograd の両方に対して、関数を不透明にするデコレーターです。関数は実行時にイーガーモードで実行されます。コンパイル時のシェイプ推論には @fn.register_fake が必要です。
• 次の場合に使用します：関数に副作用がある場合、コンパイルできないライブラリを呼び出す場合、または関数内部がコンパイルのメリットを受けない場合。

@torch._dynamo.decorators.leaf_function
def my_function(x):
    # This runs eagerly at runtime, no graph break
    ...

@my_function.register_fake
def my_function_fake(x):
    # Return a tensor with the correct shape/dtype for compile-time tracing
    return torch.empty_like(x)

• 引数を変更する関数の場合：@leaf_function(mutates_args={"buf"})

@torch._dynamo.decorators.nonstrict_trace

• Dynamo に対して関数を不透明にしつつ、AOT autograd ではトレースするデコレーターです。関数本体にグラフブレークを起こすコードを含めることができますが、AOT autograd は引き続き関数内をトレースするため、演算は Inductor によってコンパイルされます。
• 次の場合に使用します：関数本体にグラフブレークがあるものの、内部の演算には Inductor の最適化を適用したい場合。

@torch._dynamo.decorators.nonstrict_trace
def my_function(x):
    # Dynamo won't trace this, but AOT autograd will
    return x.relu()

torch.compiler.disable()

• 関数のコンパイルを完全にスキップするデコレーター／コンテキストマネージャーです。関数はイーガーモードで実行されます。呼び出し箇所でグラフブレークが発生します。
• 次の場合に使用します：関数がコンパイルと根本的に互換性がなく、コンパイルする必要もない場合。

カスタム演算子（torch.library）

• Python 関数をカスタム PyTorch 演算子として登録します。Dynamo はカスタム演算子を、グラフブレークを起こさない、型が明確な不透明ノードとして扱います。Inductor は実行時にこれらを呼び出せます。
• 次の場合に使用します：明確に定義されたテンソルの入出力シェイプを持つ、サードパーティ製の関数やトレースできない関数をラップする場合。

import torch.library

@torch.library.custom_op("mylib::my_op", mutates_args=())
def my_op(x: torch.Tensor) -> torch.Tensor:
    # arbitrary non-traceable code here
    return x.numpy().copy()  # example

@my_op.register_fake
def my_op_fake(x):
    return torch.empty_like(x)

• @leaf_function とは異なり、カスタム演算子は公開 torch.library API の一部であり、FX グラフ内に名前付きノードとして表示されます。そのため、複数のコンパイル済み関数で演算子を再利用する場合により適しています。

グラフブレークの種類ごとの一般的な修正方法

print() / ロギング（GB0059：「組み込み演算子のトレースに失敗しました」）

• 修正方法（推奨順）：
  1. torch._dynamo.config.reorderable_logging_functions.add(print) — printを保持し、コードの変更は不要です
  2. torch._dynamo.config.ignore_logging_functions = [print] — コンパイル中にprintを省略します
  3. torch._higher_order_ops.print("format {}", arg)に置き換えます — コンパイルに対応したprint HOPです
  4. デバッグ専用であれば、printを完全に削除します

データ依存の分岐（GB0035、GB0170）

• if tensor.item() > 0:やif tensor.sum(): — 制御フローがテンソルの値に依存しています。
• 修正方法：単純なif/elseにはtorch.cond()を使用するか、データ依存の分岐が発生しないように構成し直します。単純なケースでは、分岐のない代替手段としてtorch.where()を使用できます。
• 分岐の条件となる値がトレース時に実際には定数である場合は、torch._check()を使用するか、Pythonの定数にします。

Tensor.item() / tolist()（GB0124、GB0109）

• これらはテンソルからスカラー値を取り出すため、データ依存が発生します。
• 修正方法：
  1. 値をテンソルのまま扱うように構成し直します（例：h.norm().item()の後にスカラー演算を行う代わりに、torch.norm(h)を使用します）。
  2. torch._dynamo.config.capture_scalar_outputs = Trueを設定して、グラフ内で.item()をキャプチャします。
  3. .item()の呼び出しをコンパイル領域の外に移動します。

サポートされていない関数呼び出し（GB0147）

• DynamoがトレースできないPython関数です。
• 修正方法：PyTorchに相当する機能があるか確認します。サードパーティライブラリの呼び出しの場合は、以下を試してください。
  1. 形状推論が可能な場合は、register_fakeとともに@leaf_functionを使用します。
  2. 関数本体にコンパイルのメリットがある演算が含まれる場合は、@nonstrict_traceを使用します。
  3. 関数をコンパイルする必要がない場合は、torch.compiler.disable()を使用します。
  4. 呼び出しをコンパイル領域の外に移動します。

ジェネレーターのトレース試行（GB0003）

• ジェネレーターはコンパイルできません。
• 修正方法：ジェネレーターをリスト内包表記または通常のループに変換するか、コンパイル対象の関数の外からジェネレーターを呼び出します。

サポートされていないコンテキストマネージャー（GB0142）

• 修正方法：コンテキストマネージャーをコンパイル領域の外に移動するか、Dynamoに対応した代替手段があるか確認します（例：torch.inference_mode()の代わりにtorch.no_grad()を使用します）。

モジュールレベルのフック（GB0083、GB0029）

• モジュールのバックワードフックには、コンパイル済みオートグラッドが必要です。
• 修正方法：コンパイル済みオートグラッドを有効にする（torch._dynamo.config.compiled_autograd=True）か、コンパイル中にフックを削除します。

ループ内のグラフブレーク（GB7000）

• for/whileループ内でグラフブレークが発生すると、フレーム全体がスキップされます。
• 修正方法：まず、ループ内の根本的なグラフブレークを修正します。これは必ず別のグラフブレークが原因です。

一般的な戦略

• コンパイルできないコードをtorch.compileの外に移動する： 関数を分割し、コンパイル可能な中核部分を1つの関数にまとめ、コンパイルできないセットアップ／後処理をその外に配置します。
• Pythonの組み込み関数をtorchの同等の関数に置き換える： 例：max() -> torch.max()、abs() -> torch.abs()、sorted() -> torch.argsort()。
• コンパイル済みコード内でテンソルをPythonの値に変換しない： .item()、.tolist()、.numpy()、bool(tensor)、int(tensor)。
• テンソル演算はテンソルのまま扱う： スカラー値を取り出してPythonの演算を行う代わりに、テンソルを直接処理するtorchの演算（例：torch.where、torch.clamp、torch.norm）を使用します。

反復ワークフロー

修正を適用した後、TORCH_LOGS="graph_breaks"（またはfullgraph=True）を指定してスクリプトを再実行し、次を確認します。

1. グラフブレークが解消されたか確認します。
2. 新たなグラフブレークが発生した場合は、フェーズ2～3を繰り返します。
3. 目標は、fullgraph=Trueでエラーなく実行できるようにするか、影響の大きいグラフブレークを解消することです。
4. 任意：ユーザーがPyTorchの開発者であり（torchがローカルでビルドされているかどうかで判断できます）、グラフブレークのメッセージが不明瞭な場合は、エラーメッセージの変更を提案するか、issueの作成を提案します。

ベンチマークの改善

グラフブレークを修正した後は、改善効果を測定できるようユーザーを支援します。次のような、動作に影響しにくい方法を提案してください。

コンパイル時間の測定：

• バックエンドとして torch._dynamo.utils.CompileProfiler を使用し、グラフごとのコンパイル時間を確認します。

with torch._dynamo.utils.CompileProfiler() as prof:
    compiled_fn = torch.compile(fn, backend=prof)
    compiled_fn(*args)
print(prof.report())

• サブグラフが少ないほど、コンパイルのオーバーヘッドが小さくなります。グラフブレークが発生するたびに、別個のコンパイルが追加されます。

実行時パフォーマンスの測定：

• シンプルな実時間の計測（CPU または CUDA の場合）：

import time

# Warmup
for _ in range(3):
    compiled_fn(*args)
if torch.cuda.is_available():
    torch.cuda.synchronize()

start = time.perf_counter()
for _ in range(100):
    compiled_fn(*args)
if torch.cuda.is_available():
    torch.cuda.synchronize()
elapsed = (time.perf_counter() - start) / 100
print(f"Average iteration: {elapsed*1000:.2f}ms")

修正前と修正後の比較：

• グラフブレークを修正する前後で、同じベンチマークを実行します。
• 主な指標：サブグラフ数（少ないほど良い）、合計コンパイル時間、イテレーションあたりの平均実行時間。
• サブグラフ数は、TORCH_LOGS="graph_breaks" の出力（「Graph break in user code」の行数をカウント）または tlparse レポートのスタックトライから確認できます。

主な参考資料

• グラフブレークのウェブサイト：https://meta-pytorch.org/compile-graph-break-site/
• グラフブレークのヒント：torch/_dynamo/graph_break_hints.py
• コンパイルのプログラミングモデル：https://docs.pytorch.org/docs/stable/user_guide/torch_compiler/compile/programming_model.html
• leaf_function：torch/_dynamo/decorators.py
• nonstrict_trace：torch/_dynamo/decorators.py
• torch.compiler.disable()：torch/_dynamo/decorators.py
• error_on_graph_break()：torch/_dynamo/decorators.py
• torch._higher_order_ops.print：torch/_higher_order_ops/print.py
• CompileProfiler：torch/_dynamo/utils.py