کد خطا: ۰۱۰۰

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

این خطا نشان می‌دهد که تخصیص‌دهنده حافظه زمان اجرای XLA:TPU نتوانسته بلوک حافظه مناسبی را در HBM شتاب‌دهنده برای تخصیص درخواستی پیدا کند.

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

ValueError: RESOURCE_EXHAUSTED: Error allocating device buffer: Attempting to allocate 8.00M. That was not possible. There are 6.43M free.; (0x0x1_HBM0)

پایه‌های XLA: TPU

نمای کلی

این خطا روی ... اعمال می‌شود

• خرابی‌های تخصیص بافر توسط کاربر از طریق jax.device_put یا
• خرابی تخصیص خروجی برنامه زمان‌بندی شده توسط کاربر.

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

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

زمان اجرای TPU تعدادی مکانیسم برای تلاش مجدد در صورت خرابی تخصیص دارد، از جمله:

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

بنابراین، خطایی که پس از اقدامات اصلاحی فوق رخ می‌دهد، معمولاً نیاز به اقدام کاربر دارد.

اشکال‌زدایی

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

ابزارسازی

پرچم tpu_log_allocations_on_oom فعال کنید، برای این منظور، تخصیص‌دهنده هنگام وقوع یک OOM، گزارش مفصلی از تمام تخصیص‌های فعلی ارائه می‌دهد که می‌تواند برای اشکال‌زدایی بسیار ارزشمند باشد.