معیارسنجی

مبانی بنچمارک

توکامکس زیرساخت یکپارچه‌ای برای سنجش عملکرد ارائه می‌دهد. شما می‌توانید تابع f(x) عملیاتی خود را برای مجموعه‌ای از ورودی‌های مشخص به صورت زیر سنجش کنید:

# Conforming and initializing data to prepare for benchmarking
f_std, args = tokamax.standardize_function(f, kwargs={'x': x})

# Execute and measure timing
bench: tokamax.BenchmarkData = tokamax.benchmark(f_std, args)

standardize_function توابع پیچیده را ساده می‌کند، برای مثال با آرگومان‌های غیر آرایه‌ای. ابتدا یک فرم استاندارد با یک آرگومان args ایجاد می‌کند که لیستی از آرایه‌های انتزاعی یا واقعی jax.Array | jax.ShapeDtypeStruct است. سپس به طور تصادفی تمام تانسورهای انتزاعی را مقداردهی اولیه می‌کند و یک f_std(args) استاندارد را فقط با آرگومان‌های آرایه واقعی برمی‌گرداند. این تابع می‌تواند به طور تمیز و بدون نگرانی در مورد آرگومان‌های استاتیک مانند رشته‌ها، jit شود.

لطفاً برای مشاهده‌ی مجموعه‌ی کامل گزینه‌های پشتیبانی‌شده توسط هر یک از این توابع، به مستندات مربوطه مراجعه کنید؛ برخی از مباحث کلیدی در زیر مورد بحث قرار گرفته‌اند.

مباحث پیشرفته بنچمارک

تکرارها را اجرا کنید

tokamax.benchmark به شما امکان می‌دهد تعداد تکرارها را انتخاب کنید؛ تکرارهای بیشتر معمولاً منجر به کاهش نویز اندازه‌گیری می‌شوند، مثلاً tokamax.benchmark(f_std, args, iterations=num_iters). با این حال، اگر تعداد تکرارها در یک دوره زمانی کوتاه خیلی زیاد باشد، ممکن است به خصوص برای هسته‌های محاسباتی سنگین، گلوگاه حرارتی ایجاد شود و بر زمان اجرا تأثیر بگذارد. متعادل کردن این عوامل اغلب یک تمرین تجربی است. یک رویکرد پیشنهادی این است که تعداد کمی تکرار در هر آزمایش با چندین آزمایش با فاصله زمانی انجام شود، که به قیمت افزایش زمان ساعت دیواری تمام می‌شود.

روش معیارسنجی

سربار JAX پایتون اغلب بسیار بزرگتر از زمان اجرای واقعی هسته شتاب‌دهنده است. این بدان معناست که رویکرد معمول زمان‌بندی jax.block_until_ready(f(x)) مفید نخواهد بود. benchmark به شما امکان می‌دهد روش زمان‌بندی اساسی مورد استفاده برای بنچمارک را انتخاب کنید، مثلاً benchmark(f_std, args, iterations=num_iters, method=method)

برای هسته‌های TPU، ما اکیداً method=xprof_hermetic را توصیه می‌کنیم که پروفایلر XProf را فراخوانی کرده و زمان اجرا را روی سخت‌افزار اندازه‌گیری می‌کند. این روش به دلیل پشتیبانی کامل از پشته سفارشی شامل سخت‌افزار و کامپایلر، تقریباً هیچ سربار ابزار دقیق تحمیل نمی‌کند.

برای هسته‌های GPU، می‌توانید xprof_hermetic نیز استفاده کنید؛ XProf به نوبه خود از APIهای CUPTI انویدیا استفاده می‌کند. همچنین می‌توانید مستقیماً یک تایمر CUPTI را با method=cupti فراخوانی کنید. هر دو روش مقداری سربار متغیر، معمولاً تا 5٪، اعمال می‌کنند.

توزیع داده‌ها

کارهای قبلی نشان داده‌اند که عملکرد می‌تواند به دلیل تعاملات پیچیده در سطح سخت‌افزار و توان و دما، بسته به توزیع داده‌ها به طور قابل توجهی متفاوت باشد. برای رفع این مشکل، standardize_function آرایه‌های ورودی را به روشی که نمایانگر کارهای آموزشی واقعی است، مقداردهی اولیه می‌کند، به عنوان مثال، هر jax.ShapeDtypeStruct با مقدار واقعی به صورت تصادفی مقداردهی اولیه می‌شود. شما می‌توانید این را برای نیازهای خود تطبیق دهید.

حالت بنچمارک

standardize_function به شما امکان می‌دهد mode=forward برای انتخاب فقط forward، forward_res را برای انتخاب forward و محاسبه باقیمانده‌ها، vjp را برای انتخاب فقط تابع VJP و forward_and_vjp را برای محاسبه یک مسیر کامل forward و VJP برای محک‌زنی انتخاب کنید. توجه داشته باشید که محک‌زنی فقط VJP می‌تواند منجر به OOM شود، زیرا مسیر forward خارج از تابع استاندارد شده برگشتی محاسبه می‌شود و تمام واسطه‌ها در HLO که در حافظه HBM باقی می‌ماند، ذخیره می‌شوند.