کد خطا: E2001

دسته: زمان ترجمه: نوع داده سمت راست پشتیبانی‌نشده در سخت‌افزار

این خطا زمانی رخ می‌دهد که نوع داده استفاده‌شده برای عملوند سمت راست (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 نگاشت کند.

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

اشکال‌زدایی

برای رفع این خطا، باید انواع داده‌هایتان را با قابلیت‌های سخت‌افزارتان هماهنگ کنید. گزینه‌های زیر را دارید:

۱. فرستادن به انواع بومی

مطمئن‌ترین راه اصلاح این است که قبل‌از عملیات matmul، عملوندهای خود را به‌صورت دستی به نوع داده پشتیبانی‌شده سخت‌افزاری (مانند Float32 یا BFloat16 در TPU v4+) در هسته خود تبدیل کنید.

  • چرا: Float32 نوع داده جهانی است که MXU در همه نسل‌های TPU به‌طور بومی از آن پشتیبانی می‌کند.
  • مبادله: این کار با هزینه VPU (واحد پردازش برداری) همراه است - چرخه‌های موردنیاز برای انجام تبدیل، اما تضمین می‌کند که هسته شما روی سخت‌افزار فعلی اجرا خواهد شد.

۲. بررسی «حالت سازگاری»

معمولاً کامپایلر می‌تواند به‌طور خودکار این مشکلات عدم تطابق نوع را در حالت سازگاری که به‌طور پیش‌فرض فعال است مدیریت کند. پیکربندی‌های XLA را دوباره‌سنجی کنید تا مطمئن شوید --xla_mosaic_compat_mode روی «نادرست» تنظیم نشده باشد.

این ویژگی به‌عنوان «پلی‌فیل» عمل می‌کند و توالی‌های شبیه‌سازی نرم‌افزاری را برای عملیاتی که سخت‌افزار شما به‌طور بومی پشتیبانی نمی‌کند تزریق می‌کند.

«حالت سازگاری» چه چیزی را فعال می‌کند:

  • ضرب ماتریس مختلط: با درج خودکار عملیات تبدیل نوع (برای نمونه، تبدیل اعداد صحیح به Float32 قبل‌از ضرب ماتریس)، امکان ترکیب عملوندهای عدد صحیح با انباشتگرهای شناور را فراهم می‌کند.
  • شبیه‌سازی با دقت پایین: در برخی‌از نسل‌های سخت‌افزار، انواع پشتیبانی‌نشده مانند نقطه شناور 4-bit (4E2M1FN) یا نقطه شناور 8-bit (8E4M3FN) را با گسترش دادن آن‌ها به انواع پشتیبانی‌شده مانند BFloat16 یا Float32 قبل‌از اجرا شبیه‌سازی می‌کند.

توجه داشته باشید که این حالت سازگاری را بر عملکرد اوج اولویت می‌دهد زیرا شبیه‌سازی به دستورالعمل‌های اضافی برای تبدیل قالب‌های داده قبل‌از اینکه MXU بتواند روی آن‌ها عمل کند نیاز دارد.

۳. سخت‌افزار را ارتقا دهید یا درخواست پشتیبانی کنید

اگر الگوریتم شما برای انواع Int4 یا Float8 بدون سربار پخش یا شبیه‌سازی، به‌طور دقیق به عملکرد بومی نیاز دارد، باید در نسل جدیدتر TPU با پشتیبانی بومی اجرا کنید.

درخواست ویژگی: اگر فکر می‌کنید سخت‌افزارتان از این عملیات پشتیبانی می‌کند، یا اگر حتی در «حالت سازگاری» نیز مسیر شبیه‌سازی معتبری در کامپایلر وجود ندارد، لطفاً درخواست ویژگی ثبت کنید. ما معمولاً تضمین می‌کنیم که عملیات‌ها با نسخه‌های بعدی سازگار هستند. بنابراین اگر هسته شما روی نسل TPU اجرا شود، باید روی همه نسل‌های آینده نیز اجرا شود، اما تضمینی وجود ندارد که برای نسل‌های قدیمی‌تر شبیه‌سازی داشته باشد (برای برخی‌از آن‌ها تبدیل‌ها بسیار پرهزینه خواهد بود).