アップルのMLX機械学習フレームワークが、NVIDIAのCUDAプラットフォームをサポートすることになったと発表しました。これにより、アップルシリコンに最適化されたMLXが、CUDAを活用することで、NVIDIAのGPUを使用した高性能な並列計算が可能になるということです。
このプロジェクトは、GitHub上で開発者の@zcbenz氏が数か月前からCUDAのサポートのプロトタイプを開始し、その後、プロジェクトを小さな部分に分けて、徐々にMLXのメインブランチに統合してきたとしています。
このバックエンドはまだ開発途中ですが、行列の乗算やソフトマックス、リダクション、ソート、インデックス付けなど、いくつかの基本操作がすでにサポートされ、テストされているということです。
CUDAとは何かというと、NVIDIAが自社のGPUで高性能な並列計算を行うために開発したコンピューティングプラットフォームです。多くの人にとって、CUDAはNVIDIAのGPUで機械学習のワークロードを実行する標準的な方法であり、学術研究から商業展開まで、機械学習のエコシステム全体で利用されています。
MLXは元々アップルシリコンとMetalに最適化されていましたが、CUDAバックエンドの追加により、研究者や技術者がMac上でCUDAベースのモデルをプロトタイプし、その後、NVIDIAの大規模GPUクラスターで展開できるようになる方針です。
ただし、まだすべてのMLXオペレーターが実装されているわけではなく、AMD GPUのサポートもまだ先の話です。それでも、MLXとNVIDIAのGPUがより密接に連携することで、テストや実験、研究の用途が広がるとされています。
詳細はGitHubで確認することができます。