GPU busy about 30% of the time when training ACT on 8x B200 at the default settings

[TarzanZhao]

Ticket Type

💡 Feature Request / Improvement

Environment & System Info

lerobot: main at 2774d9bd, pip install -e ".[training]"
python: 3.12.14
torch: 2.11.0+cu130 (cuDNN 9.19, NCCL 2.28.9), torchvision 0.26.0, torchcodec 0.11.1
accelerate: 1.14.0
GPU: 8x NVIDIA B200, driver 580.126.20, CUDA 13.0, one node
OS: Linux

Description

I ran lerobot-train --policy.type=act (src/lerobot/scripts/lerobot_train.py, plain DDP, batch 8 per GPU, fp32) on lerobot/aloha_sim_insertion_human with 8x B200, following the 8-GPU command in docs/source/multi_gpu_training.mdx (accelerate launch --num_processes=8), and got 31.5 ms per training step. A profile showed the GPU busy about 30% of the time: the eager ACT.forward in src/lerobot/policies/act/modeling_act.py and its backward issued about 1150 small kernels per step, and on this GPU the host could not launch them as fast as they ran. With policy.model compiled under CUDA graphs (torch.compile, mode="reduce-overhead"), plus a few smaller DDP and optimizer changes listed in the PR, the same step took 13.2 ms.

Nothing in my run pointed at this. The step time looked plausible, --accelerator.compile was in the config but the train script refused it, and the docs I read described multi-GPU runs as compute-bound. A line in multi_gpu_training.mdx saying that a small policy at batch 8 is launch-bound on a fast GPU, and that compiling the model is the lever, would have saved me the profile.

Context & Reproduction

accelerate launch --num_processes=8 $(which lerobot-train) \
  --policy.type=act --policy.device=cuda --policy.push_to_hub=false \
  --dataset.repo_id=lerobot/aloha_sim_insertion_human \
  --steps=300 --seed=1000 --batch_size=8 --num_workers=4 --log_freq=1 \
  --save_checkpoint=false --env_eval_freq=0 --eval_steps=0 --wandb.enable=false

Step time is the script's step_s (max over ranks), median of steps 51 to 300. The profile covers steps 60 to 62 on all 8 ranks; the PR links the traces.

Relevant logs or stack trace

No response

Checklist

• I have searched existing tickets to ensure this isn't a duplicate.
• I am using the latest version of the main branch.
• I have verified this is not an environment-specific problem.

Additional Info / Workarounds

I opened #4607 to record my experiment: setup, measurements, traces and some potential fixes. A different fix may also well suit the codebase.

[sattyamjjain]

Two questions, since the issue also asks for a docs change. The PR review notes compiled dropout draws masks in a different order, so eager and compiled runs are not bit-reproducible: should compile stay opt-in (default off) so the docs' same-seed reproducibility statement holds? And does reduce-overhead (CUDA graphs) tolerate the shorter final batch when drop_last is False? If useful, I can run the same command on a single consumer GPU and post torch-profiler kernel counts to show the launch-bound effect is not B200-specific.