گوگل پروژه TorchTPU را برای پشتیبانی بومی پای‌تورچ از تراشه‌های TPU متن‌باز کرد

در دنیای هوش مصنوعی، انحصار سخت‌افزارهای شتاب‌دهنده و چالش‌های هماهنگی آن‌ها با فریم‌ورک‌های محبوب همواره دغدغه‌ی مهندسان یادگیری ماشین بوده است. گوگل به‌تازگی مخزن کد جدیدی به نام TorchTPU را در گیت‌هاب با شناسه google-pytorch/torch_tpur منتشر کرده است که هدف آن ارائه یک بک‌اند بومی برای اجرای عملیات PyTorch روی تراشه‌های اختصاصی گوگل یا همان Tensor Processing Units است.

پروژه TorchTPU چیست و چه مشکلی را حل می‌کند؟

پیش از این، استفاده از TPUهای گوگل با پای‌تورچ نیازمند لایه‌ها و ابزارهای واسطی بود که گاهی پیچیدگی‌های خاص خود را داشتند. پروژه TorchTPU با هدف ارائه یکپارچگی عمیق‌تر و بومی‌تر ایجاد شده است. این پروژه شامل هسته‌های سفارشی ATen برای عملیات مختلف پای‌تورچ، کش کامپایلر برای بهینه‌سازی اجرای محاسبات XLA، و مثال‌های کاربردی برای آموزش توزیع‌شده (مانند Data Parallelism و Tensor Parallelism) است.

این ابزار از معماری گراف‌مود پای‌تورچ یعنی torch.compile() پشتیبانی می‌کند که به مهندسان اجازه می‌دهد مدل‌های خود را روی سخت‌افزار TPU به بالاترین کارایی برسانند. همچنین پشتیبانی از مدل‌های بزرگ زبانی و بینایی مانند Llama، Qwen، ResNet، DLRM و UNet در این پروژه پیش‌بینی شده است.

معماری و نحوه کارکرد

هسته اصلی این پروژه با زبان C++ نوشته شده و برای مدیریت وابستگی‌ها و بیلد کردن از Bazel (و ابزار کمکی Bazelisk) استفاده می‌کند. برای توسعه‌دهندگانی که روی کدهای C++ این پروژه کار می‌کنند، تنظیمات ویژه‌ای برای استفاده از اکستنشن clangd در وی‌اس‌کد فراهم شده است تا ناوبری کدها و تکمیل خودکار با دقت بالا انجام شود.

از نظر محیط اجرایی، این ابزار نیازمند پایتون نسخه ۳.۱۱ به بالا است و برای اجرای نوت‌بوک‌های تعاملی از ابزار marimo (مشابه Jupyter) بهره می‌برد که اجرای آن‌ها روی یک TPU VM (ترجیحاً نمونه‌های v6e در Google Cloud) پیشنهاد می‌شود.

نحوه نصب و استفاده

برای شروع کار با این پکیج، ابتدا باید یک محیط مجازی پایتون با venv یا uv بسازید:

python -m venv .venv
source .venv/bin/activate

سپس آخرین نسخه پایدار آن را از PyPI نصب کنید. توجه داشته باشید که برای جلوگیری از نصب نسخه‌های مخصوص CUDA، باید نسخه CPU پای‌تورچ را به عنوان شاخص اضافی معرفی کنید:

pip install torch_tpu==0.2.0 --extra-index-url https://download.pytorch.org/whl/cpu

اگر قصد دارید پروژه را از طریق سورس‌کد و با کمک Bazel بیلد کنید، دستورات مفصلی برای ساخت ویل (Wheel) با پرچم‌های بهینه‌سازی مختلف نظیر --config=wheel_common در مستندات پروژه گیت‌هاب قرار داده شده است.

چرا این پروژه در هکرنیوز مورد توجه قرار گرفت؟

این پروژه در هکرنیوز بحث‌هایی را به همراه داشته است. برخی از کاربران به شباهت این رویکرد با اقدامات سایر شرکت‌ها اشاره کرده‌اند؛ برای مثال، آامازون نیز پروژه مشابهی به نام torch-neuronx را برای تراشه‌های AWS Trainium توسعه داده است. یکی از سوالات مطرح‌شده توسط کاربران این بود که آیا این TPUها خارج از زیرساخت گوگل کلود نیز قابل دسترس هستند یا خیر، که نشان‌دهنده کنجکاوی جامعه مهندسی نرم‌افزار درباره میزان دسترسی‌پذیری این سخت‌افزارهاست.

در یک نگاه

• مخزن: google-pytorch/torch_tpu
• زبان اصلی: C++ و Python
• لایسنس: Apache 2.0
• حداقل نسخه پایتون: ۳.۱۱ به بالا
• ویژگی‌های کلیدی: پشتیبانی از torch.compile()، هسته‌های سفارشی ATen، آموزش توزیع‌شده و کش کامپایلر XLA