دسته بندی: زمان کامپایل: HBM OOM
این خطا نشان میدهد که برنامه به حافظه با پهنای باند بالا (HBM) بیشتری نسبت به آنچه که به صورت فیزیکی در دستگاه TPU موجود است، نیاز دارد.
نمونه پیامهای خطا:
RESOURCE_EXHAUSTED: XLA:TPU compile permanent error. Ran out of memory in memory space hbm. Used 49.34G of 32.00G hbm. Exceeded hbm capacity by 17.34G.
RESOURCE_EXHAUSTED: TPU TensorCore Hbm usage: 34.82G, SparseCore Hbm usage 174.10G, exceeding available bytes: 95.74G
پایههای XLA: TPU
نمای کلی
XLA بررسیهایی را انجام میدهد تا اطمینان حاصل شود که اندازه کل تمام تخصیصهای استاتیک لازم در HBM دستگاه جای میگیرد.
کامپایلر ظرفیت ثابت HBM مربوط به TPU را برای چندین نوع تخصیص مدیریت میکند:
- ورودیها و خروجیهای برنامه: دستههای آموزشی، حالتهای بهینهساز و غیره
- حافظههای موقت TensorCore + SparseCore: حافظه پویا مورد نیاز برای محاسبات میانی (مثلاً فعالسازیها، گرادیانها و غیره).
- کامپایل دودویی: کد ماشین برای هر دو TensorCore (TC) و SparseCore (SC).
- سربار سیستم: فضای رزرو شده برای زمان اجرای XLA (مثلاً بافرهای ورودی در نسلهای قدیمیتر TPU).
- ثابتها: مقادیر ثابت تعبیهشده در HLO IR روی HBM تخصیص داده میشوند.
- داخلی کامپایلر: سطح برنامه و تخصیصهای هر HLO (مثلاً اطلاعات مسیریابی برای گرهها در شبکه)
این خطا زمانی رخ میدهد که کامپایلر XLA نتواند تمام تخصیصهای فوق را در HBM دستگاه جای دهد.
اشکالزدایی
پیام خطا و گزارشها را با دقت تجزیه و تحلیل کنید تا مشخص شود کدام دسته از خطاهای HBM OOM زیر، خطای شما را به بهترین شکل توصیف میکند:
- استفاده از حافظههای ذخیرهسازی ابری TensorCore (TC) + SparseCore (SC) بیش از حد مجاز است : اگر خطا صریحاً میزان استفاده را تجزیه و تحلیل کند، مثلاً "میزان استفاده از حافظههای ذخیرهسازی ابری TC: X، میزان استفاده از حافظههای ذخیرهسازی ابری SC: Y" . → به بخش ۱ بروید. میزان استفاده از حافظههای ذخیرهسازی ابری TC و SC را متعادل کنید .
- تخصیصهای غیرمنتظره بزرگ : اگر خطا «حافظه در فضای حافظه HBM تمام شد» را نشان میدهد، گزارشها را برای شمارش بزرگترین تخصیصها در HBM بررسی کنید. در صورتی که یک یا چند تنسور غیرمنتظره بزرگ (مثلاً > 50٪ از حد HBM) وجود داشته باشد → به بخش 2 بروید. تخصیصهای غیرمنتظره بزرگ .
- تخصیصهای تجمیعی از حد مجاز HBM فراتر میروند : اگر خطا «حافظه در فضای حافظه HBM تمام شد» را نشان میدهد اما هیچ تانسور غیرمنتظره بزرگی در گزارشها وجود ندارد → به بخش ۳ بروید. تخصیصهای تجمیعی از حد مجاز HBM فراتر میروند .
بخش ۱. ایجاد تعادل بین استفاده از HBM در TC و SC
اگر خطا صراحتاً میزان استفاده را تفکیک میکند، مثلاً «میزان استفاده از TC Hbm: X، میزان استفاده از SC Hbm: Y»، دو مقدار را برای شناسایی گلوگاه مقایسه کنید.
- استفادهی زیاد از هستهی پراکنده:
- بهینهسازی استفاده از پشته HBM: مصرف حافظه پشته HBM با
feature_width،max_unique_nz_per_rowوlogical_replica_countافزایش مییابد. میتوانید با تنظیم پرچم--xla_sc_num_serialized_tables_to_optimize_hbmکه پردازش جداول را سریالی میکند، میزان استفاده از پشته در اوج را کاهش دهید. این کار به قیمت کاهش موازیسازی تمام میشود. - بررسی سربار Padding: SparseCore جداول جاسازی را با 32B (8 floats) تراز میکند. جداول با عرض ویژگی کوچک (مثلاً کمتر از 8 floats) سربار Padding قابل توجهی ایجاد میکنند و HBM را هدر میدهند.
- کاهش استفاده از حافظه Heap: مقادیر بالای
maximum_parallel_iterationsمیزان دادههای ورودی پیشواکشی شده در حافظه HBM heap را افزایش میدهد. کاهش این مقدار میتواند حافظه قابل توجهی را آزاد کند. - تأیید تقسیمبندی: اطمینان حاصل کنید که جداول جاسازی به درستی در تمام تراشهها تقسیمبندی شدهاند. ببینید چگونه محدودیتها به جداول تبدیل میشوند .
- برای ایدههای بیشتر ، SC را بررسی کنید: گلوگاههای عملکرد و حافظه .
- بهینهسازی استفاده از پشته HBM: مصرف حافظه پشته HBM با
- استفادهی بالا از TensorCore:
- به بخش ۲ بروید.
- متعادل
- اگر هیچکدام به تنهایی بیش از حد نباشد اما مجموع آنها خیلی زیاد باشد، از حداکثر ظرفیت تراشه استفاده میکنید. باید سعی کنید میزان استفاده از هر دو جزء را کاهش دهید. توصیههای هر سه بخش را دنبال کنید.
بخش ۲. تخصیصهای غیرمنتظره بزرگ
اگر یک یا چند تخصیص غیرمنتظره و بزرگ در گزارشها وجود داشته باشد (بیش از ۵۰٪ از حد مجاز HBM)، تقریباً هرگز مشکل از ظرفیت سختافزار نیست. معمولاً یک خطای پیکربندی است. برچسب XLA (در صورت وجود) تخصیصهای بزرگ را برای نکات مربوط به کد منبع JAX آنها بررسی کنید.
- حذف مصنوعات اشکالزدایی:
- استفاده از jax.debug.print() در اجراهای بزرگ میتواند کامپایلر را مجبور کند تا کل تانسور را در HBM پیادهسازی کند تا آن را به CPU منتقل کند، که باعث اختلال در فیوژن و افزایش استفاده از حداکثر حافظه میشود. هرگونه
jax.debug.print()باقی مانده را حذف کنید.
- استفاده از jax.debug.print() در اجراهای بزرگ میتواند کامپایلر را مجبور کند تا کل تانسور را در HBM پیادهسازی کند تا آن را به CPU منتقل کند، که باعث اختلال در فیوژن و افزایش استفاده از حداکثر حافظه میشود. هرگونه
- اشکال مش یا شاردینگ ناکارآمد را اصلاح کنید:
- اشکال نادرست مش یا حاشیهنویسیهای نادرست مربوط به شاردینگ میتواند باعث شود کامپایلر به طور پیشفرض روی Replication قرار گیرد - و کامپایلر را مجبور کند تا تانسورهای بسیار بزرگ را روی یک تراشه جا دهد.
- شکل تخصیصهای بزرگ را بررسی کنید و تأیید کنید که شاردینگ به درستی توسط XLA مشخص و منتشر شده است.
بخش 3. تخصیصهای کل از حد HBM تجاوز میکنند
اگر به دلیل تجاوز مجموع تخصیصها از حد مجاز HBM، ظرفیت برنامه تمام شود، اغلب تجسم پروفایل حافظه برای شناسایی بافرهای خاصی که در اوج استفاده نقش دارند، مفید است. برای راهنمای گام به گام در مورد شناسایی مشارکتکنندگان در اوج حافظه، به بخش اشکالزدایی خطاهای OOM با XProf مراجعه کنید.
پس از شناسایی برخی از مهمترین عوامل، از مراحل زیر برای بهینهسازی ردپای حافظه استفاده کنید.
الف. بررسی میزان پر کردن و تراز کردن تانسورها
شکلهای ناکارآمد تانسور، یک علت رایج و خاموش خطاهای OOM در TPUها هستند. برای رسیدن به حداکثر عملکرد در TPUها، XLA ابعاد تانسور را - معمولاً به مضربی از ۱۲۸ برای کوچکترین بُعد و ۸ برای دومین کوچکترین بُعد - پدگذاری میکند. این پدگذاری هم بر آرایههای ورودی و هم بر تانسورهای میانی (موقتهای HLO) تأثیر میگذارد و به طور بالقوه باعث افزایش قابل توجه استفاده از حافظه، به خصوص با اندازههای کوچک بُعد میشود. به طرحبندی آرایه مراجعه کنید.
- اشکال حسابرسی بافرهای بزرگ: (روی TPU نسخه ۵ با طرحبندیهای پیشفرض)
- با نگه داشتن ماوس روی یک بافر در Xprof Memory Viewer، کارت جزئیات بافر نمایش داده میشود که حاوی جزئیات بافر از جمله اطلاعات padding است.
- مثال : شکلی با ابعاد
(129, 1024)ممکن است به(256, 1024)تبدیل شود که منجر به اتلاف تقریباً 50٪ حافظه میشود. - اصلاحیه: شکلی به شکل
(128, 1024)نیازی به padding ندارد و 0% از حافظه را اشغال میکند.
- ابعاد را تراز کنید: اطمینان حاصل کنید که تمام ابعاد بزرگ تانسور (اندازه دسته، بعد تعبیه، اندازه پنهان) مضربی از ۱۲۸ هستند.
ب. تنظیم پیکربندی
شما اغلب میتوانید خطاهای OOM را با این تنظیمات پیکربندی برطرف کنید:
- کاهش اندازه دسته: حافظه مورد نیاز برای فعالسازیهای میانی و گرادیانها مستقیماً با اندازه دسته متناسب است. کاهش اندازه دسته اغلب میتواند به کاهش استفاده از حافظه کمک کند.
- اهدا بافرهای ورودی: هنگام استفاده از
jax.jit، برای پارامترهای مدل خود donate_argnums را مشخص کنید. این به XLA اجازه میدهد تا حافظه ورودی را با خروجی بازنویسی کند. - فعال کردن دقت ترکیبی (bfloat16): در صورتی که معماری مدل و الزامات کیفی اجازه دهند، از bfloat16 یا کوانتیزاسیون (int8 و غیره) برای بزرگترین تانسورهای برنامه استفاده کنید.
ج. بهینهسازی معماری و شاردینگ
اگر تغییرات پیکربندی کافی نباشد، توپولوژی مدل ممکن است برای تنظیمات سختافزاری فعلی بیش از حد بزرگ باشد.
- از نسلهای جدیدتر TPU استفاده کنید: TPUهای جدیدتر معمولاً HBM بیشتری در هر تراشه ارائه میدهند؛ در صورت وجود، به نسلهای جدیدتر TPU روی بیاورید.
- اجرا روی توپولوژی تراشه بزرگتر: اگر وزنهای مدل برای توپولوژی موجود خیلی بزرگ هستند، میتوانید آنها را روی تراشههای بیشتری تقسیم کنید.
- تکنیکهای پیشرفتهی شاردینگ را پیادهسازی کنید:
- رویکردهای پیشرفتهتر موازیسازی داده، تانسور یا خط لوله را بررسی کنید.
- نکات مربوط به خرد کردن (sharding hints) را برای مقادیر و خروجیهای میانی مشخص کنید.
- استفاده از تخلیه بار میزبان JAX: تخلیه بار تنسورهای بزرگ به حافظه پردازنده میزبان. به عنوان مثال تخلیه بار فعالسازی و تخلیه بار وضعیت بهینهساز .
د. تنظیم حافظه کلید که بر پرچمهای XLA تأثیر میگذارد:
میتوان پرچمهای کلیدی حافظه را طوری تنظیم کرد که عملکرد را با استفاده کمتر از حافظه به خطر بیندازند. اما این موارد باید به عنوان آخرین راه حل مورد استفاده قرار گیرند زیرا میتوانند بر عملکرد تأثیر منفی بگذارند.
عبور / بازرسی دستی E. Tune XLA Rematerialization
اگر مدل به جایگیری در حافظه نزدیک است، میتوانید XLA::Rematerialization را مجبور کنید که صرفهجویی در حافظه را در اولویت قرار دهد، که احتمالاً به قیمت کندتر شدن کامپایلها تمام میشود:
| پرچم | توضیحات | تأثیر / بده بستان |
|---|---|---|
--xla_tpu_max_hbm_size_mib | محدودیت اندازه HBM مورد استفاده توسط مرحله Rematerialization را به صورت دستی تعیین میکند. | کامپایلر را مجبور میکند تا سختتر کار کند تا برنامه را در محدودهای کوچکتر از حافظه فیزیکی HBM قرار دهد. |
--xla_tpu_rematerialization_algo=PEAK_PRIORITY | تلاشها را در نقاط اوج استفاده از حافظه متمرکز میکند. | میتواند برای کاهش شدید حافظه نسبت به الگوریتم پیشفرض کارآمدتر باشد. |
--xla_tpu_rematerialization_max_block_size_limit=32 | حداکثر تعداد دستورالعملهای موجود در یک بلوک را که میتوانند به طور همزمان بازسازی شوند، کنترل میکند. | افزایش این مقدار، امکان صرفهجویی در حافظه را با هزینه افزایش قابل توجه زمان کامپایل فراهم میکند. |
--xla_tpu_rematerialization_block_effort_factor=10.0 | میزان تلاش (زمان کامپایل) صرف شده برای جستجوی بلوکها جهت بازسازی را تعریف میکند. | مقادیر بالاتر امکان جستجوی جامعتری را برای صرفهجویی در حافظه فراهم میکنند، اما به قیمت افزایش زمان کامپایل . |
--xla_tpu_pre_fusion_remat=true | قبل از مرحلهی فیوژن، یک مرحلهی Rematerialization اضافی را فعال میکند. | میتواند صرفهجویی بیشتری در حافظه ایجاد کند، اما زمان کامپایل را افزایش میدهد و ممکن است به طور بالقوه بر پایداری عددی تأثیر بگذارد . |
روش دیگر، استفاده از دکوراتور jax.checkpoint به همراه jax.grad برای کنترل دستی اینکه کدام واسطهها در مسیر رو به جلو ذخیره شوند و کدام در مسیر رو به عقب دوباره محاسبه شوند، است که در آن چرخههای محاسباتی با HBM جایگزین میشوند.
و. از ابزارهای پیشرفته پروفایلینگ استفاده کنید
«اشکالزدایی خطاهای OOM با XProf» آموزشی در مورد استفاده از نمایشگر حافظه XProf برای تجسم دیدگاه کامپایلر از میزان استفاده از HBM ارائه میدهد.
این ابزار به شما امکان میدهد حداکثر تخصیص حافظه و طول عمر بافر را مشاهده کنید، که برای درک دقیق آنچه HBM را در نقطه اوج استفاده مصرف میکند، بسیار مهم است. برای تنظیمات کلی پروفایلبندی، به شروع کار با Xprof و TensorBoard Profiling مراجعه کنید.