▶️

RAPIDS RMMでRAMをGPUメモリとして使う

2023/12/15に公開

Python

はじめに

この記事は Kaggle Advent Calendar 2023 15日目の記事です

皆さん，GPUメモリ足りてますか！？
kaggleをやっているとどうしてもGPUメモリが足りなくてout_of_memoryを見る機会が多いと思います
本記事では，RAMをGPUメモリとして扱い，cuDFとpytorchでGPUメモリ足りない問題を解決する方法をご紹介します

RAPIDS Memory Manager (RMM)

RAPIDS Memory Manager (RMM) とは，cuDFなどが有名なNVIDIAが開発しているライブラリ郡RAPIDS^[1]のうちの一つで，GPUメモリプールを使って処理を高速化したり，managed memory(=unified memory)を使えるようにするライブラリです^[2]
GPUメモリプールを使って高速化する話は，RAPIDS teamの方が書かれたこちらの記事^[3]が詳しいです，気になる方はこちらも読んでみてください

managed memoryとは，ざっくりいうとRAMとGPUメモリを同一のメモリ空間として扱えるようにする仕組みで，詳しい仕組みはこれらの記事を参照してください^[4]^[5]
本記事では，このmanaged memoryを使ってRAMをGPUメモリとして扱います

cuDFでのRMM

まずは，cuDFでRMMを使ってGPUメモリ足りない問題を解決します

環境はkaggle notebookのP100，rmmとcudfのversionはそれぞれ23.08.00を使用します

データは，4.81GBあるmeta-kaggleのUserAchievements.csvを使います^[6]

まずは何もせずにcudfでread_csvします

out_of_memoryが出ました，見たくないですね

次に，

rmm.reinitialize(managed_memory=True)

を実行し，再度read_csvします

今度はout_of_memoryになることなくcsvを読み込むことができました！

pytorchでのRMM

次は，pytorchでRMMを使ってGPUメモリ足りない問題を解決します

先程と同じ環境で，

tensor = torch.randn(size=(50000, 100000), device="cuda")

を実行し，out_of_memoryを発生させます

よく見るやつですね，15.9GBしかないGPUメモリ上に18.63GBのテンソルを作ろうとして怒られています

次に，

rmm.reinitialize(managed_memory=True)
from rmm.allocators.torch import rmm_torch_allocator
torch.cuda.memory.change_current_allocator(rmm_torch_allocator)

を実行し再度テンソルを生成します^[7]

今度は怒られることなくテンソルを作れています！

おわりに

本記事では，RMMを使ってRAMをGPUメモリとして扱い，cuDFとpytorchでout_of_memoryを回避する方法をご紹介しました．
本記事で作成したコードはこちらのnotebookで確認できます，良ければご確認ください

ちなみに，私はRAM64GB＋swap領域100GBをmanaged memoryとして使っています．

明日は@yufuinさんで
ニューラルモデルをONNXにしてブラウザ拡張で使おうです！

脚注