Entrenar un LLM con Ray en AKS

En este artículo, se envía un RayJob que ajusta Qwen2.5-7B-Instruct en el conjunto de datos viggo NLG utilizando LLaMA-Factory y el entrenamiento distribuido de Ray. El trabajo utiliza cuatro trabajadores con una GPU cada uno, lee los datos de entrenamiento desde Azure Blob Storage y carga el adaptador LoRA entrenado para la inferencia posterior.

Prerequisites

• Infraestructura implementada siguiendo Implementar la infraestructura para Ray y Kueue en AKS.
• Las colas de Kueue se configuran siguiendo Configurar colas de Kueue para cargas de trabajo de Ray en AKS.
• Al menos cuatro GPU A100 disponibles en el clúster (en este ejemplo se utilizan cuatro trabajadores con una GPU cada uno).
• Conjunto de datos Viggo cargado en el almacenamiento de blobs en llm-pipeline/data/ (se realiza automáticamente mediante el módulo de infraestructura de Terraform).
• envsubst instalado (gettext paquete en Linux, brew install gettext en macOS).

Establecimiento de variables de entorno

Vaya al ejemplo de entrenamiento de LLM en el repositorio clonado y configure las variables de entorno necesarias:

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

Enviar la carga de trabajo

Envíe el entrenamiento distribuido RayJob:

./submit.sh

El script crea un ConfigMap a partir del script de entrenamiento, genera la plantilla de manifiesto con cuatro trabajadores de GPU mediante envsubst y la aplica. Kueue admite el trabajo cuando hay cuatro GPU disponibles en la cola configurada.

Supervisión de progreso

Busca y exporta el nombre del trabajo si te encuentras en un nuevo shell:

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

Vea el estado de RayJob y la admisión de Kueue:

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

Salida esperada cuando se completa el trabajo:

NAME                      JOB STATUS   DEPLOYMENT STATUS   START TIME             END TIME               AGE
llm-training-xxxxxxxxxx   SUCCEEDED    Complete            2026-01-01T00:00:00Z   2026-01-01T00:19:00Z   19m

NAME                                   QUEUE     RESERVED IN     ADMITTED   FINISHED   AGE
rayjob-llm-training-xxxxxxxxxx-xxxxx   default   cluster-queue   True       True       19m

Sigue los registros del pod principal:

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

Comprobación de los resultados

Compruebe la subida del adaptador LoRA:

az storage blob list -c llm-pipeline --prefix lora/ \
  --account-name ${AZURE_STORAGE_ACCOUNT_NAME} --auth-mode login -o table

Resultado esperado:

Name                                       Blob Type    Blob Tier    Length    Content Type
-----------------------------------------  -----------  -----------  --------  ------------------------
lora/latest.txt                            BlockBlob    Hot          86        application/octet-stream
lora/<job-name>/rng_state_3.pth            BlockBlob    Hot          14725     application/octet-stream

Compruebe que se haya escrito el puntero latest.txt (utilizado por el ejemplo de inferencia por lotes para el autodescubrimiento):

az storage blob download -c llm-pipeline -n lora/latest.txt \
  --account-name ${AZURE_STORAGE_ACCOUNT_NAME} --auth-mode login

Resultado esperado:

azure://llm-pipeline@<storage-account>.blob.core.windows.net/lora/<job-name>

Referencia de configuración

Limpieza de recursos

Elimine RayJob y su ConfigMap:

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

Pasos siguientes