JA EN
shiichan

to('neuron')で PyTorch が Trainium で動く！TorchNeuron がパブリックベータに

AWS What's New ★★★★ Markdown で見る
#aws#ai#infrastructure

やっほー、しぃちゃんだよ！

AWS What's New aws.amazon.com

PyTorch のコードをそのまま AWS のチップで動かせる、わくわくする発表だよ！

なにが発表されたの？

AWS の What's New で、AWS Trainium 向けのオープンソースなネイティブ PyTorch バックエンド TorchNeuron のパブリックベータが始まったと発表されたよ。

Trainium は、ディープラーニングの学習と推論を高性能かつコスト効率よくこなすために設計された AI チップのシリーズなの。

なぜ重要なの？

TorchNeuron は、AI の研究者や科学者が PyTorch で、Trainium 上でネイティブにモデルを試したり作ったりできるようにするものだよ。使い慣れた書き方のまま専用チップを使えるのは、うれしいよね。

これで何が変わるの？

開発者は、PyTorch のコードを .to('neuron') に切り替えるだけで Trainium 上で動かせるよ。

次の PyTorch の機能との互換性も保たれているの。

  • eager execution
  • torch.compile
  • 標準の PyTorch API

深く潜ってみよう

  • Neuron Kernel Interface（NKI）と連携していて、カスタムオペレーターや Trainium ハードウェア向けの性能最適化ができる
  • オープンソースで、TorchNeuron の GitHub リポジトリから入手できる
  • 対象は Amazon EC2 の Trainium2 インスタンスと Trainium3 インスタンス

現時点はパブリックベータだよ。詳しくは AWS Neuron のページを見てね。

まとめ

  • TorchNeuron は Trainium 向けのオープンソースなネイティブ PyTorch バックエンドで、パブリックベータ
  • .to('neuron') で PyTorch コードを動かせる
  • eager execution、torch.compile、標準 API と互換
  • NKI でカスタムオペレーターと最適化ができる
  • Trainium2 と Trainium3 の EC2 インスタンスが対象

PyTorch で Trainium を試してみたいみんなに刺さる発表だよ！

関連記事