التدريب على AWS Trainium

يعمل AWS Trainium مع PyTorch بشكل أصلي من خلال حزمة torch-neuronx، التي تسجّل جهاز neuron. تعمل مدرّبات TRL عليه دون أي تغييرات في الكود. يقوم Neuron بتجميع رسم بياني (graph) لكل شكل إدخال (input shape) يكتشفه، ويستغرق كل عملية تجميع نحو دقيقة واحدة، لذا فإن معظم محتوى هذه الصفحة يتناول كيفية إبقاء الأشكال ثابتة. للاطلاع على التثبيت والجانب المتعلق بـ Transformers، راجع دليل Transformers الخاص بـ AWS Trainium.

التثبيت

قم بتثبيت torch-neuronx باتباع دليل إعداد AWS Neuron، ثم ثبّت TRL في نفس البيئة.

pip install trl

الضبط الدقيق الإشرافي

يعمل البدء السريع لـ SFTTrainer كما هو:

from datasets import load_dataset
from trl import SFTTrainer

trainer = SFTTrainer(
    model="Qwen/Qwen3-0.6B",
    train_dataset=load_dataset("trl-lib/Capybara", split="train"),
)
trainer.train()

على Neuron، القيمة الافتراضية لـ pad_to_multiple_of هي 128 (أي قيمتها "auto")، لذا يتم ملء التسلسلات حتى أقرب مضاعف لـ 128. وهذا يحدد عدد الأشكال، وبالتالي عدد عمليات التجميع، بـ max_length / 128. وبدون ذلك، يؤدي كل طول جديد للتسلسل إلى تشغيل عملية تجميع. على وحدة NeuronCore واحدة، مع per_device_train_batch_size=4، تستغرق 20 خطوة من المثال أعلاه 130 ثانية مع القيمة الافتراضية و260 ثانية مع pad_to_multiple_of=None.

ينطبق الأمر نفسه على كل مُدرّب يحتوي على خيار pad_to_multiple_of: DPO وKTO وReward وGRPO وRLOO وDistillation.

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

اضبط torch_compile=True لتجميع النموذج باستخدام واجهة neuron. الأشكال الديناميكية التلقائية في Dynamo تنتج رسومًا لا تستطيع واجهة Neuron تجميعها بعد (#7660)، لذا عطّلها:

import torch
from datasets import load_dataset
from trl import SFTConfig, SFTTrainer

torch._dynamo.config.automatic_dynamic_shapes = False

trainer = SFTTrainer(
    model="Qwen/Qwen3-0.6B",
    args=SFTConfig(torch_compile=True),
    train_dataset=load_dataset("trl-lib/Capybara", split="train"),
)
trainer.train()

يستغرق التجميع في الخطوة الأولى حوالي دقيقتين. بعد ذلك، مع per_device_train_batch_size=4، يتم تدريب Qwen3-0.6B بمعدل 1.15 ثانية/خطوة على NeuronCore واحد، مقابل 1.71 ثانية/خطوة دون تجميع.

الطرق عبر الإنترنت

لا يعمل vLLM على Neuron، لذا يولّد GRPO وRLOO باستخدام Transformers. مع ذاكرة التخزين المؤقت الديناميكية الافتراضية، يكون لكل رمز مولّد شكل جديد، ويستغرق كل خطوة عدة دقائق. استخدم ذاكرة التخزين المؤقت الثابتة، التي تحافظ على أشكال فك الترميز ثابتة:

from datasets import load_dataset
from trl import GRPOConfig, GRPOTrainer


# Reward completions close to 200 characters
def reward_len(completions, **kwargs):
    return [-abs(200 - len(completion)) / 100 for completion in completions]


trainer = GRPOTrainer(
    model="Qwen/Qwen3-0.6B",
    reward_funcs=reward_len,
    args=GRPOConfig(
        per_device_train_batch_size=4,
        num_generations=4,
        max_completion_length=128,
        cache_implementation="static",
        model_init_kwargs={"dtype": "bfloat16"},
    ),
    train_dataset=load_dataset("trl-lib/tldr", split="train"),
)
trainer.train()

تُملأ الأوامر (Prompts) حتى pad_to_multiple_of قبل التوليد أيضًا، لذا فإن الكاش الثابت لا يستهلك سوى بضعة أحجام. لا يزال كل حجم جديد يُترجم (compiles) لمدة حوالي 12 دقيقة على NeuronCore واحد؛ وبمجرد ترجمة حجم معين، يستغرق خطوة واحدة من هذا المثال حوالي 25 ثانية.

يتطلب الكاش الثابت نموذجًا من نوع bf16 حاليًا: فمع النموذج الافتراضي من نوع fp32، ينتهي الأمر بالمفاتيح والقيم بأنواع بيانات مختلفة (#7664).

التوازي في البيانات

على مثيل trn2.3xlarge، الذي يحتوي على شريحة Trainium2 واحدة، يجمع التكوين الافتراضي لـ Logical NeuronCore (LNC=2) بين 8 أنوية NeuronCore فيزيائية في الشريحة لتكوين 4 أنوية NeuronCore منطقية.

لتدريب النموذج على جميع الأنوية الأربعة، شغّل عملية واحدة لكل NeuronCore باستخدام torchrun:

torchrun --nproc_per_node 4 train.py

في هذا الإعداد، يعمل SFT أسرع بمقدار 3.6 مرات على 4 من NeuronCores مقارنةً بواحد.

لحصر التشغيل في مجموعة فرعية من NeuronCores، اضبط NEURON_RT_VISIBLE_CORES على 0,1 أو 2,3 أو 0-3. لا تُدعَم تركيبات أخرى مثل 1,2: إذ يفشل وقت تشغيل Neuron في إعداد العمليات الجماعية عبرها.