دسته بندی: زمان کامپایل: HBM OOM
این خطا نشان میدهد که برنامه به حافظه با پهنای باند بالا (HBM) بیشتری نسبت به آنچه که به صورت فیزیکی در دستگاه TPU موجود است، نیاز دارد.
نمونه پیامهای خطا:
RESOURCE_EXHAUSTED: TPU TensorCore Hbm usage: 34.82G, SparseCore Hbm usage 174.10G, exceeding available bytes: 95.74G
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.
پایههای XLA: TPU
نمای کلی
XLA بررسیهایی را انجام میدهد تا اطمینان حاصل شود که اندازه کل تمام تخصیصهای استاتیک لازم در HBM دستگاه جای میگیرد.
کامپایلر ظرفیت ثابت HBM مربوط به TPU را برای چندین نوع تخصیص مدیریت میکند:
- ورودیها و خروجیهای برنامه: دستههای آموزشی، حالتهای بهینهساز و غیره
- حافظههای موقت TPU: حافظه پویا که برای محاسبات میانی (مثلاً فعالسازیها، گرادیانها و غیره) مورد نیاز است.
- کد دودویی کامپایلشده: کد ماشین برای هر دو TensorCore (TC) و SparseCore (SC).
- سربار سیستم: فضای رزرو شده برای زمان اجرای XLA (مثلاً بافرهای ورودی در نسلهای قدیمیتر TPU).
- ثابتها: مقادیر ثابت تعبیهشده در HLO IR روی HBM تخصیص داده میشوند.
- داخلی کامپایلر: سطح برنامه و تخصیصهای هر HLO (مثلاً اطلاعات مسیریابی برای گرهها در شبکه).
این خطا زمانی رخ میدهد که کامپایلر XLA نتواند تمام تخصیصهای فوق را در HBM دستگاه جای دهد.
اشکالزدایی
پیام خطا و گزارشها را با دقت تجزیه و تحلیل کنید تا مشخص شود کدام دسته از خطاهای HBM OOM زیر، خطای شما را به بهترین شکل توصیف میکند:
- "میزان استفاده از Hbm در TC: X، میزان استفاده از Hbm در SC: Y": اگر خطا صراحتاً میزان استفاده از HBM را تجزیه کند، مجموع استفاده از TensorCore (TC) + SparseCore (SC) از حد مجاز HBM فراتر میرود. → پرش به سناریوی ۱. متعادل کردن میزان استفاده از HBM در TC و SC .
- "حافظه در فضای حافظه HBM تمام شد" : گزارشها را برای شمارش بزرگترین تخصیصها در HBM بررسی کنید.
- در صورتی که یک یا چند تانسور به طور غیرمنتظرهای بزرگ (مثلاً > 50٪ از حد HBM) وجود داشته باشند → به سناریوی ۲ بروید. به دلیل تخصیصهای غیرمنتظرهی بزرگ، حافظه کافی نیست .
- اگر هیچ تانسور غیرمنتظره بزرگی در لاگها وجود نداشته باشد → به سناریوی ۳ بروید. به دلیل تخصیصهای تجمیعی، حافظه کافی نیست .
سناریو ۱. ایجاد تعادل بین استفاده از HBM در TC و SC
اگر خطا به صراحت میزان استفاده را تفکیک کند، مثلاً "میزان استفاده از TC Hbm: X، میزان استفاده از SC Hbm: Y" ، این به این معنی است که مجموع استفاده از TensorCore (TC) + SparseCore (SC) از حد مجاز HBM فراتر رفته است. برای شناسایی گلوگاه، این دو مقدار را با هم مقایسه کنید:
- استفاده زیاد از SparseCore
- بهینهسازی استفاده از پشته 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میزان دادههای ورودی پیشواکشی شده در حافظه heap HBM را افزایش میدهد. کاهش این مقدار میتواند حافظه قابل توجهی را آزاد کند. - تأیید تقسیمبندی: اطمینان حاصل کنید که جداول جاسازی به درستی در تمام تراشهها تقسیمبندی شدهاند. ببینید چگونه محدودیتها به جداول تبدیل میشوند .
- برای ایدههای بیشتر، SC: Performance and Memory Bottlenecks را بررسی کنید.
- بهینهسازی استفاده از پشته HBM: مصرف حافظه پشته HBM با
- استفاده زیاد از TensorCore
- به سناریوی ۲ بروید.
- متعادل
- اگر هیچکدام به تنهایی بیش از حد نباشد اما مجموع آنها خیلی زیاد باشد، از حداکثر ظرفیت تراشه استفاده میکنید. باید سعی کنید میزان استفاده از هر دو جزء را کاهش دهید. توصیههای هر سه بخش را دنبال کنید.
سناریو ۲. کمبود حافظه به دلیل تخصیصهای غیرمنتظره و بزرگ
اگر پیام خطای "Ran out of memory in memory space HBM" را مشاهده کردید و یک یا چند تخصیص غیرمنتظره بزرگ در گزارشها وجود دارد (بیش از 50٪ از محدودیت HBM)، تقریباً هرگز مشکل ظرفیت سختافزاری نیست. معمولاً یک خطای پیکربندی است. برچسب XLA (در صورت وجود) تخصیصهای بزرگ را برای نکات مربوط به کد منبع JAX آنها بررسی کنید.
- حذف مصنوعات اشکالزدایی
- استفاده از jax.debug.print() در اجراهای بزرگ میتواند کامپایلر را مجبور کند تا کل تانسور را در HBM پیادهسازی کند تا آن را به CPU منتقل کند، که باعث اختلال در فیوژن و افزایش استفاده از حداکثر حافظه میشود. هرگونه
jax.debug.print()باقی مانده را حذف کنید.
- استفاده از jax.debug.print() در اجراهای بزرگ میتواند کامپایلر را مجبور کند تا کل تانسور را در HBM پیادهسازی کند تا آن را به CPU منتقل کند، که باعث اختلال در فیوژن و افزایش استفاده از حداکثر حافظه میشود. هرگونه
- اشکال مش یا شاردینگ ناکارآمد را اصلاح کنید
- اشکال نادرست مش یا حاشیهنویسیهای ناقصِ مربوط به شاردینگ میتواند باعث شود کامپایلر بهطور پیشفرض به حالت تکثیر (replication) برود - و کامپایلر را مجبور کند تا تانسورهای واقعاً بزرگ را روی یک تراشه جا دهد.
- شکل تخصیصهای بزرگ را بررسی کنید و تأیید کنید که شاردینگ به درستی توسط XLA مشخص و منتشر شده است.
سناریو ۳. کمبود حافظه به دلیل تخصیصهای تجمیعی
اگر پیام خطای "Ran out of memory in memory space HBM" را مشاهده کردید و هیچ تانسور غیرمنتظره بزرگی در لاگها وجود نداشت، به دلیل تجاوز مجموع تخصیصها از حد HBM، ظرفیت برنامه تمام شده است. در این حالت، اغلب تجسم پروفایل حافظه برای شناسایی بافرهای خاصی که در اوج استفاده نقش دارند، مفید است. برای راهنمای گام به گام در مورد شناسایی مشارکتکنندگان در اوج حافظه، به Debug OOM errors with XProf مراجعه کنید.
پس از شناسایی برخی از مهمترین عوامل، از مراحل زیر برای بهینهسازی ردپای حافظه استفاده کنید.
سناریو ۳.الف: تنظیم پیکربندی
شما اغلب میتوانید خطاهای OOM را با این تنظیمات پیکربندی برطرف کنید:
- کاهش اندازه دسته: حافظه مورد نیاز برای فعالسازیهای میانی و گرادیانها مستقیماً با اندازه دسته متناسب است. کاهش اندازه دسته اغلب میتواند به کاهش استفاده از حافظه کمک کند، اگرچه ممکن است لازم باشد نرخ یادگیری، تکانه یا پارامترهای بهینهساز خود را برای حفظ پایداری مدل دوباره تنظیم کنید.
- اهدای بافرهای ورودی: وقتی JAX محاسباتی را انجام میدهد، از بافرهای روی دستگاه برای همه ورودیها و خروجیها استفاده میکند. اگر میدانید که یکی از ورودیها پس از محاسبه مورد نیاز نیست و اگر با شکل و نوع عنصر یکی از خروجیها مطابقت داشته باشد، میتوانید مشخص کنید که میخواهید بافر ورودی مربوطه برای نگهداری یک خروجی اهدا شود. این کار حافظه مورد نیاز برای اجرا را به اندازه بافر اهدا شده کاهش میدهد. میتوانید با تعیین پارامتر donate_argnums به عنوان آرگومان هنگام استفاده از
jax.jitبه این هدف برسید. - فعال کردن دقت ترکیبی (bfloat16): در صورتی که معماری مدل و الزامات کیفی اجازه دهند، برای بزرگترین تانسورهای برنامه از bfloat16 یا کوانتیزاسیون (int8 و غیره) استفاده کنید. توجه داشته باشید که این تغییر میتواند بر رفتار مدل تأثیر بگذارد و باید با دقت در نظر گرفته شود.
میکرو بچینگ (اختیاری)
اگر کاهش اندازه کلی دسته یا افزایش تعداد تراشهها امکانپذیر نیست و اندازه دسته به ازای هر تراشه از قبل به حداقل نرسیده است، میتوانید یک استراتژی میکرو-بچینگ را امتحان کنید:
- هر دسته را به
nریز-دسته تقسیم کنید؛ - برای هر میکروبچ، مراحل رفت و برگشت را پردازش کنید؛
- پس از انجام این کار، گرادیانها را جمع کرده و وزن را به طور کلی بهروزرسانی کنید.
This process reduces the activation memory as we divided each batch into n micro-batches, so that the if the original batch had size M , the activation memory size becomes M/n .
مشکلات احتمالی: - این فرآیند زمان گام را افزایش میدهد زیرا ما چندین پاس رو به جلو و عقب داریم. - اگر اندازه مدل و میکروبچ خیلی متفاوت باشند، ممکن است در مدل خود با مشکلات همگرایی مواجه شوید.
سناریوی ۳.ب بهینهسازی معماری و شاردینگ
اگر تغییرات پیکربندی کافی نباشد، توپولوژی مدل ممکن است برای تنظیمات سختافزاری فعلی بیش از حد بزرگ باشد.
- از نسلهای جدیدتر TPU استفاده کنید: TPUهای جدیدتر معمولاً HBM بیشتری در هر تراشه ارائه میدهند؛ در صورت وجود، به نسلهای جدیدتر TPU روی بیاورید.
- اجرا روی توپولوژی تراشه بزرگتر: اگر وزنهای مدل برای توپولوژی موجود خیلی بزرگ هستند، میتوانید آنها را روی تراشههای بیشتری تقسیم کنید.
تکنیکهای پیشرفتهی شاردینگ را پیادهسازی کنید:
- رویکردهای پیشرفتهتر موازیسازی داده، تانسور یا خط لوله را بررسی کنید.
- نکات مربوط به خرد کردن (sharding hints) را برای مقادیر و خروجیهای میانی مشخص کنید.
توجه داشته باشید که این امر ممکن است به دلیل تقسیم تنسورها در چندین تراشه، باعث افزایش سربار ارتباطات شبکه شود.
استفاده از JAX host offloading: تکنیکهای host offloading به کاربر اجازه میدهد تا تانسورهای بزرگ را به حافظه CPU میزبان منتقل کند (مثلاً activation offloading و optimizer state offloading ). توجه داشته باشید که تکنیکهای host offloading میتوانند به شدت بر عملکرد تأثیر بگذارند، زیرا این عملیات سیستم را مجبور میکند تا تانسورهای بزرگ را دائماً بین TPU HBM و CPU RAM جابجا کند.
سناریو ۳.ج بررسی لایهگذاری و ترازبندی تانسور
شکلهای ناکارآمد تانسور، یک علت رایج و خاموش خطاهای OOM در TPUها هستند. برای رسیدن به حداکثر عملکرد در TPUها، XLA ابعاد تانسور را - معمولاً به مضربی از ۱۲۸ برای کوچکترین بُعد و ۸ برای دومین کوچکترین بُعد - پدگذاری میکند. این پدگذاری هم بر آرایههای ورودی و هم بر تانسورهای میانی (موقتهای HLO) تأثیر میگذارد و به طور بالقوه باعث افزایش قابل توجه استفاده از حافظه، به خصوص با اندازههای کوچک بُعد میشود. به طرحبندی آرایه مراجعه کنید.
- اشکال حسابرسی بافرهای بزرگ: (روی TPU نسخه ۵ با طرحبندیهای پیشفرض)
- با نگه داشتن ماوس روی یک بافر در Xprof Memory Viewer، کارت جزئیات بافر نمایش داده میشود که حاوی جزئیات بافر از جمله اطلاعات padding است.
- مثال : شکلی با ابعاد
(129, 1024)ممکن است به(256, 1024)تبدیل شود که منجر به اتلاف تقریباً 50٪ حافظه میشود. - اصلاحیه: شکلی به شکل
(128, 1024)نیازی به padding ندارد و 0% از حافظه را اشغال میکند.
- ابعاد را تراز کنید: اطمینان حاصل کنید که تمام ابعاد بزرگ تانسور (اندازه دسته، بعد تعبیه، اندازه پنهان) مضربی از ۱۲۸ هستند. توجه داشته باشید که این تغییر میتواند بر رفتار مدل تأثیر بگذارد و باید با دقت در نظر گرفته شود.
سناریوی ۳.D تنظیم حافظه کلید که بر پرچمهای XLA تأثیر میگذارد
میتوان پرچمهای کلیدی حافظه را طوری تنظیم کرد که عملکرد را با استفاده کمتر از حافظه به خطر بیندازند. با این حال، این استراتژی باید به عنوان آخرین راه حل مورد استفاده قرار گیرد زیرا میتواند بر عملکرد تأثیر منفی بگذارد.
سناریوی ۳.E: عبور از گذرگاه/چکپوینت دستی برای بازسازی Tune XLA
اگر مدل نزدیک به جا شدن در حافظه است، میتوانید از دکوراتور jax.checkpoint به همراه jax.grad برای کنترل دستی اینکه کدام واسطهها در مسیر رو به جلو ذخیره میشوند و کدام در مسیر رو به عقب دوباره محاسبه میشوند، استفاده کنید. توجه داشته باشید که این عملیات ممکن است بر عملکرد تأثیر بگذارد، زیرا شما به صراحت چرخههای محاسباتی را برای HBM مبادله میکنید. برای اطلاعات بیشتر به مستندات JAX مراجعه کنید: - بررسی گرادیان با jax.checkpoint ( jax.remat ) - کنترل مقادیر ذخیره شده autodiff با jax.checkpoint (معروف به jax.remat ) - حافظههای JAX و تخلیه میزبان
از طرف دیگر، میتوانید XLA::Rematerialization pass را مجبور کنید که صرفهجویی در حافظه را در اولویت قرار دهد، که احتمالاً به قیمت کندتر شدن کامپایلها تمام میشود:
| پرچم | توضیحات | تأثیر / بده بستان |
|---|---|---|
--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 اضافی را فعال میکند. | میتواند صرفهجویی بیشتری در حافظه ایجاد کند، اما زمان کامپایل را افزایش میدهد و ممکن است به طور بالقوه بر پایداری عددی تأثیر بگذارد . |
توجه داشته باشید که ایجاد تغییر در پرچمهای XLA باید به عنوان آخرین راه حل استفاده شود، زیرا میتواند بر عملکرد تأثیر منفی بگذارد.
سناریو ۳.و استفاده از ابزارهای پیشرفته پروفایلینگ
«اشکالزدایی خطاهای OOM با XProf» آموزشی در مورد استفاده از نمایشگر حافظه XProf برای تجسم دیدگاه کامپایلر از میزان استفاده از HBM ارائه میدهد.
این ابزار به شما امکان میدهد حداکثر تخصیص حافظه و طول عمر بافر را مشاهده کنید، که برای درک دقیق آنچه HBM را در نقطه اوج استفاده مصرف میکند، بسیار مهم است. برای تنظیمات کلی پروفایلبندی، به شروع کار با Xprof و TensorBoard Profiling مراجعه کنید.
جدول خلاصه
جدول زیر شامل خلاصهای از مداخلات بالقوه برای حل خطاهای OOM و اطلاعاتی است که به شما در تصمیمگیری در مورد اقدامات لازم کمک میکند.
| مداخله | انجامش امنه؟ (آیا رفتار برنامه رو تغییر میده؟) | سودهای بالقوه | نشانههای آشکار (آیا این واقعاً همان تنگنایی است که با آن مواجه هستید؟) |
|---|---|---|---|
| استفاده از تکنیکهای پیشرفتهی شاردینگ | بله. تقریباً هرگز صحت عددی آزمایش را تغییر نمیدهد، اگرچه میتواند به دلیل تقسیم تنسورها در چندین تراشه، باعث سربار ارتباط شبکه شود. | سود عظیم (تا ۲۵۶ برابر کاهش) | تخصیصهای انفرادی غیرمنتظره و بزرگ در نمایشگر حافظه (مثلاً، یک تانسور واحد که در تمام TPUها تکثیر شده و ۲۵۶ برابر بزرگتر از بقیه است). آرایههای فعال در هوکهای TensorBoard به صورت un-sharded نشان داده میشوند. |
| کاهش اندازه دسته | خیر. این روش دینامیک آموزش را تغییر میدهد و معمولاً نیاز به تنظیم مجدد نرخ یادگیری دارد. (توجه: میکروبچینگ یک جایگزین ایمن است که بدون تغییر رفتار، حافظه را کاهش میدهد). | سودهای کلان (میتواند هزاران برابر صرفهجویی ایجاد کند). | عدم تخصیص «موقتها» در طول محاسبات گرادیان. مشاهده «JVP» در نام عملیات، و مواجهه با تنسورهای متعدد با اندازه دستهای در پروفایل حافظه. |
| فعال کردن دقت ترکیبی (مثلاً Bfloat16) | خطرناک است. دقت عددی را تغییر میدهد، که میتواند نتایج آزمایش را تغییر دهد یا باعث شود مدل به طور کامل همگرا نشود. | افزایش متوسط (معمولاً ضریب ۲ برابر، زیرا استفاده از حافظه را نصف میکند). | نمایشگر حافظه تأیید میکند که بزرگترین تانسورها در حال حاضر از اعشار ۳۲ بیتی ( float32 ) استفاده میکنند. |
بازرسی دستی ( jax.checkpoint ) | بله. این روش رفتار را تغییر نمیدهد؛ بلکه صرفاً با محاسبه مجدد تانسورها به جای ذخیره آنها، زمان محاسبه (فلاپ) را برای صرفهجویی در حافظه مبادله میکند. | دستاوردهای بزرگ (مثلاً میتواند منجر به این شود که فقط نیمی از فعالسازیها باید همزمان در حافظه وجود داشته باشند). | چندین تانسور با اندازه دقیقاً یکسان که در طول یک حرکت رو به عقب، حافظه را پر میکنند. اغلب با "JVP" در نام عملیات همراه هستند. |
اهدای بافرهای ورودی ( donate_argnums ) | بله. یکپارچگی آزمایش را حفظ میکند. اگر به اشتباه اعمال شود، دادهها را خراب نمیکند، بلکه فقط یک پیام خطای واضح نمایش میدهد. | سود حاشیهای (حدود ۱٪ صرفهجویی در حافظه). | هیچ نشانهی خاصی برای گفتن وجود ندارد، اما این یک «برد رایگان» محسوب میشود که همیشه ارزش امتحان کردن را دارد. |
| تغییر ابعاد مدل | خیر. مستقیماً رفتار مدل را تغییر میدهد. تغییر ابعاد ورودی یا خروجی میتواند سازگاری با مجموعه دادهها را به طور کامل از بین ببرد. | بسته به اینکه ابعاد یا لایههای پنهان چقدر به شدت کاهش یافتهاند ، دستاوردهای متغیری حاصل میشود . | نمایشگر حافظه Xprof مقدار زیادی از حافظه را که صرف «padding» میشود، نشان میدهد، زیرا ابعاد آرایه توانی از دو یا مضربی از ۱۲۸ نیستند (مثلاً بُعد ۲۰۵۰ به جای ۲۰۴۸). |
| تخلیه بار میزبان (CPU) | بله (از نظر عددی) ، اما خیر (از نظر عملکرد) . اگرچه از نظر ریاضی ایمن است، اما به عنوان یک "محدوده سرعت" در نظر گرفته میشود که ممکن است به دلیل انتقال داده بین CPU و TPU باعث ایجاد گلوگاههای شدید سرعت شود. | پیشرفت جزئی (حافظه CPU فقط حدود ۳ برابر TPU است). | معمولاً آخرین راه حل برای حالتهای بهینهساز حجیم یا برای مراحل آمادهسازی/پیشپردازش دادههای سنگین از نظر حافظه است. |