OpenVLA

OpenVLA distributed training on AWS

Fine-tune OpenVLA-7B (Vision-Language-Action model) on the LIBERO robotics benchmark using LoRA, distributed across 8 GPUs on a SageMaker HyperPod Slurm cluster.

Prerequisites

• SageMaker HyperPod cluster with Slurm scheduler, Pyxis/Enroot, and P5/P5en GPU nodes
• FSx for Lustre shared filesystem mounted at /fsx
• Docker installed (for building the container image)
• SSH access to the cluster login node

1. Clone this repository

git clone https://github.com/awslabs/awsome-distributed-ai.git
cd awsome-distributed-ai/examples/use-cases/openvla

2. Build the container

docker build -t openvla-finetune -f openvla.Dockerfile .

3. Import as Enroot squashfs image

Transfer to the cluster and import:

enroot import -o /fsx/$USER/vla/openvla-finetune.sqsh dockerd://openvla-finetune:latest

4. Set up the cluster workspace

SSH into your cluster:

export VLA_HOME=/fsx/$USER/vla
mkdir -p $VLA_HOME/{models,data,checkpoints,logs,scripts}

5. Download model weights

python3 -c "
from huggingface_hub import snapshot_download
snapshot_download('openvla/openvla-7b', local_dir='$VLA_HOME/models/openvla-7b')
"

6. Download the LIBERO dataset

LIBERO RLDS datasets are distributed pre-built on the Hugging Face Hub (they are not in the stock TFDS catalog). Requires git-lfs:

git lfs install
git clone https://huggingface.co/datasets/openvla/modified_libero_rlds "$VLA_HOME/data/libero_rlds"

This yields $VLA_HOME/data/libero_rlds/libero_10_no_noops/ (the layout expected by the sbatch --data_root_dir + --dataset_name arguments).

7. Launch training

Submit from $VLA_HOME — Slurm writes logs relative to the submission directory:

cd $VLA_HOME
cp ~/awsome-distributed-ai/examples/use-cases/openvla/slurm/finetune_openvla.sbatch scripts/
sbatch scripts/finetune_openvla.sbatch

Monitor progress:

squeue -u $USER
tail -f logs/<JOB_ID>.out

Expected: ~10 minutes for 500 steps on 1 node (8x H200). Checkpoints saved to $VLA_HOME/checkpoints/run_<JOB_ID>/.

8. Verify output

ls -lh $VLA_HOME/checkpoints/run_<JOB_ID>/
# Expected: ~15 GB merged model (4 safetensors shards + config)

Configuration

Training hyperparameters can be overridden via environment variables before submitting:

export MAX_STEPS=1000
export LEARNING_RATE=2e-4
export BATCH_SIZE=8
export LORA_RANK=64
sbatch scripts/finetune_openvla.sbatch

File Structure

openvla/
├── README.md               # This file
├── openvla.Dockerfile      # Training container (pinned deps)
├── .gitignore
└── slurm/
    └── finetune_openvla.sbatch  # Slurm batch script (srun + torchrun + LoRA)

References

• OpenVLA — Vision-Language-Action model
• LIBERO — Robotic manipulation benchmark
• SageMaker HyperPod
• Pyxis + Enroot — Container runtime for Slurm