TorchTPU و اجرای بومی مدل‌های Hugging Face روی TPU

پیامد برای توسعه‌دهندگان

به گفته نویسندگان، پیش از این اجرای PyTorch روی TPU از مسیر torch_xla یا ابزار Optimum TPU انجام می‌شد. با TorchTPU، طبق متن دیگر نیازی به Optimum TPU برای اجرای مدل‌های Hugging Face نیست. برای تیم‌هایی که کدشان بر پایه PyTorch است، این یعنی امکان کاهش فورک‌ها و نیاز کمتر به یادگیری JAX. نصب هم در متن با یک دستور pip روی ماشین مجازی TPU توصیف شده و نسخه CPU بسته torch برای جلوگیری از دانلود نسخه بزرگ CUDA معرفی شده است.

عددها و ادعاهای گزارش‌شده

نویسندگان گزارش داده‌اند که استفاده از کرنل ragged-dot از کتابخانهٔ Tokamax در مدل‌های Mixture-of-Experts، با اندازهٔ دستهٔ ۶۴، توان عملیاتی مرحلهٔ تولید توکن را روی Qwen3-30B-A3B به ۲۱٫۵ برابر و روی Gemma 4 26B-A4B به ۱۲٫۸ برابر رسانده است. مبنای این مقایسه، کرنل Hub در برابر پیاده‌سازی پیش‌فرض روی همان TPU v6e-8 است. این اعداد، نتیجهٔ مقایسهٔ TPU با GPU نیستند.

این اندازه‌گیری‌ها را خود تیم گزارش کرده و منبع ارزیابی مستقلی از آن‌ها ارائه نمی‌کند. مثال‌های Diffusers جداگانه روی Cloud TPU v6e-4، با تقسیم مدل روی چهار تراشه، دقت bf16 و اندازهٔ دستهٔ یک اجرا شده‌اند. نتیجهٔ حالت کامپایل‌شده و eager به همان مدل، اندازهٔ تصویر و تنظیمات وابسته است؛ از آن نمی‌توان سرعت یکسانی برای همهٔ کاربردها نتیجه گرفت.

محدودیت‌ها و نکات باز

خود متن تصریح می‌کند که پشتیبانی TPU هنوز کامل نیست. طبق آمار زمان انتشار منبع، ۹۶٫۵ درصد از تست‌های اجراشده روی TPU پاس شده‌اند، ۳۴ مشکل در مخزن TorchTPU ثبت شده که ۱۹ مورد آن رفع شده است، و اجرای شبانه CI روی TPU هنوز در حال توسعه است. همچنین طبق توضیح متن، اجرای کامپایل‌شده برای اندازه تصویر جدید دوباره گراف می‌سازد، بنابراین برای کاربرانی که اندازه‌های متنوع دارند باید زمان کامپایل را در نظر بگیرند. پیش از استفاده در محیط تولیدی، بررسی سازگاری مدل و کرنل با نسخه‌های مورد نظر ضروری است.