گوگل TorchTPU را متن‌باز کرده است؛ لایه‌ای نرم‌افزاری که به توسعه‌دهندگان امکان می‌دهد PyTorch ــ ابزار پرکاربرد ساخت و اجرای مدل‌های هوش مصنوعی ــ را روی واحدهای پردازش تنسور گوگل یا TPU اجرا کنند. این ابزار با همکاری تیم PyTorch متا توسعه یافته و هدفش انتقال بارهای کاری موجود به سخت‌افزار TPU با تغییرات نسبتاً اندک در کد است، نه بازنویسی آن‌ها.
گوگل می‌گوید TorchTPU هم از اجرای فوری پشتیبانی می‌کند؛ روشی که عملیات را هنگام فراخوانی اجرا می‌کند و برای اشکال‌زدایی مفید است، و هم از اجرای کامپایل‌شده برای بارهای کاری عملیاتی. توسعه‌دهندگان می‌توانند تنظیم دستگاه را از «cuda» به «tpu» تغییر دهند و از بک‌اند توزیع‌شده TPU استفاده کنند؛ گوگل همچنین API تشخیص شتاب‌دهنده در PyTorch را گزینه‌ای جایگزین می‌داند. TorchTPU در GitHub در دسترس است، اما نقشه راه آن هنوز شامل کار روی شکل‌های ورودی پویا، سرعت کامپایل، پوشش عملگرها، یادگیری تقویتی و کوانتیزه‌سازی است.
برای توسعه‌دهندگان چه تغییری می‌کند؟
این عرضه، TPUها را به نرم‌افزارهایی پیوند می‌دهد که بسیاری از توسعه‌دهندگان PyTorch از قبل استفاده می‌کنند. گوگل می‌گوید TorchTPU با ابزارهای سرویس‌دهی vLLM و SGLang و همچنین پروژه‌های آموزش و ارکستراسیون، از جمله TorchTitan، Ray و TorchRec، کار می‌کند. در Transformers v5، Hugging Face می‌گوید توسعه‌دهندگان می‌توانند با تنظیم device_map="tpu" یک TPU انتخاب کنند. کتابخانه Diffusers این شرکت نیز از اجرای TPU، از جمله حالت کامپایل‌شده و تقسیم کار میان چند شتاب‌دهنده، پشتیبانی می‌کند.
این رویکرد با PyTorch/XLA، مسیر قبلی گوگل برای اجرای PyTorch روی TPUها، تفاوت دارد. گوگل می‌گوید آن سامانه به تغییرات گسترده‌تری در کد نیاز داشت، زمان کامپایلش طولانی بود و از حالت اجرای فوری PyTorch پشتیبانی نمی‌کرد. گوگل یک خط لوله آزمایشی ساخته است که به‌گزارش InfoQ، حدود 5,000 مدل Hugging Face را در بیش از 500 معماری ارزیابی می‌کند. این موضوع از آزمایش گسترده حکایت دارد، اما به این معنا نیست که اجرای همه مدل‌ها یا بارهای کاری بدون هیچ تغییری تضمین شده است.
در مقایسه با گزینه‌های سخت‌افزاری دیگر چه وضعی دارد؟
گوگل گزارش می‌دهد حالت DeferAndFuse در TorchTPU ــ ترکیبی از اجرای فوری و کامپایل‌شده ــ می‌تواند دست‌کم 50% سریع‌تر از حالت Strict Eager باشد و هم‌زمان اوج مصرف حافظه شتاب‌دهنده را کاهش دهد. در آزمایشی جداگانه، گوگل یک مدل Muse Glimmer با 30 میلیارد پارامتر را در PyTorch استاندارد اجرا کرد: به‌گزارش InfoQ، توان عملیاتی استنتاج روی یک تراشه TPUv7 با طول توالی 5,120 به 25,700 توکن در ثانیه و با طول 10,240 به 22,400 توکن در ثانیه رسید. آموزش روی چهار تراشه TPUv7 نیز برای این دو طول به‌ترتیب به 5,600 و 4,600 توکن در ثانیه رسید. این نتایج را گوگل گزارش کرده و مقایسه‌ای مستقل نیستند.
TechCrunch، TorchTPU را در کنار TorchNeuron آمازون قرار می‌دهد که در همان هفته برای اجرای بارهای کاری PyTorch روی تراشه‌های Trainium آمازون معرفی شد. هدف هر دو، آسان‌تر کردن استفاده از گزینه‌هایی به‌جای پردازنده‌های گرافیکی Nvidia است. عرضه TorchTPU گزینه سخت‌افزاری دیگری پیش روی کاربران PyTorch می‌گذارد، اما به‌تنهایی نشان نمی‌دهد که انتقال یک بار کاری هزینه را کاهش می‌دهد یا از پیکربندی فعلی توسعه‌دهنده بهتر عمل می‌کند.
چرا مهم است
برای توسعه‌دهندگانی که از PyTorch استفاده می‌کنند، TorchTPU می‌تواند آزمایش یا استقرار بارهای کاری روی TPUهای گوگل را بدون بازسازی آن‌ها از ابتدا آسان‌تر کند. هنوز روشن نیست که این انتقال برای یک مدل مشخص هزینه یا سرعت را بهبود می‌دهد؛ ارقام عملکردی ذکرشده نتایج خود گوگل هستند.
