تدريب Qwen3 على AWS Trainium باستخدام torchtitan

إن استخدام torchtitan على AWS Trainium يعني استخدامه بالطريقة التي اعتدت عليها دائمًا. يُسجّل TorchNeuron وحدة Trainium كجهاز PyTorch أصلي، مما يتيح لـ torchtitan استخدامها من خلال واجهات PyTorch القياسية. يمكنك استخدام هذه الإعدادات للتدريب المسبق من تهيئة عشوائية أو لضبط الأوزان المنشورة بدقة. يمكنك أيضًا تخصيص تنفيذ النموذج والأنوية بما يلائم عبء العمل لديك. في هذه المقالة، ستدرّب Qwen3 بثلاث طرق على مثيل واحد من نوع trn2.48xlarge.

تستخدم الطريقة الأولى Transformers v5 كخلفية نمذجة في torchtitan. وهي تشغّل نماذج Hub من خلال AutoModelForCausalLM وتقرأ ملفات safetensors المنشورة وتكتبها مباشرة. أما الطريقة الثانية فتنتقل إلى Qwen3 الأصلي في torchtitan بتغيير config.model_spec. وتضيف الطريقة الثالثة أنوية Neuron Kernel Interface (NKI) mixture-of-experts (MoE) من وصفات torchtitan-neuron، ثم تستخدم سجل التجاوز في torchtitan لإرفاق نواتك الخاصة. تستخدم الإعدادات الثلاثة جميعها نفس نمط torchrun للتشغيل وهيكل الإعدادات نفسه.

المتطلبات الأساسية

لمتابعة هذا الدليل التطبيقي، استخدم مثيلًا واحدًا من نوع trn2.48xlarge مع إما AWS Deep Learning AMI for Neuron أو حاوية Neuron PyTorch training container. توفر كلتا البيئتين مشغّل Neuron وزمن التشغيل والمترجم المطلوبين للتدريب على Trainium.

قبل تثبيت torchtitan-neuron، تحقق مما إذا كانت بيئتك تتضمنه بالفعل. يطبع الأمر التالي الإصدار المثبّت ومسار الحزمة:

python3 -c "import torchtitan, torchtitan_neuron; print(torchtitan.__version__, torchtitan_neuron.__file__)"

يستخدم هذا الدليل التطبيقي الإصدارات الموضحة في الجدول التالي. لإعادة إنتاج الأمثلة، ثبّت هذه الإصدارات. فتكامل Transformers v5 ووحدات التجاوز الخاصة بـ Trainium قد يتغيّران بين إصدارات torchtitan.

عند تثبيت torch_neuronx ، يُسجّل TorchNeuron وحدة Trainium كجهاز neuron ويكتشفها torchtitan عند بدء التشغيل. في الوضع السريع (eager mode)، يُجمّع Neuron العمليات أثناء تنفيذها. مع torch.compile، يُجمّع Neuron الرسم البياني الملتقط. يعمل نمط التشغيل نفسه عبر الإعدادات الثلاثة جميعها.

للإصدارات المستخدمة في هذا الدليل التطبيقي، ثبّت torchtitan-neuron وباقي تبعيات torchtitan باستخدام --no-deps للحفاظ على بناء Neuron من torch:

git clone https://github.com/aws-neuron/torchtitan-neuron
cd torchtitan-neuron
pip install -r requirements.txt

قبل أن تبدأ التدريب، احصل على نسخ نموذج Qwen3 من Hugging Face Hub. قم بتخزينها في دليل محلي يمكنك تمريره إلى إعدادات التدريب:

hf download Qwen/Qwen3-4B-Instruct-2507 \
  --local-dir "$PWD/assets/hf/Qwen3-4B-Instruct-2507"

كل إعداد تدريب هو دالة إعداد بلغة Python. الأمثلة الثلاثة تبدأ من إعداد أساسي تقوم بتحديثه في كل خطوة.

التدريب باستخدام Hugging Face Transformers v5

في التكوين الأول، يقوم Transformers v5 ببناء Qwen3 من خلال AutoModelForCausalLM. يوفّر تعريف النموذج على Hub التنفيذ المستخدم لتمريرات التمرير الأمامي والخلفي. تختار النموذج عن طريق تعيين معرّف Hub الخاص به في الإعدادات.

ابدأ بـ transformers_modeling_backend_full(). يترك الحقول التي تقابل مفاتيح إعداد Hugging Face دون تعيين، لذلك AutoConfig.from_pretrained() يحمّل بنية النموذج من الـ Hub.

إنشاء neuron_qwen3.py في المجلد الجذر لمستودع torchtitan-neuron، حيث تقوم بتشغيل أمر التدريب. أضف عمليات الاستيراد ومسار نموذج الأداة (model artifact) والإعدادات الأساسية:

import os
from torchtitan.experiments.transformers_modeling_backend.config_registry import (
    transformers_modeling_backend_full,
)
from torchtitan.experiments.transformers_modeling_backend.configs import (
    TransformersBackendConfig,
)

ASSETS = os.path.abspath("./assets/hf/Qwen3-4B-Instruct-2507")

def v5_qwen3_4b() -> TransformersBackendConfig:
    config = transformers_modeling_backend_full()
    config.hf_model = "Qwen/Qwen3-4B-Instruct-2507"
    config.hf_assets_path = ASSETS
    config.dataloader.dataset = "c4_test"
    config.training.steps = 10
    config.training.local_batch_size = 1
    config.parallelism.data_parallel_shard_degree = -1
    return config

torchrun --nproc_per_node=64 --rdzv_backend c10d --rdzv_endpoint localhost:0 \
  --local-ranks-filter 0 --role rank --tee 3 \
  -m torchtitan.train --module neuron_qwen3 --config v5_qwen3_4b

يطلق هذا الأمر 64 عملية تدريب، عملية واحدة لكل NeuronCore منطقي. --module neuron_qwen3 يستورد وحدة الإعدادات المحلية. --config v5_qwen3_4b يختار دالة الإعدادات. يقوم إعداد توازي البيانات بتقسيم النموذج عبر جميع الرتب الـ 64.

تحميل وحفظ نقاط التحقق من Hugging Face

لضبط النموذج بدقة من الأوزان المنشورة، فعّل تحميل نقطة التحقق ووجّه initial_load_path إلى ASSETSالمحوّل من نوع قاموس الحالة يضبط الجزء الأول model. البادئة والمقابض تربط التضمينات. وهو يحفظ نقطة التحقق النهائية بصيغة Hugging Face safetensors.

def v5_qwen3_4b_finetune() -> TransformersBackendConfig:
    config = v5_qwen3_4b()
    config.checkpoint.enable = True
    config.checkpoint.folder = "checkpoint_qwen3_4b_ft"
    config.checkpoint.initial_load_path = ASSETS # absolute path required
    config.checkpoint.initial_load_in_hf = True
    config.checkpoint.initial_load_model_only = True
    config.checkpoint.last_save_in_hf = True
    config.checkpoint.last_save_model_only = True
    return config

استخدم خسارة خطوة التدريب الأولى كفحص سريع للتأكد من أن torchtitan قام بتحميل الأوزان المنشورة. لأن Qwen/Qwen3-4B-Instruct-2507، أنتجت العملية الموثّقة خسارة في الخطوة 1 بقيمة 6.54 من الأوزان المنشورة و12.43 من التهيئة العشوائية.

بعد التدريب، قم بتحميل نقطة التحقق المُصدَّرة باستخدام AutoModelForCausalLM و load_state_dict. الرحلة الذهاب والعودة الناجحة لا تُبلّغ عن أي مفاتيح مفقودة أو غير متوقعة بخلاف lm_head.weight، وهو ما يربطه Qwen3-4B بتضمينات الرموز (token embeddings). يؤكد هذا أن الملفات تظل نقطة فحص (checkpoint) قياسية لـ Hugging Face.

الشكل 1: نقطة التحقق المنشورة تبدأ بخسارة قدرها 6.54، وتُحمَّل الأوزان المُصدَّرة دون مفاتيح مفقودة أو غير متوقعة

التدريب باستخدام النماذج الأصلية في torchtitan

في التكوين الثاني، يبني torchtitan نموذجه الأصلي Qwen3Model. أضف استيرادات Qwen3 الأصلية ودالة الإعداد إلى neuron_qwen3.py:

يبدأ هذا المثال من qwen3_32b() لأن نسخة torchtitan المُختبَرة لا توفر منشئ إعدادات عامًا لـ qwen3_4b() . السطر التالي يستبدل مواصفة نموذج 32B بمواصفة 4B. أما الحقول المتبقية فتعيد إنتاج التشغيل المُختبَر.

from torchtitan.models.qwen3 import model_registry
from torchtitan.models.qwen3.config_registry import qwen3_32b
from torchtitan.trainer import Trainer

def native_qwen3_4b() -> Trainer.Config:
    config = qwen3_32b()
    config.model_spec = model_registry("4B", attn_backend="sdpa")
    config.hf_assets_path = ASSETS
    config.dataloader.dataset = "c4_test"
    config.training.steps = 10
    config.training.local_batch_size = 1
    config.training.global_batch_size = 64
    config.parallelism.data_parallel_shard_degree = -1
    return config

torchrun --nproc_per_node=64 --rdzv_backend c10d --rdzv_endpoint localhost:0 \
  --local-ranks-filter 0 --role rank --tee 3 \
  -m torchtitan.train --module neuron_qwen3 --config native_qwen3_4b

config.model_spec يختار فئة النموذج ويحدد كيفية تقسيم torchtitan للنموذج (sharding) وتعيين مفاتيح نقطة التحقق الخاصة به. هنا، model_registry("4B", attn_backend="sdpa") يختار تنفيذ Qwen3 الأصلي في torchtitan مع الانتباه بضرب النقاط المُقاس (SDPA).

إضافة نوى NKI لمزيج الخبراء

في التكوين الثالث، يشغّل torchtitan طبقات مزيج الخبراء بنوى مكتوبة بلغة NKI، لغة نوى Neuron. أضف الاستيرادات المطلوبة ودالة إعداد Qwen3 MoE إلى neuron_qwen3.py:

from torchtitan.config import CompileConfig
from torchtitan_neuron.models.qwen3.config_registry import qwen3_nki_moe_config

def nki_qwen3_30b_a3b() -> Trainer.Config:
    config = qwen3_nki_moe_config(
        "30B-A3B",
        compile_config=CompileConfig(enable=False, backend="neuron"),
        n_layers=4,
        num_experts=64,
        top_k=4,
    )
    config.hf_assets_path = ASSETS
    return config

torchrun --nproc_per_node=64 --rdzv_backend c10d --rdzv_endpoint localhost:0 \
  --local-ranks-filter 0 --role rank --tee 3 \
  -m torchtitan.train --module neuron_qwen3 --config nki_qwen3_30b_a3b

يستخدم هذا التكوين تقسيم البيانات المتوازي عبر 64 رتبة (rank) مع درجة توازٍ موتري قدرها 1. ولا يستخدم توازي الخبراء. num_experts=64 يضبط عدد الخبراء في النموذج، وليس درجة توازي الخبراء.

ضبط CompileConfig(enable=False) يشغّل هذا المثال في الوضع المتلهف (eager). يدعم TorchNeuron معاملات NKI المسجّلة في كل من التنفيذ المتلهف و torch.compile.

في نسخة torchtitan المستخدمة في هذا الشرح، qwen3_nki_moe_config() يستبدل كتلة MoE بتنفيذ الخبراء المجمَّعين عبر NKI. كما يختار SDPA وتضمين الدوران الخالي من التجميع (gather-free). عند تمكين torch.compile، أضف "grad_clip" إلى CompileConfig.components لتجميع قصّ معيار التدرج (gradient-norm clipping).

تدخل نوى NKI التدريب كمعاملات PyTorch. wrap_nki تجعل النواة NKI قابلة للاستدعاء من PyTorch، و nki_op يسجّل المعامل. register_autograd يُرفق تمريرته الخلفية باستخدام بيانات تعريف التوجيه المحفوظة أثناء التمريرة الأمامية.

يُختزل هذا المثال Qwen3-30B-A3B إلى 4 طبقات و64 خبيرًا حتى يتسع على trn2.48xlarge. يوجّه كل رمز (token) إلى 4 خبراء باستخدام top_k=4.

أرفق نواة NKI الخاصة بك

في بعض الحالات، قد ترغب في تحسين عملية لتناسب أشكال الموترات في نموذجك. يمكن للنواة المتخصصة أن تُكيّف تقسيمها (tiling) وحركة بياناتها وفق تلك الأشكال. كما يمكنها استبدال تسلسل حرج للأداء من معاملات PyTorch بمعامل واحد محسَّن. يمكنك كتابة النواة بنفسك أو البدء بواحدة من nkilib. يستخدم هذا المثال rope_hf من nkilib.core.embeddings، والتي تتوقع تخطيط موترات Hugging Face وتدعم وضعها الخلفي الخاص. يوضح هذا المثال تسجيل النواة وتكاملها مع autograd، وليس تحسينًا للأداء.

إنشاء my_neuron_kernels.py في جذر المستودع، حيث تُطلق التدريب. يخزّن torchtitan الاستعلام (query) والمفتاح (key) بصيغة [batch, seq, heads, head_dim]بينما rope_hf يتوقع [batch, heads, seq, head_dim]. حوِّل التخطيط داخل المشغّل وخصص موترات المخرجات باستخدام torch.empty_like():

import torch
from nkilib.core.embeddings.rope_hf import rope_hf
from torch_neuronx import nki_op
from torch_neuronx.nki_hop import wrap_nki
from torch_neuronx.utils import get_logical_neuron_cores

LNC = (int(get_logical_neuron_cores()),)
wrapped_rope = wrap_nki(rope_hf)

@nki_op("my_kernels::rope_fwd")
def rope_fwd(query, key, rope_cache):
    q = query.permute(0, 2, 1, 3).contiguous()
    k = key.permute(0, 2, 1, 3).contiguous()
    cache = rope_cache.reshape(-1, rope_cache.shape[-1])
    q_out, k_out = wrapped_rope[LNC](
        q,
        k,
        torch.empty_like(q),
        torch.empty_like(k),
        rope_cache=cache,
    )
    return (
        q_out.permute(0, 2, 1, 3).contiguous(),
        k_out.permute(0, 2, 1, 3).contiguous(),
    )

rope_hf يحسب التدرجات عند استدعائه بـ backward=True. سجِّل مشغّلاً ثانيًا لذلك الوضع. احفظ ذاكرة rotary المؤقتة أثناء التمرير الأمامي حتى يتمكن التمرير الخلفي من إعادة استخدامها. أعد None لتدرج الذاكرة المؤقتة لأن الذاكرة المؤقتة غير قابلة للاشتقاق.

@nki_op("my_kernels::rope_bwd")
def rope_bwd(grad_query, grad_key, rope_cache):
    gq = grad_query.permute(0, 2, 1, 3).contiguous()
    gk = grad_key.permute(0, 2, 1, 3).contiguous()
    cache = rope_cache.reshape(-1, rope_cache.shape[-1])
    out_q, out_k = wrapped_rope[LNC](
        gq,
        gk,
        torch.empty_like(gq),
        torch.empty_like(gk),
        rope_cache=cache,
        backward=True,
    )
    return (
        out_q.permute(0, 2, 1, 3).contiguous(),
        out_k.permute(0, 2, 1, 3).contiguous(),
    )

def _setup_context(ctx, inputs, output):
    ctx.save_for_backward(inputs[2])

def _backward(ctx, grad_query, grad_key):
    (rope_cache,) = ctx.saved_tensors
    grad_q, grad_k = rope_bwd(grad_query, grad_key, rope_cache)
    return grad_q, grad_k, None

rope_fwd.register_autograd(_backward, setup_context=_setup_context)

لأن rope_hf يتطلب طول تسلسل مدعومً، تحقّق من ذلك القيد قبل استدعاء النواة (kernel). اشتق فئة من CosSinRoPE واستدعِ rope_fwd من apply_rotary_emb:

from dataclasses import dataclass
from torchtitan.config import derive, override
from torchtitan.models.common.rope import CosSinRoPE

class NKIRoPE(CosSinRoPE):
    @dataclass(kw_only=True, slots=True)
    class Config(CosSinRoPE.Config):
        pass

    @staticmethod
    def apply_rotary_emb(query, key, rope_cache):
        stride = 128 * LNC[0]
        if query.shape[1] % stride:
            raise ValueError(
                f"rope_hf requires a sequence length that is a multiple of "
                f"128 * LNC={stride}, got {query.shape[1]}"
            )
        q_out, k_out = rope_fwd(query, key, rope_cache)
        return q_out.type_as(query), k_out.type_as(key)

سجِّل NKIRoPE بديلاً عن CosSinRoPE. derive() ينسخ حقول الإعداد المشتركة إلى إعداد البديل:

@override(
    "nki_rope",
    target=CosSinRoPE.Config,
    exact=True,
    description="nkilib rope_hf kernel in place of the PyTorch rotation",
)
def _nki_rope(cfg: CosSinRoPE.Config) -> NKIRoPE.Config:
    return derive(cfg, NKIRoPE.Config)

لتحميل التجاوز المخصص، مرِّر اسم وحدته إلى --override.imports. يستورد torchtitan إدخال السجل قبل بناء النموذج:

torchrun --nproc_per_node=64 --rdzv_backend c10d --rdzv_endpoint localhost:0 \
  --local-ranks-filter 0 --role rank --tee 3 \
  -m torchtitan.train --module neuron_qwen3 --config native_qwen3_4b \
  --override.imports my_neuron_kernels

قبل بدء التدريب، يسجّل torchtitan كل عقدة إعداد يستبدلها التجاوز. تأكّد من أن المخرجات تحتوي على إدخال nki_rope واحد لكل طبقة من طبقات Qwen3:

[Override] nki_rope: model_spec.model.layers.0.attention.rope CosSinRoPEConfig -> NKIRoPE.Config
[Override] nki_rope: model_spec.model.layers.1.attention.rope CosSinRoPEConfig -> NKIRoPE.Config

اضبط طول التسلسل ليكون مضاعفاً من 128 * LNC، وهو 256 عندما تكون LNC تساوي 2. أعد تشكيل ذاكرة rotary المؤقتة إلى [seq, head_dim * 2]، مع قيم جيب التمام (cosine) في النصف الأول وقيم الجيب (sine) في النصف الثاني. تتحقق النواة من هذه القيود قبل بدء التدريب.

لمثال كامل على torchtitan، انظر torchtitan_neuron/models/gpt_oss/nki_attention.py. يربط نوى nkilib flash-attention عبر نفس آليات المشغّل والتجاوز. كما يعالج أيضاً مدخلات DTensor للتوازي الموتري.

قبل التدريب بنواة مخصصة، قارن مخرجاتها بالتنفيذ الذي تحل محله. في الاختبار العددي بمدخلات float32، rope_hf وأنتج التضمين الدوّار (rotary embedding) في PyTorch الخاص بـ torchtitan فرقاً مطلقاً أقصى قدره 0.0 لكل من الاستعلام والمفتاح.

عند تسجيل عدة تجاوزات، حدِّد كل واحد منها بمثيلات الوحدة المقصودة باستخدام fqns. على سبيل المثال، fqns=["layers.1[0-9].attention.rope"] يستهدف التضمين الدوّار في الطبقات من 10 إلى 19. يتوقف torchtitan عندما تطالب تجاوزان بنفس العقدة، ويحذّر عندما لا يطابق تجاوز أي عقد.

راجع مخرجات التدريب

عند بدء التشغيل، تأكد من أن torchtitan يُبلغ عن 64 رانكًا وعن النموذج المتوقع. تتضمن الخطوة الأولى الترجمة (ال compilation) وتستغرق وقتًا أطول من الخطوات اللاحقة. تُبلغ كل خطوة تدريب عن الخسارة ومعيار التدرج (gradient norm) واستخدام ذاكرة النطاق الترددي العالي (HBM) ومعدل النقل وقيمة model FLOPs utilization (MFU). في تشغيل نقطة الحفظ (checkpoint)، تأكد من أن خسارة الخطوة الأولى تقترب من 6.54. وفي تشغيل النواة المخصصة (custom kernel)، تأكد من أن nki_rope يظهر التجاوز (override) مرة واحدة لكل طبقة من طبقات Qwen3.

التوسع عبر المثيل

تنطبق إعدادات التوازي نفسها على تهيئات Transformers والتهيئة الأصلية للنموذج. يُعرِّض A trn2.48xlarge ‏64 نواة NeuronCore منطقية عند LNC=2. اضبط data_parallel_shard_degree = -1 لتقسيم معاملات النموذج والتدرجات وحالة المُحسِّن عبر كل رتبة (rank) يبدأها torchrun ‏.

اضبط config.parallelism.tensor_parallel_degree لتقسيم أوزان كل طبقة عبر مجموعة من الرتب. تحسب كل رتبة شريحة من كل عملية ضرب مصفوفات. مع تمكين توازي التسلسل افتراضيًا، تعمل طبقات التطبيع على التنشيطات المُقسَّمة على طول بُعد التسلسل. يجمع torchtitan بين توازي الموترات (tensor parallelism) وتوازي البيانات المُقسَّم بالكامل (FSDP) على DeviceMesh.

راقب سعة HBM لكل رتبة أثناء زيادة حجم النموذج أو طول التسلسل أو حجم الدفعة المحلي. تحدد هذه القيم مقدار الذاكرة التي تحتاجها كل رتبة.

التنظيف

إذا أردت الاحتفاظ بنقاط التحقق (checkpoints) الخاصة بك، فانسخها إلى وحدة تخزين دائمة قبل التنظيف. إذا كنت تخطط لإعادة استخدام المثيل، أوقفه. وعندما لم تعد بحاجة إلى المثيل، أنهِه.

الخلاصة

يعمل AWS Trainium كجهاز مدعوم من torchtitan عبر التهيئات الثلاث في هذا الدليل التوضيحي. يمكنك تدريب Qwen3 من خلال Hugging Face Transformers v5 وتحميل أو حفظ ملفات safetensors المنشورة مباشرة. يمكنك أيضًا التدريب باستخدام تنفيذ torchtitan الأصلي لـ Qwen3 عن طريق تعيين config.model_spec. أما بالنسبة لنماذج mixture-of-experts، فاستخدم وصفات NKI الخاصة بـ torchtitan-neuron. وللنوى المخصصة، استخدم سجل التجاوزات (override registry) لإرفاق عوامل NKI عبر مسار الإطلاق نفسه.

للحصول على إرشادات الإعداد، راجع وثائق AWS Neuron. أما للاطلاع على المصدر الحالي ووصفات التدريب، فراجع مستودع torchtitan.

عن المؤلفين