اشکال‌زدایی خطاهای OOM با XProf

خطاهای کمبود حافظه (OOM) زمانی رخ می‌دهند که ظرفیت حافظه پهنای باند بالای (HBM) شتاب‌دهنده (GPU یا TPU) به پایان رسیده باشد. برخی از دلایل رایج مشکلات OOM و تکنیک‌های اشکال‌زدایی در مستندات E1000 - Compile Time HBM OOM و مستندات JAX در مورد تخصیص حافظه GPU به تفصیل شرح داده شده است.

این صفحه نحوه استفاده از ابزار نمایشگر حافظه XProf را برای نمایش میزان استفاده از حافظه برنامه JAX شما، شناسایی موارد اوج استفاده و اشکال‌زدایی خطاهای OOM شرح می‌دهد. این شامل مراحل زیر است:

  1. برنامه خود را با jax.profiler.trace اجرا کنید تا پروفایل را ثبت کنید.
  2. XProv را در پس‌زمینه اجرا کنید و از ابزار Memory Viewer برای مشاهده جزئیات استفاده از حافظه استفاده کنید.

برنامه نمونه

برنامه 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>

(در یک دستگاه GPU، خطا به شکل زیر است: 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 ) بروید. در منوی کشویی Tools ، گزینه Memory Viewer را انتخاب کنید و در پنجره ابزار Memory Viewer، در منوی کشویی Memory Types ، گزینه HBM را انتخاب کنید (که معمولاً به طور پیش‌فرض انتخاب شده است).

صفحه نمایشگر حافظه XProf برای برنامه مثال بالا

مستندات ابزار XProf: Memory Viewer اجزای این ابزار و اطلاعات ارائه شده را شرح می‌دهد.

روی بخش HLO Ops at Peak Memory Allocation تمرکز کنید که سه نمودار بافر را در نقطه اوج استفاده از حافظه نشان می‌دهد. بافر شامل موارد زیر است: * ورودی‌ها و خروجی‌های برنامه: دسته‌های آموزشی، حالت‌های بهینه‌ساز و غیره. * TensorCore و SparseCore Temporaries: حافظه پویا مورد نیاز برای محاسبات میانی (مانند فعال‌سازی‌ها، گرادیان‌ها و غیره)

می‌توانید روی نمودارهای بافر حرکت کنید تا جزئیات بیشتری در مورد Op مانند اندازه، شکل، نوع تخصیص و موارد دیگر به دست آورید. این می‌تواند به شما در شناسایی و ارزیابی Opهایی که ممکن است موقت‌های بالا یا طولانی‌مدت داشته باشند، هرگونه تانسور ورودی/میانی/خروجی بزرگ که دارای padding ناکارآمد هستند و غیره که در اوج حافظه نقش دارند و نیاز به تنظیم یا بهینه‌سازی دارند، کمک کند.

تکنیک‌های خاص اشکال‌زدایی را در E1000 بیاموزید: اشکال‌زدایی .