Amazon、PyTorchのコードを自社チップで動かす技術を発表
Amazonは、PyTorchで書いたコードを自社のAIチップAWS Trainium上で動かせるソフトウェアTorchNeuronを発表しました。設定を1か所変えるだけで、コードの書き直しは不要です。10月20日のPyTorch Conference North Americaで、オープンソースとして公開するとしています。
ARCHETYP Staffingでは現在クリエイターを募集しています。
エンジニア、デザイナー、ディレクター以外に、生成AI人材など幅広い職種を募集していますのでぜひチェックしてみてください！
PyTorchのコードを書き直さずにTrainiumで動かす仕組み
PyTorchは、AIモデルの開発や学習に使われるフレームワーク（開発の土台となるソフトウェア群）です。Amazonによると、現在AI研究者と開発者の85%がPyTorchでコードを書いています。一方で、コードを動かせるハードウェアは限られており、新しいチップに対応させるには、数か月かけて仕上げたコードを書き直し、改めて検証する必要がありました。
TorchNeuronを使うと、開発者が普段書いているPyTorchのコードが、AWS Trainium上でネイティブに（コードを書き直さず、そのまま）動作します。Amazonは、TorchNeuronを車のエンジンの載せ替えにたとえています。手頃で高出力な新しいエンジンに替えるには、通常は車を分解して組み直す必要がありますが、TorchNeuronなら分解せずに載せ替えられるといいます。
参照：Amazon News
実際に、Amazonのエンジニアは同じAI学習用スクリプトを、従来のチップとTrainiumで実行するデモを行いました。違いは、どのチップを使うかを指定する設定だけです。結果として、AIモデルは両方のチップで同じ内容を同じペースで学習したとしています。
仕組みの面では、PyTorchにもともと備わっている、新しいハードウェアを登録するための標準的な機能を利用しています。TorchNeuronはPyTorchの標準機能に組み込まれる形で動くため、開発者がAIモデルのテスト、デバッグ、規模の拡大に使っているツールは、変更せずにTrainium上でも使えます。品質を保つため、Amazonはビルドのたびに数十万のAIモデルでTorchNeuronをテストし、ユーザーの処理が動く前に問題を見つけられるようにしています。対応するチップはTrainium2とTrainium3です。TorchNeuronは、10月20日のPyTorch Conference North Americaで、Apache 2.0ライセンス（オープンソースライセンスの一種）のもとで公開すると発表されています。
開発コストの削減と、導入企業・業界の対応
Amazonは、AIにかかる実際のコストはチップの価格だけではないと指摘しています。多くの企業、特にスタートアップにとって最も大きな出費は、エンジニアの作業時間です。新しいハードウェアに合わせてソフトウェアを調整する1週間は、製品の改善に使えない1週間になり、人員や資金が限られる小規模なチームでは、事業の成否に影響することもあるとしています。
Amazonによると、大企業にとっては、同じハードウェアを学習と推論の両方に使えることが重要です。利用が少ない時間帯には、推論に使っていた計算資源を学習に切り替え、専用のハードウェアを遊ばせずに済むといいます。Amazonは、開発者が不利益なくチップを移行できるようになれば、チップメーカーは開発者に選ばれる努力を求められ、開発者はコスト削減や開発サイクルの短縮といった恩恵を受けられると説明しています。
TorchNeuronはすでに、ロボティクス、動画、世界モデル（現実世界の仕組みを学習・予測するAIモデル）、自動運転車、音声、創薬、大規模言語モデル、金融などの分野で利用されています。会話を起点に、人間のように世界から学び関わるAIを開発するCartesiaは、TrainiumのPyTorchネイティブ対応で、最小限のコード変更により標準の推論処理を実行しました。また、精神神経疾患の新たな治療法の開発に取り組む米国のリーバー脳発達研究所は、計算コストが非常に高い分子動力学シミュレーションをAIで高速化・代替する手法を、TorchNeuronを使ってTrainium上で検討しています。
業界の対応としては、AmazonがPyTorchのコミュニティと協力し、PyTorch本体から特定のチップに依存するコードを取り除く変更を提供しています。Trainiumの利用者に限らず、PyTorchを特定のハードウェアに縛られないものにする狙いです。Hugging Faceも、自社のAI開発ツールでTrainiumへのネイティブ対応を有効にしており、開発者はAmazon固有のコードを書かずにTrainium上で作業できます。
まとめ
いかがだったでしょうか？
Amazonが発表したTorchNeuronは、PyTorchで書いたコードを書き直さずにAWS Trainium上で動かせるソフトウェアで、Apache 2.0ライセンスのオープンソースとして公開するとしています。AmazonによるPyTorch本体への変更の提供に加え、Hugging FaceもTrainiumに対応するなど、周辺の環境も整いつつあります。
開発者が用途やコストに応じてチップを選びやすくなれば、AI開発にかかる時間や費用の見直しにつながると考えられます。