Optimieren Sie das Aurora-Wettermodell mit Ray auf AKS

In diesem Artikel übermitteln Sie einen RayJob, der das Microsoft Aurora Weather Foundation-Modell auf regionalen WeatherBench2-Daten mit LoRA fein einstimmt. Der Einzelvorgang wird auf einer einzelnen A100-GPU ausgeführt, Kueue steuert die Zulassung, und der Einzelvorgang schreibt den Adapter-Checkpoint und die Trainingsmetriken nach Azure Blob Storage.

Voraussetzungen

• Infrastruktur, die gemäß Infrastruktur für Ray und Kueue auf AKS bereitstellen bereitgestellt wurde.
• Kueue-Warteschlangen, die nach dem Konfigurieren von Kueue-Warteschlangen für Ray-Workloads auf AKS konfiguriert wurden.
• Mindestens eine A100-GPU im Cluster (in diesem Beispiel wird eine GPU verwendet).
• Aurora WeatherBench2-Daten, die in Blob Storage hochgeladen wurden (das Infrastruktur-Terraform-Modul lädt diese Daten automatisch hoch).
• envsubst installiert (gettext Paket unter Linux, brew install gettext unter macOS).

Festlegen von Umgebungsvariablen

Navigieren Sie zum Aurora-Feinabstimmungsbeispiel im geklonten Repository, und konfigurieren Sie die erforderlichen Umgebungsvariablen:

cd <path-to-cloned-repo>/AKS/examples/kueue-and-ray-on-aks/3-workloads/aurora-finetune
export AZURE_STORAGE_ACCOUNT_NAME=$(terraform -chdir=../../1-infrastructure/terraform output -raw storage_account_name)
source env.example

Die env.example Datei legt Standardwerte fest, einschließlich des Ray-Bilds, des Warteschlangennamens und der Schulungsparameter. Die JOB_NAME Variable wird automatisch mit einem Zeitstempel generiert.

Workload übermitteln

Übermitteln Sie den RayJob an den Cluster:

./submit.sh

Das Skript erstellt eine ConfigMap aus dem Python Schulungsskript, rendert die Manifestvorlage über envsubstund wendet sie auf den Cluster an. Kueue lässt den Job zu, wenn in der konfigurierten Warteschlange ein GPU-Kontingent verfügbar ist.

Fortschritt überwachen

Suchen und exportieren Sie den Auftragsnamen, wenn Sie sich in einer neuen Shell befinden:

export JOB_NAME=$(kubectl -n ray get rayjob --no-headers -o custom-columns=":metadata.name" | grep aurora-finetune)

Überwachen Sie den RayJob-Status und die Kueue-Zulassung:

kubectl -n ray get rayjob ${JOB_NAME} -w
kubectl -n ray get workload -w

Erwartete Ausgabe, wenn der Einzelvorgang abgeschlossen ist:

NAME                         JOB STATUS   DEPLOYMENT STATUS   START TIME             END TIME               AGE
aurora-finetune-xxxxxxxxxx   SUCCEEDED    Complete            2026-01-01T00:00:00Z   2026-01-01T00:08:00Z   8m

NAME                                      QUEUE     RESERVED IN     ADMITTED   FINISHED   AGE
rayjob-aurora-finetune-xxxxxxxxxx-xxxxx   default   cluster-queue   True       True       8m

Arbeitsprotokolle verfolgen:

kubectl -n ray logs -l ray.io/cluster=$(kubectl -n ray get rayjob ${JOB_NAME} -o jsonpath='{.status.rayClusterName}') -f

Überprüfen der Ergebnisse

Überprüfen Sie die hochgeladenen Artefakte im BLOB-Speicher:

az storage blob list -c aurora --prefix checkpoints/${JOB_NAME}/ \
  --account-name ${AZURE_STORAGE_ACCOUNT_NAME} --auth-mode login -o table

Erwartete Ausgabe:

Name                                                    Blob Type    Blob Tier    Length    Content Type
------------------------------------------------------  -----------  -----------  --------  ------------------------
checkpoints/<job-name>/last.safetensors                 BlockBlob    Hot          11432328  application/octet-stream
checkpoints/<job-name>/train-metrics.json               BlockBlob    Hot          985       application/json

Laden Sie die Schulungsmetriken herunter, und überprüfen Sie sie:

az storage blob download -c aurora \
  -n checkpoints/${JOB_NAME}/train-metrics.json \
  --account-name ${AZURE_STORAGE_ACCOUNT_NAME} --auth-mode login \
  --file /tmp/train-metrics.json
cat /tmp/train-metrics.json | python3 -m json.tool

Erwartete Ausgabe:

{
    "final_loss": 24051.97,
    "initial_loss": 24183.10,
    "loss_history": [24183.10],
    "loss_improvement": 131.13,
    "max_steps": 1,
    "trainable_parameters": 2850816,
    "gpu_name": "NVIDIA A100-SXM4-80GB",
    ...
}

Das loss_history Array sollte endliche Werte (nicht NaN) enthalten, um sicherzustellen, dass der Datenpfad und die Modellschulung ordnungsgemäß funktionieren.

Konfigurationsreferenz

Bereinigen von Ressourcen

Löschen Sie den RayJob und dessen ConfigMap:

kubectl -n ray delete rayjob ${JOB_NAME}
kubectl -n ray delete configmap ${CONFIGMAP_NAME}

Nächste Schritte

• Trainieren eines LLM mit Ray
• Ausführen von Batch-Rückschlüssen mit Ray