TorchTPU؛ اجرای بومی PyTorch روی تراشه‌های TPU گوگل

اگر تا امروز اجرای PyTorch روی TPU یعنی سر و کله زدن با torch_xla یا کوچ‌کردن به JAX بود، TorchTPU این معادله را عوض می‌کند: TPU حالا یک دستگاه دیگر در PyTorch است. این بک‌اند بومی، پشتیبانی از TPU را به قلب اکوسیستم هاگینگ فیس می‌آورد تا Transformers، Diffusers و Kernels بدون واسطه روی آن اجرا شوند.

TPU دقیقاً چیست؟

TPU مخفف Tensor Processing Unit است؛ شتاب‌دهنده‌های هوش مصنوعی ساخت گوگل. این تراشه‌ها برای یک کار ساخته شده‌اند: همان ضرب‌های ماتریسی که شبکه‌های عصبی بیشترین وقت خود را صرف آن می‌کنند. گوگل حدود ده سال است روی طراحی آن‌ها کار می‌کند و از آن‌ها برای آموزش و سرویس‌دهی مدل‌های خودش، از جمله Gemini، استفاده می‌کند؛ مسیری که نمونه‌ای از خروجی‌هایش را در نوشته معرفی Gemini 4 Argon دیده‌اید.

TPUها روی Google Cloud و از یک تراشه تا پادهایی با هزاران تراشه در دسترس‌اند. جایی که به توان محاسباتی زیاد برای آموزش یا سرویس‌دهی در مقیاس نیاز دارید، به کار می‌آیند. گلوگاه اما همیشه نرم‌افزار بوده، نه سخت‌افزار: پشته TPU حول TensorFlow و JAX رشد کرد و کاربران PyTorch راه دشوارتری برای ورود داشتند.

راه قدیمی: torch_xla و Optimum TPU

تا پیش از این، اجرای PyTorch روی TPU از مسیر PyTorch/XLA یا همان torch_xla می‌گذشت. کار می‌کرد، اما رفتارهای عجیب خودش را داشت: تنزورها lazy بودند، یعنی کد همیشه در زمانی که فکر می‌کردید اجرا نمی‌شد؛ تغییرهای کوچک می‌توانست بازکامپایل گراف را به راه بیندازد و دیباگ‌کردن هم سخت بود.

برای راحت‌ترشدن کار، Optimum TPU ساخته شد تا مدل‌های هاگینگ فیس روی TPU هم استنتاج و هم آموزش بگیرند. حالا با بومی‌شدن پشتیبانی TorchTPU، دیگر برای اجرای مدل‌های هاگینگ فیس با Transformers و Diffusers به Optimum TPU نیازی نیست. اگر پیگیری این جنس مهاجرت‌ها برایتان جالب است، تجربه مشابه اجرای مدل‌های هاگینگ فیس روی شتاب‌دهنده Trainium شرکت AWS را در نوشته TorchNeuron خوانده‌اید.

TorchTPU؛ TPU به‌عنوان یک بک‌اند تازه

TorchTPU یک بک‌اند تازه به PyTorch اضافه می‌کند: کد معمولی PyTorch می‌نویسید و TPU فقط یک دستگاه مقصد دیگر است. همراه آن، کرنل‌های Pallas تنظیم‌شده برای TPU، پشتیبانی از آموزش توزیع‌شده و یکپارچگی با torch.compile برای جاهایی که سرعت بیشتری می‌خواهید، در دسترس است. نصب روی یک TPU VM با یک فرمان انجام می‌شود:

pip install torch-tpu --extra-index-url https://download.pytorch.org/whl/cpu

ایندکس افزوده در این فرمان، نسخه CPU پکیج torch را می‌آورد. بدون آن، pip سراغ نسخه بسیار بزرگ‌تر CUDA می‌رود که روی TPU به آن نیازی ندارید.

کتابخانه‌های محبوب، حالا روی TPU

Transformers

Transformers v5 حول یک ایده ساخته شد: مرجع تعریف مدل‌ها باشد تا هر جا اجرا شوند؛ و «هر جا» حالا TPU را هم شامل می‌شود. برای بارگذاری مدل روی TPU کافی است در from_pretrained مقدار device_map="tpu" را بدهید. مدل به‌طور پیش‌فرض در حالت eager اجرا می‌شود و در صورت نیاز می‌توانید از torch.compile برای کارایی بهتر یا از یک چیدمان توزیع‌شده استفاده کنید. در مستندات، نمونه‌ای با مدل Qwen/Qwen3.5-9B آمده که متن ورودی را توکنایز می‌کند و با ۶۴ توکن جدید خروجی می‌گیرد.

Diffusers؛ تصویر و ویدیو روی TPU

TPU حالا در Diffusers یک دستگاه درجه‌یک است. برای نمونه‌سازی یا دیباگ، در حالت eager فقط pipe.to("tpu") لازم است و برای استقرار، حالت کامپایل‌شده یک فراخوانی compile اضافه می‌کند. پایپ‌لاین‌ها و کد شما ثابت می‌ماند و لازم نیست فورک مخصوص TPU را نگه‌داری کنید. نمونه‌اش با پایپ‌لاین StableDiffusion3Pipeline و مدل stabilityai/stable-diffusion-3.5-medium است که با dtype=bfloat16 روی TPU می‌رود و در حالت اختیاری کامپایل‌شده، این خط به آن اضافه می‌شود:

pipe.transformer.compile(backend="tpu", fullgraph=True, dynamic=False)

اولین فراخوانی کامپایل می‌کند و هر فراخوانی بعدی با همان اندازه تصویر، گراف را بازاستفاده می‌کند. اگر مدلی از یک تراشه بزرگ‌تر باشد، کافی است یک TensorParallelConfig به from_pretrained بدهید تا ترنسفورمر همان لحظه بارگذاری بین چند تراشه تقسیم شود و لازم نباشد کل مدل در یک جا جا بگیرد. بسیاری از پایپ‌لاین‌های پرکاربرد این قابلیت را از همان ابتدا دارند و بقیه هم فقط یک tp_plan فاصله دارند؛ راهنمای Diffusers برای TPU دستور کامل را دارد.

نمونه‌هایی از اجرای این پایپ‌لاین‌ها روی Cloud TPU v6e-4 با TP=4 (تقسیم هر مدل روی چهار تراشه)، دقت bf16 و batch برابر ۱:

Kernels

در Kernels Hub کرنل‌های آماده‌ای برای TPU وجود دارد. کافی است یکی را با کتابخانه Kernels بارگذاری کنید و به مدل‌تان سرعت بدهید. کرنل‌های TPU با Pallas نوشته می‌شوند که کنترل سطح پایین روی سخت‌افزار می‌دهد؛ می‌توانید کرنل‌ها را از کتابخانه Tokamax هم بیاورید. در مستندات، یک کرنل ragged-dot از Tokamax به یک مدل Mixture-of-Experts وصل شده است: کرنل tengomucho/ragged-dot-tpu با get_kernel بارگذاری و با نام tokamax_ragged_dot در ALL_EXPERTS_FUNCTIONS ثبت می‌شود، سپس مدل Qwen/Qwen3-30B-A3B-Instruct-2507 با experts_implementation="tokamax_ragged_dot" بارگذاری می‌شود.

هرچه batch بزرگ‌تر باشد، برد بیشتر است: در batch برابر ۶۴، رمزگشایی روی Qwen3-30B-A3B حدود ۲۱.۵ برابر و روی Gemma 4 26B-A4B حدود ۱۲.۸ برابر سریع‌تر از پیاده‌سازی پیش‌فرض شده است.

تست می‌کنیم تا شما درگیر نشوید

پشتیبانی TPU در Transformers به‌طور پیوسته با اجرای مجموعه آزمون‌های کتابخانه بهتر می‌شود. تیم هاگینگ فیس روی افزودن یک CI شبانه روی TPU کار می‌کند تا نتایج در داشبورد CI ترنسفورمرز، کنار دیگر شتاب‌دهنده‌ها، دیده شود.

همه‌چیز سبز نیست و قرار هم نیست باشد: از آزمون‌هایی که روی TPU اجرا می‌شوند ۹۶.۵ درصد قبول می‌شوند. هر شکستی به یک اسکریپت بازتولید کوچک و یک ایشو در مخزن TorchTPU تبدیل می‌شود؛ تا حالا ۳۴ ایشو گزارش شده و ۱۹ مورد آن رفع شده است. اجرا روی TPU چند باگ در خود Transformers را هم رو کرده، مثل StaticCache و نقشه‌های tensor-parallel مدل Qwen3، که بالادست اصلاح شده‌اند.

شروع کنید

TPU دیگر یک ماجراجویی حاشیه‌ای برای کاربران PyTorch نیست. TorchTPU را نصب کنید، مدلی از هاب بردارید و آن را به "tpu" بفرستید. برای قدم‌های بعدی، مخزن TorchTPU روی گیت‌هاب، کرنل‌های TPU در هاب و مستندات Cloud TPU سه مرجع اصلی‌اند.

پرسش‌های متداول

منبع: Hugging Face