Hata kodu: 0101

Kategori: Çalışma zamanı: Program ayırma hatası

Bu hata, bir TPU cihazındaki XLA çalışma zamanının, derlenmiş bir XLA programı yürütülebilir dosyasını TPU'nun HBM'sine yükleyemediğini gösterir.

Örnek Hata Mesajı:

XlaRuntimeError: RESOURCE_EXHAUSTED: Error loading program 'jit_embedding_pipeline_step_fn': Attempting to reserve 29.49G at the bottom of memory. That was not possible. There are 147.64M free, 0B reserved, and 147.64M reservable. Scope: unknown..: while running replica 0 and partition 34 of a replicated computation (other replicas may have failed as well).

XLA arka uçları: TPU

Genel Bakış

Bu hata genellikle aşağıdaki nedenlerden biriyle oluşur:

• Program Boyutu, Kullanılabilir HBM'yi Aşıyor: Talimatları, statik verileri ve yerleştirilmiş sabitleri de dahil olmak üzere derlenmiş XLA programı, programın yüklendiği belirli TPU çekirdeklerinde şu anda kullanılabilen toplam ücretsiz HBM miktarından daha büyük.
• HBM Parçalanması: Cihazdaki toplam boş HBM, toplu olarak yeterli olsa da programın tamamını sığdıracak kadar büyük tek ve bitişik bir blokta kullanılamaz.

TPU çalışma zamanının belleğe nasıl öncelik verdiğini anlamak önemlidir. Arabellek ayrımları, yüklü programlara göre önceliklidir. Arabellek ayırma işlemi başarısız olursa çalışma zamanı, yer açmak için önceden yüklenmiş programları HBM'den çıkarır. Bu durum, daha önce başarıyla yüklenen bir programın artık HBM daha fazla veri arabelleğiyle dolu olduğundan OOM hatasıyla başarısız olmasına neden olabilir.

Hata ayıklama

• Arabellek Belleğinde Kaplanan Yeri Azaltma: Veri arabellekleri tarafından kullanılan belleği boşaltmak, programın kendisi için daha fazla alan bırakır:
  • Toplu İş Boyutunu Küçültme: Bu, etkinleştirmeler için kullanılan bellek miktarını azaltmanın en etkili yollarından biridir.
  • Parametre Parçalama: Çok büyük modellerde, modelin parametrelerini ve hesaplamasını birden fazla TPU çekirdeğine veya ana makineye dağıtmak için model paralelliği ya da parçalama teknikleri (ör. FSDP veya Megascale) kullanın.
  • Sıra/Bağlam Uzunluğunu Kısaltma: Sıralı verileri işleyen modeller için (ör. NLP modelleri), dizinin uzunluğunu azaltmak bellek kullanımını önemli ölçüde azaltabilir.
  • Arabellek bağışı: Çerçeve özelliklerini kullanın (ör. jax.jit(..., donate_argnums=...)) XLA'nın, çıkışı depolamak için giriş arabelleklerinin belleğini yeniden kullanmasına izin vererek en yüksek bellek kullanımını azaltır.
• Geçici dosyalar için programın bellek gereksinimlerini azaltın:
  • tpu_shared_memory_percent işaretini kullanarak geçici dosyalar için programların bellek kullanımını azaltın. Bu durumun performansı olumsuz etkileyebileceğini unutmayın.
• Yürütme stratejisini optimize etme/Yayınlama yükünü azaltma:
  • Program Yüklemeyi Yönetme: Birden fazla işlevi JIT derliyorsanız her işlevin bir programın yüklenmesine neden olabileceğini unutmayın. Eşzamanlı olarak yüklenen program sayısını en aza indirmek için iş yükünüzü yapılandırmaya çalışın.
• Bellek sızıntısı olmadığından emin olun:
  • jax.Array nesnelerine yapılan referansların amaçlanandan daha uzun süre tutulmadığından emin olun. jax.Array nesnelerini tutmak, program derlemesi tamamlandıktan sonra bile otomatik olarak bellekten kaldırmayı engelleyebilir.

Takım

• Allocator'ın, bellek yetersizliği oluştuğunda mevcut tüm ayırmaların ayrıntılı bir raporunu dökeceği tpu_log_allocations_on_oom işaretini etkinleştirin. Bu, hata ayıklama için çok değerli olabilir.
• Programınızın Profilini Oluşturma: Programınızın zaman içindeki bellek kullanımına ilişkin ayrıntılı bir görünüm elde etmek için JAX bellek profilleme aracını veya TensorFlow profilleme aracını kullanın. Bu, bellek tüketimindeki beklenmedik artışları belirlemeye yardımcı olabilir.