Qwen3 auf AWS Trainium mit torchtitan trainieren

Die Verwendung von torchtitan auf AWS Trainium bedeutet, es so zu nutzen, wie Sie es immer getan haben. TorchNeuron registriert Trainium als natives PyTorch-Gerät, wodurch torchtitan es über Standard-PyTorch-Schnittstellen verwenden kann. Sie können diese Konfigurationen verwenden, um von einer zufälligen Initialisierung aus vorzutrainieren oder veröffentlichte Gewichte feinabzustimmen. Sie können auch die Modellimplementierung und die Kernel für Ihre Workload anpassen. In diesem Beitrag trainieren Sie Qwen3 auf drei Arten auf einer einzelnen trn2.48xlarge.

Die erste Methode verwendet Transformers v5 als Modeling-Backend in torchtitan. Sie führt Hub-Modelle aus über AutoModelForCausalLM und liest und schreibt veröffentlichte safetensors direkt. Die zweite Methode wechselt zu torchtitan’s nativem Qwen3 durch Änderung von config.model_spec. Die dritte fügt Neuron Kernel Interface (NKI) Mixture-of-Experts-Kernel (MoE) aus den torchtitan-neuron-Rezepten hinzu und verwendet dann torchtitan’s Override-Registry, um Ihren eigenen Kernel anzuhängen. Alle drei Konfigurationen verwenden dasselbe torchrun Launch-Muster und dieselbe Konfigurationsstruktur.

Voraussetzungen

Um diesem Walkthrough zu folgen, verwenden Sie eine einzelne trn2.48xlarge mit entweder dem AWS Deep Learning AMI für Neuron oder einem Neuron PyTorch Training Container. Beide Umgebungen bieten den Neuron-Treiber, die Laufzeit und den Compiler, die für das Training auf Trainium erforderlich sind.

Bevor Sie torchtitan-neuroninstallieren, prüfen Sie, ob Ihre Umgebung es bereits enthält. Der folgende Befehl gibt die installierte Version und den Paketpfad aus:

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

Dieses Walkthrough verwendet die Versionen in der folgenden Tabelle. Um die Beispiele zu reproduzieren, fixieren Sie diese Versionen. Die Transformers v5-Integration und die Trainium-spezifischen Override-Module können sich zwischen torchtitan-Versionen ändern.

Wenn torch_neuronx installiert ist, registriert TorchNeuron Trainium als neuron Gerät, und torchtitan erkennt es beim Start. Im Eager-Modus kompiliert Neuron Operationen, während sie ausgeführt werden. Mit torch.compilekompiliert Neuron den erfassten Graphen. Dasselbe Launch-Muster funktioniert über alle drei Konfigurationen hinweg.

Für die in diesem Walkthrough verwendeten Versionen installieren Sie torchtitan-neuron und torchtitan’s übrige Abhängigkeiten mit --no-deps um den Neuron-Build von torch zu erhalten:

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

Bevor Sie mit dem Training beginnen, holen Sie die Qwen3-Modellartefakte aus dem Hugging Face Hub. Speichern Sie sie in einem lokalen Verzeichnis, das Sie an die Trainingskonfiguration übergeben können:

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

Jede Trainingskonfiguration ist eine Python-Konfigurationsfunktion. Die drei Beispiele gehen von einer Basiskonfiguration aus, die Sie für jeden Schritt anpassen.

Trainieren mit Hugging Face Transformers v5

In der ersten Konfiguration baut Transformers v5 Qwen3 durch AutoModelForCausalLMauf. Die Hub-Modelldefinition liefert die Implementierung für die Forward- und Backward-Pässe. Sie wählen das Modell aus, indem Sie seine Hub-ID in der Konfiguration setzen.

Beginnen Sie mit transformers_modeling_backend_full(). Es lässt die Felder, die den Hugging-Face-Konfigurationsschlüsseln entsprechen, ungesetzt, sodass AutoConfig.from_pretrained() die Modellarchitektur aus dem Hub lädt.

Erstellen Sie neuron_qwen3.py im Stammverzeichnis des torchtitan-neuron-Repositorys, in dem Sie den Trainingsbefehl ausführen. Fügen Sie die Importe, den Pfad zu den Modellartefakten und die Basiskonfiguration hinzu:

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

Dieser Befehl startet 64 Trainingsprozesse, einen für jeden logischen NeuronCore. --module neuron_qwen3 importiert das lokale Konfigurationsmodul. --config v5_qwen3_4b wählt die Konfigurationsfunktion aus. Die Data-Parallel-Einstellung shardet das Modell über alle 64 Ranks.

Hugging-Face-Checkpoints laden und speichern

Um von veröffentlichten Gewichten aus zu fine-tunen, aktivieren Sie das Laden von Checkpoints und verweisen initial_load_path auf ASSETS. Der State-Dictionary-Adapter passt das führende model. -Präfix an und behandelt geteilte Embeddings. Er speichert den finalen Checkpoint als 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

Verwenden Sie den Trainingsverlust des ersten Schritts als schnelle Prüfung, ob torchtitan die veröffentlichten Gewichte geladen hat. Für Qwen/Qwen3-4B-Instruct-2507erzeugte der verifizierte Lauf einen Loss von 6.54 in Schritt 1 mit veröffentlichten Gewichten und 12.43 bei zufälliger Initialisierung.

Laden Sie nach dem Training den exportierten Checkpoint mit AutoModelForCausalLM und load_state_dict. Ein erfolgreicher Roundtrip meldet keine fehlenden oder unerwarteten Schlüssel außer lm_head.weight, das Qwen3-4B an die Token-Embeddings koppelt. Dies bestätigt, dass die Dateien ein standardmäßiger Hugging-Face-Checkpoint bleiben.

Abbildung 1: Der veröffentlichte Checkpoint startet mit einem Loss von 6.54, und die exportierten Gewichte laden ohne fehlende oder unerwartete Keys

Trainieren Sie mit torchtitans nativen Modellen

In der zweiten Konfiguration baut torchtitan sein natives Qwen3Model. Fügen Sie die nativen Qwen3-Importe und die Config-Funktion hinzu zu neuron_qwen3.py:

Dieses Beispiel beginnt bei qwen3_32b() weil der getestete torchtitan-Build keinen generischen qwen3_4b() Config-Builder bereitstellt. Die nächste Zeile ersetzt die 32B-Modellspezifikation durch die 4B-Spezifikation. Die übrigen Felder reproduzieren den getesteten Lauf.

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 wählt die Modellklasse und definiert, wie torchtitan das Modell shardet und seine Checkpoint-Keys abbildet. Hier, model_registry("4B", attn_backend="sdpa") wählt torchtitans native Qwen3-Implementierung mit scaled dot product attention (SDPA).

Fügen Sie NKI-Kernel für Mixture of Experts hinzu

In der dritten Konfiguration führt torchtitan Mixture-of-Experts-Schichten mit Kerneln aus, die in NKI geschrieben sind, der Neuron-Kernelsprache. Fügen Sie die erforderlichen Importe und die Qwen3-MoE-Config-Funktion hinzu zu 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

Diese Konfiguration verwendet Data-Parallel-Sharding über 64 Ranks mit einem Tensor-Parallel-Grad von 1. Sie verwendet keine Expert Parallelism. num_experts=64 legt die Anzahl der Experten im Modell fest, nicht den Expert-Parallel-Grad.

Das Setzen von CompileConfig(enable=False) führt dieses Beispiel im Eager-Modus aus. TorchNeuron unterstützt registrierte NKI-Operatoren sowohl in der Eager-Ausführung als auch torch.compile.

Im torchtitan-Build, der für diese Anleitung verwendet wurde, qwen3_nki_moe_config() ersetzt den MoE-Block durch die NKI-Grouped-Experts-Implementierung. Er wählt außerdem SDPA und das gather-freie Rotary Embedding. Wenn Sie torch.compileaktivieren, fügen Sie "grad_clip" zu CompileConfig.components hinzu, um Gradient-Norm-Clipping zu kompilieren.

NKI-Kernel gelangen als PyTorch-Operatoren ins Training. wrap_nki macht einen NKI-Kernel aus PyTorch aufrufbar, und nki_op registriert den Operator. register_autograd verknüpft seinen Backward-Pass unter Verwendung von Routing-Metadaten, die während des Forward-Pass gespeichert wurden.

Das Beispiel reduziert Qwen3-30B-A3B auf 4 Schichten und 64 Experten, damit es auf einen trn2.48xlargepasst. Es routet jeden Token an 4 Experten mit top_k=4.

Binden Sie Ihren eigenen NKI-Kernel ein

In einigen Fällen möchten Sie möglicherweise eine Operation für die Tensorformen in Ihrem Modell optimieren. Ein spezialisierter Kernel kann sein Tiling und seine Datenbewegung auf diese Formen zuschneiden. Er kann außerdem eine leistungskritische Sequenz von PyTorch-Operationen durch einen einzigen optimierten Operator ersetzen. Sie können den Kernel selbst schreiben oder mit einem aus nkilibbeginnen. Dieses Beispiel verwendet rope_hf aus nkilib.core.embeddings, das das Hugging Face Tensor-Layout erwartet und seinen eigenen Backward-Modus unterstützt. Dieses Beispiel demonstriert die Kernel-Registrierung und die Autograd-Integration, keine Leistungsverbesserung.

Erstellen Sie my_neuron_kernels.py im Repository-Stammverzeichnis, von wo aus du das Training startest. torchtitan speichert query und key als [batch, seq, heads, head_dim], während rope_hf Folgendes erwartet [batch, heads, seq, head_dim]. Konvertiere das Layout innerhalb des Operators und alloziere die Ausgabe-Tensoren mit 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 berechnet Gradienten, wenn du es mit backward=Trueaufrufst. Registriere einen zweiten Operator für diesen Modus. Speichere den rotary-Cache während des Forward-Pass, damit der Backward-Pass ihn wiederverwenden kann. Gib None für den Cache-Gradienten zurück, da der Cache nicht differenzierbar ist.

@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)

Da rope_hf eine unterstützte Sequenzlänge benötigt, prüfe diese Einschränkung, bevor du den Kernel aufrufst. Leite von CosSinRoPE ab und rufe rope_fwd aus 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)

auf. Registriere NKIRoPE als Ersatz für CosSinRoPE. derive() kopiert die gemeinsamen Konfigurationsfelder in die Ersatzkonfiguration:

@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)

Um den benutzerdefinierten Override zu laden, übergib seinen Modulnamen an --override.imports. torchtitan importiert den Registry-Eintrag, bevor es das Modell erstellt:

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

Bevor das Training beginnt, protokolliert torchtitan jeden Konfigurationsknoten, den der Override ersetzt. Stelle sicher, dass die Ausgabe einen nki_rope -Eintrag für jede Qwen3-Schicht enthält:

[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

Setze die Sequenzlänge auf ein Vielfaches von 128 * LNC, was 256 ist, wenn LNC 2 ist. Forme den rotary-Cache zu [seq, head_dim * 2]um, mit Kosinuswerten in der ersten Hälfte und Sinuswerten in der zweiten. Der Kernel prüft diese Einschränkungen, bevor das Training beginnt.

Für ein vollständiges torchtitan-Beispiel siehe torchtitan_neuron/models/gpt_oss/nki_attention.py. Es verbindet die nkilib -flash-attention-Kernels über dieselben Mechanismen für Operatoren und Overrides. Es behandelt außerdem DTensor -Eingaben für Tensorparallelität.

Bevor du mit einem benutzerdefinierten Kernel trainierst, vergleiche seine Ausgaben mit der Implementierung, die er ersetzt. Im numerischen Test mit float32-Eingaben erzeugten rope_hf und torchtitan’s PyTorch-rotary-embedding eine maximale absolute Differenz von 0.0 sowohl für query als auch für key.

Wenn du mehrere Overrides registrierst, beschränke jeden einzelnen auf die beabsichtigten Modulinstanzen mit fqns. Zum Beispiel zielt fqns=["layers.1[0-9].attention.rope"] auf das rotary-embedding in den Schichten 10 bis 19. torchtitan bricht ab, wenn zwei Overrides denselben Knoten beanspruchen, und warnt, wenn ein Override keinen Knoten trifft.

Überprüfe die Trainingsausgabe

Bestätigen Sie beim Start, dass torchtitan 64 Ranks und das erwartete Modell meldet. Der erste Schritt umfasst die Kompilierung und dauert länger als die späteren Schritte. Jeder Trainingsschritt meldet den Loss, die Gradientennorm, die Auslastung des High-Bandwidth-Memory (HBM), den Durchsatz und die Model-FLOPs-Utilization (MFU). Bestätigen Sie beim Checkpoint-Run, dass der Loss des ersten Schritts nahe 6.54 liegt. Bestätigen Sie beim Custom-Kernel-Run, dass das nki_rope Override einmal für jede Qwen3-Schicht erscheint.

Skalierung über die Instanz hinweg

Dieselben Parallelisierungseinstellungen gelten für die Transformers- und die native Modellkonfiguration. Eine trn2.48xlarge stellt 64 logische NeuronCores bei LNC=2 bereit. Setzen Sie data_parallel_shard_degree = -1 um die Modellparameter, Gradienten und den Optimizer-State über jeden Rank zu sharden, den torchrun startet.

Setzen Sie config.parallelism.tensor_parallel_degree um die Gewichte jeder Schicht über eine Gruppe von Ranks zu sharden. Jeder Rank berechnet ein Segment jeder Matrixmultiplikation. Mit standardmäßig aktivierter Sequence Parallelism laufen Normalisierungsschichten auf Aktivierungen, die entlang der Sequenzdimension geshardet sind. torchtitan kombiniert Tensor Parallelism mit Fully Sharded Data Parallel (FSDP) auf einem mehrdimensionalen DeviceMesh.

Beobachten Sie das HBM pro Rank, wenn Sie die Modellgröße, die Sequenzlänge oder die lokale Batch-Größe erhöhen. Diese Werte bestimmen, wie viel Speicher jeder Rank benötigt.

Aufräumen

Wenn Sie Ihre Checkpoints behalten möchten, kopieren Sie sie in einen persistenten Speicher, bevor Sie aufräumen. Wenn Sie die Instanz wiederverwenden möchten, stoppen Sie sie. Wenn Sie die Instanz nicht mehr benötigen, beenden Sie sie.

Fazit

AWS Trainium fungiert als unterstütztes torchtitan-Gerät in allen drei Konfigurationen dieser Anleitung. Sie können Qwen3 über Hugging Face Transformers v5 trainieren und veröffentlichte safetensors direkt laden oder speichern. Sie können auch mit torchtitans nativer Qwen3-Implementierung trainieren, indem Sie config.model_specsetzen. Verwenden Sie für Mixture-of-Experts-Modelle die NKI-Rezepte von torchtitan-neuron. Verwenden Sie für Custom Kernels die Override-Registry, um NKI-Operatoren über denselben Start-Ablauf anzuhängen.

Hinweise zur Einrichtung finden Sie in der AWS Neuron documentation. Den aktuellen Quellcode und Trainingsrezepte finden Sie im torchtitan-Repository.

Über die Autoren