آیا مدل شما روی Trainium اجرا می‌شود؟ این ابزار جواب می‌دهد

ابزار torch-neuronx model support در هاگینگ‌فیس (Hugging Face) به یک پرسش مشخص جواب می‌دهد: آیا یک مدل بدون هیچ تغییری روی ترینوم (Trainium) اجرا می‌شود یا نه. این ابزار برای مدل‌های محبوب هاب، نتیجهٔ تست، تعداد هسته‌های درگیر، زمان اجرا و اسکریپت لازم را کنار هم می‌گذارد و پشتیبانی مدل‌ها از Trainium را به‌جای حدس زدن، قابل بررسی می‌کند.

از TorchNeuron تا فهرست پشتیبانی

پیش‌تر تورچ‌نورون (TorchNeuron) معرفی شد؛ بک‌اند بومی PyTorch برای AWS Trainium که در قالب بستهٔ torch-neuronx منتشر می‌شود. با آن، کتابخانهٔ Transformers کد مدل‌سازی خودش را مثل هر دستگاه دیگری روی Trainium اجرا می‌کند؛ یعنی به‌طور نظری هر مدل Transformers روی هاب می‌تواند اجرا شود. پرسش واقعی این است: کدام‌ها در عمل اجرا می‌شوند؟

این Space دقیقاً همان پرسش را برای پرکاربردترین مدل‌های هاب پاسخ می‌دهد. خروجی آن دو چیز است: اینکه مدل بدون تغییر روی Trainium اجرا می‌شود یا نه، و اسکریپتی که می‌توانید خودتان با آن اجرا کنید. اگر با مسیر کلی اجرای مدل‌های هاگینگ‌فیس روی Trainium آشنا نیستید، نوشتهٔ TorchNeuron: اجرای مدل‌های هوگینگ‌فیس روی AWS Trainium ساده‌تر از همیشه نقطهٔ شروع خوبی است.

جست‌وجوی یک مدل در فهرست

Space با نمای «Selected» باز می‌شود؛ مجموعه‌ای دستچین‌شده از مدل‌های برتر هاب بر اساس تعداد دانلود و امتیاز ترند. هر ردیف چهار قلم اطلاعات می‌دهد: نتیجهٔ تست، تعداد هسته‌های Trainium که مدل روی آن‌ها شارد شده، زمان واقعی اجرا (شامل کامپایل) و اوج حافظهٔ دستگاه.

برای پیدا کردن یک مدل مشخص، به نمای Trending یا Downloads بروید و بخشی از نام مدل را تایپ کنید. این دو نما از رتبه‌بندی خود هاب پیروی می‌کنند، هر بار برای یک وظیفه، و فرقی نمی‌کند مدل تست شده باشد یا نه. اگر مدلی عامدانه کنار گذاشته شده باشد، Space دلیلش را می‌گوید؛ در بیشتر موارد اولویت‌بندی نشده و گاهی هم واقعاً ناسازگار است. امکان انتخاب نسخهٔ کامپایلر Neuron (neuronx-cc) و نسخهٔ Transformers هم وجود دارد تا نتیجهٔ متناسب با پشتهٔ خودتان را ببینید.

در زمان نوشتن این مطلب، فقط مدل‌های Transformers تست می‌شوند و دامنهٔ انتخاب به دو دستهٔ تولید متن و متن-به-متن تصویری محدود است. برنامهٔ آینده گسترش به diffusers، sentence transformers و وظیفه‌های دیگر Transformers است.

خودتان اجرا کنید

با انتخاب یک مدل، ابزار نشان می‌دهد چطور آن را اجرا کنید: پشتهٔ Neuron که مدل روی آن تست شده، بسته‌هایی که باید روی آن نصب شوند و فایل run.py که اجرا شده است. این اسکریپت همان دستورالعمل معرفی TorchNeuron را دنبال می‌کند؛ مدل را با from_pretrained بارگذاری می‌کند، در صورت نیاز شارد می‌کند و سپس با کش ایستا و torch.compile(backend="neuron") تولید متن را انجام می‌دهد.

روی یک نمونهٔ Trainium2 (trn2) با همان پشته — مثلاً در کانتینر AWS Neuron Deep Learning — کافی است بسته‌ها را نصب کنید و اجرا کنید:

python3 run.py

نیازی به torchrun نیست؛ وقتی مدل شارد شده باشد، خود اسکریپت رنک‌های torch.distributed را راه می‌اندازد. اجرای اول با کامپایل گراف‌های مدل شروع می‌شود. برای پریدن از این مرحله می‌توانید گراف‌هایی را که همان کامپایلر در جریان تست ساخته، از باکت عمومی کش NEFF بردارید. Space دستور مربوطه را با نسخهٔ درست کامپایلر در اختیارتان می‌گذارد:

hf buckets sync hf://buckets/aws-neuron/torch-neuronx-neff-cache/neuronxcc-<version> /tmp/neff_cache

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

«اجرا می‌شود» یعنی چه؟

هر تست، مدل را روی یک نمونهٔ trn2 و روی کمترین تعداد هستهٔ Trainium که می‌تواند وزن‌هایش را نگه دارد بارگذاری می‌کند. بیشتر مدل‌ها با موازی‌سازی تانسوری شارد می‌شوند. مدل‌های ترکیب خبرگان (MoE) که برنامهٔ موازی‌سازی تانسوری ندارند، فقط با موازی‌سازی خبره شارد می‌شوند. سپس تست ۳۲ توکن تولید و بررسی می‌کند.

روی یک هستهٔ تنها، توکن‌ها با خروجی اجرای CPU همان چک‌پوینت مقایسه می‌شوند. وقتی مدل شارد شده باشد، یک مدل داور کوچک بررسی می‌کند که متن، انگلیسی روان باشد.

در زمان انتشار، فقط مدل‌هایی تست شده‌اند که روی چهار دستگاه Trainium2 جا می‌شوند؛ یعنی حداکثر ۱۶ هسته. این نکته را باید در تفسیر نتیجه در نظر گرفت: یک PASS هیچ چیزی دربارهٔ توان عملیاتی یا تولید متن‌های طولانی‌تر نمی‌گوید.

نتایج یک دیتاست است

هر چیزی که در Space می‌بینید از دیتاست عمومی aws-neuron/torch-neuronx-model-support می‌آید؛ برای هر جفت از نسخه‌های neuronx-cc و کتابخانه‌ها یک فایل Parquet وجود دارد. هر ردیف نتیجهٔ مدل، خطا در صورت شکست، نسخه‌های کل پشته و همان run.py اجراشده را نگه می‌دارد. کارت این دیتاست معنای دقیق هر نتیجه را توضیح می‌دهد.

می‌توانید مستقیماً از آن پرس‌وجو بگیرید؛ مثلاً با pandas فایل را بخوانید و سطرهای مربوط به یک مدل را با ستون‌های نسخهٔ کامپایلر، کتابخانه، تعداد هسته و نتیجه فیلتر کنید:

import pandas as pd
df = pd.read_parquet("hf://datasets/aws-neuron/torch-neuronx-model-support/data")
df[df.model_id == "Qwen/Qwen3-0.6B"][["neuronx_cc", "library", "library_version", "neuron_cores", "result"]]

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

پرسش‌های متداول

منبع: Hugging Face