AKS で Ray を使用して LLM をトレーニングする

この記事では、 LLaMA-Factory と分散 Ray Train を使用して viggo NLG データセットに対して Qwen2.5-7B-Instruct を微調整する RayJob を送信します。 このジョブでは、それぞれ 1 つの GPU を持つ 4 つのワーカーを使用し、Azure Blob Storageからトレーニング データを読み取り、ダウンストリーム推論用にトレーニング済みの LoRA アダプターをアップロードします。

重要

オープンソース ソフトウェアは、AKS のドキュメントとサンプル全体で説明されています。 デプロイするソフトウェアは、AKS サービス レベル アグリーメント、限定保証、Azure サポートから除外されます。 AKS と共にオープンソース テクノロジを使用する場合は、それぞれのコミュニティとプロジェクト保守担当者から受けられるサポート オプションを調べ、計画を策定してください。

Microsoft は、AKS 上に展開するオープンソース パッケージを構築する責任を負います。 その責任には、ビルド、スキャン、署名、検証、修正プログラム プロセスの完全な所有権と、コンテナー イメージ内のバイナリの制御権が伴います。 詳細については、「AKS の脆弱性の管理」と「AKS のサポート範囲」を参照してください。

前提条件

環境変数の設定

複製されたリポジトリの LLM トレーニングの例に移動し、必要な環境変数を構成します。

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

チーム キュー (オプション B) を構成した場合は、export QUEUE_NAME=team-aを実行する前にexport QUEUE_NAME=team-bまたはsource env.exampleを設定します。 既定の QUEUE_NAME=default は、単一キュー構成 (オプション A) でのみ機能します。

ワークロードを送信する

分散トレーニング RayJob を送信します。

./submit.sh

このスクリプトは、トレーニング スクリプトから ConfigMap を作成し、 envsubstを介して 4 つの GPU ワーカーでマニフェスト テンプレートをレンダリングし、それを適用します。 構成されたキューで 4 つの GPU が使用可能な場合、Kueue はジョブを許可します。

Tip

./submit.sh --dry-runを実行して、レンダリングされたマニフェストをクラスターに適用せずに検証します。

進行状況の監視

新しいシェルを使用している場合は、ジョブ名を検索してエクスポートします。

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

RayJob の状態と Kueue の受付状況を確認します。

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

ジョブの完了時に予想される出力:

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

ヘッド ポッドのログを追跡する:

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

結果を確認する

LoRA アダプターのアップロードを確認します。

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

予想される出力:

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

latest.txt ポインターが書き込まれた (自動検出のためにバッチ推論の例で使用) を確認します。

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

予想される出力:

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

構成のリファレンス

Variable Default Description
AZURE_STORAGE_ACCOUNT_NAME (必須) モジュール 1 のストレージ アカウント
NUM_WORKERS 4 GPU ワーカー レプリカ (各 4 x 1 GPU)
QUEUE_NAME default Kueue LocalQueue の名前
LLM_DATA_CONTAINER llm-pipeline 入力データ用の BLOB コンテナー
LLM_LORA_CONTAINER llm-pipeline LoRA アップロード用の BLOB コンテナー
CONFIGMAP_NAME llm-training-scripts トレーニング スクリプトを保持する ConfigMap の名前

リソースをクリーンアップする

RayJob とその ConfigMap を削除します。

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

次のステップ