01 / UNDERSTAND
仕組みと使いどころ
TorchDynamoは、PyTorch 2のtorch.compileを構成するPythonレベルのJITコンパイラです。Pythonの柔軟性を保ちながら、実行前にPythonバイトコードを動的に変更し、連続するPyTorch演算をFXグラフとして抽出します。取得したグラフは、拡張可能なコンパイラバックエンドへ渡されてJITコンパイルされます。TorchDynamoの主な役割はグラフ取得であり、GPUやCPU向けコードへの変換はバックエンドが担います。既定のバックエンドはTorchInductorです。
02 / IN PRACTICE
TorchDynamoの活用例
たとえば、こんな場面で
PyTorchプログラムをtorch.compileの対象にし、演算列をFXグラフとして取得してコンパイラによる最適化へ渡す用途が考えられます。バックエンドを拡張し、取得したグラフを独自のコンパイル処理に接続することも可能です。
03 / GOOD TO KNOW
使うときのポイント
TorchDynamoとバックエンドの役割を分けて確認してください。TorchDynamoはグラフ取得を担当し、実際のコード生成や実行先への変換はTorchInductorなどのバックエンドが担当します。
参考資料
PyTorch 2 paper and tutorial @ ASPLOS 2024 - PyTorch （新しいタブで開きます）用語の説明や活用例に関連する参考資料です。
READ & LEARN
記事一覧を見る 用語を知ったら、活用を知る。
AI LABで、AIの動向や業務に取り入れるヒントを読む。