知り合いと新しいMacBookProを買う買わないの話をしていたのですが、その際に「Apple SiliconだからGPUもCPUと同じメモリを共有している」という話をしました。
そうしたら、「GPUでどこまでメモリ使えるの？」と聞かれ、気になったので調べてみました。
Apple Siliconのアーキテクチャ
2020年に発表された、Apple SiliconはCPU、GPUなど、全てのコンポーネントでメモリを共有しています。
今までだと、GPUはGPUのメモリを、CPUはCPUのメモリを使い、必要に応じてデータをコピーしてそれぞれと共有していました。
Apple Siliconでは、全てのユニットが同じメモリを使用できる「ユニファイドメモリアーキテクチャ」と呼ばれる構造になっています。
これにより、データのやり取りを行う必要がなく、オーバーヘッドを減らすことができます。
実験
とはいえ、実際はメモリ全ては使い切れないでしょ？と思い、どのぐらいまでメモリを使用できるのか調べてみました。
検証に使用したMacはM2 Pro、16GBメモリ搭載です。
virtualenvで仮想環境を用意し、pytorchとnumpyをインストールしました。
とりあえず、諸々インストールし、テスト用の配列を用意します。
Python 3.11.4 (main, Jun 20 2023, 17:23:00) [Clang 14.0.3 (clang-1403.0.22.14.1)] on darwin
Type "help", "copyright", "credits" or "license" for more information.
>>> import torch
>>> import numpy
>>> import torch.mps
>>> device=torch.device("mps")
>>> x=list(range(250000000))
250,000,000個の配列を用意し、それをfloatのtensorに入れます。
floatなので1要素あたり4byteとなり、xは約1GB分となります。
>>> a=torch.tensor(x,dtype=torch.float).to(device)
>>> torch.mps.current_allocated_memory()
1000000000
こんな感じで、用意した配列をtensorに変換し、GPUに載せていきます。
その後に、メモリの使用量をbyteで表示しています。
上記から想定通り、1,000,000,000byte = 1GBのメモリが使用されていることがわかります。
こんな感じで変数名を変えてどんどん代入していきましょう。
>>> b=torch.tensor(x,dtype=torch.float).to(device)
>>> torch.mps.current_allocated_memory()
2000000000
>>> c=torch.tensor(x,dtype=torch.float).to(device)
>>> torch.mps.current_allocated_memory()
3000000000
>>> d=torch.tensor(x,dtype=torch.float).to(device)
>>> torch.mps.current_allocated_memory()
4000000000
>>> e=torch.tensor(x,dtype=torch.float).to(device)
>>> torch.mps.current_allocated_memory()
5000000000
>>> f=torch.tensor(x,dtype=torch.float).to(device)
>>> torch.mps.current_allocated_memory()
6000000000
>>> g=torch.tensor(x,dtype=torch.float).to(device)
>>> torch.mps.current_allocated_memory()
7000000000
>>> h=torch.tensor(x,dtype=torch.float).to(device)
>>> torch.mps.current_allocated_memory()
8000000000
ここまでは予定通りです。
8GBのメモリがGPUで使用できていることがわかります。
>>> i=torch.tensor(x,dtype=torch.float).to(device)
>>> torch.mps.current_allocated_memory()
9000000000
>>> j=torch.tensor(x,dtype=torch.float).to(device)
>>> torch.mps.current_allocated_memory()
10000000000
>>> k=torch.tensor(x,dtype=torch.float).to(device)
>>> torch.mps.current_allocated_memory()
11000000000
>>> l=torch.tensor(x,dtype=torch.float).to(device)
>>> torch.mps.current_allocated_memory()
12000000000
搭載メモリの3/4もGPUで使用することができました。
限界まで調べてみましょう。
>>> m=torch.tensor(x,dtype=torch.float).to(device)
>>> torch.mps.current_allocated_memory()
13000000000
>>> n=torch.tensor(x,dtype=torch.float).to(device)
>>> torch.mps.current_allocated_memory()
14000000000
>>> o=torch.tensor(x,dtype=torch.float).to(device)
>>> torch.mps.current_allocated_memory()
15000000000
おお。かなりギリギリです。
せっかくなので使い切ってエラーを出してみましょう。
>>> p=torch.tensor(x,dtype=torch.float).to(device)
>>> torch.mps.current_allocated_memory()
16000000000
？
まさかの搭載メモリ全てを使い切れました。
まあ、多少誤差はありますので、ギリギリ空いているのでしょう。
>>> q=torch.tensor(x,dtype=torch.float).to(device)
>>> torch.mps.current_allocated_memory()
17000000000
？？
あれ、物理メモリの容量を超えた？
>>> r=torch.tensor(x,dtype=torch.float).to(device)
>>> torch.mps.current_allocated_memory()
18000000000
？？？
>>> s=torch.tensor(x,dtype=torch.float).to(device)
>>> torch.mps.current_allocated_memory()
19000000000
？？？？
>>> t=torch.tensor(x,dtype=torch.float).to(device)
Traceback (most recent call last):
File "<stdin>", line 1, in <module>
RuntimeError: MPS backend out of memory (MPS allocated: 17.70 GB, other allocations: 7.00 MB, max allowed: 18.13 GB). Tried to allocate 953.67 MB on private pool. Use PYTORCH_MPS_HIGH_WATERMARK_RATIO=0.0 to disable upper limit for memory allocations (may cause system failure).
ようやくエラーで落ちました。
どういう理屈か分かりませんが、物理メモリの容量を超えてGPUに割り当てることができてしまいました。
16GB以上の物理メモリを搭載していたか、Swapを使うことができるのか、理由は定かではありませんが、とりあえずGPUでもメモリをギリギリまで使うことができることがわかりました。
あとがき
今回は思いつきで実験したら、予想以上の結果になったため、雑ですがまとめてみました。
今回の結果で知り合いがメモリマシマシのMacを買う気になったようなので今回の実験をぜひやってもらいたいですね。（笑）
もし何か間違い等ございましたら、コメントいただけますと幸いです。
この記事が何かの役に立てば幸いです。