JavaScriptを有効にしてください

【k8s・DDP】クラスタ上でのtorch.loadが遅い

 ·  ☕ 3 min read

概要

  • 巨大な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が明確なボトルネックである場合は,圧縮等も検討する

「RAM展開→VRAM展開」の場合

  • DDP等がうまく設定されていない可能性あり
    • ちゃんと各GPUで処理されているか?DataloaderはGPUごとに分割されているか?
  • map_locationで適切なGPUを指定しているか確認
    • 特に,すべてをcpurank=0に任せていないか?
  • そもそもDataloaderに任せられないか?
    • num_workersを増やしすぎると不用意にRAMを消費し,スワップ領域まで食いつぶす可能性があるので注意
      • スワップに手を付け始めると,過度な書き込みでディスクが消耗する可能性があるのでやめる
        • SSDの寿命を擦り減らさらないためにも,スワップ領域への信頼は避けましょう.
共有

YuWd (Yuiga Wada)
著者
YuWd (Yuiga Wada)
機械学習・競プロ・iOS・Web