ハックされた Nvidia 4090 GPUドライバ、P2Pを有効化

NVIDIA Linux Open GPUドライバにP2Pサポートを追加

このプロジェクトはNVIDIAドライバをフォークし、4090 GPUにP2Pサポートを追加したもの。

• ./install.sh を実行するだけで簡単にインストール可能
• DKMSで既存のドライバを先に削除する必要がある場合がある
• システムにLarge BARサポートが必要で、IOMMUは無効化する必要がある
• すべてのキャッシュフラッシュが正しいかは確信できないため、問題を見つけた場合は報告してほしい
• これはハックではなくPCIe仕様に沿って実装したものであり、整理を進めればアップストリームに反映される可能性がある

仕組み

従来のNVIDIA P2Pの問題点

• 従来はMAILBOXP2Pというハードウェアインターフェースを使ってGPU間のメモリ転送を行っていた
• 4090ではこのハードウェアが存在しないか無効化されており、P2Pが動作しない
• 初期ドライバでは正常に動作していると報告していたが、実際にはPCIeバス経由で転送していた
• しかしメールボックス用ハードウェアがないため、コピーした内容が正しい場所に届かず、システムクラッシュを引き起こす可能性があった

Large BARサポートの追加

• 一部の3090とすべての4090でNVIDIAがLarge BARサポートを追加した
• H100ではメールボックスの代わりにBARを直接使うBAR1P2PというPCIeモードが追加されている
• 4090でこれを有効化するにはHALを迂回し、GH100メソッドを直接呼び出す必要があった
  • kbusEnableStaticBar1Mapping_GH100 のようなメソッドでVRAM全体をBAR1にマップ
  • MapAperture 関数でその領域の使用を無効化する必要があった

P2P有効化の難しさ

• VRAMをマッピングした後でも、cuda-samples の ./simpleP2P 実行時にMMUエラーが発生
  • GMMU_APERTURE_PEER をマッピングタイプとして使うが、4090ではサポートされていない
  • 4090でサポートされるタイプは GMMU_APERTURE_VIDEO、GMMU_APERTURE_SYS_NONCOH、GMMU_APERTURE_SYS_COH のみ
• GMMU_APERTURE_PEER を GMMU_APERTURE_SYS_NONCOH に変更
  • CPU L2キャッシュとの一貫性は不要だが、PCIeバスに出る必要があるため
• ピアアドレスフィールド fldAddrPeer を fldAddrSysmem に変更
• fabricBaseAddress フィールドに BAR1 基準アドレスを設定

動作確認

• ./simpleP2P が正常に動作することを確認
  • GPU0とGPU1の間で24GB/sの速度でP2Pが動作
• p2pBandwidthLatencyTest で双方向帯域幅を確認
  • 920GB/sのローカル帯域幅と51GB/sのP2P帯域幅を計測
• NCCLテストで互換性を確認
  • 6基の4090 GPUで平均バス帯域幅24.5GB/sを達成

GN⁺の見解

• NVIDIAドライバの大部分がオープンソースとして公開されているため、開発者コミュニティでこのような試みが可能になったのだと思われる。今後さらに多くの部分がオープンソース化されることを期待したい。
• 4090の強力な性能を複数枚接続して活用できるようになれば、個人開発者や小規模ラボでも巨大なAIモデルを学習させられるようになるだろう。
• ただし、このようなハードウェア依存で扱いの難しい部分を開発者が直接触らざるを得ないということは、NVIDIAがまだ4090サポートを完全には仕上げられていないことの裏返しでもある。
• また、これはLinuxドライバに限った話であり、Windowsでの商用利用はまだ先が長そうだ。NVIDIAによる公式サポートが一日も早く実現してほしい。
• 4090は非常に新しいハードウェアであるだけに、CUDA、PyTorch、TensorflowのようなライブラリやMLフレームワークとの完全な互換性を期待するのは難しそうだ。安定化するまで待つ必要がありそう。