PFNでは、PFVM向けのONNXエクスポータとしてFX2ONNXを開発しています。PFVMはMN-Coreで深層学習モデルを実行するためのコンパイラであり、その入力形式としてONNXを採用しています。したがって、PyTorchで構築されたモデルをPFVMでコンパイルするには、まずはPythonのコードをONNXに変換するONNXエクスポータが必要になります（図1）。

図1: PyTorchで構築されたモデルをMN-Core上で実行するまでのフロー

ONNXエクスポータの詳細に関しては PFVM向けのONNXエクスポータの開発 – Preferred Networks Research & Development もご覧ください。

FX2ONNXは図2のような工程を経てモデルをONNXへとエクスポートします。まず、FakeTensorと呼ばれる仮想的な値を入力に与えてモデルを実行し、実行されたテンソル演算をトレースします。このトレースをもとにしてPyTorchの計算グラフであるfx.Graphを構築します。そして、fx.Graphの各ノードを読み替えてONNXを構築し、エクスポート結果として出力します。

図2: FX2ONNXがエクスポート時に行う処理

問題点: エクスポートによる挙動の変化

前述の通り、FX2ONNXはエクスポート時に仮想的にモデルを実行し、そのトレースに基づいてfx.Graphを構築します。この方式には次のような欠点があります。

実行されなかったテンソル演算はエクスポートされない
テンソル以外の値に対する演算はトレース対象でなく、エクスポートされない

これらの欠点により、Pythonのグローバルオブジェクトに依存したモデルをFX2ONNXでエクスポートした際に、元のモデルとエクスポート後のONNXの挙動が一致しない場合があります。

例として、図3に示した関数fをエクスポートする場合を考えます。関数fはグローバル変数aの値に応じてx + 1かx – 1を返し、呼び出しのたびにaの値をインクリメントします。この場合、エクスポート後のONNXに含まれるのはx + 1かx – 1のうち実行された方の演算のみです。どちらが含まれるかはエクスポート時のaの値に依存して決まります。また、aのインクリメントはintに対する演算であり、これはエクスポート後のONNXに含まれません。

図3: エクスポートによって挙動が変化するモデルの例

元のモデルでは関数呼び出し時のaの値によって処理が分岐する一方で、エクスポート後のONNXでは分岐先が固定されてしまいました。このような原理で、エクスポートによってモデルの挙動が変化することがあります。

現状のトレースの方式では、この種の「グローバルオブジェクトに依存した演算や分岐」がモデルに存在することを検出するのは難しいという課題がありました。またそれによって、サイレントにモデルの挙動差が発生する可能性があり、その発見や修正も難しいという状況がありました。

改良案: エクスポートによって挙動が変化しうるコードの検出

この問題に対する改良案として、「実行された分岐」と「その分岐の影響を受けうる、実行されたテンソル演算」をエクスポート時に検出する機能の実装に取り組みました。グローバルオブジェクトへの依存は検出が難しいため、スコープを限定し、実行された分岐に着目して検出を試みる方針です。

本機能の動作イメージとして、先ほど図3に示した関数fをエクスポートする際の検出例を図4に示します。

図4: エクスポート時に検出されるコードの例

この例では、まずグローバル変数aに依存した分岐が実行されたこと、そしてその後に実行されたx – 1の演算が分岐の影響を受ける可能性があることが検出されています。ユーザーがこれらの検出を確認し、手動で該当箇所のコードを確認、必要に応じて修正するようなユースケースを期待しています。

実装にあたっては、動的解析と静的解析の両手法を用いました。概要としては、動的解析で実行経路を追跡し、静的解析で関数の処理やデータの依存関係の情報を補完することによって検出を行っています（図5）。

図5: 検出機能の処理概要

動的解析による検出の実装

まず、動的解析にはPython 3.12で導入されたsys.monitoringモジュールの機能を用いました。このモジュールを用いることで、関数の呼び出しや分岐などの実行イベントを高速にフックできます。フック対象の分岐はifに限らず、for/while/and/or/assertなどによって発生する分岐も含まれます。

動的解析で行っている処理は大まかには以下の通りになります。

分岐が実行されたら分岐フラグを立てる
分岐フラグが立っている間に、テンソル演算が行われたら検出
関数からreturnしたら分岐フラグを降ろす

この処理では分岐の影響範囲を「分岐が実行されてから関数のreturnまで」と仮定し、その範囲で行われたテンソル演算を検出します。

実装上の工夫としては、スタックトレースの記録を自前で行うことでオーバーヘッドを小さくしています。スタックトレースはフック関数が呼ばれた時点でのコンテキストの把握や、検出結果の表示のために必要です。今回の場合フック関数が高頻度で呼ばれるため、その度にtracebackやinspectモジュールなどでスタックトレースを取得するとオーバーヘッドが大きくなってしまいます。そのため、sys.monitoringで関数呼び出しとreturnもフックし、自前でスタックトレースを記録するようにしました。

一方で、「分岐が実行されてからreturnするまで」という仮定の後の範囲でも、テンソル演算が分岐の影響を受ける可能性があります。例えば図6に示すような、分岐によって異なるスカラー値を返す関数の戻り値が、その後のテンソル演算で使われるようなケースは検出できません。このほかにも、主に分岐の影響範囲の仮定に起因して、動的解析のみでは網羅できないケースが存在します。

図6: 動的解析のみでは検出が難しいケース

静的解析による検出の実装

そこで、そのようなケースを検出するためにLibCSTによる静的解析を併せて導入しました。LibCSTではPythonのソースコードをパース・解析し、シンボルのスコープ情報や完全修飾名等のメタデータを取得できます。これらのメタデータをもとに、動的解析では追跡が難しい参照関係の情報を補強しています。

より具体的には、動的解析のステップ3において、静的解析で「関数の戻り値が分岐の影響を受ける可能性がある」と判定された場合、分岐フラグを降ろさずに関数の呼び出し元へ伝搬します。これによって、関数をまたぐような分岐の影響を一定程度追跡できます。

「関数の戻り値が分岐の影響を受ける可能性がある」条件は以下のいずれかを満たす場合と定めました。

分岐内でreturnが行われる
返り値を構成する変数が、分岐内で代入される
返り値を構成する関数について、過去に「関数の戻り値が分岐の影響を受ける可能性がある」と判定された

実装面では、動的解析で関数returnをフックしたタイミングでLibCSTによる静的解析を行い、その関数が上記の条件を満たしているか判定します。その結果を分岐フラグの伝搬条件として利用しています。

結果

一般的な公開モデルやテストケースとして用意したモデルをエクスポートし、検出されたコードの例を図7に示します。図中の赤線はそれぞれ以下の意味を持ちます。

赤の実線：実行された分岐
赤の波線：その分岐の影響を受けうる、実行されたテンソル演算

図7: 検出されたコードの例

図7では複数のケースで意図どおりの検出結果が得られることを確認しました。特に興味深いケースとして、図7の右側中央に示したStepLr.get_lr関数の例が挙げられます。この関数は一定のエポック数ごとに学習率を更新するために用いられます。こうした学習率スケジューラのような典型的なコードにおいても、本検出機能が有効に機能することを実証できました。

現状の課題点

今回実装した検出機能に関してはいくつかの点で改善の余地があります。

1点目に、誤検出・過検出がともに発生する可能性があります。実行トレースに依存して検出ロジックを実装しているため、実行されなかった分岐・テンソル操作は検出できません。また、静的解析でも分岐とテンソル演算の依存関係を追いきれない場合があり、その際にも検出もれが発生する可能性があります。対して、分岐とテンソル演算の依存関係が実際には無視できる場合においても、分岐の影響範囲を関数単位の粒度で見ていることによって、誤って検出してしまう可能性があります。

2点目に、検出機能を有効化した際にエクスポートにかかる時間が増加します。図8に、モデルからfx.Graphを構築する際の所要時間を計測した結果を示します。計測にはモデルとしてEfficientNet B0とGPT-2を使用しました。左から順に検出機能無効時、検出機能有効時（動的解析のみ）、検出機能有効時（動的解析と静的解析を併用）の場合です。