AWS Trainium

تعمل وحدات AWS Trainium على PyTorch بشكل أصلي من خلال حزمة torch-neuronx، التي تسجّل جهاز neuron. يقوم Neuron بتجميع نموذجك إلى رسم بياني واحد أو أكثر لكل شكل إدخال يراه، لذا فإن إبقاء أشكال الإدخال ثابتة يتجنب إعادة التجميع.

التثبيت

قم بتثبيت torch-neuronx باتباع دليل إعداد AWS Neuron. ثم تأكد من أن PyTorch يرى أجهزة Neuron.

import torch

print(torch.neuron.is_available(), torch.neuron.device_count())

يطبع True وعدد أجهزة Neuron، على سبيل المثال True 4.

تشغيل الاستدلال

device_map="auto" لا يكتشف Neuron. يعتمد التوزيع التلقائي على Accelerate لذكر الأجهزة المتاحة، ولا يذكر Accelerate أجهزة Neuron، لذا يتم تحميل النموذج على وحدة المعالجة المركزية (CPU) بصمت. مرّر device_map="neuron" إلى from_pretrained() بشكل صريح، أو مرّر device="neuron" إلى Pipeline.

مرّر مدخلات التوليد على وحدة المعالجة المركزية (CPU) بدلاً من استدعاء .to(model.device). تنقل generate() فقط المصفوفات التي تستهلكها forward إلى جهاز Neuron. تبقى المصفوفات التي تنمو في كل خطوة (تسلسل الإخراج، attention_mask) على وحدة المعالجة المركزية، مما يحافظ على استقرار الرسم المُجمّع. تُرجع generate الإخراج على نفس الجهاز الذي توجد عليه المدخلات. راجع إبقاء مصفوفات التوليد على وحدة المعالجة المركزية لمزيد من التفاصيل.

from transformers import AutoModelForCausalLM, AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-0.6B")
model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-0.6B", dtype="bfloat16", device_map="neuron"
)

# leave the inputs on CPU
inputs = tokenizer("The French Bread Law states", return_tensors="pt")
output = model.generate(
    **inputs, do_sample=False, max_new_tokens=50, cache_implementation="static"
)
print(tokenizer.decode(output[0], skip_special_tokens=True))

The French Bread Law states that the French government must provide a certain amount of bread for the citizens of the country. The law is passed by the French parliament. The French parliament is a body of the French government. The French government is a body of the French state. The

مع ذاكرة مؤقتة ثابتة، يجمع generate() مسار التقدّم تلقائيًا باستخدام torch.compile على Neuron. مرّر disable_compile=True إلى generate() للاستثناء من ذلك.

التدريب باستخدام Trainer

يتعرّف TrainingArguments على Neuron ويضبط الجهاز على neuron تلقائيًا. تدعم جميع أجيال Trainium نمط bf16.

from transformers import TrainingArguments

training_args = TrainingArguments(
    output_dir="./outputs",
    bf16=True,
)

تجميع النموذج

اضبط torch_compile=True لتغليف النموذج بـ torch.compile. يقوم هذا بتجميع المسار الأمامي للنموذج، بينما يتم تنفيذ خطوة المُحسِّن بشكل فوري. يضبط TrainingArguments قيمة torch_compile_backend على "neuron" تلقائيًا عند اكتشافه لـ Neuron.

from transformers import TrainingArguments

training_args = TrainingArguments(
    output_dir="./outputs",
    bf16=True,
    torch_compile=True,
)

يُطلق كل شكل إدخال جديد عملية تجميع جديدة. قم بتعبئة المدخلات إلى طول ثابت (padding="max_length" في مُجزّئ الرموز) و اضبط dataloader_drop_last=True حتى لا يختلف حجم الدفعة الأخيرة.

التدريب الموزّع

يدعم Neuron واجهة Transformers الموزعة. أنشئ DistributedConfig بأحجام التوازي التي تريدها وقم بتمريره إلى from_pretrained(). تقوم Transformers بتقسيم النموذج عند التحميل، ويستخدم Trainer ذلك كما هو.

import torch
from transformers import AutoModelForCausalLM, DistributedConfig

model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-1.7B",
    dtype=torch.bfloat16,
    distributed_config=DistributedConfig(tp_size=2, fsdp_size=2),
)

شغّل السكربت باستخدام torchrun.

torchrun --nproc_per_node=4 train.py

راجع توازي الموترات، وتوازي الخبراء، وFSDP لمعرفة الخيارات المتاحة.

القيود

• device_map="auto" يحمّل النموذج على وحدة المعالجة المركزية (CPU) بدلاً من Neuron. مرّر device_map="neuron" بشكل صريح.
• تُحفّز الأشكال الديناميكية إعادة التجميع. يتباطأ كلٌّ من التوليد والتدريب عندما تتغيّر أشكال المدخلات من خطوة إلى أخرى. استخدم ذاكرة مؤقتة ثابتة للتوليد، والتعبئة بطول ثابت للتدريب.