PyTorch از قفل CUDA فاصله می‌گیرد؛ TPU و Trainium بی‌واسطه می‌شوند

گوگل در ۲۰ اکتبر ۲۰۲۶ و در کنفرانس PyTorch، TorchTPU را به‌صورت متن‌باز منتشر کرد و TPU را در کنار CPU و GPU به مقصدی بومی برای PyTorch تبدیل کرد؛ جانا ون گرونن، مدیر مهندسی PyTorch در Meta، در همان اعلامیه گفت: «We cannot be constrained by a single hardware architecture.» این انتشار به توسعه‌دهنده اجازه می‌دهد بخش‌هایی از کد آشنای PyTorch را با تغییر کمتر روی شتاب‌دهندهٔ گوگل اجرا کند، اما میزان تغییر به وابستگی‌های همان برنامه بستگی دارد.

AWS نیز در ۲۰ اکتبر ۲۰۲۶، بتای عمومی TorchNeuron را برای Trainium2 و Trainium3 عرضه کرد. این بک‌اند متن‌باز با مجوز Apache 2.0، اجرای eager، ‏torch.compile و APIهای استاندارد PyTorch را پشتیبانی می‌کند. برای کدی که دستگاه را صریح تعیین می‌کند، تغییر آن به .to('neuron') نقطهٔ شروع است؛ سازگاری عملگرها، کتابخانه‌ها و مسیر کامپایل هنوز تعیین می‌کند چه مقدار کار باقی می‌ماند.

اجرای بومی در PyTorch چه چیزی را تغییر می‌دهد؟

هر دو بک‌اند، شتاب‌دهندهٔ مقصد را به الگوی آشنای دستگاه‌ها در PyTorch نزدیک می‌کنند. در این الگو، تعریف مدل و بسیاری از APIهای رایج همان‌هایی هستند که توسعه‌دهنده برای CPU یا GPU می‌شناسد. حالت eager اجرای مرحله‌به‌مرحله و عیب‌یابی تعاملی را ممکن می‌کند؛ ‏torch.compile مسیر دیگری برای تبدیل بخش‌های مناسب مدل به اجرای بهینه‌شده روی سخت‌افزار مقصد فراهم می‌آورد. بنابراین انتخاب تراشه کمتر از گذشته به نگهداری نسخه‌ای جداگانه از تعریف مدل گره می‌خورد.

این تغییر به‌ویژه برای کدی مؤثر است که از APIهای مستقل از دستگاه، مانند torch.accelerator، استفاده می‌کند. در مقابل، نام‌گذاری صریح CUDA در برنامه، فراخوانی کرنل اختصاصی آن یا فرض دربارهٔ ارتباط و حافظهٔ GPU با عوض‌کردن نام دستگاه از میان نمی‌رود. TorchNeuron برای نوشتن کرنل اختصاصی Trainium مسیر Neuron Kernel Interface یا NKI را دارد؛ در سمت TPU نیز برخی مسیرهای بهینه‌سازی از کرنل‌های Pallas استفاده می‌کنند. چنین بخش‌هایی محل واقعی هزینهٔ انتقال‌اند، حتی اگر حلقهٔ اصلی آموزش تقریباً یکسان بماند.

ماتریس سازگاری: وضعیت، مجوز و کتابخانه‌ها

کد TorchTPU و مجوز Apache 2.0 آن در مخزن عمومی پروژه ثبت شده است. مقایسهٔ دو مسیر نشان می‌دهد «تغییر کم» بیشتر دربارهٔ کد متکی به APIهای استاندارد صدق می‌کند؛ کتابخانه و قابلیت موردنیاز مدل نیز باید در بک‌اند مقصد پشتیبانی شوند.

• TorchTPU برای Google TPU: کد آن منتشر شده و تحت مجوز Apache 2.0 است. مسیر معمول، قرار دادن مدل و داده روی TPU و استفاده از eager یا torch.compile است. در مجموعهٔ ابزارهای معرفی‌شده، Transformers و Diffusers برای کار با مدل‌ها، vLLM و SGLang برای سرویس‌دهی، و TorchTitan برای آموزش توزیع‌شده حضور دارند. برای نمونه، مسیر Transformers از device_map="tpu" استفاده می‌کند. اتصال این ابزارها به TPU به معنای سازگاری خودکار هر مدل، افزونه یا کرنل وابسته به CUDA نیست.
• TorchNeuron برای AWS Trainium: انتشار آن در وضعیت بتای عمومی و با مجوز Apache 2.0 است. دستگاه neuron برای کد معمول PyTorch در دسترس قرار می‌گیرد و اجرای کامپایل‌شده می‌تواند بک‌اند neuron را برای torch.compile انتخاب کند. Transformers، Diffusers، TRL و TorchTitan از مسیرهای پشتیبانی‌شده‌اند. TorchAO و FP8 هنوز در بتای خصوصی‌اند؛ یکپارچه‌سازی گسترده‌تر vLLM و SGLang با مدل‌های Hugging Face نیز در حال فعال‌سازی و اعتبارسنجی است.

تفاوت وضعیت انتشار اهمیت دارد: کد TorchTPU اکنون در مخزن عمومی قابل دریافت است، در حالی که TorchNeuron با وجود متن‌باز بودن، برچسب بتای عمومی دارد. در هر دو مورد، متن‌باز بودن نرم‌افزار هزینهٔ ظرفیت ابری را حذف نمی‌کند. همچنین نام یک کتابخانه در فهرست پشتیبانی، جای بررسی نسخهٔ همان کتابخانه، عملگرهای مدل و تنظیمات اجرای توزیع‌شده را نمی‌گیرد.

مرز وعدهٔ «تغییر کم» در نسخهٔ نخست

محدودیت‌های TorchNeuron نمونهٔ روشنی از فاصلهٔ میان شناختن دستگاه و اجرای کامل یک برنامه است. بخشی از عملگرهای PyTorch هنوز پیاده‌سازی بومی روی Neuron ندارند. تنسورهای sparse، ‏NestedTensor، مشتق‌گیری رو‌به‌جلو و شکل‌های پویا یا نمادین نیز در فهرست قابلیت‌های پشتیبانی‌نشدهٔ انتشار نخست هستند. برنامه‌ای که به یکی از این ویژگی‌ها متکی باشد، ممکن است پس از تغییر دستگاه هم برای آموزش یا استنتاج به اصلاح نیاز داشته باشد.

مسیر eager و مسیر کامپایل‌شده نیز یکسان رفتار نمی‌کنند. eager عملگرها را هنگام اجرا پیش می‌برد و برای دیدن جریان برنامه و یافتن خطا مناسب است؛ در مسیر torch.compile، گراف محاسباتی برای Trainium آماده می‌شود. بخشی از مدل می‌تواند کامپایل شود و بخش دیگر در eager بماند، اما همین ترکیب ممکن است مرزهای تازه‌ای برای حافظه، انتقال داده و عیب‌یابی بسازد. اجرای موفق مدل در حالت تعاملی، نتیجهٔ عملکرد آن در حالت کامپایل‌شده را مشخص نمی‌کند.

در سمت TPU، پشتیبانی از چند ابزار شناخته‌شده مسیر آغاز را کوتاه می‌کند، اما هر اتصال کار متفاوتی انجام می‌دهد. vLLM و SGLang به سرویس‌دهی مربوط‌اند؛ TorchTitan برای آموزش در مقیاس بزرگ به کار می‌رود؛ Transformers و Diffusers نیز خانواده‌های متفاوتی از مدل‌ها را پوشش می‌دهند. سازگاری یک مسیر سرویس‌دهی را نمی‌توان به همهٔ مراحل آموزش، همهٔ مدل‌ها یا هر پیکربندی چندتراشه‌ای تعمیم داد. این تمایز برای تیمی که کدش میان آموزش و استقرار مشترک است، مستقیماً بر مقدار بازنویسی اثر می‌گذارد.

پیامد این دو انتشار برای وابستگی به CUDA

اکنون توسعه‌دهندهٔ PyTorch برای آزمایش TPU و Trainium مسیر بومی و متن‌باز در اختیار دارد؛ این مسیرها وابستگی پیش‌فرض به ابزارهای ویژهٔ Nvidia را کمتر می‌کنند. نتیجهٔ فوری، گسترش انتخاب سخت‌افزار در سطح چارچوب و کتابخانه‌های اطراف آن است، نه تضمین انتقال بی‌هزینهٔ هر برنامه. کدی که بر عملگرهای عمومی PyTorch بنا شده آسان‌تر جابه‌جا می‌شود، ولی بخش‌های اختصاصی CUDA همچنان به جایگزین یا بازنویسی نیاز دارند.

از این انتشارها نمی‌توان برتری سرعت یا هزینهٔ یکی از شتاب‌دهنده‌ها را برای یک مدل مشخص نتیجه گرفت. نتیجه به پوشش عملگرها، حالت اجرا، اندازه و ساختار مدل، ارتباط میان تراشه‌ها و ظرفیت ابری مورد استفاده بستگی دارد. برای تیمی که میان این سخت‌افزارها انتخاب می‌کند، محدودیت‌های ثبت‌شدهٔ بتای TorchNeuron و پوشش واقعی ابزارهای متصل به TorchTPU اکنون همان‌قدر تعیین‌کننده‌اند که سادگی تغییر نام دستگاه در کد.

بیشتر بخوانید:

• رزومه یا GitHub؛ کد واقعی جای روایت شغلی را نمی‌گیرد
• Cloudflare تیم Deno را می‌گیرد؛ workerd و celld یکی می‌شوند
• عامل‌های OpenAI داخل AWS اجرا می‌شوند؛ پیش‌نمایش فقط در سه منطقه است