کد خطا: ۱۰۰۱

دسته بندی: زمان کامپایل: محدود شده Vmem OOM

این خطا نشان می‌دهد که برنامه به حافظه برداری محدود (Vmem) بیشتری نسبت به آنچه اختصاص داده شده است، نیاز دارد.

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

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

نمای کلی

TPUها دارای حافظه برداری (VMEM) هستند که یک حافظه محلی برای یادداشت‌برداری است که منحصراً توسط TensorCore (TC) استفاده می‌شود. کامپایلر، Vmem را برای انواع مختلف تخصیص‌ها مدیریت می‌کند:

  • تخصیص‌های محدود به دستورالعمل: ذخیره‌سازی موقت در Vmem هنگام اجرای یک دستورالعمل HLO. این شامل بافر دهانه عملوند (مثلاً برای بافرینگ دوگانه) و نشت‌های ثبات می‌شود.
  • تخصیص‌های محدود به برنامه: تخصیص‌هایی که فراتر از محدوده یک دستورالعمل HLO واحد هستند. اینها معمولاً HLOهای موقت و نتایج میانی هستند که ورودی‌ها و/یا خروجی‌های دستورالعمل‌های HLO می‌باشند.

یک Vmem OOM با محدوده زمانی کامپایل زمانی رخ می‌دهد که تخصیص‌های محدوده دستورالعمل از محدودیت تخصیص برای آن دستورالعمل فراتر رود. این محدودیت کنترل می‌شود.

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

اشکال‌زدایی

پیام خطا را با دقت تجزیه و تحلیل کنید تا مشخص شود که آیا خطا از یک هسته سفارشی یا یک HLO استاندارد ناشی می‌شود. خطایی که به دلیل یک هسته سفارشی ایجاد می‌شود باید امضای زیر را داشته باشد:

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" ...
  • کرنل سفارشی در محدوده Vmem OOM : اگر خطا به یک کرنل سفارشی اشاره دارد → به تنظیم مجدد کرنل بروید.
  • مشکلات Vmem غیر هسته ای : اگر خطای OOM Vmem به دلیل یک عملیات غیر هسته ای سفارشی رخ دهد، احتمالاً یک اشکال کامپایلر داخلی است. لطفاً یک اشکال در XLA را با یک فایل HLO dump ثبت کنید.

تنظیم مجدد هسته

اگر خطا از یک هسته سفارشی سرچشمه می‌گیرد، از تکنیک‌های زیر برای کاهش نیاز هسته به حافظه استفاده کنید:

  • تنظیم اندازه بلوک‌ها: اندازه بلوک‌ها (اندازه کاشی‌ها) را در پیکربندی هسته خود کاهش دهید تا میزان استفاده از Scoped Vmem کمتر شود.
  • تنظیم محدودیت‌های Vmem به ازای هر هسته: با استفاده از پارامتر vmem_limit_bytes، به طور صریح مقدار حافظه مورد نیاز برای آن هسته خاص را درخواست کنید.
  • تغییر رنگ حافظه: ورودی‌ها/خروجی‌های هسته به VMEM را با استفاده از pallas.tpu.with_memory_space_constraint به طور صریح رنگ/محدود کنید. اما مراقب باشید که ورودی‌ها و خروجی‌های زیادی به Vmem رنگ نکنید، زیرا این امر ممکن است باعث ایجاد یک خطای کلی VMEM OOM شود.
  • اگر تنظیم مجدد هسته خاص دشوار است یا مشکل بر بسیاری از هسته‌ها تأثیر می‌گذارد، می‌توانید محدودیت Vmem سراسری را با استفاده از پرچم --xla_tpu_scoped_vmem_limit_kib تنظیم کنید.