XProf ile OOM hatalarını ayıklama

Yetersiz bellek (OOM) hataları, hızlandırıcının (GPU veya TPU) yüksek bant genişlikli bellek (HBM) kapasitesi tükendiğinde oluşur. OOM sorunlarının bazı yaygın nedenleri ve hata ayıklama teknikleri E1000 - Compile Time HBM OOM dokümanında ve GPU bellek ayırma ile ilgili JAX dokümanında ayrıntılı olarak açıklanmaktadır.

Bu sayfada, XProf'un bellek görüntüleyici aracını kullanarak JAX programınızın bellek kullanımını görselleştirme, en yüksek kullanım örneklerini belirleme ve OOM hatalarını ayıklama hakkında bilgi verilmektedir. Bu işlem aşağıdaki adımları içerir:

  1. Profili yakalamak için programınızı jax.profiler.trace ile çalıştırın.
  2. XProf'u arka planda başlatın ve bellek kullanım ayrıntılarını görüntülemek için Bellek Görüntüleyici aracını kullanın.

Örnek program

Aşağıdaki JAX programı, bellek yetersizliği hatasına neden oluyor:

import jax
from jax import random
import jax.numpy as jnp


@jax.profiler.trace("/tmp/xprof")
@jax.jit
def oom():
    a = random.normal(random.PRNGKey(1), (327680, 327680), dtype=jnp.bfloat16)
    return a @ a


if __name__ == "__main__":
    oom()

Bu program, TPU makinesinde şu hatayla başarısız olur:

XlaRuntimeError: RESOURCE_EXHAUSTED: Allocation (size=107374182400) would exceed memory (size=17179869184) :: #allocation7 [shape = 'u8[327680,327680]{1,0:T(8,128)(4,1)}', space=hbm, size = 0xffffffffffffffff, tag = 'output of xor_convert_fusion@{}'] :: <no-hlo-instruction>

(GPU makinesinde hata şu şekilde görünür: XlaRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 214748364800 bytes.)

XProf'u çalıştırma

xprof (pip install xprof) uygulamasını yükleyin ve profilin depolandığı dizini belirterek bir XProf örneği başlatın:

xprof --logdir=/tmp/xprof/ --port=6006

Örneğe gidin (yerel bir makinede, http://localhost:6006 konumunda). Araçlar açılır listesinde Bellek Görüntüleyici'yi seçin ve Bellek Görüntüleyici araç penceresinde Bellek Türleri açılır listesinde HBM'yi seçin (genellikle varsayılan olarak seçilir).

Yukarıdaki örnek program için XProf Bellek Görüntüleyici sayfası

XProf: Memory Viewer aracı dokümanında, aracın bileşenleri ve sunulan bilgiler açıklanmaktadır.

En yüksek bellek kullanımı noktasında üç arabellek grafiğini gösteren En Yüksek Bellek Ayırımında HLO İşlemleri bölümüne odaklanın. Arabellek şunları içerir: * Program Girişleri ve Çıkışları: Eğitim grupları, optimize edici durumları vb. * TensorCore ve SparseCore Geçicileri: Ara hesaplamalar (etkinleştirmeler, gradyanlar vb.) için gereken dinamik bellek

İşlem hakkında daha fazla ayrıntı (ör. boyutu, şekli, ayırma türü) almak için arabellek grafiklerinin üzerine gelebilirsiniz. Bu, yüksek veya uzun süreli geçici değerlere sahip olabilecek, verimsiz dolguya sahip büyük giriş/ara/çıkış tensörleri gibi, en yüksek belleğe katkıda bulunan ve ayarlanması ya da optimize edilmesi gereken işlemleri belirlemenize ve değerlendirmenize yardımcı olabilir.

E1000: Hata Ayıklama başlıklı makalede hata ayıklamayla ilgili özel teknikler hakkında bilgi edinin.