AWSのPhysical AI Scaffolding Kit (PASK)を試す② — HyperPodでNVIDIA Isaac GR00Tをファインチューニング

連載第2回。第1回で構築したAmazon SageMaker HyperPodのSlurmクラスタにGPUワーカーを追加し、学習サンプルsamples/gr00tでNVIDIA Isaac GR00T-N1.6-3Bをファインチューニングします。手順どおり進めれば素直に動きますが、最初に押さえておくと詰まらないポイントが3つあるので、そこを中心にまとめます。

前回はこちら
https://dev.classmethod.jp/articles/try-aws-pask-01-build-slurm-cluster-with-hyperpod/

はじめに

第1回ではHyperPodのSlurmクラスタ(GPUなし)を構築しました。第2回ではGPUワーカーを足し、samples/gr00tでGR00T-N1.6-3Bをcube_to_bowlタスクのデモデータでファインチューニングします。GR00Tは「カメラ画像 + 言語指示」から「ロボットの動作」を出力するVLA(Vision-Language-Action)モデルです。

最初に押さえる3つのポイント

ここだけ先に押さえておけば、あとは手順どおりで通ります。

① GPUは「単一GPUで40GB+VRAM」を選ぶ

GR00T公式の要件は 「40GB+ VRAM推奨、H100 or L40 nodes」(1枚あたり)。手頃なのは L40S 48GBのml.g6e.4xlarge。24GB級(A10G/L4)だと学習時にoptimizer stateが載らずGPUのOOMになります。なお g6(L4 24GB)とg6e(L40S 48GB)は別物なので、選ぶのは「e付き」のg6eです。

② Isaac-GR00Tはn1.6-releaseにcheckoutする

公式のgit clone手順だとmainブランチ(=N1.7のコード)が落ちてきますが、サンプルの既定モデルはN1.6です。そのまま学習すると:

ValueError: model type `Gr00tN1d6` but Transformers does not recognize this architecture

となるので、cloneのあとにgit checkout n1.6-releaseでコードをN1.6に合わせます。

③ 出力先の書き込み権限を付けておく

学習のOUTPUT_DIR既定が/fsx/s3link/so100(S3リンク領域)なので、ubuntuが書けるよう事前にsudo chmod -R a+w /fsx/s3linkしておきます(公式手順にも記載あり)。

全体像: 4つのフェーズ

このサンプルは、HPC環境でよく使われるコンテナ実行の流れ(Dockerでイメージを作り、Enrootで変換し、Slurmジョブとして動かす)です。

[Login Node]──docker build──►[ECR]      ① イメージを作りECRに保管
[Login Node]◄─enroot import──[ECR]      ② .sqsh形式に変換
[Login Node]──sbatch────────►[GPU Node]  ③ Slurmジョブとして学習
        すべて [FSx /fsx] を共有ストレージとして利用

なぜEnroot/Pyxisか: 通常のDockerはroot権限が必要でHPC共有環境では使いにくいため、Enroot(rootlessコンテナ)で.sqsh(SquashFS)に変換し、Pyxis(Slurmプラグイン)経由でsrun --container-image=...として計算ノードで動かします。.sqshもデータもログも全ノード共有のFSxに置くことで、どのノードからでも同じものが見えます。

手順

Phase 0: GPUワーカーを追加

hyperpod/cdk.jsonのWorkerGroupにL40S 48GBのg6e.4xlargeを追加してcdk deployします(事前にクオータの確認と、必要に応じて申請をしておきましょう)。

"WorkerGroup": [
  { "Name": "worker-group-1", "Count": 1, "InstanceType": "ml.g6e.4xlarge" }
]

export AWS_PROFILE=pask
cd hyperpod
cdk deploy

sinfo(ログインノード上)でGPUノードがidleになればOKです。

Phase 1: ログインノードでの準備

ssh pask-clusterで入り、ubuntuユーザーで実行します。

cd
git clone https://github.com/aws-samples/sample-physical-ai-scaffolding-kit.git
# Git LFS
curl -s https://packagecloud.io/install/repositories/github/git-lfs/script.deb.sh | sudo bash
sudo apt-get update && sudo apt-get install -y git-lfs
git lfs install
# ジョブログ
mkdir -p /fsx/ubuntu/joblog
# Isaac-GR00T(サブモジュール込み) → N1.6に合わせる(ポイント②)
cd
git clone --recurse-submodules https://github.com/NVIDIA/Isaac-GR00T
echo 'export GR00T_HOME="$HOME/Isaac-GR00T"' >> ~/.bashrc
source ~/.bashrc
cd ~/Isaac-GR00T
git checkout n1.6-release        # gr00t/model/ に gr00t_n1d6 が現れる

Phase 2: Dockerビルド & ECRプッシュ

cd ~/sample-physical-ai-scaffolding-kit/samples/gr00t/training
sbatch slurm_build_docker.sh
squeue   # gr00t_do が R なら実行中

pytorch3d/flash-attnのソースビルドで30分〜1hかかります。完了の合図はdocker_build_<JOBID>.out末尾のDocker build and push completed successfully。

Phase 3: Enrootインポート

この処理はフォアグラウンドで動き続けるため、SSHが切れても中断しないようtmuxの中で実行します(あとからtmux attachで戻れます)。

tmux new -s import
cd ~/sample-physical-ai-scaffolding-kit/samples/gr00t/training
bash ./hyperpod_import_container.sh
# 完了確認
export ENROOT_DATA_PATH=/fsx/enroot/data
enroot list   # gr00t-train+latest.sqsh が出ればOK

Phase 4: 学習

sudo chmod -R a+w /fsx/s3link          # ポイント③: 出力先の書き込み権限
cd ~/sample-physical-ai-scaffolding-kit/samples/gr00t/training
export GR00T_HOME="$HOME/Isaac-GR00T"
sbatch slurm_finetune_container.sh     # NUM_GPUS指定なし=1。L40S 48GBで3Bが載る

進捗とログ:

squeue                                            # gr00t_fi が R
tail -f $(ls -t /fsx/ubuntu/joblog/finetune_*.out | head -1)

ログに{'loss': 0.03..., 'learning_rate': ...}が回り始めれば学習中です。step 2000/2000で完了し、/fsx/s3link/so100/checkpoint-2000/にモデル(model-*.safetensors)が保存されます。

チェックポイントはS3に残る

OUTPUT_DIRの/fsx/s3linkはS3リンク領域なので、保存したチェックポイントはS3に自動同期されます。

aws s3 ls s3://<データバケット名>/so100/checkpoint-2000/ --recursive --profile pask
# model-00001-of-00002.safetensors / model-00002-of-00002.safetensors / ...

クラスタをcdk destroyしてもFSxは消えますが、**S3データバケットはRETAIN(保持)**なので、チェックポイントは残ります。学習(高VRAM GPUが必要)と評価(Isaac Sim環境)を、S3を介して別構成・別タイミングでやれるのがPASKの設計の良いところです。

後片付け(課金停止)

export AWS_PROFILE=pask     # プロファイルを通したシェルで
cd hyperpod
cdk destroy

GPUワーカーは稼働している限り課金されるので、学習が終わったら忘れずにdestroyしておきましょう。チェックポイントはS3に残るので、安心してdestroyできます。

まとめと次回

GR00T-N1.6-3Bのファインチューニングを完走し、チェックポイントをS3に保存しました。要点は冒頭の3つ(単一40GB+のGPU / n1.6-release / OUTPUT_DIR権限)だけ押さえれば、あとは公式の手順どおりで通ります。

次回は、このS3のチェックポイントを使って、isaacsim-workstationでロボットの挙動を確認します。