PyTorch
TorchDynamo (torch._dynamo) は任意のPythonコードをコンパイルのためのサブシステムである。JITコンパイルするためのtorch.compile API、事前コンパイルするためのtorch.export APIもある(後者は#事前コンパイルと量子化参照)。コンパイルできないコードは自動的にPythonインタプリタへとフォールバックされる。2023年にリリースされたPyTorch 2.0で本体へと統合された。
狭義のTorchDynamoはPythonのバイトコードから前述のFX Graphを作るサブシステムであり、FX Graphを各実行環境向けIRへと変換するための各バックエンド(TorchInductor、onnxrtバックエンドなど)と組み合わせてPythonコードのコンパイルを行うことができる。バックエンドを実装しやすくするためのものとして、FX Graphの複合的なAten命令を単純なPrim命令へと置き換えるPrimTorchも存在する。
以下は前向き (forward) 関数の最適化の例である:
また派生プロジェクトとしてバックエンドにTensorFlow XLAコンパイラを使うためのPyTorch/XLAも開発中となっているが、こちらも試験的にTorchDynamoへと対応しており、場合によってはTorchDynamo+Inductorよりも高速であるとされる。
PyTorchは以下のパイプライン並列処理に対応している:
パイプライン並列処理 (torch.distributed.pipeline) - nn.Sequentialにのみ対応。FairScaleの実装 (Fairscale.nn.Pipe) を統合したもので、元々は torchgpipe の GPipe実装に由来する。
PyTorchモデルの自動パイプライン並列処理に向けては PiPPy が開発中となっている。また DeepSpeed もパイプライン並列処理に対応している (deepspeed.pipe)。