AWS、Trainium向けネイティブPyTorchバックエンドTorchNeuronを公開ベータで提供

AWS、Trainium向けネイティブPyTorchバックエンドTorchNeuronを公開ベータで提供

AWSは10月20日、AIチップTrainium向けのオープンソースPyTorchバックエンドTorchNeuronが公開ベータに入ったと発表した。AWSの技術者による長めの記事がPyTorchブログにあり、日付は10月19日、北米のPyTorch Conferenceの当日である10月20日に更新されている。AWSの発表は、コードがAmazon EC2のTrainium2とTrainium3インスタンス向けにオープンソースで提供されると述べる。記事はライセンスをApache 2.0とし、aws-neuron配下のGitHubリポジトリを挙げている。10月27日号は元の日付から1週間後に伝えるものだ。どちらもAWS側の説明だが、PyTorchの記事はDeepgramのCTOの発言として、同社が評価したベンチマークではTrainiumの入力音声秒数が1ドル当たり同社の既存GPU基準より約60%多かったと記す。これは顧客の自己申告であり、独立した測定でもTorchNeuron全般の性能保証でもない。

開発者にとって何が変わるか

これまでTrainiumでPyTorchのコードを動かすには、AWS独自の別のツールを使う必要があった。TorchNeuronは、PyTorchがカスタムハードウェア向けに用意する公式の拡張点PrivateUse1と`torch.accelerator`インターフェースを通じて、Trainiumを通常のPyTorchデバイスとして登録する。最も簡単な場合、開発者は`.to('neuron')`でモデルを移し、学習ループのほかの部分はそのままにできる。実行方式は二つある。イーガーモードは演算を一つずつ実行し、デバッグに向く。記事は実験的な「アダプティブ・イーガー」モードも加えており、コードを変えずに起動コストを減らし、演算を融合することを狙う。`torch.compile`モードはモデルをグラフとして捉えてNeuronコンパイラーに渡し、AWSはこれを本番向けと位置づける。記事はコンパイルの経路も記している。TorchDynamo、AOTAutograd、Torch-MLIR、StableHLOを経て、Neuronコンパイラーがチップ上で実行されるファイルを生成する。

対応範囲とその限界

AWSは、テキスト、視覚、動画などの多様なモデルに向けた幅広い演算子対応、BF16、FP32、FP16のデータ型、標準の分散API(DDP、FSDP2、テンソル並列、パイプライン並列、DTensor、DeviceMesh)を挙げる。通信はノード内ではTrainiumの集合通信コア、ノード間ではElastic Fabric Adapterを通る。記事は不足している点を隠していない。TorchAOとFP8は非公開ベータだ。デフォルト以外のストリームは単一のハードウェアキューを共有するので、順序は保たれるが、それぞれに専用のエンジンは割り当てられない。エラーは問題の演算子ではなく次の同期点で表面化し、TorchNeuronは最も早いエラーを文脈つきで再送出する。カスタムカーネルはNeuron Kernel Interfaceで書き、AWSはアテンションや混合専門家(MoE)のルーチンを含む、そうしたカーネルのオープンソースライブラリを出している。プロファイラーとの連携で、フレームワークのイベントとハードウェアの動きを結びつけられる。

発表の読み方

製品説明の中心はPyTorchとの互換性だ。AWSは、開発者がPyTorchのAPIを保ったままTrainiumを利用できると述べ、Metaのエンジニアリング責任者がこの対応を歓迎する言葉も引用されている。アダプティブ・イーガーがコンパイル済みグラフの性能に近づくという狙いは、記事自身が述べる目標である。PyTorchの記事には上記のDeepgramによる入力音声の費用効率の数字があるが、独立検証や一般的な性能結論はない。また、10万回のモデル実行はHugging Faceモデルを縮小した代表的構成で動かし、読み込み、実行、コンパイル、演算子の配置を確認する網羅テストであり、GPUとの性能比較を10万回行ったという意味ではない。Trainiumを検討するチームは、自分のモデルをイーガーとコンパイルの両方で動かし、必要な演算子が対応済みかを確かめ、本番の作業を移す前にベータ版の注意書きを読むべきだ。