کد خطا: ۰۱۰۱

دسته بندی: زمان اجرا: شکست در تخصیص برنامه

این خطا نشان می‌دهد که زمان اجرای XLA روی یک دستگاه TPU نتوانسته یک برنامه XLA کامپایل شده قابل اجرا را در HBM مربوط به TPU بارگذاری کند.

نمونه پیام خطا:

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: TPU

نمای کلی

این خطا معمولاً به یکی از دلایل زیر ایجاد می‌شود:

• حجم برنامه از HBM موجود بیشتر است: برنامه XLA کامپایل شده، شامل دستورالعمل‌ها، داده‌های استاتیک و هرگونه ثابت تعبیه شده در آن، بزرگتر از کل مقدار HBM آزاد موجود در هسته(های) TPU خاص است که برنامه در آن بارگذاری می‌شود.
• قطعه قطعه شدن HBM: اگرچه کل HBM آزاد روی دستگاه ممکن است در مجموع کافی باشد، اما در یک بلوک واحد و پیوسته که به اندازه کافی بزرگ باشد تا کل برنامه را در خود جای دهد، در دسترس نیست.

درک چگونگی اولویت‌بندی حافظه توسط TPU runtime بسیار مهم است. تخصیص بافر نسبت به برنامه‌های بارگذاری‌شده امتیاز بیشتری دارد. اگر تخصیص بافر با شکست مواجه شود، runtime برنامه‌های بارگذاری‌شده قبلی را از HBM حذف می‌کند تا فضا را آزاد کند. این می‌تواند به وضعیتی منجر شود که برنامه‌ای که قبلاً با موفقیت بارگذاری شده بود، اکنون با خطای OOM از کار بیفتد، زیرا HBM اکنون با بافرهای داده بیشتری اشغال شده است.

اشکال‌زدایی

• کاهش ردپای حافظه بافر: آزاد کردن حافظه مورد استفاده توسط بافرهای داده، فضای بیشتری برای خود برنامه باقی می‌گذارد:
  • کاهش اندازه دسته: این یکی از موثرترین راه‌ها برای کاهش میزان حافظه مورد استفاده برای فعال‌سازی‌ها است.
  • تقسیم‌بندی پارامترها: برای مدل‌های بسیار بزرگ، از تکنیک‌های موازی‌سازی مدل یا تقسیم‌بندی (مانند FSDP یا Megascale) برای توزیع پارامترها و محاسبات مدل در چندین هسته یا میزبان TPU استفاده کنید.
  • کوتاه کردن طول توالی/متن: برای مدل‌هایی که داده‌های ترتیبی را پردازش می‌کنند (مثلاً مدل‌های NLP)، کاهش طول توالی می‌تواند به طور قابل توجهی استفاده از حافظه را کاهش دهد.
  • اهدای بافر: از ویژگی‌های چارچوب (مثلاً jax.jit(..., donate_argnums=...) ) استفاده کنید تا به XLA اجازه دهید از حافظه بافرهای ورودی برای ذخیره خروجی استفاده مجدد کند و میزان استفاده از حافظه در اوج مصرف را کاهش دهد.
• کاهش نیاز برنامه به حافظه برای فایل‌های موقت:
  • با استفاده از پرچم tpu_shared_memory_percent میزان استفاده برنامه‌ها از حافظه برای فایل‌های موقت را کاهش دهید. توجه داشته باشید که این ممکن است بر عملکرد تأثیر منفی بگذارد.
• بهینه سازی استراتژی اجرا/کاهش بار سرویس:
  • مدیریت بارگذاری برنامه: اگر چندین تابع را به صورت JIT کامپایل می‌کنید، توجه داشته باشید که هر تابع می‌تواند منجر به بارگذاری یک برنامه شود. سعی کنید حجم کار خود را طوری ساختار دهید که تعداد برنامه‌های بارگذاری شده همزمان را به حداقل برساند.
• اطمینان حاصل کنید که هیچ نشت حافظه‌ای وجود ندارد:
  • مطمئن شوید که ارجاعات به اشیاء jax.Array بیش از مدت زمان مورد نظر نگهداری نمی‌شوند. نگهداری اشیاء jax.Array ممکن است از حذف خودکار تخصیص حتی پس از تکمیل کامپایل برنامه جلوگیری کند.

ابزارسازی

• پرچم tpu_log_allocations_on_oom فعال کنید، برای این منظور، تخصیص‌دهنده هنگام وقوع یک OOM، گزارش مفصلی از تمام تخصیص‌های فعلی ارائه می‌دهد که می‌تواند برای اشکال‌زدایی بسیار ارزشمند باشد.
• برنامه خود را پروفایل کنید: از پروفایل حافظه JAX یا پروفایل TensorFlow برای مشاهده دقیق میزان استفاده از حافظه برنامه خود در طول زمان استفاده کنید. این می‌تواند به شناسایی اوج‌های غیرمنتظره در مصرف حافظه کمک کند.