مقیاس‌بندی و ارائه مدل‌های JAX روی پردازنده‌های گرافیکی NVIDIA

1. بهترین بررسی اولیه قبل از اشکال‌زدایی عملکرد JAX GPU چیست؟

   افزایش اندازه دسته

   jax.devices() را اجرا کنید و تأیید کنید که JAX یک GPU را می‌بیند.

   تبدیل تمام آرایه‌های JAX به NumPy

   غیرفعال کردن jax.jit

2. اگر jax.default_backend() cpu را چاپ کند و jax.devices() فقط CpuDevice نشان دهد، محتمل‌ترین علت چیست؟

   این مدل برای پردازنده گرافیکی (GPU) خیلی بزرگ است.

   یک jaxlib فقط برای CPU، عدم تطابق درایور یا کانتینری که بدون دسترسی به GPU شروع شده است

   اندازه دسته خیلی زیاد تنظیم شده است

   jax.jit هنوز فراخوانی نشده است

3. چرا اولین فراخوانی یک تابع jax.jit می‌تواند بسیار کندتر از فراخوانی‌های بعدی باشد؟

   JAX همیشه اولین فراخوانی را روی CPU اجرا می‌کند.

   اولین فراخوانی ممکن است شامل ردیابی و کامپایل باشد

   JAX در اولین فراخوانی، حافظه GPU را غیرفعال می‌کند.

   گرادیان‌ها فقط در اولین فراخوانی محاسبه می‌شوند

4. علت رایج کامپایل مجدد تصادفی در JAX چیست؟

   استفاده مجدد از همان شکل ورودی و نوع داده (dtype)

   نگهداری دسته‌ها روی پردازنده گرافیکی (GPU)

   تغییر شکل ورودی بین فراخوانی‌ها

   فراخوانی تابع ()block_until_ready قبل از زمان‌بندی

5. چرا درون یک تابع @jax.jit ، if x > 0: روی یک آرایه ردیابی شده باشد، خطایی رخ می‌دهد؟

   JAX از عملگر > پشتیبانی نمی‌کند.

   در طول ردیابی، x یک مقدار انتزاعی با شکل و نوع داده شناخته شده است اما مقدار مشخصی ندارد، بنابراین پایتون نمی‌تواند روی آن شاخه ایجاد کند.

   برای مقایسه باید به جای > از jnp.greater استفاده شود.

   jit تمام جریان کنترل پایتون را برای همیشه غیرفعال می‌کند.

6. چرا باید قبل از متوقف کردن تایمر در حین کار GPU، تابع block_until_ready() را فراخوانی کنید؟

   قبل از مرحله بعدی، حافظه پردازنده گرافیکی را آزاد می‌کند

   JAX کار GPU را به صورت ناهمزمان ارسال می‌کند، بنابراین بدون مسدود کردن، فقط می‌توانید مدت زمانی را که پایتون برای قرار دادن کار در صف صرف کرده است، اندازه‌گیری کنید.

   تابع را برای زمان‌بندی دقیق، دوباره کامپایل می‌کند.

   نتیجه را به CPU برمی‌گرداند

7. فراخوانی تابع float(loss) در هر مرحله از حلقه‌ی آموزشی می‌تواند آن را کند کند، زیرا:

   تابع زیان را در هر مرحله مجدداً کامپایل می‌کند

   انتقال دستگاه میزبان را اجباری می‌کند و می‌تواند دستگاه را در هر مرحله همگام‌سازی کند.

   نوع داده‌ی loss را به float64 تغییر می‌دهد.

   حافظه نهان کامپایل JAX را غیرفعال می‌کند.

8. چرا حلقه آموزش، داده‌ها را به دسته‌هایی با اندازه ثابت تغییر شکل می‌دهد و دسته نهایی جزئی را حذف می‌کند؟

   برای اینکه داده‌ها را به طور مؤثرتری تغییر دهید

   برای ثابت نگه داشتن شکل ورودی تا مرحله‌ی jitted برای دسته‌ی آخرِ کوچک‌تر، دوباره کامپایل نشود.

   زیرا JAX نمی‌تواند دسته‌های بزرگتر از ۵۱۲ را پردازش کند.

   برای کاهش تعداد پارامترهای مدل

9. در یک مرحله آموزش JAX، پارامترها چگونه به‌روزرسانی می‌شوند؟

   وزن‌ها در داخل مدل به صورت جهش‌یافته قرار می‌گیرند.

   یک تابع خالص پارامترها، حالت و یک دسته را دریافت می‌کند و پارامترها و حالت‌های جدید را برمی‌گرداند - بدون جهش درجا

   jax.jit وزن‌ها را به‌طور خودکار به‌روزرسانی می‌کند

   گرادیان‌ها مستقیماً توسط درایور GPU اعمال می‌شوند

10. تنظیم implementation="cudnn" در jax.nn.dot_product_attention چه چیزی را می‌طلبد؟

    ورودی‌های جایگزین CPU و float64

    ورودی‌های bfloat16 یا float16 و قابلیت محاسبه >= 8.0

    حداقل دو پردازنده گرافیکی

    غیرفعال کردن پوشش سببی

11. چرا توجه ساده‌لوحانه در توالی‌های طولانی از حافظه بیشتری نسبت به هسته cuDNN ادغام‌شده استفاده می‌کند؟

    وزن‌های مدل را دو بار ذخیره می‌کند

    این ماتریس توجه کامل (B, N, T, S) را که با مربع طول توالی رشد می‌کند، محقق می‌سازد.

    به طور پیش‌فرض با float64 اجرا می‌شود

    داده‌ها را روی CPU نگه می‌دارد

12. در آموزش موازی داده‌ها با JAX، چه چیزی در واقع موازی‌سازی را ایجاد می‌کند؟

    بازنویسی مرحله آموزش با حلقه‌های صریح روی GPUها

    نحوه قرارگیری آرایه‌ها - داده‌های خرد شده و پارامترهای تکثیر شده - در حالی که کد مرحله آموزش ثابت می‌ماند

    دو بار فراخوانی jax.jit

    افزایش نرخ یادگیری به ازای هر پردازنده گرافیکی

13. PartitionSpec('data', None) برای یک آرایه دوبعدی روی یک مش با محور 'data' به چه معناست؟

    کل آرایه را روی هر پردازنده گرافیکی (GPU) تکثیر کنید

    بُعد ۰ را در امتداد محور 'data' خرد کنید و بُعد ۱ را تکثیر کنید

    هر دو بُعد را در امتداد محور 'data' خرد کنید

    آرایه را به CPU منتقل کنید

14. چگونه مبدل درس ۷ توجه علّی (فقط رو به عقب) را دریافت می‌کند؟

    در هر مرحله، یک ماتریس ماسک مثلثی را به صورت دستی می‌سازد.

    این causal_sdpa (که dot_product_attention با is_causal=True فراخوانی می‌کند) را به عنوان attention_fn از nnx.MultiHeadAttention ارسال می‌کند.

    توکن‌ها را به ترتیب معکوس مرتب می‌کند

    از یک کتابخانه رمزگشای جداگانه استفاده می‌کند

15. هنگام ذخیره یک مدل NNX با Orbax StandardCheckpointer ، چه چیزی ذخیره می‌شود و چه چیزی ذخیره نمی‌شود؟

    کل کلاس پایتون، شامل متدها، pickled شده است.

    پارامترها ( nnx.state(model, nnx.Param) ) ذخیره می‌شوند؛ ساختار گراف ذخیره نمی‌شود — از کلاس پایتون در restore می‌آید.

    فقط وضعیت بهینه‌ساز ذخیره می‌شود

    تا پایان آموزش هیچ چیز ذخیره نمی‌شود

16. کامپایل AOT ( lower() → compile() ) چه مشکلی را برای سرویس‌دهی حل می‌کند؟

    تعداد پارامترهای مدل را کاهش می‌دهد

    کامپایل را از اجرا جدا می‌کند، بنابراین شما یک بار در هنگام راه‌اندازی کامپایل می‌کنید و درخواست‌ها را بدون تأخیر در شروع سرد (cold-start) ارائه می‌دهید.

    این مدل را به طور خودکار به float16 تبدیل می‌کند.

    این مدل را در چندین پردازنده گرافیکی تقسیم می‌کند.

17. کدام مسیر سرویس‌دهی، یک مصنوع قابل حمل تولید می‌کند که در هر زمان اجرای JAX و بدون کد مدل اصلی اجرا می‌شود؟

    jax.jit در حال پردازش

    jax.export — بایت‌های StableHLO سریالیزه شده

    فراخوانی مدل درون یک حلقه پایتون

    کامپایل AOT با lower().compile()

ارسال پاسخ‌ها