perf(act): training on 8x B200, 28.0 ms to 13.0 ms per step, mostly from the eager model's kernel launches

[TarzanZhao]

Summary / Motivation

When I trained ACT with lerobot-train under 8-GPU DDP on B200s, a step took 28 ms, but the GPU was busy for only about 30 % of it. The eager step launched about 1,150 small kernels, every forward waited for a buffer broadcast across all ranks, the forward read two losses back to the host, and the backward sent its gradients in nine all-reduces that each waited for the slowest rank. This PR compiles the ACT model with CUDA graphs, sends the gradients in one all-reduce, drops the broadcast and the loss reads, and uses fused AdamW. On those 8 B200 GPUs the step goes from 28.0 ms to 13.0 ms (2.15x); I have not measured other GPUs, single-process training or other policies.

1. Compile the ACT model with CUDA graphs. --accelerator.compile was a placeholder on main; it now compiles the regions a policy declares (policy.model for ACT) with mode="reduce-overhead", so each step replays a captured graph instead of launching about 1,150 kernels. It is on by default for policies that declare regions, which today is only ACT, and --accelerator.compile.enabled=false turns it off. Step 1 pays a one-time compile of 27 s with a warm inductor cache and 260 s cold.
2. One all-reduce per backward. A new bucket_cap_mb field on DDPConfig defaults to 1024, so the compiled backward's gradients go out in one all-reduce instead of nine that each waited for the slowest rank. --accelerator.ddp.bucket_cap_mb=25 restores torch's default.
3. No buffer broadcast when there is nothing to sync. DDP broadcast 286 KB of constant buffers at the start of every forward and blocked all ranks until each arrived. A new broadcast_buffers field on DDPConfig is switched off after prepare() when the policy has no BatchNorm layer; a policy with trainable BatchNorm keeps the broadcast.
4. Losses stay on the device. ACTPolicy.forward called .item() on l1_loss and kld_loss, which held the host until the forward finished. It now returns them as detached 0-d tensors, and MetricsTracker reads them back after the optimizer step, where the loss is read anyway.
5. Fused AdamW. When every parameter is a CUDA float tensor, AdamW uses fused=True instead of one update per parameter; any other parameter set keeps the old path. On this job it did not change the step time, because the per-parameter work was off the critical path.

Speedup Result

Measured on 8 B200 GPUs.

	baseline 2774d9bd	this PR	change
step time, slowest rank, median of steps 51 to 300	28.0 ms (27.0 / 29.0)	13.0 ms (13.0 / 13.0)	2.15x
step time, p90	32 to 33 ms	15 ms
slowest step in steps 51 to 300	168 ms	17 ms
peak allocated memory per rank	3.37 GB	3.37 GB
nvidia-smi memory per GPU	5.1 GiB	10.6 GiB	+5.5 GiB for CUDA-graph pools
step 1	12.5 / 12.8 s	26.8 / 26.9 s warm cache, 260 s cold	one-time compile
300-step wall	52.1 / 52.8 s	62.8 / 64.4 s

Per commit, measured cumulatively in an earlier round against the unmodified main without the two setup items listed under Details (two runs per commit, same step-time definition):

commit	change	step time
	unmodified 2774d9bd	31.5 ms
978b1485	fused AdamW	31.4 ms
73bb3232	buffer broadcast off	28.6 ms
663535e8, 26680519	compile policy.model, reduce-overhead	17.0 ms
d37994a1	losses stay on the device	16.0 ms
5cc82b53	one all-reduce, bucket_cap_mb=1024	13.8 ms
e5621af6, ae4af846	type annotation for mypy, ruff format	no change
bc51b2d0, 710ae0f7, 23a91bd2	optimizer before compile, tests, comments	not re-measured; none changes what runs

Correctness Verification

I trained the baseline and this branch with the same seed on the same 300 batches on all 8 ranks, and compared 1,248 recorded values per rank (9,984 per run): the loss, l1_loss, kld_loss and gradient norm at every step, and at steps 1 to 3 the batch shapes, the action values and the mean magnitudes of the ACT outputs. The tolerances were set before the runs and checked against three runs of the unmodified code. The recording code is on the perf/act-ddp-compile-verify branch of my fork, not in this PR.

172 of the 9,984 checks fail the tolerances set before the runs, so numerical equivalence has not been shown. The 172 are 122 per-step loss values (up to 8.1 % against a 3 % + 0.01 limit) and 50 ACT output means at steps 1 to 3 (0.3 to 0.6 % against 0.1 %). The other 9,812 are within tolerance, including every gradient norm, batch shape and action value, and the all-rank loss at step 300 is 2.442 on the baseline and 2.452 on this branch in both runs of each.

I guess the failures come from the order in which the dropout masks are drawn. ACT applies dropout in every transformer layer and inside nn.MultiheadAttention, and the compiled graph draws those masks at different points inside its fused regions, so the two runs part at step 1 and follow different trajectories from there.

recorded	baseline vs this branch	tolerance
loss, all-rank mean, step 300	2.442 vs 2.452 (+0.4 %)	3 %
gradient norm, steps 1 to 5 / 6 to 300	0.63 % / 4.8 %	1 % / 10 %
loss per rank, steps 1 to 5 / 6 to 300	1.7 % / 8.1 %; above the limit on some steps	0.01 / 3 % + 0.01
l1_loss and kld_loss per rank, steps 1 to 5 / 6 to 300	0.12 % and 1.8 % / 2.0 % and 8.8 %; above the limit on some steps	same as loss
ACT output means (actions, mu, log sigma), steps 1 to 3	0.27 %, 0.60 %, 0.58 %; above the limit	0.1 %
same outputs with --policy.dropout=0, earlier run, not repeated	0.06 %, 0.005 %, 0.007 %	0.1 %
batch shapes, dtypes, devices, action values, padded-step count, steps 1 to 3	identical	exact

Related issues

• Related: GPU busy about 30% of the time when training ACT on 8x B200 at the default settings #4612

How was this tested (or how to run locally)

• pre-commit run -a: all 19 hooks pass.
• pytest tests/scripts tests/distributed tests/configs/test_accelerator_config.py tests/processor/test_act_processor.py (CPU): 101 passed, 10 skipped.
• pytest tests/configs/test_train_config_distributed.py: 14 passed. test_compile_placeholder asserted that compile is always rejected; it is replaced by one test per side of the new gate, rejected when sharded and accepted otherwise.
• New tests in tests/distributed/test_policy_surface.py for apply_torch_compile: the enable logic including the auto default, that regional compile wraps only the declared region, that the wrapper keeps the same parameter objects, that a parameter group selected by name is lost once the wrapper is in place, and that the train script builds the optimizer before compiling. The last one fails if the two calls are swapped back.
• Still untested: the tensor path in MetricsTracker.update_metrics.
• Speed and correctness runs as above; setup and full command under Details.

Checklist (required before merge)

• Linting/formatting run (pre-commit run -a)
• All tests pass locally (pytest): 115 passed, 10 skipped on CPU over the files listed above
• Documentation updated
• CI is green
• Community Review: I have reviewed another contributor's open PR and linked it here: #

Reviewer notes

• Defaults that change for every policy: bucket_cap_mb=1024 also applies to eager policies, where one bucket delays the all-reduce until the whole backward is done; fused AdamW and the broadcast skip have no config key. I can scope bucket_cap_mb to compiled runs if that is preferred.
• Fixed on this head, from the Copilot review: the optimizer is now built before the compile, so ACT's model.backbone prefix still matches and optimizer_lr_backbone is no longer dropped when it differs from optimizer_lr. The measurements above predate that commit; the two learning rates are equal by default, so the grouping does not change any number in them.
• Also fixed: the fallback_random comment claimed eager-equivalent numerics, and the CompileConfig docstring still called compile an unwired placeholder.
• Open, from the same review: _all_cuda_float consumes a generator passed as a group's params; compile.enabled=None bypasses the sharded check.
• Not tested: single-process runs, FSDP, checkpoint save and resume, other policies, other GPUs.
• The compile pays for itself on these B200s after about 1,000 steps; the default is 100,000.
• Left out of this PR: the same .item() change for the other eight policies whose forward reads losses back, and the launcher-level NUMA pinning and gradient_as_bucket_view findings in GPU busy about 30% of the time when training ACT on 8x B200 at the default settings #4612.

Details: hardware, model, full command, traces

Hardware and environment. One node, 8 × NVIDIA B200 (sm_100, about 180 GB each), host CPUs on 2 NUMA nodes (GPUs 0 to 3 on node 0, 4 to 7 on node 1), driver 580.126.20, CUDA 13.0; node empty before every timed run. Python 3.12.14, torch 2.11.0+cu130 from PyPI (cuDNN 9.19, NCCL 2.28.9, triton 3.6.0), torchvision 0.26.0, torchcodec 0.11.1 with conda-forge ffmpeg 8.0.1, accelerate 1.14.0; lerobot main at 2774d9bd, pip install -e ".[training]", nothing built from source. The dataset is copied to node-local disk before each run; the conda environment and the inductor cache are on NFS.

Model and job. ACT as configuration_act.py defaults it: ResNet18 backbone with torchvision ImageNet weights; transformer width 512, 8 heads, feed-forward 3200 with ReLU, 4 encoder layers, 1 decoder layer; VAE encoder with 4 layers and latent size 32; chunk size 100; dropout 0.1; kl_weight 10; 51 613 582 parameters, all trained. AdamW, lr 1e-5, weight decay 1e-4, gradient clipping at norm 10. fp32 with torch.backends.cuda.matmul.allow_tf32 = True and cudnn.benchmark = True (the script's own settings, lerobot_train.py:438-439). 8-process DDP, 4 loader workers per process. Per process per step: 8 samples of observation.images.top (3×480×640 float32), observation.state (14) and action (100×14) with an action_is_pad mask; 64 samples per step. Loss: L1 over the unpadded action steps plus 10 × the KL of the VAE latent to a standard normal. --seed=1000 is the script's default and goes through its set_seed; the two baseline runs reproduce each other's losses (78.861 to 2.442 / 2.442), as do the two branch runs (78.707 to 2.452 / 2.452).

Measurement. Against the unmodified 2774d9bd, both versions with the same setup, which includes two items that are not part of this PR: --accelerator.ddp.gradient_as_bucket_view=true, an existing flag worth 2 ms on the eager step, and per-rank numactl --cpunodebind to the NUMA node of the rank's GPU, worth 3 to 4 ms on the eager step. Without them the branch is 2.42x against the unmodified code. Step time is step_s from the script's per-step log (the slowest rank, 1 ms resolution), median over steps 51 to 300; steps 1 to 50 are excluded because step 1 carries worker start, cuDNN autotune and, on the branch, the compile. Runs alternate between the versions on an otherwise empty node.

Data (once): lerobot/aloha_sim_insertion_human (50 episodes, 25 000 frames, 88 MB, v3 format) copied to node-local disk and passed as --dataset.root; ResNet18 weights under TORCH_HOME; HF_HUB_OFFLINE=1, nothing downloads.

The measured job, both versions; only the checkout on PYTHONPATH differs:

export LD_LIBRARY_PATH=$CONDA_PREFIX/lib          # torchcodec needs the env's libstdc++
export TMPDIR=<node-local dir>                    # loader workers die on NFS
export HF_HUB_OFFLINE=1 TORCH_HOME=<ResNet18 weights dir>
export TORCHINDUCTOR_CACHE_DIR=<scratch> TRITON_CACHE_DIR=<scratch>
export PYTHONPATH=<checkout>/src                  # baseline: 2774d9bd; branch: ae4af846

accelerate launch --num_processes=8 --no_python bash numa_wrap.sh \
  --policy.type=act --policy.device=cuda --policy.push_to_hub=false \
  --dataset.repo_id=lerobot/aloha_sim_insertion_human --dataset.root=<local copy> \
  --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 \
  --output_dir=<out> --job_name=exp0914_act_bench \
  --accelerator.ddp.gradient_as_bucket_view=true

numa_wrap.sh is the per-rank program accelerate launch --no_python starts; it reads the rank's GPU from LOCAL_RANK, looks up its NUMA node in sysfs, and execs lerobot-train under numactl --cpunodebind (CPU affinity only; falls through unpinned when anything is missing):

#!/usr/bin/env bash
set -u
node=""
if [ -n "${LOCAL_RANK:-}" ] && command -v numactl >/dev/null; then
  bus=$(nvidia-smi --query-gpu=pci.bus_id --format=csv,noheader -i "$LOCAL_RANK" 2>/dev/null | tr 'A-F' 'a-f')
  bus=${bus#0000}                      # 00000000:1B:00.0 -> 0000:1b:00.0
  f=/sys/bus/pci/devices/$bus/numa_node
  [ -r "$f" ] && node=$(cat "$f")
fi
if [ -n "$node" ] && [ "$node" -ge 0 ] 2>/dev/null; then
  exec numactl --cpunodebind="$node" lerobot-train "$@"
fi
exec lerobot-train "$@"

Rejected in the setup sweep on the unmodified code, within the 2 ms noise: find_unused_parameters=false, static_graph=true, OMP_NUM_THREADS=8, NCCL_CUMEM_ENABLE=0 (and step 1 +5 s), PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True. Accelerate's --dynamo_backend inductor --dynamo_mode reduce-overhead compiled for 303 s and segfaulted on all 8 ranks at step 2. Rejected in an earlier round: inductor's default mode (no gain, 273 s compile), bf16 autocast for the ACT core (step unchanged, final loss +17 %), gradient_as_bucket_view=false under compile (worse). Not tried because they change the numerics or the job: mixed precision, batch size, step count, --policy.dropout=0, --cudnn_deterministic, gradient accumulation.

Traces (torch.profiler, steps 60 to 62, all 8 ranks; open in https://ui.perfetto.dev; the README next to them lists code, script, hardware, env and capture window; the baseline trace is the unmodified main):

• baseline: https://drive.google.com/open?id=1V9ViJQBTGt8kB1XzJnXOkJ42MUJ9wGIy
• optimized: https://drive.google.com/open?id=1nVYRbFm5_s0kSgV_9yjZOB0P3tjhiAwa

[TarzanZhao]

Ready for review. This PR speeds up the ACT training step in lerobot-train under DDP through six changes (compile, fused optimizer, and related fixes): 31.5 ms to 13.2 ms per step on 8 B200 GPUs, 2.4x.

cc @HaomingSong (author of the distributed and accelerator modules) @pkooij (most recent changes to lerobot_train.py) @imstevenpmwork (merged the last changes to optimizers.py and modeling_act.py)

Could a maintainer also approve the first-time contributor workflow runs so CI can start? Thanks.

[Copilot]

🟡 Changes recommended

Compilation currently breaks parameter and checkpoint names, bypasses sharded validation in auto mode, and leaves existing tests failing.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Pull request overview

Optimizes ACT training throughput, especially for multi-GPU CUDA workloads.

Changes:

• Adds regional torch.compile support and device-side loss metrics.
• Tunes DDP, AdamW, and NUMA affinity for lower overhead.
• Integrates the optimizations into lerobot-train.

File summaries

File	Description
src/lerobot/utils/logging_utils.py	Converts scalar tensor metrics during logging.
src/lerobot/scripts/lerobot_train.py	Integrates compilation, NUMA binding, and DDP tuning.
src/lerobot/policies/act/modeling_act.py	Declares compile regions and avoids early synchronization.
src/lerobot/optim/optimizers.py	Enables fused CUDA AdamW.
src/lerobot/distributed/utils.py	Adds compilation, broadcast, and NUMA helpers.
src/lerobot/distributed/__init__.py	Exports new distributed utilities.
src/lerobot/configs/train.py	Permits compilation for non-sharded training.
src/lerobot/configs/accelerator.py	Updates DDP and compilation defaults.

Review details

Suppressed comments (2)

src/lerobot/configs/train.py:358

• The existing tests/configs/test_train_config_distributed.py::test_compile_placeholder still expects enabled=True on a non-sharded config to raise, so this change makes that test fail. Update it to cover the new contract: single-process/DDP is accepted, while explicit and auto-enabled sharded compilation are rejected.

        if self.accelerator.compile.enabled and self.parallelism.is_sharded:
            raise ValueError("--accelerator.compile is wired for DDP/single-process runs only.")

src/lerobot/configs/accelerator.py:189

• This claim contradicts the measured behavior in the PR: compiled ACT consumes dropout RNG in a different order, and same-seed losses exceed the stated tolerance. fallback_random uses eager RNG operators but does not preserve eager draw ordering, so document that same-seed compiled and eager trajectories can differ.

    # Keep eager RNG semantics inside compiled regions (dropout, the VAE's randn_like), so a
    # compiled run reproduces the eager one to fp32 rounding at the same seed.
    fallback_random: bool = True

• Files reviewed: 8/8 changed files
• Comments generated: 7
• Review effort level: Balanced

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[HaomingSong]

@TarzanZhao Thanks so much for your contribution; we will review it asap.