ハクソク

世界を動かす技術を、日本語で。

AppleのMLXがCUDAサポートを追加

315日前原文(github.com)

概要

CUDAバックエンドの最適化によって、トレーニング速度が大幅に向上。 cudaEventやprefetchコール削減、cudaLaunchHostFuncの呼び出し最適化を実施。 prefetchの削除により、カーネル実行の遅延が解消し、さらなる高速化を実現。 今後はバッファキャッシュやカーネル最適化などが課題。 現状でPyTorchを超える性能を達成。

CUDAバックエンド最適化の経緯

  • Event実装 をcuda::std::atomicから cudaEvent へ切り替え、トレーニング速度が 500 it/sから900 it/s に向上
  • prefetchコール の削減で、速度が 900 it/sから1100 it/s へ
  • Metalバックエンドでは、各op評価後に オペランドや一時変数 をカーネル終了まで保持
    • コマンドバッファの必要時にcommitし、完了ハンドラで解放・エラーチェック
  • CUDAではcudaLaunchHostFunc を用いて類似実装
    • ただし、 20µs以上のストリーム遅延 が発生し、各カーネルが待たされる課題
    • 対策として、 mx::eval_impl()完了時 までオペランド・一時変数を保持
      • これにより cudaLaunchHostFuncの呼び出し回数を大幅削減
      • 2カーネル間の遅延が 1µs未満 となり、PyTorch以上のパフォーマンス
      • デメリットは 配列の解放遅延によるメモリ使用量増加
      • MAX_ACTIVE_TASKSを超えるタスク時の待機も廃止
  • この最適化で速度が 1100 it/sから1600 it/s へ

カーネル遅延の原因と解決

  • 一部カーネルの 遅延共通点 は、実行前に cudaMemPrefetch API を呼んでいた点
  • CUDAバックエンドは Unified Memory API を利用し、自動でデータ転送
    • 明示的な cudaMemPrefetch でデバイス側に転送指示をしていた
  • プロファイリングで prefetchがカーネル実行を大幅遅延 させていることを発見
  • prefetch削除 により、速度が 1600 it/sから2100 it/s に向上
    • プロファイラのタイムラインも大幅に改善

今後の最適化課題

  • バッファキャッシュ の実装
    • 主要opの実装とサードパーティライブラリ統合後に着手予定
  • 残る課題は主に カーネル最適化 と CPUコードの高速化
    • すべてのop実装完了後に本格対応予定

まとめ

  • CUDAバックエンドの最適化により、 PyTorchを超える性能 を実現
  • 今後は バッファキャッシュ・カーネル最適化・CPU高速化 が焦点
  • 継続的なプロファイリングとボトルネック解消が重要

Hackerたちの意見

「待って、AppleのプラットフォームにはファーストパーティのCUDAサポートがない！」って思うかもしれないけど、このパッチセットは「CUDA 12とSM 7.0（Volta）以上のLinuxプラットフォーム」もサポートしてるからね。

MLXの重要な機能の一つが統一メモリアーキテクチャを使うことなんだけど、これってどうなるの？（リポジトリのREADMEに書いてある箇条書きを見てみて: https://github.com/ml-explore/mlx）全てのUMA APU（どのベンダーでも）にそれを持ってくるのは面白いと思うけど、ディスクリートGPUは別のアプローチが必要だよね？ 編集: PRのコメントを読んでみたら、CUDAはUMA APIを直接サポートしていて、必要に応じて自動的にコピーするみたいだね。

└

うん、そうだけど、私の経験から言うと、プリフェッチがないとメモリの待機時間がかなり長くなって、コピーを待つことになるんだ。もしデータセット全体がVRAMに収まるなら「手動プリフェッチ」をすれば適してるかもしれないけど、私のアプリケーション（MLトレーニング）ではパフォーマンスがめちゃくちゃ悪くなって、ストリーミングロードに移行する時間ができたくらいだよ。

これはワクワクするね。これってCUDAの統一メモリを使ってるの？どれくらい上手く機能するんだろう。CUDAの統一メモリの動作はAppleシリコンのと同じなのかな？Appleシリコンの場合、GPUとCPUの間でメモリが共有されてるって理解してるけど、CUDAではそうじゃないよね。だから、CPUにテンソルがあるとき、どうやってGPUに送るの？何らかの形でコピーが必要だよね。それとも、CUDAが全部隠してくれてるのかな？

└

これは私の推測だけど、彼らが販売している高性能なハードウェア、例えばAI用のサーバーラックのやつには、統一メモリがあるのかな？標準のGPUにはないって知ってるけど。パッチによると、その理由の一つはMacでの開発を簡単にして、スパコンで動かせるようにするためだったみたい。だから、統一メモリを持ってるハードウェアはそのクラスに入るかもしれないね。

└

ハードウェアの統一メモリがない場合、CUDAはページフォルトが発生したときに自動的にCPUとGPUの間でデータをコピーするよ。

これはElectronを作ったzcbenzから来てるよ、他にも色々やってるね。https://zcbenz.com/ いいね。

これってStrix Haloの影響がどれくらいあるんだろう。仕事用のコンピュータのために結構標準的な手当をもらってたんだけど、しばらく使わなかったから最近EVO-X2に使ったんだ。マジで、あれは中堅のznver5 EPYCマシンと簡単に競争できるよ。普通のEC2やGCEインスタンスなんて、r1337.xxxxlarge.metal.metalとかがちょっと優位かもしれないけど、z1d.metalやc6.2xlargeみたいなやつ（速いコア、いいNIC、基本的な条件）はそれを圧倒してる。あれは月に3-10Kかかるけど、こっちは本物のNVMEで1800ドルだった。ローカル推論はあまりやってないけど、いろんなYouTuberがDGX SparkはStrix Haloに比べて過剰/高すぎるって言い始めてる。もちろん、ROCｍはまだそこにないけど（今は真剣に取り組んでるみたいだから、時間の問題だね）。Appleの機材で完璧なCUDAが使えたら、Strixが安くて良いのとは違った魅力があるよね。

└

PRの説明では、クラウドクラスターのトレーニングを明確にターゲットにしてるね。

└

もちろん、問題はROCmがまだ整ってないってことだね（でも今は真剣に取り組んでるみたいだから、時間の問題だと思う）。競争力のあるAMDのGPUニューラルコンピュートは、もう10年以上前から「いつでも出る」って言われてる。

└

初めての人に説明すると、Strix HaloはAMD Ryzen AI Max+ 395と同じもので、Framework Desktopに搭載される予定で、最近はミニPCにも出てきてるよ。そのメモリ帯域幅は200GB/sなんだ。ほとんどの消費者向けx86プラットフォームと比べるとすごく良いけど、NvidiaのGPU（5090は1792GB/sだし、プロレベルのカードはどうかは知らないけど）やAppleの最高モデル（M3 Ultraは800GB/s）には遠く及ばないね。確かにお得感はあるけど、LLMのようなメモリ帯域幅を重視するアプリケーションには、やっと「十分なレベル」に入ってきたって感じ。

└

M4 Mac Miniと比べてどうなの？

このPRは、MLXにCUDAバックエンドを追加するための進行中の取り組みで、MLXコードがx86 + GeForceハードウェアでコンパイルして実行できるようにするみたいだね。逆はダメってこと。

Appleは、アプリ開発やテスト、外部サービスをホストするためにMシリーズのチップを使ったデータセンターを作る計画を立ててるよ！

つまり、私が理解している限りでは、これはこういうことだよね：1. MLXに対して構築されたプログラムは、CUDA対応チップを利用できるけど、2. CUDAプログラムはApple Siliconで実行できない。これが正しいなら、#2は著作権侵害になるから（特にNvidiaの有名な防壁に関して）。合ってる？

└

いや、2をやるのはかなり難しいからだよ。

└

#2はGoogle対Oracleってことだよね？

└

それは1だよ。つまり、開発者は比較的低性能なAppleデバイス（UMA付き）を使って、NVIDIAの比較的高性能なシステム向けに開発できるってこと。これはいろんな理由で便利だよね。

└

2はCUDAを事実上のターゲットAPIとしてさらに確立するし、誰もMLX向けのコードを書くことはないだろうね。こうすることで、みんながMLXを書いてどこでも動かせるようになるから、みんなにとってウィンウィンな状況だよ。特にAppleにとっては、自社のプラットフォーム向けにさらに最適化できるからね。

└

#3ではNVIDIAのGPUをApple Siliconと接続してCUDAを動かせるようにしたいな。Apple Silicon + 統一メモリ + GPU + CUDAを使って、PyTorch、JAX、またはTensorFlowを活用したい。MLXについてはあまり探ってないから、詳しくは言えないけど。

余談だけど、MLXに関わってる人たちって、少なくとも表面的にはAppleと正式に関係があるようには見えないんだよね。例えば、これを見てみて: https://x.com/prince_canuma ふと思ったんだけど、Appleがこれに資金を出してるけど、表に出さないようにしてるのかな？チームが一時期移動を考えてたって噂もあったし。

MLXモデルにはすごく感心してるよ。家の中の皆にローカルモデルを開放できるし、Nvidiaのコンピュータでは家が燃えるリスクがあるから、そんなことは絶対にできないんだ。Apple SiliconがNvidiaチップの本格的な競争相手になることを期待してるけど、CUDAサポートはただの「抱きしめて、拡張して、消す（EEE）」なのかな。