الفئة: وقت التشغيل: تعذُّر تخصيص البرنامج
يشير هذا الخطأ إلى تعذُّر تحميل وقت تشغيل XLA على جهاز TPU لملف تنفيذي لبرنامج XLA تم تجميعه إلى ذاكرة النطاق الترددي العالي (HBM) في وحدة TPU.
نموذج رسالة خطأ:
XlaRuntimeError: RESOURCE_EXHAUSTED: Error loading program 'jit_embedding_pipeline_step_fn': Attempting to reserve 29.49G at the bottom of memory. That was not possible. There are 147.64M free, 0B reserved, and 147.64M reservable. Scope: unknown..: while running replica 0 and partition 34 of a replicated computation (other replicas may have failed as well).
البرامج الخلفية في XLA: وحدة معالجة الموتّرات (TPU)
نظرة عامة
يحدث هذا الخطأ عادةً لأحد الأسباب التالية:
- حجم البرنامج يتجاوز سعة ذاكرة النطاق الترددي العالي (HBM) المتاحة: حجم برنامج XLA الذي تم تجميعه، بما في ذلك التعليمات والبيانات الثابتة وأي ثوابت مضمّنة، أكبر من إجمالي سعة ذاكرة النطاق الترددي العالي (HBM) المتاحة حاليًا على نوى وحدات معالجة Tensor المحدّدة التي يتم تحميل البرنامج عليها.
- تجزئة الذاكرة العالية النطاق (HBM): على الرغم من أنّ إجمالي الذاكرة العالية النطاق (HBM) المتاحة على الجهاز قد يكون كافيًا بشكل إجمالي، إلا أنّه غير متاح في كتلة واحدة متجاورة وكبيرة بما يكفي لاستيعاب البرنامج بأكمله.
من المهم فهم الطريقة التي يحدّد بها وقت تشغيل وحدة TPU أولوية الذاكرة. تكون عمليات تخصيص المخزن المؤقت أكثر أهمية من البرامج التي تم تحميلها. في حال تعذُّر تخصيص المخزن المؤقت، سيزيل وقت التشغيل البرامج التي تم تحميلها من ذاكرة النطاق الترددي العالي (HBM) لإخلاء بعض المساحة. وقد يؤدي ذلك إلى حدوث خطأ OOM في برنامج كان يتم تحميله بنجاح في السابق، لأنّ ذاكرة النطاق الترددي العالي أصبحت مشغولة بمخازن مؤقتة أكبر للبيانات.
تصحيح الأخطاء
- تقليل استهلاك الذاكرة المؤقتة: سيؤدي إخلاء الذاكرة المستخدَمة في مخازن البيانات المؤقتة إلى توفير مساحة أكبر للبرنامج نفسه:
- تقليل حجم الدفعة: هذه إحدى أكثر الطرق فعالية لتقليل مقدار الذاكرة المستخدَمة في عمليات التفعيل.
- تقسيم المعلمات: بالنسبة إلى النماذج الكبيرة جدًا، استخدِم التوازي على مستوى النموذج أو تقنيات التقسيم (مثل FSDP أو Megascale) لتوزيع معلمات النموذج وعمليات الحساب على عدة نوى أو مضيفات لوحدة معالجة Tensor.
- تقصير طول التسلسل/السياق: للنماذج التي تعالج البيانات التسلسلية (مثل نماذج معالجة اللغة الطبيعية)، يمكن أن يؤدي تقليل طول التسلسل إلى خفض استخدام الذاكرة بشكل كبير.
- Buffer Donation: استخدام ميزات إطار العمل (مثل
jax.jit(..., donate_argnums=...)) للسماح لـ XLA بإعادة استخدام ذاكرة مخازن الإدخال المؤقتة لتخزين الإخراج، ما يقلّل من الحد الأقصى لاستخدام الذاكرة.
- تقليل متطلبات ذاكرة البرنامج للمتغيرات المؤقتة:
- يمكنك تقليل استخدام الذاكرة المؤقتة للبرامج من خلال استخدام العلامة
tpu_shared_memory_percent. يُرجى العِلم أنّ ذلك قد يؤثّر سلبًا في الأداء.
- يمكنك تقليل استخدام الذاكرة المؤقتة للبرامج من خلال استخدام العلامة
- تحسين استراتيجية التنفيذ/تقليل عبء العرض:
- إدارة تحميل البرامج: إذا كنت بصدد تجميع عدة دوال برمجية في الوقت الفعلي، يجب أن تعلم أنّ كل دالة برمجية يمكن أن تؤدي إلى تحميل برنامج. حاوِل تنظيم عبء العمل لتقليل عدد البرامج التي يتم تحميلها في الوقت نفسه.
- التأكّد من عدم حدوث تسرّبات في الذاكرة:
- تأكَّد من عدم الاحتفاظ بإشارات إلى عناصر
jax.Arrayلفترة أطول من المدة المحدّدة. قد يؤدي الاحتفاظ بكائناتjax.Arrayإلى منع إلغاء التخصيص التلقائي حتى بعد اكتمال تجميع البرنامج.
- تأكَّد من عدم الاحتفاظ بإشارات إلى عناصر
الأدوات
- فعِّل العلامة
tpu_log_allocations_on_oomالتي سيعرض فيها برنامج التخصيص تقريرًا تفصيليًا عن جميع عمليات التخصيص الحالية عند حدوث خطأ OOM، وهو ما قد يكون مفيدًا جدًا في تصحيح الأخطاء. - تحديد المشاكل في برنامجك: استخدِم أداة JAX لتحديد المشاكل المتعلّقة بالذاكرة أو أداة TensorFlow لتحديد المشاكل للحصول على عرض تفصيلي لاستخدام برنامجك للذاكرة بمرور الوقت. يمكن أن يساعد ذلك في تحديد الارتفاعات غير المتوقّعة في استهلاك الذاكرة.