Skip to content

[vla-fine-tuning] perf: ~5× per-step speedup; zero data spillage #705

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
irradiantlife wants to merge 2 commits into
base: main
Choose a base branch
from
Open

[vla-fine-tuning] perf: ~5× per-step speedup; zero data spillage #705

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
a27ac57
[vla-fine-tuning] perf: ~5× per-step speedup; zero data spillage
irradiantlife May 19, 2026
bcaebac
Merge branch 'main' into perf/vla-gpu-cpu-bottlenecks
Aydin-ab May 21, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
94 changes: 61 additions & 33 deletions templates/vla-fine-tuning/README.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
{
"cell_type": "markdown",
"metadata": {},
"source": "# Distributed VLA Fine-Tuning with Ray\n\n**⏱️ Time to complete:** 30 min\n\nThis notebook fine-tunes the **PI0.5 Vision-Language-Action (VLA)** model on a\n[LeRobot](https://github.com/huggingface/lerobot) robotics dataset stored in S3.\n\n## Why Ray?\n\nFine-tuning a VLA model involves two fundamentally different workloads:\n\n1. **DATA** — Decoding mp4 video, renaming/transposing images, normalizing\n actions/states. This is CPU-heavy, I/O-heavy, and embarrassingly parallel.\n\n2. **TRAINING** — Forward/backward passes on a large transformer with mixed\n precision, gradient accumulation, and distributed data parallelism.\n This is GPU-heavy.\n\nWithout Ray, you'd stitch these together yourself: write a PyTorch Dataset\nthat blocks GPU workers while they decode video, or build a separate\npreprocessing pipeline and serialize to disk. Either way, the GPUs sit idle\nwaiting for data, and you own all the plumbing.\n\nRay splits this cleanly:\n\n- **[Ray Data](https://docs.ray.io/en/latest/data/data.html)** handles (1): it streams, decodes, and preprocesses data on\n auto-scaled CPU workers, feeding batches to GPU workers through a\n pipelined, backpressure-aware channel. No GPU ever stalls waiting for\n a video decode.\n\n- **[Ray Train](https://docs.ray.io/en/latest/train/train.html)** handles (2): it launches distributed PyTorch workers across\n GPUs, manages DDP synchronization, checkpointing, and fault recovery.\n If a worker dies, training resumes from the last checkpoint — you\n don't re-run from scratch.\n\nThe result: you write a short, single-file script that scales from 1 GPU\non a laptop to 32 GPUs across a cluster, with zero changes to your code.\n\n## Architecture Overview\n\n```\n+-----------+ +------------+ +-----------+\n| S3 Bucket | ----> | Ray Data | ----> | Ray Train |\n| (LeRobot | | (CPU pool) | | (N GPUs) |\n| mp4+pqt) | | | | |\n+-----------+ +------------+ +-----------+\n | |\n | - read parquet | - load PI0.5\n | - decode mp4 | - freeze backbone\n | - rename cameras | - train action heads\n | - transpose HWC | - mixed-precision\n | -> CHW float32 | - gradient accum\n | - stream batches | - checkpoint & resume\n +--------------------+---------------------\n```"
"source": "# Distributed VLA Fine-Tuning with Ray\n\n**\u23f1\ufe0f Time to complete:** 30 min\n\nThis notebook fine-tunes the **PI0.5 Vision-Language-Action (VLA)** model on a\n[LeRobot](https://github.com/huggingface/lerobot) robotics dataset stored in S3.\n\n## Why Ray?\n\nFine-tuning a VLA model involves two fundamentally different workloads:\n\n1. **DATA** \u2014 Decoding mp4 video, renaming/transposing images, normalizing\n actions/states. This is CPU-heavy, I/O-heavy, and embarrassingly parallel.\n\n2. **TRAINING** \u2014 Forward/backward passes on a large transformer with mixed\n precision, gradient accumulation, and distributed data parallelism.\n This is GPU-heavy.\n\nWithout Ray, you'd stitch these together yourself: write a PyTorch Dataset\nthat blocks GPU workers while they decode video, or build a separate\npreprocessing pipeline and serialize to disk. Either way, the GPUs sit idle\nwaiting for data, and you own all the plumbing.\n\nRay splits this cleanly:\n\n- **[Ray Data](https://docs.ray.io/en/latest/data/data.html)** handles (1): it streams, decodes, and preprocesses data on\n auto-scaled CPU workers, feeding batches to GPU workers through a\n pipelined, backpressure-aware channel. No GPU ever stalls waiting for\n a video decode.\n\n- **[Ray Train](https://docs.ray.io/en/latest/train/train.html)** handles (2): it launches distributed PyTorch workers across\n GPUs, manages DDP synchronization, checkpointing, and fault recovery.\n If a worker dies, training resumes from the last checkpoint \u2014 you\n don't re-run from scratch.\n\nThe result: you write a short, single-file script that scales from 1 GPU\non a laptop to 32 GPUs across a cluster, with zero changes to your code.\n\n## Architecture Overview\n\n```\n+-----------+ +------------+ +-----------+\n| S3 Bucket | ----> | Ray Data | ----> | Ray Train |\n| (LeRobot | | (CPU pool) | | (N GPUs) |\n| mp4+pqt) | | | | |\n+-----------+ +------------+ +-----------+\n | |\n | - read parquet | - load PI0.5\n | - decode mp4 | - freeze backbone\n | - rename cameras | - train action heads\n | - transpose HWC | - mixed-precision\n | -> CHW float32 | - gradient accum\n | - stream batches | - checkpoint & resume\n +--------------------+---------------------\n```"
},
{
"cell_type": "markdown",
Expand All @@ -13,9 +13,9 @@
"\n",
"| File | Description |\n",
"|------|-------------|\n",
"| `README.ipynb` (this notebook) | Interactive walkthrough — open in Jupyter and run cells top-to-bottom |\n",
"| `vla.py` | Job script — same pipeline, submittable with `uv run python vla.py` |\n",
"| `util.py` | Training utilities — model loading, checkpointing, collation, training step helpers |\n",
"| `README.ipynb` (this notebook) | Interactive walkthrough \u2014 open in Jupyter and run cells top-to-bottom |\n",
"| `vla.py` | Job script \u2014 same pipeline, submittable with `uv run python vla.py` |\n",
"| `util.py` | Training utilities \u2014 model loading, checkpointing, collation, training step helpers |\n",
"| `lerobot_datasource.py` | Custom Ray Data datasource for LeRobot v3 datasets |"
]
},
Expand Down Expand Up @@ -53,7 +53,7 @@
{
"cell_type": "markdown",
"metadata": {},
"source": "## GPU Requirements\n\nThis template supports both **A100** and **L4** GPUs. Pick one and set\n`accelerator_type` plus matching hyperparameters in\n[`ScalingConfig`](https://docs.ray.io/en/latest/train/api/doc/ray.train.ScalingConfig.html)\nand `train_loop_config` (see [Section 5](#5-launch-distributed-training)).\n\n| | A100 (80 GB) | L4 (24 GB) |\n|---------------|--------------|------------|\n| `batch_size` | 4 | 1 |\n| `grad_accum` | 2 | 8 |\n| `num_workers` | 4 | 4 |\n\nPick instances with **≥ 64 GB RAM** to avoid OOM during model checkpointing.\n\n**A100s** have enough VRAM to run larger batch sizes with minimal gradient\naccumulation. This keeps GPU utilization high and the data pipeline straightforward.\n\n**L4s** require `batch_size=1` to fit in 24 GB VRAM. To compensate, increase\n`grad_accum` so the effective batch size stays reasonable. Smaller per-step batches\nmean faster consumption, which can cause the Ray Data pipeline to fall behind and\nspill to disk. If you see frequent object store spillage, reduce the data pipeline's\nthroughput to match training speed — for example, lower `map_batches` concurrency\nor decrease the number of CPU data workers so the producer and consumer stay in balance."
"source": "## GPU Requirements\n\nThis template supports both **A100** and **L4** GPUs. Pick one and set\n`accelerator_type` plus matching hyperparameters in\n[`ScalingConfig`](https://docs.ray.io/en/latest/train/api/doc/ray.train.ScalingConfig.html)\nand `train_loop_config` (see [Section 5](#5-launch-distributed-training)).\n\n| | A100 (80 GB) | L4 (24 GB) |\n|---------------|--------------|------------|\n| `batch_size` | 4 | 1 |\n| `grad_accum` | 2 | 8 |\n| `num_workers` | 4 | 4 |\n\nPick instances with **\u2265 64 GB RAM** to avoid OOM during model checkpointing.\n\n**A100s** have enough VRAM to run larger batch sizes with minimal gradient\naccumulation. This keeps GPU utilization high and the data pipeline straightforward.\n\n**L4s** require `batch_size=1` to fit in 24 GB VRAM. To compensate, increase\n`grad_accum` so the effective batch size stays reasonable. Smaller per-step batches\nmean faster consumption, which can cause the Ray Data pipeline to fall behind and\nspill to disk. If you see frequent object store spillage, reduce the data pipeline's\nthroughput to match training speed \u2014 for example, lower `map_batches` concurrency\nor decrease the number of CPU data workers so the producer and consumer stay in balance."
},
{
"cell_type": "markdown",
Expand Down Expand Up @@ -133,9 +133,9 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"## 3. Ray Data — Build the preprocessing pipeline\n",
"## 3. Ray Data \u2014 Build the preprocessing pipeline\n",
"\n",
"These functions run on Ray Data's **CPU worker pool** — NOT on the GPU training\n",
"These functions run on Ray Data's **CPU worker pool** \u2014 NOT on the GPU training\n",
"workers. Ray Data calls them automatically as it streams data from S3,\n",
"keeping GPU workers fed without blocking them on I/O or image decoding.\n",
"\n",
Expand All @@ -162,11 +162,22 @@
"\n",
" PI0.5 (like most vision models) expects (batch, channels, height, width).\n",
" The raw dataset stores images as (height, width, channels) uint8.\n",
"\n",
" Implementation note: writing into a pre-allocated float32 buffer fuses the\n",
" stack + transpose + dtype cast into a single pass, avoiding the two\n",
" intermediate allocations that ``np.stack`` + ``np.transpose`` + ``.astype``\n",
" would produce. Cuts CPU time and host memory bandwidth on this stage.\n",
" \"\"\"\n",
" result = dict(batch)\n",
" for key in camera_keys:\n",
" result[key] = np.transpose(np.stack(batch[key]), (0, 3, 1, 2)).astype(np.float32)\n",
" return result"
" frames = batch[key]\n",
" n = len(frames)\n",
" h, w, c = frames[0].shape\n",
" out = np.empty((n, c, h, w), dtype=np.float32)\n",
" for i, f in enumerate(frames):\n",
" np.copyto(out[i], np.moveaxis(f, -1, 0), casting=\"unsafe\")\n",
" result[key] = out\n",
" return result\n"
]
},
{
Expand All @@ -175,9 +186,9 @@
"source": [
"### Build the pipeline\n",
"\n",
"This is **lazy** — no data is read until training starts.\n",
"This is **lazy** \u2014 no data is read until training starts.\n",
"Ray Data will stream from S3, decode mp4 video frames, rename camera\n",
"columns, and transpose images — all on CPU workers, in parallel, with\n",
"columns, and transpose images \u2014 all on CPU workers, in parallel, with\n",
"backpressure so it never overwhelms GPU memory.\n",
"\n",
"Docs: [Loading Data](https://docs.ray.io/en/latest/data/loading-data.html)"
Expand Down Expand Up @@ -207,7 +218,7 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"## 4. Ray Train — Distributed training loop\n",
"## 4. Ray Train \u2014 Distributed training loop\n",
"\n",
"`train_loop_per_worker` runs inside each GPU worker. Ray Train launches N\n",
"copies of it (one per GPU), each receiving its own shard of the streaming\n",
Expand All @@ -222,7 +233,7 @@
" last checkpoint, not from scratch\n",
"\n",
"All of this happens transparently. The code below reads like **single-GPU\n",
"training code** — Ray handles the distribution.\n",
"training code** \u2014 Ray handles the distribution.\n",
"\n",
"Docs: [Getting Started with PyTorch](https://docs.ray.io/en/latest/train/getting-started-pytorch.html)"
]
Expand All @@ -244,6 +255,12 @@
"\n",
" device = torch.device(\"cuda\")\n",
"\n",
" # -- Enable GPU perf knobs --------------------------------------------------\n",
" # TF32 + cudnn.benchmark, set inside the worker (must NOT reference\n",
" # torch.backends.cudnn from the driver's notebook globals -- cloudpickle\n",
" # then refuses to ship this fn to workers).\n",
" util.enable_gpu_perf_flags()\n",
"\n",
" # -- Load model and freeze backbone ----------------------------------------\n",
" #\n",
" # load_pi05_policy() loads the pretrained PI0.5 model, applies the\n",
Expand All @@ -258,10 +275,10 @@
" policy = util.load_pi05_policy()\n",
" policy = ray.train.torch.prepare_model(policy) # <-- RAY TRAIN: wrap in DDP\n",
"\n",
" # AdamW optimizer -- only updates the unfrozen action/time projection heads.\n",
" optimizer = torch.optim.AdamW(\n",
" [p for p in policy.parameters() if p.requires_grad], lr=config.get(\"lr\", 1e-4),\n",
" )\n",
" # AdamW optimizer (fused CUDA kernel) -- only updates the unfrozen\n",
" # action/time projection heads. Caches the trainable-param list on the\n",
" # policy so clip_grad_norm_ doesn't re-walk every step.\n",
" optimizer = util.make_trainable_optimizer(policy, lr=config.get(\"lr\", 1e-4))\n",
"\n",
" # GradScaler for mixed-precision training (fp16 forward, fp32 gradients).\n",
" scaler = torch.amp.GradScaler(\"cuda\")\n",
Expand Down Expand Up @@ -315,31 +332,42 @@
"\n",
" for epoch in range(start_epoch, num_epochs):\n",
" optimizer.zero_grad(set_to_none=True)\n",
" epoch_loss_sum, epoch_loss_count = 0.0, 0\n",
" # Keep the running loss on-device so we don't block on a host sync\n",
" # every micro-step. Materialize a scalar only at log/epoch boundaries.\n",
" epoch_loss_sum = torch.zeros((), device=device)\n",
" epoch_loss_count = 0\n",
" accum_count = 0\n",
"\n",
" # iter_torch_batches() streams pre-processed batches from the Ray Data\n",
" # CPU worker pool directly into GPU memory.\n",
" # Docs: https://docs.ray.io/en/latest/data/api/doc/ray.data.DataIterator.iter_torch_batches.html\n",
" for batch in shard.iter_torch_batches( # <-- RAY DATA: streaming batches\n",
" batch_size=batch_size,\n",
" collate_fn=util.NumpyToTorchCollate(device),\n",
" ):\n",
" # Standard PyTorch: forward + scaled backward (see util.py)\n",
" loss_val = util.train_step(policy, batch, preprocessor, max_len, grad_accum, scaler)\n",
" # cuda_prefetcher wraps the Ray Data iterator with a 1-batch GPU\n",
" # prefetch on a separate CUDA stream. Batch N+1's H2D copy overlaps\n",
" # with batch N's forward/backward -- no stalls on data transfer.\n",
" batch_stream = util.cuda_prefetcher(\n",
" shard.iter_torch_batches( # <-- RAY DATA: streaming batches\n",
" batch_size=batch_size,\n",
" collate_fn=util.NumpyToTorchCollate(device),\n",
" ),\n",
" device,\n",
" )\n",
" for batch in batch_stream:\n",
" # train_step returns a 0-d device tensor -- no implicit host sync.\n",
" loss_t = util.train_step(policy, batch, preprocessor, max_len, grad_accum, scaler)\n",
" step += 1\n",
" accum_count += 1\n",
" epoch_loss_sum += loss_val\n",
" epoch_loss_sum += loss_t\n",
" epoch_loss_count += 1\n",
"\n",
" if accum_count % grad_accum == 0:\n",
" # Standard PyTorch: unscale, clip, step, zero_grad (see util.py)\n",
" util.optimizer_step(policy, optimizer, scaler, scheduler)\n",
" accum_count = 0\n",
"\n",
" # Log every 10 steps so you can watch training progress.\n",
" # Log every 10 steps. .item() is the only host sync in this loop.\n",
" if step % 10 == 0:\n",
" log.info(\"epoch=%d step=%d loss=%.4f lr=%.2e\", epoch, step, loss_val, scheduler.get_last_lr()[0])\n",
" log.info(\"epoch=%d step=%d loss=%.4f lr=%.2e\",\n",
" epoch, step, loss_t.item(), scheduler.get_last_lr()[0])\n",
"\n",
" # smoke-test cap (unset MAX_TRAIN_STEPS for full training)\n",
" if max_train_steps and step >= max_train_steps:\n",
Expand All @@ -354,7 +382,7 @@
" # ray.train.report() is a synchronization barrier -- every worker must\n",
" # call it. Only rank 0 creates the actual checkpoint.\n",
" # Docs: https://docs.ray.io/en/latest/train/api/doc/ray.train.report.html\n",
" avg_loss = epoch_loss_sum / max(epoch_loss_count, 1)\n",
" avg_loss = (epoch_loss_sum / max(epoch_loss_count, 1)).item()\n",
" metrics = {\"epoch\": epoch, \"steps\": step, \"loss\": avg_loss, \"lr\": scheduler.get_last_lr()[0]}\n",
"\n",
" if ray.train.get_context().get_world_rank() == 0: # <-- RAY TRAIN\n",
Expand Down Expand Up @@ -398,17 +426,17 @@
"\n",
"**What was covered:**\n",
"\n",
"1. **Configuration** — Dataset path, camera column rename map, and HuggingFace token.\n",
"1. **Configuration** \u2014 Dataset path, camera column rename map, and HuggingFace token.\n",
"\n",
"2. **Ray Data preprocessing** — A lazy, streaming pipeline that reads from S3,\n",
" decodes mp4 video, renames camera columns, and transposes images from HWC to CHW —\n",
"2. **Ray Data preprocessing** \u2014 A lazy, streaming pipeline that reads from S3,\n",
" decodes mp4 video, renames camera columns, and transposes images from HWC to CHW \u2014\n",
" all on CPU workers, in parallel, with backpressure.\n",
"\n",
"3. **Ray Train distributed training** — `train_loop_per_worker` runs on each GPU\n",
"3. **Ray Train distributed training** \u2014 `train_loop_per_worker` runs on each GPU\n",
" with automatic DDP wrapping (`prepare_model`), data sharding (`get_dataset_shard`),\n",
" and fault-tolerant checkpointing (`get_checkpoint` / `report`).\n",
"\n",
"4. **Launch** — `TorchTrainer` ties it all together. To scale from 1 to N GPUs,\n",
"4. **Launch** \u2014 `TorchTrainer` ties it all together. To scale from 1 to N GPUs,\n",
" change `ScalingConfig.num_workers`. Everything else adapts automatically.\n",
"\n",
"### Job submission\n",
Expand Down
Loading
Loading