admin/lerobot
1
0
Fork 0
mirror of https://github.com/huggingface/lerobot.git synced 2026-08-08 17:39:44 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(train): split phase timing metrics (#4344)

Browse Source
This commit is contained in:
Pepijn
2026-08-06 17:20:35 +02:00
committed by GitHub
parent ef88d4e52b
commit b1bf24f565
2 changed files with 14 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/source/hardware_guide.mdx
+4 -1
View File
@@ -62,7 +62,10 @@ Reference data points on a 4×H100 80 GB cluster (`accelerate launch --num_proce
| `smolvla` | 27m 49s | 0.312 | 0.011 | ~80% | `--policy.path=lerobot/smolvla_base`, `freeze_vision_encoder=false`, `train_expert_only=false` |
| `pi05` | 3h 41m | 2.548 | 0.014 | ~95% | `--policy.pretrained_path=lerobot/pi05_base`, `gradient_checkpointing=true`, `dtype=bfloat16`, vision encoder + expert trained |
The `dataloading_s` vs. `update_s` ratio is the diagnostic that matters: when `dataloading_s` approaches `update_s`, more GPUs stop helping — your dataloader is the bottleneck and you should look at `--num_workers`, image resolution, and disk speed before adding compute.
Training logs separate the full iteration into `dataloading_s` (`next(dl_iter)`), `preprocessing_s`
(image conversion and the policy pipeline), and `update_s` (the optimizer update). `step_s` covers all
three and drives `samples_per_s`. The benchmark above predates this split, so its `dataloading_s` includes
preprocessing.
### Schedule and checkpoints
src/lerobot/scripts/lerobot_train.py
+10 -9
View File
@@ -602,9 +602,10 @@ def train(cfg: TrainPipelineConfig):
"lr": AverageMeter("lr", ":0.1e"),
# Report the slowest rank for bottleneck-style timings so multi-GPU runs surface the
# true straggler instead of rank 0's view.
"update_s": AverageMeter("updt_s", ":.3f", reduction="max"),
"dataloading_s": AverageMeter("data_s", ":.3f", reduction="max"),
# Derived from the post-reduce max step time; set once per log window on the main rank.
"preprocessing_s": AverageMeter("prep_s", ":.3f", reduction="max"),
"update_s": AverageMeter("updt_s", ":.3f", reduction="max"),
"step_s": AverageMeter("step_s", ":.3f", reduction="max"),
"samples_per_s": AverageMeter("smp/s", ":.0f"),
}
if torch.cuda.is_available():
@@ -634,13 +635,15 @@ def train(cfg: TrainPipelineConfig):
)
for _ in range(step, cfg.steps):
start_time = time.perf_counter()
step_start = time.perf_counter()
batch = next(dl_iter)
preprocessing_start = time.perf_counter()
train_tracker.dataloading_s = preprocessing_start - step_start
for cam_key in dataset.meta.camera_keys:
if cam_key in batch and batch[cam_key].dtype == torch.uint8:
batch[cam_key] = batch[cam_key].to(dtype=torch.float32) / 255.0
batch = preprocessor(batch)
train_tracker.dataloading_s = time.perf_counter() - start_time
train_tracker.preprocessing_s = time.perf_counter() - preprocessing_start
train_tracker, _ = update_policy(
train_tracker,
@@ -652,6 +655,7 @@ def train(cfg: TrainPipelineConfig):
lr_scheduler=lr_scheduler,
sample_weighter=sample_weighter,
)
train_tracker.step_s = time.perf_counter() - step_start
# Note: eval and checkpoint happens *after* the `step`th training update has completed, so we
# increment `step` here.
@@ -668,11 +672,8 @@ def train(cfg: TrainPipelineConfig):
# Collective reduce must run on every rank, before the main-process gate below.
train_tracker.reduce_across_ranks()
if is_main_process():
# Cluster-wide throughput, derived from the already-reduced (max) step time so it
# reflects the slowest rank — which is what actually gates the next iteration.
step_time = train_tracker.update_s.avg + train_tracker.dataloading_s.avg
if step_time > 0:
train_tracker.samples_per_s = samples_per_step / step_time
if train_tracker.step_s.avg > 0:
train_tracker.samples_per_s = samples_per_step / train_tracker.step_s.avg
logging.info(train_tracker)
if wandb_logger:
# Policy sub-losses (latent_loss, action_loss, ...) are aggregated into the