گوگل و Meta پشتیبانی بومی PyTorch روی TPU را متن‌باز کردند
حدود ۴ دقیقه مطالعه
گوگل و Meta روز ۲۰ اکتبر ۲۰۲۶ در PyTorch Conference پروژه‌ی TorchTPU را معرفی کردند؛ یک backend بومی و متن‌باز برای اجرای PyTorch روی TPU های Google Cloud. طبق بلاگ رسمی PyTorch، این backend یک لایه‌ی روی JAX نیست و کد موجود با تغییرهای کوچک، مثلاً عوض کردن cuda با tpu در انتقال مدل به دستگاه، روی TPU اجرا می‌شود.
TorchTPU چطور کار می‌کند؟
طبق بلاگ، TorchTPU به‌عنوان دستگاه PrivateUse1 در dispatcher پایتورچ ثبت می‌شود و runtime وابستگی به JAX ندارد. سه حالت eager وجود دارد: Strict (پیش‌فرض و ناهمزمان با CPU)، Debug (همزمان و برای اشکال‌زدایی با بررسی مرزها) و DeferAndFuse که عملیات را عقب می‌اندازد و در تکه‌های بزرگ‌تر ترکیب می‌کند؛ بلاگ ادعا می‌کند این حالت ۵۰ درصد یا بیشتر از Strict سریع‌تر است و HBM کمتری مصرف می‌کند. عملیات به StableHLO تبدیل می‌شوند، یک backend برای torch.compile با نام tpu ثبت شده، و collective ها در سطح c10d پیاده‌سازی شده‌اند تا DDP، FSDP و DTensor روی آن کار کنند.
برای امتحان، MarkTechPost می‌نویسد Python 3.11 یا بالاتر روی یک TPU VM لازم است و باید ایندکس wheel مخصوص CPU برای PyTorch را بدهی، وگرنه pip نسخه‌ی CUDA پکیج torch را نصب می‌کند. بلاگ رسمی نسخه‌ی Python یا PyTorch را مشخص نکرده است. دستور نصب و استفاده این است:
pip install torch_tpu
model = model.to("tpu")
compiled = torch.compile(model, backend="tpu")
مخزن پروژه در این نشانی است و مجوز Apache 2.0 را فقط MarkTechPost نوشته؛ بلاگ رسمی مجوز را ذکر نکرده:
github.com/google-pytorch/torch_tpu
عددهای ادعایی
همه‌ی عددهای زیر ادعای نویسنده‌های پروژه (گوگل و Meta) هستند و ارزیابی مستقل نیستند. طبق بلاگ، ۴٬۹۱۷ مدل Hugging Face به‌صورت end-to-end تست شده‌اند و ۹۸ درصد مدل‌های محبوب بدون crash اجرا شده‌اند. در ۶۵۷ نماینده‌ی معماری، ۹۷ درصد با نتیجه‌ی عددی GPU می‌خواند. برای آموزش مدل Qwen3 235B MoE با TorchTitan روی ۶۴ چیپ TPU v7، بلاگ ۱٬۸۲۸ TPS برای PyTorch در برابر ۱٬۹۳۷ TPS برای JAX/MaxText گزارش می‌کند و weak scaling تا ۸٬۱۹۲ چیپ با MFU حدود ۲۷ تا ۲۸ درصد را ادعا می‌کند. MarkTechPost عدد دیگری هم می‌آورد: ۵۲ درصد MFU برای Qwen 3 8B روی Ironwood. در inference، گوگل می‌گوید در benchmark ای از SemiAnalysis به نام InferenceX کمترین هزینه به‌ازای هر میلیون توکن را داشته است.
محدودیت‌ها
کامپایلر XLA به shape های ثابت نیاز دارد و dynamic shape هنوز کامل نشده؛ زیر torch.compile فعلاً مقدار dynamic برابر False اجباری است. دقت پیش‌فرض matmul از نوع bfloat16 با جمع fp32 است و حالت‌های دقیق‌تر ۳ تا ۶ برابر محاسبه‌ی MXU می‌خواهند. هر TPU را فقط یک process می‌تواند استفاده کند، پس colocate کردن آموزش و inference در RL بدون کار اضافه ممکن نیست. FlexAttention هنوز پشتیبانی نمی‌شود و Helion آزمایشی است. در کد MPMD ممکن است collective ها deadlock شوند و حالت پیش‌فرض ایمن آن طبق بلاگ ۳۰ تا ۴۰ درصد از عملکرد کل را می‌گیرد. MarkTechPost هم می‌نویسد eager بدون تنظیم خیلی کندتر از اجرای کامپایل‌شده است (۱۱٫۴۹ ثانیه در برابر ۲٫۷۸ ثانیه برای هر گام در یک مثال).
بلاگ رسمی می‌گوید چهار مخزن TorchTPU، vLLM-TorchTPU، SGLang-TorchTPU و TorchTitan در همین کنفرانس open-source می‌شوند، در حالی که MarkTechPost vLLM و SGLang را همین الان پشتیبانی‌شده می‌نویسد؛ قبل از استفاده وضعیت هر مخزن را چک کن. TechCrunch هم می‌نویسد آمازون یک روز قبل، ۱۹ اکتبر، TorchNeuron را برای Trainium معرفی کرده است.
منابع:
