تصور کنید می‌خواهید در یک مدل یادگیری عمیق، داده‌ها را به‌صورت فشرده ذخیره کنید اما هم‌زمان نیاز دارید که تمام محاسبات ریاضی و مشتقات مدل به‌طور دقیق و بدون خطا باقی بماند. تا پیش از این، توسعه‌دهندگان JAX برای رسیدن به این هدف مجبور بودند بین ساختارهای سخت‌گیرانه یا مدل‌های منعطف اما مبهمی به نام pytree یکی را انتخاب کنند. پیاده‌سازی یک آرایه کوانتیده معمولاً کاربر را مجبور می‌کرد تا بین یک ساختار صلب و یک pytree منعطف اما نامشخص، یکی را برگزیند.
طبق مستندات فنی این پروژه، کتابخانه آزمایشی hijax این تضاد را برطرف کرده است. این ابزار به برنامه‌نویسان اجازه می‌دهد «انواع سطح بالا» (hi types) تعریف کنند که به‌عنوان مقادیر واحد و شناسنامه‌دار در نمایش‌های داخلی JAX (jaxprs) وجود دارند.
تا ۱۲ ژوئیه ۲۰۲۶، اکثر توسعه‌دهندگان برای دسته‌بندی آرایه‌ها از پایتری (pytree) — که شبیه به یک جعبه ابزار است که وسایل مختلف را کنار هم می‌گذارد اما هویت هر وسیله را به تنهایی نمی‌بیند — استفاده می‌کردند. مشکل این بود که پایتری‌ها شفاف هستند؛ یعنی JAX فقط مجموعه‌ای از برگ‌های آرایه‌ای مستقل را می‌بیند. بنابراین، اعمال قواعد ثابت (invariants) دشوار بود؛ مثلاً نمی‌شد به‌راحتی تضمین کرد که یک ضریب مقیاس (scale factor) با ابعاد یک مقدار کوانتیده مطابقت دارد یا خیر. همچنین تعریف رفتارهای سفارشی برای مشتق‌گیری (differentiation) و دسته‌بندی (batching) در این ساختار ممکن نبود. در حالی که پایتری صرفاً یک ظرف است، یک hi type هویت مستقل خود را در jaxprs دارد، مفهوم خاص خود را از دسته‌بندی تحت vmap دارد و می‌تواند اطلاعات مربوط به توزیع داده (sharding) را در نوع داده خود برای حالت توزیع صریح حمل کند.
نکته قابل توجه این است که hi typeها در jaxprs به‌عنوان یک مقدار واحد از یک نوع واحد ظاهر می‌شوند، نه مجموعه‌ای از برگ‌ها. آن‌ها اجازه می‌دهند قواعد داخلی حفظ شوند تا کاربران تنها از طریق عملیات‌های تعریف‌شده و ثابت، آن‌ها را تولید یا مصرف کنند. حیاتی‌ترین نکته این است که «نوع مماس» (tangent type) آن‌ها می‌تواند با ساختار اصلی (primal) متفاوت باشد؛ یعنی مشتقات لزوماً نباید «همان پایتری، اما برای مماس‌ها» باشند.
همان‌طور که در بحث‌های گذشته‌ی ما درباره‌ی بهینه‌سازی حافظه در مدل‌های زبانی اشاره کردیم، مدیریت دقیق نحوه نمایش داده‌ها در حافظه، کلید افزایش سرعت استنتاج است. در hijax، این انواع داده در jaxprs به‌عنوان یک مقدار واحد ظاهر می‌شوند و کاربران تنها از طریق عملیات‌های تعریف‌شده با آن‌ها تعامل دارند. این موضوع مانع از پراکندگی محاسبات در لایه‌های پایین‌تر می‌شود.
سازوکار انواع سطح بالا
در hijax، توسعه‌دهندگان با ارث‌بری از کلاس HiType یک هویت جدید تعریف می‌کنند. از آن‌جهت که این نوع باید قابل هش (hashable) و قابل مقایسه برای برابری باشد، توسعه‌دهندگان معمولاً از frozen dataclasses برای این منظور استفاده می‌کنند. این فرآیند در سه مرحله اصلی رخ می‌دهد:
- Lowering (پایین‌افکنی): شما متد
lo_tyرا پیاده‌سازی می‌کنید تا مشخص کنید کدام آرایه‌های استاندارد JAX (lojax) این نوع را تشکیل می‌دهند. این کار به JAX اجازه می‌دهد انواع پایین‌افکنی‌شده را بدون نیاز به داشتن یک مقدار فیزیکی در دست، محاسبه کند. همچنین متدهایlower_valوraise_valرا برای تبدیل بین شیء سطح بالا و لیست اجزای آرایه‌ای آن پیاده می‌کنید. برای مثال، پیاده‌سازیQArrayTyازlo_tyاستفاده می‌کند تا لیستی شامل یکShapedArrayبرای مقادیرint8و یکShapedArrayدیگر برای مقیاسfloat32برگرداند. - Registration (ثبت): شما از
register_hitypeاستفاده می‌کنید تا یک کلاس مقدار پایتونی (مانند یک dataclass) را با تعریف hi type مرتبط کنید. این فراخوانی شامل یک lambda است که نوع داده (مثلاً خواندن شکل و توزیع) را از هر مقدار داده شده محاسبه می‌کند، مشابه روش کارjax.typeofبرای آرایه‌ها. پس از ثبت،jax.typeofمستقیماً روی اشیاء سفارشی کار می‌کند و تبدیل‌های JAX آن‌ها را در هر جایی که یک مقدار انتظار می‌رود، می‌پذیرند. - Primitives (پریمیتیوها): شما زیرکلاس‌های
VJPHiPrimitiveرا می‌نویسید. این‌ها تنها روش‌های مجاز و sanctioned برای تولید یا مصرف نوع جدید هستند. با تعریفin_avalsوout_avalخاص، این پریمیتیوها تضمین می‌کنند که قواعد داخلی هرگز توسط کدهای خارجی نقض نشوند. این پریمیتیوها متدexpandرا فراهم می‌کنند، جایی که نوع سطح بالا در نهایت برای اجرا به آرایه‌های lojax سازنده خود تجزیه می‌شود.
حل چالش کوانتش
در یک آرایه‌ی کوانتیده (Quantization) — که شامل یک مقدار qvalue (از نوع int8) و یک مقیاس scale (از نوع f32) است — کوانتش در امتداد آخرین محور اتفاق می‌افتد و برای هر ردیف یک مقیاس وجود دارد. در یک پایتری استاندارد، نوع مماس (که برای گرادیان‌ها استفاده می‌شود) مجبور است با انواع برگ‌ها مطابقت داشته باشد. از آن‌جا که آرایه‌های int8 مماس‌های بدیهی float0 دارند، یک آرایه کوانتیده مبتنی بر پایتری نمی‌تواند به‌طور مؤثر گرادیان‌ها را منتقل کند.
اما با hijax، شما می‌توانید to_tangent_aval را تعریف کنید تا مشخص شود مماسِ یک آرایه کوانتیده، در واقع یک آرایه پیوسته float32 است. این کار ساختار اصلی را از ساختار مماس جدا می‌کند (decouple). این جداسازی اجازه می‌دهد گرادیان‌ها از طریق یک «تخمین‌گر مستقیم» (Straight-Through Estimator یا STE) جریان یابند. با پیاده‌سازی قواعد خاص vjp_fwd و vjp_bwd_retval روی پریمیتیوهای Quantize و Dequantize، سیستم گام گسسته کوانتش را در مسیر بازگشتی (backward pass) به‌صورت یک تابع همانی (identity function) در نظر می‌گیرد.
برای مثال، پریمیتیو Quantize تضمین می‌کند که x_aval.dtype در هنگام مقداردهی اولیه float32 باشد. متد expand آن منطق واقعی کوانتش را پیاده می‌کند: محاسبه مقیاس به صورت jnp.max(jnp.abs(x), axis=-1) / 127.0 و سپس گرد کردن ورودی مقیاس‌بندی‌شده به int8. پریمیتیو متناظر Dequantize صرفاً عملیات معکوس را انجام می‌دهد: qx.qvalue.astype('float32') * qx.scale[..., None].
عملکرد و اجرا
انواع سفارشی امکان اجرای عملیاتی را می‌دهند که بسیار سریع‌تر از تبدیل دائمی داده‌های کوانتیده به حالت عادی (dequantizing) برای هر مرحله است. نمونه بارز آن پریمیتیو MatmulQ است که ضرب یک آرایه متراکم f32 (به نام x) در یک ماتریس وزنی کوانتیده (qw) را مدیریت می‌کند.
جزئیات پیاده‌سازی MatmulQ:
- ادغام مقیاس‌ها (Folding Scales): به‌جای تبدیل کامل ماتریس وزنی به حالت متراکم، متد
expandمقیاس‌های هر ردیف را مستقیماً در عملوند متراکم ادغام می‌کند:(x * qw.scale) @ qw.qvalue.astype(jnp.float32). این روش از این حقیقت بهره می‌برد که مقیاس‌ها در امتداد محور انقباض به‌صورت ردیفی اعمال می‌شوند. - بررسی نوع (Type Checking): این پریمیتیو در زمان ساخت بررسی می‌کند که محورهای انقباض توزیع‌نشده (unsharded) باشند و در صورت عدم تطابق ابعاد یا مشخصات توزیع، خطای
TypeErrorصادر می‌کند. به‌طور خاص بررسی می‌کند کهx_aval.shape[1] == q_aval.shape[0]باشد و اطمینان حاصل می‌کند که هیچ‌کدام از محورهای انقباض sharded نباشند. - بهره‌وری: انجام ضرب ماتریسی سنگین مستقیماً روی محموله
int8از سربار حافظه ناشی از ایجاد یک کپی کامل و متراکم (dequantized copy) جلوگیری می‌کند. خروجی یکShapedArrayبا ابعاد حاصل و نوع دادهfloat32است.
به دلیل شناسایی این انواع در jaxprs، کد ردیابی‌شده (traced code) پاکیزه می‌ماند. یک عملیات کوانتیده به‌عنوان یک فراخوانی واحد به یک پریمیتیو سطح بالا (مثلاً call_hi_primitive[_prim=Quantize[]{}]) ظاهر می‌شود، نه مجموعه‌ای از دست‌کاری‌های پراکنده آرایه‌ای. این موضوع از خطاهای «نشت ردیاب» (leaked tracer) جلوگیری می‌کند. برای مثال، اگر کاربر سعی کند درون یک تابع jit به qx.qvalue دسترسی پیدا کند، JAX خطای AttributeError می‌دهد زیرا شیء در آن لحظه در واقع یک DynamicJaxprTracer است (به‌طور دقیق: DynamicJaxprTracer has no attribute qvalue).
دسترسی مستقیم به ویژگی‌ها (attributes) تنها در متد expand مجاز است، جایی که JAX قبلاً تصمیم به استفاده از اجزای lojax گرفته است. تلاش برای فراخوانی یک سازنده (constructor) روی آرایه‌های ردیابی‌شده به‌جای استفاده از یک پریمیتیو، منجر به خطای TypeError: No constant handler for type: <class 'jax._src.interpreters.partial_eval.DynamicJaxprTracer'> می‌شود.
تبدیل‌های پیشرفته: vmap و scan
مدیریت دسته‌ها (batches) در hijax نیازمند یک MappingSpec است. از آن‌جا که JAX نمی‌تواند حدس بزند یک نوع سفارشی چگونه باید نگاشت شود، توسعه‌دهنده باید تعریف کند که رتبه‌ها چگونه از طریق متدهای inc_rank و dec_rank افزایش یا کاهش یابند. این متدها اندازه محور و یک spec را می‌گیرند تا به‌ترتیب نوع عنصر و نوع دسته‌بندی‌شده را برگردانند.
برای QArrayTy در حالت vmap (بردارسازی)، یک QArraySpec یک زیرکلاس ساده از MappingSpec است که داده داخلی ندارد. یک دسته از آرایه‌های q8[2,3] که روی یک محور پیشرو جدید پشته شده‌اند، تبدیل به q8[n,2,3] می‌شوند. متد dec_rank محور نگاشت‌شده را حذف کرده و مشخصات توزیع (sharding spec) را با برش دادن پارتیشن توزیع به‌روز می‌کند.
در سمت پریمیتیو، قواعد دسته‌بندی (batch rules) باید برای مدیریت ترکیباتی از آرگومان‌های دسته‌بندی‌شده و نشده پیاده‌سازی شوند. برای مثال، quantize_batch بررسی می‌کند که آیا بعد ورودی d برابر با None است یا خیر. اگر نباشد، محور را به موقعیت ۰ منتقل کرده، کوانتش را اعمال می‌کند و نتیجه را با یک QArraySpec() برمی‌گرداند.
در مورد jax.lax.scan (حلقه‌ها)، سیستم بر leading_axis_spec تکیه دارد. این اجازه می‌دهد حلقه در هر گام یک برش (slice) از نوع سطح بالا را مصرف کند. چون scan همیشه روی محور پیشرو حرکت می‌کند، اگر تمام مقادیر اسکن‌شده از نوع hi type باشند، توسعه‌دهنده باید length را به‌طور صریح ارائه دهد، زیرا هیچ آرگومان آرایه‌ای برای استنتاج اندازه توسط JAX وجود ندارد. بنابراین یک عملیات sum_dequantized می‌تواند روی پشته‌ای از آرایه‌های q8 پیمایش کرده و مجموع مقادیر تبدیل‌شده آن‌ها را جمع کند.
توزیع داده و محاسبات پراکنده
در حالت توزیع صریح JAX، اطلاعات توزیع (sharding) بخشی از نوع داده هستند. hijax اجازه می‌دهد داده‌های توزیع، مانند یک فیلد NamedSharding را مستقیماً روی HiType ثبت کنید.
- انتشار (Propagation): پریمیتیوها توزیع را از ورودی به خروجی از طریق قواعد تایپینگ (typing rules) منتقل می‌کنند. برای
QArrayTyمقدارqvalueتوزیع نوع را حمل می‌کند، در حالی که مقیاس (scale) با حذف آخرین محور از spec با استفاده ازself.sharding.update(spec=jax.P(*self.sharding.spec[:-1]))به‌دست می‌آید. - تعامل با Lojax: متد
lo_tyآرایه‌های سازنده را با توزیع‌های مربوطه مهر می‌زند تا JAX بتواند پارتیشن‌بندی را در سراسر مش (mesh) ردیابی کند. هنگام چاپ، این‌ها به‌عنوان نشانگرهای@ظاهر می‌شوند (مثلاًq8[8@i,3]). - shard_map: برای عبور از مرز دیدهای جهانی (global views) به دیدهای هر-دستگاه (per-device views)، hijax از
HiPspecاستفاده می‌کند. این یک معادلMappingSpecبرای مشخصات پارتیشن است. متدto_loآن، یک spec سطح بالا را به مشخصات پارتیشنjax.Pمجزا برای هر جزء پایین‌افکنی‌شده تبدیل می‌کند. برایQArrayPمتدto_loمقدار(self.spec, self.spec)را برمی‌گرداند زیراqvalueوscaleباید با هم توزیع شوند. - تقارن (Symmetry): چون کوانتش ردیفی صرف‌نظر از رتبه اعمال می‌شود، کوانتش تکه‌به‌تکه (shard-by-shard) در یک
shard_mapدقیقاً با کوانتش جهانی مطابقت دارد. این موضوع با بررسی اینکهqrows.qvalueوqrows.scaleدقیقاً با کوانتش جهانی یکی هستند، تایید می‌شود.
تطبیق‌پذیری: از آرایه‌های رتبه-۱ تا تاپل‌ها
این چارچوب فراتر از کوانتش کاربرد دارد. یک مورد استفاده، آرایه‌های رتبه-۱ (Rank-1 arrays) است که در آن یک ماتریس $m \times n$ به‌صورت حاصل‌ضرب خارجی دو بردار: col (با ابعاد f32[m]) و row (با ابعاد f32[n]) نمایش داده می‌شود.
ویژگی‌های آرایه رتبه-۱:
- جلوگیری از تجسم (Avoiding Materialization): پریمیتیو
MatmulR1ضرب را به‌صورتjnp.outer(x @ r1.col, r1.row)انجام می‌دهد تا تضمین شود که حاصل‌ضرب متراکم کامل $m \times n$ هرگز در حافظه ساخته (materialize) نشود. - فضای مماس (Tangent Space): برخلاف آرایه‌های کوانتیده، آرایه‌های رتبه-۱ تحت عمل جمع بسته نیستند (جمع دو ماتریس رتبه-۱ عموماً رتبه-۲ است). بنابراین، نوع مماس برای
Rank1Tyبه‌صورت یکShapedArrayمتراکم از f32 تعریف شده است. اختلالات از روی منیفولد خارج می‌شوند، به این معنی که مماس‌ها و کوتانس‌ها (cotangents) باید متراکم باشند. - دسترسی پریمیتیو: برای حفظ مرز انتزاع، از پریمیتیو
Factorsبرای بازیابی اجزایcolوrowدر طول قواعد ردیابی‌شده استفاده می‌شود، نه دسترسی مستقیم به ویژگی‌ها. پریمیتیوOuterبه‌عنوان سازنده عمل کرده و دو بردار را گرفته و یک شیءRank1برمی‌گرداند.
علاوه بر این، TupTy اجازه ایجاد کانتینرهای عمومی را می‌دهد. برخلاف ساختار ثابت یک آرایه کوانتیده، یک hi-tuple می‌تواند هر ترکیبی از انواع آرایه و سایر انواع سطح بالا، از جمله تاپل‌های تو در تو، را در خود جای دهد.
پیاده‌سازی تاپل و نگاشت:
- تفویض (Delegation):
TupTyمتدهایlo_ty،lower_valوto_tangent_avalرا به انواع سازنده‌اش تفویض می‌کند. این کلاس ازitertools.isliceدرraise_valاستفاده می‌کند تا لیست تخت‌شده‌ی آرایه‌های lojax را دوباره به انواع اجزای سطح بالای مربوطه تقسیم کند. - اندیس‌گذاری پویا: پریمیتیوهایی مانند
GetTupEltاز یک پارامتر استاتیکself.idxکه از طریقparamsارسال شده، برای دسترسی به عناصر خاص استفاده می‌کنند. قاعدهbatchبرایGetTupEltباید هر دو حالت ورودی‌های دسته‌بندی‌شده و نشده را مدیریت کند، زیرا ورودی‌هایin_dimsمی‌توانندNoneباشند. - دسته‌بندی منعطف:
TupSpecبرای هر جزء یک ورودی محور حمل می‌کند. این اجازه می‌دهد هر عنصر یک تاپل در یک عملیاتvmapروی محورهای مختلف نگاشت شود (یا اصلاً نشود)، همان‌طور که در تابعswapمشاهده می‌شود که در آن عنصر اول در ورودی و عنصر دوم در خروجی نگاشت شده است. این کار اجازه ایجاد امضاهای پیچیده‌ای مانندTup{float32[3],float32[]} -> Tup{float32[],float32[3]}را می‌دهد.
این انعطاف‌پذیری ساختاری به این معنی است که توسعه‌دهندگان اکنون می‌توانند انواع داده‌های متناسب با دامنه تخصصی خود را طراحی کنند که به‌طور کامل در اکوسیستم تبدیل‌های JAX — شامل jit ،grad ،vmap ،scan و shard_map — مشارکت کنند، بدون اینکه کارایی لایه‌ی پایین‌افکنی XLA را فدا کنند.
گفتگو