تصحيح أخطاء نفاد الذاكرة باستخدام XProf

تحدث أخطاء نفاد الذاكرة (OOM) عندما تنتهي سعة ذاكرة النطاق الترددي العالي (HBM) الخاصة بأداة التسريع (وحدة معالجة الرسومات أو وحدة معالجة Tensor). يمكنك الاطّلاع على بعض الأسباب الشائعة لمشاكل OOM وتقنيات تصحيح الأخطاء بالتفصيل في مستند E1000 - Compile Time HBM OOM ومستند JAX حول تخصيص ذاكرة وحدة معالجة الرسومات.

توضّح هذه الصفحة كيفية استخدام أداة Memory Viewer في XProf لتصوُّر استخدام الذاكرة في برنامج JAX، وتحديد حالات الاستخدام القصوى، وتصحيح أخطاء نفاد الذاكرة. يتضمّن ذلك الخطوات التالية:

  1. شغِّل برنامجك باستخدام jax.profiler.trace لتسجيل بيانات الأداء.
  2. ابدأ XProf في الخلفية، واستخدِم أداة "عارض الذاكرة" لعرض تفاصيل استخدام الذاكرة.

مثال على برنامج

يؤدي برنامج JAX التالي إلى حدوث خطأ OOM:

import jax
from jax import random
import jax.numpy as jnp


@jax.profiler.trace("/tmp/xprof")
@jax.jit
def oom():
    a = random.normal(random.PRNGKey(1), (327680, 327680), dtype=jnp.bfloat16)
    return a @ a


if __name__ == "__main__":
    oom()

على جهاز TPU، يتعذّر تنفيذ هذا البرنامج مع ظهور الخطأ التالي:

XlaRuntimeError: RESOURCE_EXHAUSTED: Allocation (size=107374182400) would exceed memory (size=17179869184) :: #allocation7 [shape = 'u8[327680,327680]{1,0:T(8,128)(4,1)}', space=hbm, size = 0xffffffffffffffff, tag = 'output of xor_convert_fusion@{}'] :: <no-hlo-instruction>

(على جهاز مزوّد بوحدة معالجة الرسومات، يبدو الخطأ على النحو التالي: XlaRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 214748364800 bytes.)

تشغيل XProf

ثبِّت xprof (pip install xprof)، وابدأ مثيلاً من XProf مع تحديد الدليل الذي يتم فيه تخزين الملف الشخصي:

xprof --logdir=/tmp/xprof/ --port=6006

انتقِل إلى الجهاز الظاهري (على جهاز محلي، على http://localhost:6006). في القائمة المنسدلة الأدوات، اختَر عارض الذاكرة، وفي نافذة أداة "عارض الذاكرة"، اختَر ذاكرة النطاق الترددي العالي في القائمة المنسدلة أنواع الذاكرة (يتم اختيارها تلقائيًا عادةً).

صفحة &quot;أداة فحص الذاكرة&quot; في XProf لبرنامج المثال أعلاه

توضّح مستندات أداة &quot;عارض الذاكرة&quot; في XProf مكوّنات الأداة والمعلومات المعروضة.

ركِّز على قسم عمليات HLO عند ذروة تخصيص الذاكرة الذي يعرض ثلاثة رسومات بيانية للمخزن المؤقت عند ذروة استخدام الذاكرة. تتضمّن المخزن المؤقت ما يلي: * مدخلات ومخرجات البرنامج: دفعات التدريب، وحالات المحسّن، وما إلى ذلك * بيانات TensorCore وSparseCore المؤقتة: الذاكرة الديناميكية المطلوبة لإجراء العمليات الحسابية الوسيطة (مثل عمليات التنشيط والتدرجات وما إلى ذلك)

يمكنك تمرير مؤشر الماوس فوق الرسومات البيانية الخاصة بالمخزن المؤقت للحصول على مزيد من التفاصيل حول العملية، مثل حجمها وشكلها ونوع التخصيص وغير ذلك. يمكن أن يساعدك ذلك في تحديد وتقييم العمليات التي قد تتضمّن قيمًا مؤقتة عالية أو طويلة الأمد، وأي موترات كبيرة للإدخال أو الوسيطة أو الإخراج تتضمّن مساحة غير مستخدَمة بشكل فعّال، وما إلى ذلك، والتي تساهم في زيادة استخدام الذاكرة إلى الحد الأقصى ويجب تعديلها أو تحسينها.

يمكنك التعرّف على تقنيات تصحيح الأخطاء المحدّدة في E1000: تصحيح الأخطاء.