Kategori: Derleme Zamanı: Kapsamlı Vmem OOM
Bu hata, programın ayrılan miktardan daha fazla Kapsamlı Vektör Belleği (Vmem) gerektirdiğini gösterir.
Örnek Hata Mesajları:
RESOURCE_EXHAUSTED: Ran out of memory in memory space vmem while allocating on stack for %my-custom-kernel = bf16[2048,4096]{1,0:T(8,128)(2,1)} custom-call(...) ...
XLA Arka Uçları: TPU
Genel Bakış
TPU'lar, yalnızca TensorCore (TC) tarafından kullanılan yerel bir geçici bellek olan Vector Memory'ye (VMEM) sahiptir. Derleyici, farklı türlerdeki ayırmalar için Vmem'i yönetir:
- Talimat kapsamlı ayırmalar: Tek bir HLO talimatı yürütülürken Vmem'deki geçici depolama alanı. Buna, operand aralığı arabelleği (ör. çift arabellek için) ve kayıt dökülmeleri dahildir.
- Program kapsamlı ayırmalar: Tek bir HLO talimatının kapsamı dışında kalan ayırmalar. Bunlar genellikle HLO talimatlarının girişleri ve/veya çıkışları olan HLO geçicileri ve ara sonuçlardır.
Derleme zamanı kapsamlı Vmem OOM, talimat kapsamlı ayırmalar söz konusu talimatın ayırma sınırını aştığında meydana gelir. Bu sınır kontrol ediliyor
- --xla_tpu_scoped_vmem_limit_kib işaretiyle programın tamamı için genel olarak ve
- vmem_limit_bytes param aracılığıyla özel çekirdek başına.
Bu hatalar genellikle dahili bir derleyici hatasından veya tahsis sınırını aşan özel bir çekirdekten kaynaklanır.
Hata ayıklama
Hatayı özel bir çekirdekten mi yoksa standart bir HLO'dan mı kaynaklandığını belirlemek için hata mesajını dikkatlice analiz edin. Özel çekirdekten kaynaklanan bir hata şu imzaya sahip olmalıdır:
Ran out of memory in memory space vmem while allocating on stack for %my-custom-call = <output-shape> custom-call(<params>), custom_call_target="tpu_custom_call" ...
- Özel çekirdek kapsamlı Vmem OOM: Hata özel bir çekirdeği işaret ediyorsa → Çekirdeği yeniden ayarlama bölümüne gidin.
- Çekirdek dışı Vmem sorunları: Vmem OOM, özel olmayan bir çekirdek işlemi nedeniyle oluşuyorsa büyük olasılıkla dahili bir derleyici hatasıdır. Lütfen bir HLO dökümüyle birlikte XLA'da hata bildiriminde bulunun.
Çekirdeği yeniden ayarlama
Hata özel bir çekirdekten kaynaklanıyorsa çekirdeğin bellek gereksinimini azaltmak için aşağıdaki teknikleri kullanın:
- Blok Boyutlarını Ayarlama: Kapsamlı Vmem kullanımını azaltmak için çekirdek yapılandırmanızdaki blok boyutlarını (döşeme boyutları) küçültün.
- Çekirdek Kapsamlı Vmem Sınırları Ayarlama: vmem_limit_bytes parametresini kullanarak söz konusu çekirdek için gereken bellek miktarını açıkça isteyin.
- Bellek Renklendirmesini Değiştirme: pallas.tpu.with_memory_space_constraint kullanarak çekirdeğin girişlerini/çıkışlarını VMEM'e açıkça renklendirin/sınırlandırın. Ancak genel bir VMEM OOM'ye neden olabileceği için çok fazla giriş/çıkışa Vmem rengi vermemeye dikkat edin.
- Çekirdeğe özgü yeniden ayarlama zor veya sorun birçok çekirdeği etkiliyorsa --xla_tpu_scoped_vmem_limit_kib işaretini kullanarak genel Vmem sınırını ayarlayabilirsiniz.