概要
- 巨大なembeddingをチャンクで外部に保存し,DDP(Distributed Data Parallel)を使った学習時に各GPUで読み込みたい
- そんなとき
torch.load(path, map_location=f"cuda:{rank}")
にかかる時間の分散が大きい場合がある- 前提:
torch.load
はまずRAM(CPU)にデータをallocateし,その後各GPUに転送する - なので,ボトルネックは 「I/O」または「RAM展開→VRAM展開」のどちらかであることが多い
- 前提:
I/Oの場合
-
ネットワーク越しにアクセスしている場合はもちろん遅い
-
HDD / SSDの違いによる差も割と顕著に現れる
- なので,そのpodに対して一番高速であろうNFSにデータを格納しておくと良い
- ただし,ディスククォータ等で,1ユーザ当たりのディスク制限が掛けられていることがあるので注意
df -h
等で見ているのはディスク全体の容量であって,クォータではないので留意すべしquota
コマンドは使えないことが多いので,クラスタの内製ドキュメントを参照する
- 書き込みすぎると
OSError: [Errno 122] Disk quota exceeded
が出るので注意- 適宜
du -sh
で当該ディレクトリを参照すること
- 適宜
- I/Oが極端なボトルネックになっているならば,クォータぎりぎりまでembeddingデータを格納して,NFS AとNFS Bの二刀流でやると良い
- ファイルの存在確認は速いので,AとBとでファイルが存在している方を読み出すようにする
-
I/Oが明確なボトルネックである場合は,圧縮等も検討する
- そもそも階層構造が多すぎる場合は,
.pth
ではなくhdf5
を使っても良い? - numpyの
memmap
が優秀という説もある
- そもそも階層構造が多すぎる場合は,
「RAM展開→VRAM展開」の場合
- DDP等がうまく設定されていない可能性あり
- ちゃんと各GPUで処理されているか?DataloaderはGPUごとに分割されているか?
map_location
で適切なGPUを指定しているか確認- 特に,すべてを
cpu
やrank=0
に任せていないか?
- 特に,すべてを
- そもそもDataloaderに任せられないか?
num_workers
を増やしすぎると不用意にRAMを消費し,スワップ領域まで食いつぶす可能性があるので注意- スワップに手を付け始めると,過度な書き込みでディスクが消耗する可能性があるのでやめる
- SSDの寿命を擦り減らさらないためにも,スワップ領域への信頼は避けましょう.
- スワップに手を付け始めると,過度な書き込みでディスクが消耗する可能性があるのでやめる