Hata kodu: 1000

Kategori: Derleme Süresi: HBM OOM

Bu hata, programın TPU cihazda fiziksel olarak mevcut olandan daha fazla yüksek bant genişlikli bellek (HBM) gerektirdiğini gösterir.

Örnek Hata Mesajları:

RESOURCE_EXHAUSTED: XLA:TPU compile permanent error. Ran out of memory in memory space hbm. Used 49.34G of 32.00G hbm. Exceeded hbm capacity by 17.34G.

RESOURCE_EXHAUSTED: TPU TensorCore Hbm usage: 34.82G, SparseCore Hbm usage 174.10G, exceeding available bytes: 95.74G

XLA Arka Uçları: TPU

Genel Bakış

XLA, gerekli tüm statik ayırmaların toplam boyutunun cihazın HBM'sine sığdığından emin olmak için kontroller gerçekleştirir.

Derleyici, TPU'nun sabit HBM kapasitesini çeşitli tahsis türleri için yönetir:

• Program Girişleri ve Çıkışları: Eğitim grupları, optimize edici durumları vb.
• TensorCore + SparseCore Geçicileri: Ara hesaplamalar (ör.etkinleştirmeler, gradyanlar vb.) için dinamik bellek gerekir.
• Derlenmiş İkili: Hem TensorCore (TC) hem de SparseCore (SC) için makine kodu.
• Sistem ek yükü: XLA çalışma zamanı için ayrılmış alan (ör. eski TPU nesillerinde feed içi arabellekler).
• Sabitler: HLO IR'ye yerleştirilmiş sabit değerler HBM'ye ayrılır.
• Derleyici İç İşleyişi: Program düzeyi ve HLO başına tahsisler (ör. ağdaki düğümler için yönlendirme bilgileri)

Bu hata, XLA derleyicisi yukarıdaki tüm tahsisleri cihaz HBM'sine sığdıramadığında oluşur.

Hata ayıklama

Aşağıdaki HBM OOM kategorilerinden hangisinin hatanızı en iyi şekilde tanımladığını belirlemek için hata mesajını ve günlükleri dikkatlice analiz edin:

• TensorCore (TC) + SparseCore (SC) HBM Kullanımı Sınırı Aşıyor: Hata, kullanımı açıkça açıklıyorsa (ör. "TC Hbm usage: X, SC Hbm usage Y". → 1. Bölüm TC ve SC HBM kullanımını dengeleyin.
• Beklenmedik Şekilde Büyük Tahsisler: Hata "HBM bellek alanında bellek tükendi" şeklinde görünüyorsa HBM'deki en büyük tahsislerin listesi için günlükleri kontrol edin. Bir veya daha fazla beklenmedik şekilde büyük tensör (ör.HBM sınırının% 50'sinden fazlası) varsa → 2. Bölüm'e gidin. Beklenmedik derecede büyük tahsisler.
• Aggregate Allocations Exceed HBM Limit (Toplam Tahsisler HBM Sınırını Aşıyor): Hata "Ran out of memory in memory space HBM" (HBM bellek alanında bellek tükendi) şeklinde okunuyorsa ancak günlüklerde beklenmedik şekilde büyük tensörler yoksa → 3. Bölüm'e gidin. Toplam tahsisler HBM sınırını aşıyor.

1. Bölüm TC ve SC HBM kullanımını dengeleme

Hata, kullanımı açıkça açıklıyorsa (ör. "TC Hbm kullanımı: X, SC Hbm kullanımı Y" değerlerini karşılaştırarak darboğazı belirleyin.

• Yüksek SparseCore Kullanımı:
  • HBM yığın kullanımını optimize etme: HBM yığın bellek tüketimi feature_width, max_unique_nz_per_row ve logical_replica_count ile ölçeklenir. Tabloların işlenmesini sıralayan --xla_sc_num_serialized_tables_to_optimize_hbm işaretini ayarlayarak en yüksek yığın kullanımını azaltabilirsiniz. Bu durum, paralelliğin azalması pahasına gerçekleşir.
  • Dolgu ek yükünü kontrol edin: SparseCore, yerleştirme tablolarını 32 B'ye (8 kayan nokta) göre hizalar. Küçük özellik genişliklerine sahip tablolar (ör. < 8 kayan nokta) önemli bir dolgu ek yüküne neden olarak HBM'yi boşa harcar.
  • Yığın Kullanımını Azaltma: maximum_parallel_iterations için yüksek değerler, HBM yığınına önceden getirilen giriş verilerinin miktarını artırır. Bu değeri düşürmek önemli miktarda bellek alanı açabilir.
  • Parçalama işlemini doğrulayın: Yerleştirme tablolarının tüm çiplerde düzgün şekilde mod-parçalandığından emin olun. Sınırlar tablolara nasıl yansır? başlıklı makaleyi inceleyin.
  • Daha fazla fikir için SC: Performance and memory bottlenecks (SC: Performans ve bellek darboğazları) başlıklı makaleyi inceleyin.
• Yüksek Tensor Çekirdeği Kullanımı:
  • 2. Bölüm'e geçin.
• Dengeli
  • Hiçbiri ayrı ayrı aşırı olmasa da toplam çok yüksekse çipin kapasitesine ulaşmışsınızdır. Her iki bileşenin de kullanımını azaltmayı denemeniz gerekir. Üç bölümdeki önerileri uygulayın.

Bölüm 2. Beklenmedik Şekilde Büyük Tahsisler

Günlüklerde beklenmedik şekilde büyük bir veya daha fazla tahsis varsa (> HBM sınırının% 50'si) bu neredeyse hiçbir zaman donanım kapasitesi sorunu değildir. Bu hata genellikle yapılandırma hatasıdır. Büyük ayırmaların JAX kaynak koduyla ilgili ipuçları için XLA etiketini (varsa) kontrol edin.

• Hata Ayıklama Yapılarını Kaldırma:
  • Büyük ölçekli çalıştırmalarda jax.debug.print() kullanmak, derleyiciyi tensörün tamamını HBM'de somutlaştırmaya zorlayarak CPU'ya aktarılmasına, birleştirmenin bozulmasına ve bellek kullanımının zirveye çıkmasına neden olabilir. Kalan jax.debug.print() karakterlerini kaldırın.
• Verimsiz Mesh Şekillerini veya Parçalama İşlemini Düzeltme:
  • Yanlış örgü şekilleri veya eksik parçalama ek açıklamaları, derleyicinin varsayılan olarak Replication'ı kullanmasına neden olabilir. Bu durumda derleyici, gerçekten büyük tensörleri tek bir çipe sığdırmaya çalışır.
  • Büyük ayırmaların şekillerini kontrol edin ve parçalama işleminin XLA tarafından doğru şekilde belirtilip yayılıp yayılmadığını doğrulayın.

Bölüm 3. Toplam ayırmalar HBM sınırını aşıyor

Program, tahsislerin toplam tutarının HBM sınırını aşması nedeniyle kapasite dışı kalırsa, en yüksek kullanıma katkıda bulunan belirli arabellekleri belirlemek için bellek profilini görselleştirmek genellikle yararlıdır. En yüksek bellek kullanımına neden olan öğeleri belirleme hakkında adım adım talimatlar için XProf ile OOM hatalarını ayıklama başlıklı makaleyi inceleyin.

En çok katkıda bulunanlardan bazılarını belirledikten sonra, bellekte kaplanan yeri optimize etmek için aşağıdaki adımları uygulayın.

C. Tensor dolgusunu ve hizalamasını kontrol etme

Verimsiz tensör şekilleri, TPU'larda OOM'lerin yaygın ve sessiz bir nedenidir. TPU'larda en iyi performansı elde etmek için XLA, tensör boyutlarını doldurur. Bu işlem genellikle en küçük boyut için 128'in, ikinci en küçük boyut için ise 8'in katları olacak şekilde yapılır. Bu dolgu, hem giriş dizilerini hem de ara tensörleri (HLO geçicileri) etkiler ve özellikle küçük boyutlarda bellek kullanımını önemli ölçüde artırabilir. Dizi Düzenleri'ne bakın.

• Büyük arabelleklerin şekillerini denetleme: (Varsayılan düzenlere sahip TPU v5'te)
  • Xprof Memory Viewer'da bir arabelleğin üzerine gelindiğinde, dolgu bilgileri de dahil olmak üzere arabellek ayrıntılarını içeren arabellek ayrıntıları kartı gösterilir.
  • Örnek: (129, 1024) şekli (256, 1024) şeklinde doldurulabilir ve bu da yaklaşık% 50 bellek israfına neden olur.
  • Düzeltme: (128, 1024) şekli dolgu gerektirmez ve %0 bellek israfına neden olur.
• Boyutları hizalama: Tüm büyük tensör boyutlarının (toplu iş boyutu, yerleştirme boyutu, gizli boyut) 128'in katları olduğundan emin olun.

B. Yapılandırmayı ayarlama

Aşağıdaki yapılandırma düzenlemeleriyle genellikle OOM'leri çözebilirsiniz:

• Toplu İş Boyutunu Küçültün: Ara etkinleştirmeler ve gradyanlar için gereken bellek, toplu iş boyutuyla doğru orantılıdır. Toplu iş boyutunu küçültmek genellikle bellek kullanımını azaltmaya yardımcı olabilir.
• Giriş arabelleklerine bağış yapma: jax.jit kullanılırken model parametreleriniz için donate_argnums değerini belirtin. Bu, XLA'nın giriş belleğinin üzerine çıkışı yazmasına olanak tanır.
• Karışık duyarlılığı (bfloat16) etkinleştirin: Model mimarisi ve kalite koşulları izin veriyorsa programdaki en büyük tensörler için bfloat16 veya nicemleme (int8 vb.) kullanın.

C. Mimarinin ve parçalamanın optimize edilmesi

Yapılandırma değişiklikleri yetersizse model topolojisi, mevcut donanım kurulumu için çok büyük olabilir.

• Daha Yeni TPU Nesillerini Kullanın: Daha yeni TPU'lar genellikle çip başına daha fazla HBM sunar. Mümkünse daha yeni TPU nesillerine geçin.
• Daha büyük bir çip topolojisinde çalıştırma: Model ağırlıkları mevcut topoloji için çok büyükse bunları daha fazla çipte parçalamayı deneyebilirsiniz.
• Gelişmiş parçalama tekniklerini uygulayın:
  • Daha gelişmiş veri, tensör veya ardışık düzen paralelliği yaklaşımlarını keşfedin.
  • Ara değerler ve çıkışlar için parçalama ipuçları belirtin.
• JAX Host Offloading'i kullanın: Büyük tensörleri ana makine CPU belleğine boşaltın. Örneğin, etkinleştirme boşaltma ve optimizasyon aracı durumu boşaltma.

D. Belleği etkileyen önemli XLA işaretlerini ayarlama:

Temel bellek işaretleri, bellek kullanımını azaltmak için performansla dengeleyecek şekilde ayarlanabilir. Ancak bu yöntemler performansı olumsuz etkileyebileceğinden son çare olarak kullanılmalıdır.

E. Tune XLA Rematerialization Pass / Manual Checkpointing

Model belleğe sığmaya yakınsa XLA::Rematerialization geçişini, derlemelerin daha yavaş olması pahasına da olsa bellek tasarrufuna öncelik verecek şekilde zorlayabilirsiniz:

Alternatif olarak, jax.grad ile jax.checkpoint dekoratörünü kullanarak hangi ara değerlerin ileri geçişte kaydedileceğini, hangilerinin geri geçişte yeniden hesaplanacağını manuel olarak kontrol edebilir, HBM için işlem döngülerini değiştirebilirsiniz.

F. Gelişmiş profilleme araçlarını kullanma

XProf ile OOM hatalarını ayıklama, derleyicinin HBM kullanımıyla ilgili görünümünü görselleştirmek için XProf Memory Viewer'ı kullanmayla ilgili bir eğitim sunar.

Bu araç, en yüksek bellek ayırma ve arabellek ömrünü görmenize olanak tanır. Bu, en yüksek kullanım noktasında HBM'yi tam olarak neyin tükettiğini anlamak için çok önemlidir. Genel profil oluşturma kurulumu için Xprof'u kullanmaya başlama ve TensorBoard Profiling başlıklı makaleleri inceleyin.