‏TorchNeuron: اجرای مستقیم PyTorch روی تراشه‌های Trainium آمازون

AWS از TorchNeuron رونمایی کرده؛ یه backend متن‌باز و native برای PyTorch که تراشه‌های Trainium رو مثل یه device استاندارد در دسترس می‌ذاره. Trainium خانواده‌ی تراشه‌هاییه که آمازون برای آموزش و اجرای مدل‌های هوش مصنوعی ساخته. تا قبل از این، استفاده ازش نیاز به ابزار و تنظیمات مخصوص Neuron داشت؛ حالا به گفته‌ی تیم AWS، کد مدل، کتابخونه‌ها و ابزارهایی که الان دارید، همون‌طور روی Trainium اجرا می‌شن.

این کار روی PrivateUse1 سوار شده، یعنی نقطه‌ی رسمی‌ای که PyTorch برای اضافه کردن سخت‌افزار سفارشی گذاشته. نتیجه اینه که torch.nn، torch.autocast، torch.distributed و torch.profiler همه روی device به اسم neuron کار می‌کنن. برای شروع فقط کافیه مدل رو به این device بفرستید:

model = MyModel().to('neuron')   # the only change
optimizer = torch.optim.AdamW(model.parameters())
for batch in dataloader:
    optimizer.zero_grad()
    loss = criterion(model(batch), targets)
    loss.backward()

دو حالت اجرا پشتیبانی می‌شه. Eager (یعنی هر عملیات همون لحظه اجرا می‌شه و breakpoint و print کار می‌کنن) برای تحقیق و دیباگه؛ یه حالت آزمایشی به اسم adaptive eager هم اضافه شده که با ترکیب عملیات‌ها، سربار اجرا رو کم می‌کنه. torch.compile هم برای production هست و مدل رو به یه گراف تبدیل می‌کنه که کامپایلر Neuron بهینه‌ش می‌کنه. برای کرنل‌های سفارشی هم NKI یا Neuron Kernel Interface هست.

بخش توزیع‌شده هم کامله: DDP، FSDP2، tensor parallelism و pipeline parallelism بدون تغییر کار می‌کنن. ارتباط بین نودها از طریق EFA انجام می‌شه. پروفایل‌گیری هم با همون torch.profiler انجام می‌شه و خروجی‌ش تو ابزار Neuron Explorer تا خط کد اصلی قابل ردگیریه.

سازگاری با اکوسیستم هم یکی از هدف‌های اصلیه: Hugging Face Transformers v5، Diffusers و TRL حالا Trainium رو به‌عنوان device پشتیبانی می‌کنن و TorchTitan هم روش آموزش می‌ده. AWS می‌گه تست رو روی ۱۰۰ هزار مدل Hugging Face انجام داده که ۹۹.۹ درصد استفاده‌ی واقعی رو پوشش می‌دن. Deepgram هم گزارش داده Trainium حدود ۶۰ درصد ثانیه‌ی صوتی بیشتری به ازای هر دلار نسبت به GPUهای فعلی‌ش پردازش می‌کنه.

البته محدودیت‌هایی هم هست. حدود ۲۲ هزار تست upstream به‌خاطر عملیات‌هایی که هنوز پیاده‌سازی native ندارن رد شدن. چیزهایی مثل sparse tensor، NestedTensor و dynamic shapes هم تو نسخه‌ی اول پشتیبانی نمی‌شن. FP8 و TorchAO هم فعلاً تو بتای خصوصی‌ان.

نکات کلیدی:

• TorchNeuron از ۲۰ اکتبر ۲۰۲۶ به‌صورت بتای عمومی و با لایسنس Apache 2.0 منتشر می‌شه.
• روی Trainium2 و Trainium3 کار می‌کنه و از حالت eager و torch.compile پشتیبانی می‌کنه.
• نصبش با پکیج torch-neuronx از مخزن pip خود AWS Neuron هست.
• Transformers v5 با device_map="neuron" مدل رو مستقیم روی Trainium لود می‌کنه.
• sparse tensor، NestedTensor و dynamic shapes تو نسخه‌ی اول پشتیبانی نمی‌شن.