feat(training): support gradient accumulation

This commit is contained in:
Pepijn
2026-07-15 23:34:10 +02:00
parent 441c7e4bea
commit 663971aa5d
6 changed files with 180 additions and 44 deletions
+13
View File
@@ -113,6 +113,19 @@ accelerate launch --num_processes=2 $(which lerobot-train) \
--policy.type=act
```
When the desired global batch is larger than the per-GPU batch that fits in memory, use gradient
accumulation. `steps` continues to count optimizer updates:
```bash
# 8 samples/GPU × 4 GPUs × 2 microbatches = effective global batch 64.
accelerate launch --num_processes=4 $(which lerobot-train) \
--batch_size=8 \
--gradient_accumulation_steps=2 \
--steps=80000 \
--dataset.repo_id=lerobot/pusht \
--policy.type=act
```
## Training Large Models with FSDP
DDP replicates the full model on every GPU, so a model that doesn't fit on one GPU won't fit under
+5 -3
View File
@@ -149,18 +149,20 @@ For a sample-matched SmolVLA visual-memory ablation, use
`examples/robomme/smolvla_visual_memory_ablation.sh`. `TARGET_SAMPLES` counts examples across all
GPUs, and `NUM_PROCESSES` is included when the script converts that target into optimizer steps. For
example, the following reproduces 5.12 million example exposures (the exposure of RoboMME's
80,000-step, global-batch-64 memory-policy recipe) with four GPUs and a global batch of 48:
80,000-step, global-batch-64 memory-policy recipe) with four GPUs, two accumulated microbatches,
and an effective global batch of 64:
```bash
TARGET_SAMPLES=5120000 \
NUM_PROCESSES=4 \
BATCH_SIZE=12 \
BATCH_SIZE=8 \
GRADIENT_ACCUMULATION_STEPS=2 \
VARIANT=visual-memory \
RUN_EVAL=false \
bash examples/robomme/smolvla_visual_memory_ablation.sh
```
This becomes 106,667 optimizer steps, or about 10.74 execution-target epochs. Run the baseline with
This becomes 80,000 optimizer steps, or about 10.74 execution-target epochs. Run the baseline with
the same `TARGET_SAMPLES`, effective batch size, seed, and scheduler settings to isolate the visual
memory change. RoboMME's released baseline uses a different global batch (128), so its nominal
80,000-step recipe is not sample-matched to its global-batch-64 memory recipe.