MSN-007 · 2026 · robotics, vla, distributed, lora, docker
Distributed fine-tuning framework for VLA robot policies built on HuggingFace LeRobot. Fine-tune π₀.₅ on custom tasks from 12 GB VRAM and deploy to a Jetson AGX Orin at 30 Hz with a single command.
VLASH-Forge is the open-source evolution of VLASH — the research framework from *"VLASH: Real-Time VLAs via Future-state-aware Asynchronous Inference"*. The paper established Temporal Delay Augmentation as a path to real-time robot control; this project packages that into a fully reproducible distributed training pipeline that anyone can run.
The core problem: fine-tuning a state-of-the-art VLA model on custom robot tasks requires multi-GPU distributed training, a carefully managed Python environment, and infrastructure that was never released alongside the original work. VLASH-Forge removes those barriers with a single container image that runs identically on a local workstation, HPC cluster, or cloud VM.
Lab setup — Piper arm, depth camera on tripod, two coloured target pedestals, and the ball on a staging platform
Demonstrations were collected on a Piper robotic arm in a tabletop pick-and-place configuration. A depth camera mounted on a tripod provides the visual observation stream. The task: pick a ball from a staging platform and place it on the matching-colour target pedestal. Simple to describe, but the 1.3B-parameter model running on the Jetson makes this non-trivial to do at robot-control frequency.
Second angle — Piper arm homed with both targets and the ball visible; the black WD external drive doubles as the edge compute node
The intended pipeline is three commands:
1. Collect demonstrations on your robot → upload to HuggingFace Hub
2. Run: ./scripts/train.sh config.yaml → fine-tuned checkpoint on HF Hub
3. Load checkpoint on your inference hardware → deploy
The container handles dataset download, base model weight acquisition, distributed process setup, and checkpoint upload. The only inputs you provide are a HuggingFace token, a dataset repo ID, and a YAML config file.
The same Docker/Singularity image supports five deployment modes with no environment changes:
- SLURM (`sbatch scripts/train_slurm.sbatch`) — HPC clusters, including NSCC ASPIRE where the VLASH models were originally trained - PBS (`qsub scripts/train_pbs.pbs`) — PBSpro clusters; the container entrypoint normalises GPU UUID assignment automatically - Docker Compose — single-node multi-GPU, any cloud VM with NVIDIA Container Runtime - Kubernetes — cloud-native deployments; PersistentVolumeClaims replace the `/scratch` bind-mount - Plain Docker — local workstation, single or multi-GPU via `scripts/train.sh`
`scripts/train.sh` detects whether `vlash.sif` exists and selects Singularity or Docker automatically. `$SCRATCH` is always bind-mounted to `/scratch` inside the container — model weights, datasets, and checkpoints persist across runs without touching the image.
Nothing important is written inside the container itself. Base model weights (~10 GB, cached after first run), dataset files, and training checkpoints all land in `$SCRATCH`. This means the image can be rebuilt or updated without losing any training state.
LoRA with shared observation encoding brings the VRAM floor for π₀.₅ (1.3B parameters) down to 12 GB — an RTX 3090 or A10G is enough. A QLoRA path (`lora.use_qlora: true`) cuts it further to 8 GB using a 4-bit base model, which fits most consumer cards.
Mode · Backend · Min VRAM · Typical hardware
LoRA · DeepSpeed ZeRO-2 · 12 GB · RTX 3090 / A10G / T4 QLoRA · DeepSpeed ZeRO-2 · 8 GB · RTX 3080 / consumer GPUs Full fine-tuning · PyTorch FSDP · 40 GB per GPU · 4× A100 40GB
DeepSpeed ZeRO-2 is the default for LoRA — it shards optimiser states and gradients across GPUs while keeping model parameters replicated, halving per-GPU memory versus standard DDP without the all-gather overhead of ZeRO-3. FSDP (`TRAIN_BACKEND=fsdp`) is used for full fine-tuning runs where the full parameter shard is needed.
Sharing the visual observation encoder across LoRA adapter steps eliminates the need to track gradients through the entire vision backbone for every observation. This is what gets the VRAM floor low enough to run on a 12 GB card without quantization, while still updating the action head and attention projections.
The harder problem was deployment. The Jetson AGX Orin is the target inference platform — capable edge hardware, but a 1.3B-parameter diffusion-policy model does not run at robot-control frequency synchronously. A synchronous inference call takes ~5444 ms. A 30 Hz control loop has 33 ms between steps. The gap is three orders of magnitude.
Temporal Delay Augmentation closes it by training the model to predict actions with a fixed time offset. At inference time, the robot executes actions from the previous chunk while the model computes the next one in parallel. The `VLASHAsyncManager` in `vlash/run.py` manages this overlap — it pre-fetches the next action chunk starting when the current chunk is partially through execution, using the predicted end state as the conditioning observation.
Side-by-side comparison at t=02s — async (left) has the ball in the gripper and is moving to target; sync (right) has barely begun its first approach
On the ball sorting task, the numbers:
- Async inference latency: 184.5 ms per call - Synchronous baseline: 5444.1 ms per call — 29.5× slower - Task success rate: 65% async vs. 5% sync - Task completion time: 52 s async vs. 153 s sync — 2.9× faster - Control frequency on Jetson AGX Orin: 30 Hz
Demonstrations collected on a Piper arm → uploaded to HuggingFace Hub → fine-tuned on NSCC ASPIRE (4× A100) → checkpoint deployed on a Jetson AGX Orin. The full pipeline ran without modification across three different hardware environments.
The demo video (from `assets/piper_demo.mp4`) shows this live. The async run completes ball placement in under a minute; the synchronous run is still executing its first reach by that point.
Demo title card — 15 Hz async inference achieved on a laptop RTX 5090; Jetson achieves 30 Hz with action chunk overlap
Training loss, gradient norm, and learning rate are streamed to Weights & Biases every `log_freq` steps when `wandb.enable: true` is set in the config. This is the primary tool for catching training failures early — if a run OOMs or diverges mid-way, the dashboard shows the exact step and the loss curve shape before it died. The container ships W&B pre-installed; only the API key needs to be provided.
Container portability sounds straightforward until you are debugging why NCCL hangs on an HPC cluster that assigns GPU UUIDs instead of integer indices — a problem that does not show up locally at all. The entrypoint UUID normalisation fix was one of the first things that had to go in before SLURM runs were reliable.
I had underestimated how much of distributed training is environment management rather than ML. Getting DeepSpeed ZeRO-2 and FSDP to initialise cleanly across five different execution contexts (SLURM, PBS, Docker Compose, Kubernetes, plain Docker) with the same image was the majority of the engineering work. The ML configuration — LoRA rank, learning rate schedule, shared encoder — was comparatively straightforward once the scaffolding was solid.
The async inference result surprised me. The 29.5× latency improvement is real and repeatable, but the 5% → 65% success rate jump is the more important number — it is what makes the technique worth using at all. A synchronous VLA on a Jetson is not a slow version of the same thing; it is effectively non-functional for pick-and-place because the arm has moved past the target by the time the model says to close the gripper. Temporal Delay Augmentation is not a speed optimisation. It is what makes VLA deployment on embedded hardware possible in the first place.
I also learned that running the full pipeline end-to-end across three hardware environments (Piper arm → NSCC ASPIRE A100s → Jetson AGX Orin) surfaces a category of bugs that unit tests cannot catch — dataset format assumptions baked into the video decoder path, checkpoint schema mismatches between LoRA and full fine-tuning runs, HuggingFace Hub rate limits during the first weight download on a cold cluster node. None of these are interesting problems. All of them cost time. Writing the pipeline to fail loudly with actionable error messages rather than silently at a CUDA kernel was worth the effort.
- π₀.₅ (1.3B params) via HuggingFace LeRobot, built on PaliGemma - LoRA / QLoRA fine-tuning with shared observation encoding (PEFT backend) - DeepSpeed ZeRO-2 (LoRA) and PyTorch FSDP FULL_SHARD (full fine-tuning) - Docker + Singularity container, Apache 2.0 licensed - SLURM, PBS, Docker Compose, Kubernetes, and plain Docker deployment paths - Temporal Delay Augmentation for embedded real-time deployment - W&B experiment tracking, bfloat16 mixed precision (bf16 on A100/H100/RTX 40xx)