کد خطا: ۰۱۰۱

دسته بندی: خرابی تخصیص برنامه

نوع: زمان اجرا

مثال گزارش خطا

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 نتوانسته است یک برنامه XLA کامپایل شده قابل اجرا را در HBM مربوط به 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 برای مشاهده دقیق میزان استفاده از حافظه برنامه خود در طول زمان استفاده کنید. این می‌تواند به شناسایی اوج‌های غیرمنتظره در مصرف حافظه کمک کند.