کد خطا: ۰۱۰۰

دسته بندی: خرابی تخصیص بافر - TPU

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

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

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