‏TorchTPU: اجرای مستقیم PyTorch روی TPUهای گوگل

تیم گوگل تو این مطلب از TorchTPU رونمایی کرده؛ یه backend متن‌باز که از اواسط ۲۰۲۵ با همکاری تیم PyTorch تو متا شروع شده. TPUها (چیپ‌های اختصاصی گوگل برای محاسبات یادگیری عمیق) تا حالا بیشتر با TensorFlow و JAX شناخته می‌شدن. نسخه قبلی یعنی PyTorch/XLA مشکلات زیادی داشت: باید کد رو دستی تغییر می‌دادی، کامپایلش کند بود و از حالت eager (یعنی اجرای خط‌به‌خط و عادی PyTorch) پشتیبانی نمی‌کرد.

به گفته نویسنده‌ها، هدف اینه که منتقل کردن یه مدل به TPU به سادگی عوض کردن cuda به tpu باشه. البته خودشون می‌گن با این تغییر ساده به حداکثر سرعت نمی‌رسی، چون shapeها و ops مدل ممکنه برای GPU تنظیم شده باشن؛ ولی مدل بدون دردسر اجرا می‌شه:

device = torch.device("tpu")
dist.init_process_group(backend="tpu_dist")
model = model.to(device)

TorchTPU سه حالت eager داره. Strict (پیش‌فرض) هر op رو جدا و async می‌فرسته، مثل PyTorch روی GPU. Debug با TORCH_TPU_LAUNCH_BLOCKING=1 همه‌چی رو همگام اجرا می‌کنه تا خطایابی راحت‌تر بشه. حالت جدید DeferAndFuse چند op رو کنار هم جمع می‌کنه و یکجا کامپایل می‌کنه؛ طبق مقاله تا ۵۰ درصد یا بیشتر از Strict سریع‌تره و مصرف حافظه HBM رو هم پایین میاره.

برای اینکه کامپایل‌های تکراری وقت نگیرن، یه Compilation Cache سه‌لایه هست: لایه اول تو حافظه همون process، لایه دوم حافظه مشترک روی یه ماشین، و لایه سوم یه فضای ذخیره‌سازی راه دور مثل Google Cloud Storage که بین ماشین‌ها و بعد از ری‌استارت هم قابل استفاده‌ست. برای حالت کامپایل‌شده هم کافیه torch.compile(model, backend="tpu") بنویسی که روی Dynamo سواره.

دقت محاسبات هم فرق داره. فرمت بومی ضرب ماتریس روی TPU هست bfloat16 و TorchTPU به‌صورت پیش‌فرض از حالت medium استفاده می‌کنه، یعنی ضرب با BF16 و جمع‌زدن با FP32. حالت‌های high و highest دقت رو بالا می‌برن ولی حدود ۳ و ۶ برابر هزینه محاسبه دارن. یه محدودیت مهم دیگه اینه که هر TPU فقط به یه process دسترسی می‌ده، پس اجرای هم‌زمان چند بار کاری روی یه دستگاه (مثلاً برای RL) فعلاً ممکن نیست.

نکات کلیدی:

• پشتیبانی از حدود ۴۵۰ op از Aten و تست روی ۵۰۰۰ مدل پرطرفدار Hugging Face مثل Qwen، Llama، Whisper و BERT
• هیچ وابستگی به JAX تو runtime نداره و از DDP، FSDP و DTensor پشتیبانی می‌کنه
• حالت DeferAndFuse می‌تونه ۵۰ درصد یا بیشتر سریع‌تر از Strict eager باشه
• با vLLM، SGLang، TorchTitan، Ray و TorchRec تست شده
• kernelهای سفارشی رو می‌شه با Pallas، Helion یا کتابخونه Tokamax نوشت
• Bounded Dynamism (پشتیبانی از shapeهای متغیر بدون کامپایل دوباره) هنوز در حال توسعه‌ست