【ITニュース解説】MPS backend out of memory on a Mac: what the numbers mean and how to set the watermark

Macで機械学習のプログラムを作る際、PyTorchというツールはApple Siliconチップの「MPS（Metal Performance Shaders）」機能を活用し、高速な計算を可能にする。しかし、このMPSを使っていると、「MPS backend out of memory」というメモリ不足エラーに遭遇することがある。これは、MPSが使用可能なメモリ容量を超過したことを示すエラーで、システムエンジニアを目指す上でメモリ管理の重要性を学ぶ良い機会となる。

エラーメッセージには「MPS allocated: 18.50 GiB」「other allocations: 448.00 KiB」「max allowed: 18.72 GiB」といった数値が含まれる。これらは、それぞれ現在MPSが機械学習計算のために確保しているメモリ量、PyTorchプログラムが他の用途で使っているメモリ量、そしてMPSがこのプログラムに許可しているメモリ使用量の上限を意味する。GiB（ギガバイト）、MiB（メガバイト）、KiB（キロバイト）はデータ量を表す単位だ。この例では、MPSが上限近くまでメモリを使い切り、新たに512 MiBを確保しようとしたが失敗した、という状況が示されている。つまり、MPSのメモリが満杯で、新しいデータのための場所がない状態である。

この「max allowed」（許可される最大メモリ量）は、「PYTORCH_MPS_HIGH_WATERMARK_RATIO」という比率と、MacのMPSが推奨する最大メモリ量を掛け合わせて計算される。例えば、推奨される最大メモリが37.44 GiBの場合、比率が0.5なら上限は18.72 GiBとなる。自分のMacの推奨値はimport torchの後、torch.mps.recommended_max_memory()で確認できる。

しかし、多くのMacユーザーは、このメモリ不足エラーを見る前に、Mac全体の動作が著しく遅くなる経験をするかもしれない。これは「スワップ」という現象が起きているためだ。Macの物理メモリ（RAM）が不足すると、macOSは一時的にRAMの内容をSSD（ストレージ）に書き出してRAMを空ける。SSDへのアクセスはRAMよりはるかに遅いため、Mac全体の処理速度が低下する。なぜPyTorchのエラーの前にスワップが発生するかというと、PyTorchのデフォルト設定における「PYTORCH_MPS_HIGH_WATERMARK_RATIO」が1.7と非常に高く設定されているためだ。筆者の48GBのMacでは、この比率で計算されるMPSのメモリ上限は物理RAM容量を超える約63.6 GiBにもなる。結果的に、PyTorchが自身のメモリ上限に達する前にMacの物理RAMが枯渇し、スワップが先に発生してしまうのだ。スワップは開発効率を大きく損なうため、避けるべきだ。

このスワップを避けるためには、MPSのメモリ上限を適切に設定することが重要になる。エラーメッセージにある「PYTORCH_MPS_HIGH_WATERMARK_RATIO=0.0」で上限を無効にする方法は、結局スワップを招きやすいため推奨されない。MPSのメモリ上限を設定するには、「PYTORCH_MPS_HIGH_WATERMARK_RATIO」と「PYTORCH_MPS_LOW_WATERMARK_RATIO」という二つの環境変数を同時に設定する必要がある。高い方の比率だけを設定しようとすると、「invalid low watermark ratio」というエラーが出る場合がある。これは、低い方の比率がデフォルトの1.4で、高い方の比率より大きくなってしまうためだ。これら二つの比率は、例えば以下のように同じ値に設定するのが一般的だ。

PYTORCH_MPS_HIGH_WATERMARK_RATIO=0.8 PYTORCH_MPS_LOW_WATERMARK_RATIO=0.8 python run.py

これらの設定は、PyTorchプログラムを実行する前に、コマンドラインで環境変数として指定するか、Pythonコード内でimport torchする前にos.environを使って設定する。PyTorchはMPSの初期化時に値を読み込むため、torchモジュールインポート後の変更は無効である。筆者の検証では、低い方の比率を高い方と同じ値に設定することで、スワップを発生させることなくPyTorchからのクリーンなメモリ不足エラーを受け取ることができた。これにより、Mac全体の動作を遅らせることなく、問題の原因を明確に特定できる。

では、適切な上限値はどれくらいだろうか。理想は、普段使っているアプリケーションを起動した状態でMacが利用できる空きメモリ量を把握し、それよりも低い値を設定することだ。筆者のMacでは、PYTORCH_MPS_HIGH_WATERMARK_RATIOを1.0（37.44 GiB）に設定するとまだスワップが発生したが、0.8（29.95 GiB）に設定したところ、スワップなしでPyTorchから直接エラーを受け取ることができた。このように物理メモリ容量より少し低い値を設定することで、スワップを防ぎ、メモリ不足時はPyTorchがエラーとして知らせてくれる。クリーンなエラーを受け取れば、バッチサイズ削減や解像度調整といった具体的な対策を立てやすくなる。

メモリ使用量をテストする際には注意が必要だ。torch.onesのようにすべて同じ値で埋められたデータを使うと、macOSはそのようなデータを効率的に圧縮するため、実際のメモリ使用量より少なく表示される可能性がある。機械学習モデルの重みのような複雑なデータは圧縮されにくいため、正確なテストにはtorch.randのようなランダムなデータを使ってシミュレーションすることが重要だ。

もしメモリ制限を調整してもなお、機械学習モデルがMacのメモリに収まらない場合もある。その際は、単に比率を調整するだけでなく、より根本的な対策が必要だ。例えば、複数のモデルを連携させるような処理では、一度にすべてのモデルをメモリにロードするのではなく、現在処理に必要なモデルだけをロードし、不要になったら解放するという方法が有効だ。このようにメモリのロード・アンロードを効率的に管理することで、Macの限られたリソースを最大限に活用し、大規模なモデルでも安定して動かすことが可能になる。メモリ管理は、高性能なアプリケーションを開発する上で非常に重要なスキルであり、この経験を通じてその基礎を学ぶことができるだろう。