دسته بندی: زمان کامپایل: نوع داده RHS پشتیبانی نشده روی سخت افزار
این خطا زمانی رخ میدهد که نوع داده مورد استفاده برای عملوند سمت راست (RHS) در ضرب ماتریس (مثلاً jax.lax.dot_general ، jax.lax.conv ، jax.numpy.matmul یا عملگر @ ) به طور بومی توسط نسل خاص TPU مورد استفاده پشتیبانی نمیشود.
نمونه پیامهای خطا:
INTERNAL: Mosaic failed to compile TPU kernel: Unsupported matmul RHS type on target: 'vector<256x256xi8>'
...
The MLIR operation involved:
%13440 = "tpu.matmul"(%13435, %13437, %13439) <dimension_numbers = #tpu.dot_dimension_numbers<...>
پایههای XLA: TPU
نمای کلی
واحد ضرب ماتریس (MXU) در TPU به صورت بومی از عملیات Float32 در تمام نسلهای سختافزار پشتیبانی میکند.
با این حال، پشتیبانی بومی از BFloat16 و سایر انواع دادههای کوانتیزه شده (مثلاً Int4، Int8 یا Float8) بسته به نسل سختافزار متفاوت است. این خطا زمانی رخ میدهد که هسته شما تلاش میکند یک ضرب ماتریسی را با استفاده از نوع دادهای که نسل خاص TPU شما مدار فیزیکی لازم برای اجرای آن را ندارد، به MXU نگاشت کند.
این خطا معمولاً نشان میدهد که مسیر Canonicalization کامپایلر - که سعی میکند به طور خودکار انواع پشتیبانی نشده را به انواع پشتیبانی شده تبدیل کند (مثلاً از طریق شبیهسازی نرمافزاری) - قادر به یافتن یک قانون تبدیل معتبر نبوده یا به دلیل غیرفعال بودن حالت سازگاری ، از انجام این کار جلوگیری شده است.
اشکالزدایی
برای رفع این خطا، باید انواع دادههای خود را با قابلیتهای سختافزارتان هماهنگ کنید. گزینههای زیر را دارید:
۱. تبدیل به انواع بومی
قابل اعتمادترین راه حل این است که قبل از عملیات matmul، عملوندهای خود را به صورت دستی به یک نوع داده پشتیبانی شده توسط سختافزار (مانند Float32 یا BFloat16 در TPU نسخه ۴+) در هسته خود تبدیل کنید.
- دلیل:
Float32نوع دادهی جهانی است که به صورت بومی توسط MXU در تمام نسلهای TPU پشتیبانی میشود. - معامله: این با هزینه VPU (واحد پردازش برداری) همراه است - چرخههای مورد نیاز برای انجام تبدیل، اما تضمین میکند که هسته شما روی سختافزار فعلی اجرا خواهد شد.
۲. حالت سازگاری را بررسی کنید
معمولاً کامپایلر میتواند به طور خودکار این مشکلات عدم تطابق نوع را در حالت سازگاری که به طور پیشفرض فعال است، مدیریت کند. پیکربندیهای XLA را دوباره بررسی کنید تا مطمئن شوید که --xla_mosaic_compat_mode روی false تنظیم نشده باشد.
این به عنوان یک "polyfill" عمل میکند و توالیهای شبیهسازی نرمافزاری را برای عملیاتی که سختافزار شما به طور طبیعی از آنها پشتیبانی نمیکند، تزریق میکند.
حالت سازگاری چه چیزهایی را فعال میکند:
- MatMulهای با دقت مختلط: امکان ترکیب عملوندهای عدد صحیح با انبارههای عدد اعشاری را با وارد کردن خودکار عملیات تبدیل نوع (مثلاً بسط اعداد صحیح به
Float32قبل از matmul) فراهم میکند. - شبیهسازی با دقت پایین: در نسلهای خاصی از سختافزار، انواع پشتیبانینشده مانند ممیز شناور
4-bit(4E2M1FN) یا ممیز شناور8-bit(8E4M3FN) را با گسترش آنها به انواع پشتیبانیشده مانندBFloat16یاFloat32قبل از اجرا شبیهسازی میکند.
توجه داشته باشید که این حالت، سازگاری را بر عملکرد اوج اولویت میدهد، زیرا شبیهسازی به دستورالعملهای اضافی برای تبدیل فرمتهای داده نیاز دارد تا MXU بتواند روی آنها کار کند.
۳. سختافزار را ارتقا دهید یا درخواست پشتیبانی کنید
اگر الگوریتم شما اکیداً به عملکرد بومی برای انواعی مانند Int4 یا Float8 بدون سربار تبدیل یا شبیهسازی نیاز دارد، باید آن را روی نسل جدیدتر TPU با پشتیبانی بومی اجرا کنید.
درخواست ویژگی: اگر معتقدید سختافزار شما از این عملیات پشتیبانی میکند، یا اگر کامپایلر حتی در حالت سازگاری، مسیر شبیهسازی معتبری را از دست داده است، لطفاً یک درخواست ویژگی ثبت کنید. ما معمولاً تضمین میکنیم که عملیات با نسخههای قبلی سازگار هستند. بنابراین اگر هسته شما روی یک نسل TPU اجرا میشود، باید روی تمام نسلهای آینده نیز اجرا شود. اما تضمینی وجود ندارد که شبیهسازی برای نسلهای قدیمیتر وجود داشته باشد (برای برخی از آنها، تبدیلها بسیار گران خواهند بود).