TorchNeuron اجرای مدل‌های PyTorch روی AWS Trainium را ساده‌تر می‌کند

Hugging Face می‌گوید TorchNeuron، افزونه‌ای برای PyTorch، امکان اجرای مدل‌ها روی تراشه‌های AWS Trainium را با قابلیت‌هایی مانند اجرای eager، ‏torch.compile و APIهای توزیع‌شدهٔ بومی PyTorch فراهم می‌کند. این افزونه قرار است از ۲۰ اکتبر ۲۰۲۶ به‌صورت بتای عمومی عرضه شود و نصب آن با دستور pip install torch-neuronx انجام می‌شود.

پیش‌تر، پشتیبانی Hugging Face از Trainium و Inferentia از مسیر کتابخانهٔ Optimum Neuron ارائه می‌شد؛ کتابخانه‌ای متکی بر PyTorch XLA که به گفتهٔ Hugging Face، توسعه‌دهندگان را با محدودیت‌هایی مانند اجرای تنبل، شکل‌های ثابت و کامپایل مجدد پس از شکستن گراف روبه‌رو می‌کرد. TorchNeuron، Neuron را به‌عنوان دستگاهی بومی در PyTorch معرفی می‌کند و از قابلیت‌هایی مانند DTensor و FSDP2 پشتیبانی می‌کند.

در نتیجه، کتابخانه‌های Transformers، Diffusers، Accelerate، PEFT و TRL می‌توانند مستقیماً روی Trainium اجرا شوند. Hugging Face توصیه می‌کند کاربران Neuron به‌جای Optimum Neuron از این کتابخانه‌ها استفاده کنند. در تولید متن، طبق اعلام این شرکت، TorchNeuron از ۸۹ درصد ۱۰۰ معماری پرکاربرد پشتیبانی می‌کند؛ این رقم برای Optimum Neuron نسخهٔ ۰.۴.۵ برابر با ۴۶ درصد است. برای انتقال مدل به دستگاه‌های Neuron، کافی است آرگومان device_map="neuron" را در کد به‌کار برد. Hugging Face همچنین استفاده از static cache را برای کاهش کامپایل‌های مجدد توصیه می‌کند.