TorchTPU پشتیبانی بومی از TPU را به PyTorch و ابزارهای Hugging Face می‌آورد

Hugging Face می‌گوید TorchTPU، یک backend بومی برای PyTorch، پشتیبانی از شتاب‌دهنده‌های TPU را به اکوسیستم این شرکت می‌آورد. کاربران می‌توانند مدل‌های Transformers و Diffusers را روی TPU اجرا کنند و برای بارگذاری مدل در Transformers کافی است `device_map="tpu"` را مشخص کنند. در Diffusers نیز انتقال خط لوله به TPU با `pipe.to("tpu")` انجام می‌شود.

TorchTPU امکان اجرای کد معمول PyTorch را فراهم می‌کند و از Pallas kernels بهینه‌شده برای TPU، آموزش توزیع‌شده و `torch.compile` پشتیبانی می‌کند. مدل‌ها به‌طور پیش‌فرض در حالت eager اجرا می‌شوند؛ `torch.compile` نیز گزینه‌ای برای بهبود عملکرد است. Hugging Face می‌گوید با اضافه‌شدن این پشتیبانی بومی، Optimum TPU دیگر برای اجرای مدل‌های این شرکت با Transformers و Diffusers لازم نیست.

پیش از TorchTPU، اجرای PyTorch روی TPU به PyTorch/XLA وابسته بود. به‌گفتهٔ Hugging Face، اجرای تنبل tensorها، بازکامپایل گراف در پی بعضی تغییرات کوچک و دشواری اشکال‌زدایی از چالش‌های آن روش بود. TPUهای گوگل از طریق Google Cloud، از یک تراشه تا podهای چند‌هزار‌تراشه‌ای، در دسترس‌اند.