vla / docs /cluster.md
anhtld's picture
Initial commit: DoVLA-CIL codebase (h=16 breakthrough)
adc02fa verified
|
Raw
History Blame
9.28 kB

Cluster Usage

This page describes generic Slurm launchers for large-scale DoVLA-CIL generation, training, scaling, and evaluation. Templates live under scripts/slurm/ and contain placeholders only.

Environment Variables

Common runtime variables:

export PROJECT_DIR="/path/to/dovla-cil"
export VENV_PATH="$PROJECT_DIR/.venv"
export DOVLA_LOG_DIR="$PROJECT_DIR/logs/slurm"
export DOVLA_PARTITION="<partition>"
export DOVLA_CPUS_PER_TASK="8"
export DOVLA_GPUS_PER_TASK="1"
export DOVLA_MEM="64G"
export DOVLA_TIME="24:00:00"

Some Slurm installations do not expand shell variables in #SBATCH headers. If yours does not, pass those values with sbatch --partition ... --gres ... or edit the template header before submitting.

Python Environment

Venv:

python -m venv "$PROJECT_DIR/.venv"
source "$PROJECT_DIR/.venv/bin/activate"
pip install -e ".[dev]"

Conda:

conda create -n dovla-cil python=3.10
conda activate dovla-cil
cd "$PROJECT_DIR"
pip install -e ".[dev]"

Optional distributed generation:

pip install -e ".[ray]"

Install ManiSkill3, Genesis, CUDA-specific wheels, and cluster modules separately. They are not required by the base install.

Secure VLM Configuration

Set OpenClaude-compatible variables in the job environment or scheduler secret store:

export OPENCLAUDE_BASE_URL="https://open-claude.com/v1"
export OPENCLAUDE_API_KEY="<your-key>"
export OPENCLAUDE_MODEL="<model>"

Do not put real keys in Slurm scripts, Git-tracked files, .env, command lines, job names, or shell traces. Avoid set -x in jobs that touch secrets. The VLM client redacts configured API keys from logs, but scheduler logs can still capture environment or command-line mistakes.

For no-network smoke jobs:

export OPENCLAUDE_MOCK=1

Recommended Directory Layout

$PROJECT_DIR/
  configs/
  data/
    tasks/
    cil_array/
    cil_merged/
  logs/
    slurm/
  runs/
    dovla_base/
    scaling/
    baselines/
  reports/
  paper_artifacts/

Use scratch storage for large shards when possible. Copy manifests, indices, checkpoints, reports, and paper artifacts back to persistent storage.

Generation Arrays

export PROJECT_DIR="/path/to/dovla-cil"
export TASKS_PATH="$PROJECT_DIR/data/tasks/tasks.jsonl"
export OUT_ROOT="$PROJECT_DIR/data/cil_array"
export DOVLA_ARRAY="0-31"
export NUM_WORKERS="8"
export STATES_PER_TASK="1000"
export K="32"
export SHARD_SIZE="10000"

sbatch scripts/slurm/generate_cil_array.sbatch

Each array task writes one dataset part:

$OUT_ROOT/part_${SLURM_ARRAY_TASK_ID}

Resume Generation

export RESUME_FLAG=1
sbatch scripts/slurm/generate_cil_array.sbatch

Resume mode reads existing group_index.jsonl, skips deterministic completed group_ids, and writes distributed_manifest.json with planned, skipped, generated, and completed counts.

Aggregate Shards

Inspect individual parts:

python scripts/inspect_shard.py "$OUT_ROOT/part_0"
python scripts/report_dataset.py --dataset "$OUT_ROOT/part_0" --out reports/part_0

Merge parts through the sharding API:

python - <<'PY'
from pathlib import Path
from dovla_cil.data.sharding import ShardReader, write_cil_shards

parts = sorted(Path("data/cil_array").glob("part_*"))
records = []
for part in parts:
    records.extend(ShardReader(part).iterate_records())

write_cil_shards(
    records,
    output_dir="data/cil_merged",
    max_records_per_shard=10000,
    dataset_name="cil_merged",
    backend="toy",
    k=32,
    task_count=0,
    seed=0,
)
PY

Training

export DATASET="$PROJECT_DIR/data/cil_merged"
export OUT_DIR="$PROJECT_DIR/runs/dovla_base"
export EPOCHS="5"
export BATCH_GROUPS="8"
export RECORDS_PER_GROUP="8"
export HIDDEN_DIM="256"
export LR="0.001"

sbatch scripts/slurm/train_dovla.sbatch

Scaling

export OUT_DIR="$PROJECT_DIR/runs/scaling_toy"
export TOTAL_RECORDS="4096"
export K_VALUES="1,2,4,8,16,32"
export EPOCHS="3"

sbatch scripts/slurm/run_scaling.sbatch

External VLA Baseline Bridge

Full SmolVLA/OpenVLA policy baselines should run in a separate environment or container because their dependency stacks can conflict with the pinned ManiSkill/DoVLA stack. First export one expert action per CIL group with deterministic task-balanced sampling:

export PROJECT_DIR="/path/to/dovla-cil"
export DATASET="/scratch/$USER/dovla/experiments/maniskill_presuccess_six_task_collection"
export OUT="$PROJECT_DIR/runs/external_vla/lerobot_export"
export SELECTION="expert"
export GROUP_SAMPLING="task_balanced"
export SEED="0"
sbatch scripts/slurm/export_lerobot_dataset.sbatch

Then write a dry-run plan for the external VLA environment:

export PROJECT_DIR="/path/to/dovla-cil"
export DATASET="$PROJECT_DIR/runs/external_vla/lerobot_export"
export OUT="$PROJECT_DIR/runs/external_vla/smolvla_plan"
export MODEL_FAMILY="smolvla"
export DRY_RUN=1
sbatch scripts/slurm/run_external_vla_baseline.sbatch

The generated external_vla_baseline_plan.json contains secret-free commands for creating the isolated env, downloading the pinned public checkpoint, and running the adapter. The repository ships a SmolVLA expert-only candidate-selection adapter; other model families can provide the same entrypoint contract.

If the pinned SmolVLA directory only contains config files, download the public weights through the containerized Hugging Face CLI before the measured run:

export PROJECT_DIR="/path/to/dovla-cil"
export LOCAL_DIR="/scratch/$USER/dovla/models/smolvla_base-c83c316"
export REVISION="c83c3163b8ca9b7e67c509fffd9121e66cb96205"
export DRY_RUN=1
sbatch scripts/slurm/download_smolvla_checkpoint.sbatch

# After checking the dry-run log:
export DRY_RUN=0
sbatch scripts/slurm/download_smolvla_checkpoint.sbatch

The downloader writes dovla_download_manifest.json with file sizes and SHA256 digests. It does not pass Hugging Face tokens on the command line. For gated/private repos, authenticate through the cluster secret store or an interactive login inside the isolated environment.

If the dry-run log reports Network is unreachable, the current compute partition cannot reach the Hub. Use a network-enabled login/data-transfer node to stage the public checkpoint into LOCAL_DIR, or copy a verified local snapshot there, then rerun the downloader with DRY_RUN=0 only to write the manifest and verify file digests.

On the reference cluster, compute-node Hub access is unavailable. The pinned checkpoint was staged from the login node and verified at revision c83c3163b8ca9b7e67c509fffd9121e66cb96205. Its model.safetensors SHA256 is 7cd549ac2351fb069c0ddb3c34ad2d09cfc92b56a15dccdfc2e41467aaca01eb. Keep LeRobot in a separate environment because stable 0.4.3 requires transformers>=4.57.1,<5 and huggingface-hub>=0.34.2,<0.36.

The validated aligned run uses configs/external/smolvla_cil_aligned.json. Export job 14555244 created 3,500 expert episodes in 39 seconds; GPU job 14555245 loaded the pinned checkpoint, fine-tuned for 1,000 steps, and evaluated 700 held-out groups in 4 minutes 42 seconds on an H100 40 GB MIG slice. The run peaked at about 3.7 GiB host RSS. Its metrics are copied to outputs/external_vla/; no network access or API secret is required at runtime.

After installing that isolated runtime, run a local-only GPU load test:

export LEROBOT_WHEEL="/scratch/$USER/dovla/wheels/lerobot-0.4.3-py3-none-any.whl"
sbatch scripts/slurm/install_smolvla_env.sbatch

export CHECKPOINT="/scratch/$USER/dovla/models/smolvla_base-c83c316"
export PYTHON="/scratch/$USER/dovla/envs/smolvla/bin/python"
sbatch scripts/slurm/smoke_smolvla_checkpoint.sbatch

The installer is intentionally offline: it uses the staged LeRobot wheel plus pinned compatible packages from the Compute Canada CVMFS wheelhouse and passes --no-index. This prevents compute jobs from silently changing dependency versions or hanging on unavailable egress.

The job sets HF_HUB_OFFLINE=1 and TRANSFORMERS_OFFLINE=1 and writes a JSON smoke artifact with the resolved device, policy class, parameter count, and load time.

export CHECKPOINT="/scratch/$USER/dovla/models/smolvla_base-c83c316"
sbatch scripts/slurm/run_smolvla_cil_baseline.sbatch

The included adapter returns measured same-state candidate-selection metrics, not online rollout metrics. A custom adapter function receives (spec_dict, plan_dict) and must return JSON. Do not place Hugging Face tokens or API keys in the Slurm script, job name, or command line; use your cluster secret store or an interactive hf auth login in the isolated environment if needed.

Evaluation and Reports

export CHECKPOINT="$PROJECT_DIR/runs/dovla_base/best.pt"
export OUT_PATH="$PROJECT_DIR/runs/dovla_base/causalstress.json"
export NUM_TASKS="20"
export K="16"

sbatch scripts/slurm/eval_causalstress.sbatch

Aggregate and prepare artifacts:

python scripts/report_eval.py \
  --inputs "$PROJECT_DIR/runs/scaling_toy/*/metrics.json" \
  --out "$PROJECT_DIR/reports/scaling_toy"

python scripts/make_paper_artifacts.py \
  --runs "$PROJECT_DIR/runs" \
  --out "$PROJECT_DIR/paper_artifacts"