Benchmarking and Profiling

OpenTau ships four diagnostic scripts under src/opentau/scripts/. Three pin down where a training run is spending its time; the fourth (benchmark_inference.py) measures per-call inference latency. All read the same TrainPipelineConfig as opentau-train, so you can point them at any config JSON. The three training profilers reproduce the exact model / dataset / batch size the real run uses; benchmark_inference.py instead builds a dummy observation and (for a random-init policy) forces IDENTITY normalization, since it is measuring shape / dtype / compile-bound latency, which is weight- and data-independent.

Script	When to reach for it
`profile_step.py`	Measure per-step wall-clock, broken down into forward / backward / optimizer / sync phases. Use this first when asking “why is training slow?”
`profile_dataloader.py`	Measure the dataloader-only throughput ceiling (no model, no collective). Use this to rule out input-pipeline starvation before looking at the GPU.
`find_unused_params.py`	List parameters that receive no gradient during a real forward+backward. Use before setting `FIND_UNUSED_PARAMS=false` on a new policy.
`benchmark_inference.py`	Measure per-call `policy.sample_actions` latency under deployment-like conditions (random-init expert, bf16, batch 1). Use to characterize serving / on-robot inference speed.

A worked example end-to-end — diagnosing a real “low GPU utilization” issue, ruling out dataloader starvation, and confirming the bottleneck is DeepSpeed’s per-parameter hook overhead — is documented in issue #177.

profile_step.py — per-step timing breakdown

Mirrors opentau-train’s setup (Accelerator, dataset mixture, policy, optimizer, LR scheduler) and runs a short loop. Splits per-step wall-clock into eight phases and reports mean / median / p95 for each.

Basic usage — same launch incantation as opentau-train:

accelerate launch \
    --config_file configs/examples/accelerate_ddp_config.yaml \
    src/opentau/scripts/profile_step.py \
    --config_path=configs/libero/reproduce_pi05_libero.json \
    --batch_size=12

Example output (8×A100, DDP, pi05 / LIBERO at bs=12):

=========== profile_step results (rank 0) ===========
warmup=20 measured=200 ranks=8
batch_size=12 num_workers=16 prefetch_factor=8
wall-clock over full loop: 265.89s

phase            stats                                                           share
------------------------------------------------------------------------------------------
dataload_wait    mean=   1.24ms  median=   1.18ms  p95=   1.60ms                 0.1%
forward          mean= 378.86ms  median= 378.65ms  p95= 381.75ms                32.8%
bwd              mean= 706.32ms  median= 704.74ms  p95= 710.24ms                61.1%
unscale_clip     mean=  19.55ms  median=  19.43ms  p95=  21.12ms                 1.7%
optim_step       mean=  39.22ms  median=  39.17ms  p95=  39.81ms                 3.4%
zero_grad_sched  mean=   1.79ms  median=   1.74ms  p95=   2.74ms                 0.2%
backward_step    mean= 766.89ms  median= 765.58ms  p95= 770.31ms                66.3%
sync_gather      mean=   9.64ms  median=   6.47ms  p95=   9.78ms                 0.8%
total            mean=1156.63ms  median=1151.73ms  p95=1158.28ms                <-- total

throughput: 0.86 steps/s, 83.0 global samples/s
=====================================================

Reading the output:

High share in dataload_wait (> ~5%) means the dataloader is not keeping up with the GPUs. Run profile_dataloader.py to confirm and raise num_workers / set persistent_workers=True / enable_cpu_affinity=True as appropriate.
High share in bwd with a large gap between 1-GPU and N-GPU typically indicates distributed-backend overhead. Try a single-GPU run (PROFILE_NO_OPTIM=1) for comparison — if single-GPU bwd is close to ~2× forward and N-GPU is much larger, the delta is host-side work, not compute.
High share in optim_step is normal for large models; if > 10% and you’re on CUDA, make sure AdamConfig.fused is True (the default since PR #176).
High share in sync_gather points at the accelerator.gather_for_metrics(...).item() calls in update_policy. They run every step (not just every log_freq) and can be gated behind log_freq if they become a bottleneck.

Environment variables (all optional):

Variable	Default	Effect
`PROFILE_STEPS`	200	Number of measured steps (after 20 warmup). `cfg.steps` from the training config is intentionally ignored because production configs typically set it to 1M.
`PROFILE_NO_OPTIM`	0	When `1`, skip optimizer creation and `optimizer.step` / `zero_grad` entirely. Useful for isolating raw forward + backward compute on a single GPU (no Adam state allocated, so a large bf16 model fits without ZeRO partitioning).
`FIND_UNUSED_PARAMS`	true	Toggles DDP’s `find_unused_parameters` kwarg. Silently ignored under DeepSpeed. Set to `false` after auditing the policy with `find_unused_params.py`.
`FUSED_ADAMW`	(unset)	When `true` / `false`, force-toggle `torch.optim.AdamW(fused=...)` for an A/B comparison without editing the optimizer config JSON. Unset leaves the factory default alone.
`PROFILE_STEP_JSON`	(unset)	When set to a file path, rank 0 writes a JSON summary of phase means and medians after the loop. Convenient for scripted A/B sweeps.

profile_dataloader.py — dataloader throughput ceiling

Builds the exact same WeightedDatasetMixture.get_dataloader() the training loop uses (num_workers, prefetch_factor, pin_memory, HierarchicalSampler) and iterates batches with no model, no optimizer, no collective. Any slowdown here is pure input-pipeline cost.

Run under the same launcher as training so the host CPU sees the real multi-rank × N-worker pressure:

accelerate launch \
    --config_file configs/examples/accelerate_ddp_config.yaml \
    src/opentau/scripts/profile_dataloader.py \
    --config_path=configs/libero/reproduce_pi05_libero.json

Example output (8 ranks):

[rank 0/8] fetch=mean=  80.76ms median=   0.20ms p95= 605.39ms | h2d=mean= 0.40ms ...
[rank 1/8] fetch=mean=  83.40ms median=   0.16ms p95= 754.82ms | h2d=mean= 0.49ms ...
... one line per rank ...

=========== profile_dataloader summary (rank 0) ===========
world_size=8 batch_size=12 num_workers=16 prefetch_factor=8
wall-clock over full loop: 27.23s
per-rank batches/s (min / mean / max): 11.99 / 12.45 / 12.93
cluster-wide samples/s (ceiling, no model): 597.6
===========================================================

Reading the output:

Compare cluster-wide samples/s against the samples/s from profile_step.py. If dataloader throughput is at or below the training step rate, the input pipeline is your bottleneck. If dataloader throughput is comfortably ahead, the bottleneck is GPU-side (forward / backward / optim).
Bimodal fetch distribution (median ≈ 0 ms, p95 ≈ hundreds of ms) means you’re alternately hitting the prefetch buffer and blocking on worker decode. That’s normal — only the mean matters for long-run throughput.

Environment variables:

Variable	Default	Effect
`PROFILE_BATCHES`	300	Number of measured batches after 20 warmup. Raise if p95 is high and you want a more stable mean.

find_unused_params.py — list parameters DDP would reject

Runs one forward + backward on a real batch on a single GPU (no DDP, no DeepSpeed) and prints every parameter where param.requires_grad is True but param.grad is None after backward. Those are exactly the parameters DDP would refuse to sync with find_unused_parameters=False.

Run as a plain Python invocation — no accelerate launch needed:

python src/opentau/scripts/find_unused_params.py \
    --config_path=configs/libero/reproduce_pi05_libero.json

Example output (pi05, after the PR that dropped gemma_expert.lm_head):

#==============================================================================
# pi05 parameter audit — single forward + backward, single GPU
# include_zero_grad=False
#==============================================================================

========== UNUSED (requires_grad=True, grad is None) — DDP will refuse without
                  find_unused_parameters=True (0 tensors, 0 params) ==========

========== FROZEN (requires_grad=False) — context (8 tensors, 256 params) ==========
  [normalize_discrete_actions.buffer_actions.max]  (1 tensors, 32 params)
    - normalize_discrete_actions.buffer_actions.max  shape=(32,)
  ...

# USED (requires_grad=True, grad is non-trivial): 814 tensors
# Tip: if UNUSED list is empty, you can flip
#  DistributedDataParallelKwargs(find_unused_parameters=False) safely.

Recommended workflow:

Run find_unused_params.py on your policy.
If UNUSED is empty, set FIND_UNUSED_PARAMS=false when launching opentau-train (or drop the DistributedDataParallelKwargs kwarg in train.py for your fork) to reclaim the per-step graph-walk cost.
If UNUSED is non-empty, each reported tensor is either an orphan in the model graph (fix by freezing it or deleting the module) or a parameter that’s only conditionally reached (fix by adding an unconditional graph edge, e.g. + 0 * unused_param.sum() in the loss).

Environment variables:

Variable	Default	Effect
`INCLUDE_ZERO_GRAD`	false	When `true`, also list parameters whose grad tensor exists but is all-zero. These are typically paths touched on some batches but producing no learning signal on this one. Usually safe to ignore; useful for auditing conditionally-active heads.

Example: a typical benchmarking session

Given low samples/s in a training run, a sensible sequence to rule out candidates in order of likelihood:

# 1. Is the dataloader keeping up? (~2-10 minutes)
accelerate launch \
    --config_file configs/examples/accelerate_ddp_config.yaml \
    src/opentau/scripts/profile_dataloader.py \
    --config_path=<your_config.json>

# 2. Where does per-step time go? (~4 minutes at 200 steps)
accelerate launch \
    --config_file configs/examples/accelerate_ddp_config.yaml \
    src/opentau/scripts/profile_step.py \
    --config_path=<your_config.json> \
    --batch_size=<your_bs>

# 3. Is DDP's find_unused_parameters=True costing you? Only if
#    backward_step looks unusually high. (~1 minute, single GPU)
python src/opentau/scripts/find_unused_params.py \
    --config_path=<your_config.json>

# 4. A/B an optimizer or distributed-backend change without
#    touching any config file:
FUSED_ADAMW=false accelerate launch ... profile_step.py ...
FUSED_ADAMW=true  accelerate launch ... profile_step.py ...

DeepSpeed ZeRO-2 vs ZeRO-3 for pi05 full fine-tuning

ZeRO-3 shards the model parameters across ranks on top of the gradient and optimizer-state sharding that ZeRO-2 already does. That extra sharding pays off only when a single replica of the model does not fit in one GPU’s memory — it adds a per-layer parameter all-gather in the forward (and a matching reduce-scatter in the backward) that ZeRO-2 does not need. pi05 is ≈3.3B parameters and fits comfortably replicated on an 80 GB GPU, so ZeRO-3 has nothing to gain and pays the all-gather cost.

Measured on 8×A100-80GB, full fine-tuning (no frozen weights; ``freeze_vision_encoder=false``, ``train_expert_only=false``), bf16, sdpa attention, ``use_torch_compile=false``, ``gradient_accumulation_steps=1``, 8 ranks, with the configs/examples/accelerate_deepspeed* configs and the pi05 reference policy (2 cameras at 224×224, chunk_size=10, predict_response=true) on TensorAuto/libero:

Backend	Per-rank batch	Global batch	sec/step	samples/s	Peak GPU mem
ZeRO-2	8	64	5.64	11.4	53.3 GiB
ZeRO-3	8	64	7.71	8.3	62.8 GiB
ZeRO-2	16	128	5.29	24.2	78.7 GiB
ZeRO-3	16	128	9.86	13.0	79.2 GiB

Both backends OOM at the same per-rank batch size on this hardware (16 fits, 18 OOMs): ZeRO-3 frees ~5 GB of replicated parameters per rank, but its parameter all-gather/prefetch buffers plus the extra allocator fragmentation consume a comparable amount, so the maximum batch is unchanged. At a matched batch size ZeRO-2 is ~1.4× faster at batch 8 and ~1.9× faster at batch 16 (the per-step parameter all-gather is the difference; both keep fp32 master weights and step the optimizer identically).

Recommendation: use plain DDP (fastest) or ZeRO-2 for pi05 and similarly sized policies. Reach for ZeRO-3 only when a single replica no longer fits per GPU (much larger backbones / many-billion-parameter experts). ZeRO-3 is fully supported and validated for pi05 — training, checkpoint save, resume, offline checkpoint consolidation (convert_checkpoint.sh), and in-training validation all work — it is simply not the throughput-optimal choice at this model size.

To reproduce, run the same config under each accelerate file (both at 8 ranks / gradient_accumulation_steps=1) and read the per-step time from the logs; samples/s = per_rank_batch × num_ranks ÷ sec_per_step:

COMMON="--policy.freeze_vision_encoder=false --policy.train_expert_only=false \
    --policy.use_torch_compile=false --policy.attention_implementation=sdpa \
    --batch_size=16 --dataloader_batch_size=16 --gradient_accumulation_steps=1"

# ZeRO-2
accelerate launch --config_file configs/examples/accelerate_deepspeed_config.yaml --num_processes 8 \
    src/opentau/scripts/train.py --config_path=configs/examples/pi05_training_config.json $COMMON

# ZeRO-3
accelerate launch --config_file configs/examples/accelerate_deepspeed_zero3_config.yaml \
    src/opentau/scripts/train.py --config_path=configs/examples/pi05_training_config.json $COMMON

Note

If a ZeRO-3 run OOMs from fragmentation (the error mentions “reserved but unallocated” memory), set PYTORCH_ALLOC_CONF=expandable_segments:True in the environment to recover the fragmented blocks.

benchmark_inference.py — single-policy inference latency

The three scripts above dissect a training step. benchmark_inference.py answers a different question — how long a single policy.sample_actions call takes at deployment time — and so is set up to look like inference, not training: it needs no policy checkpoint (the action expert and its projections are randomly initialized because pretrained_path is null), while still loading the real pretrained backbone the config specifies (load_pretrained_backbone=true); it runs in bfloat16 at batch_size=1 with N cameras and a short language prompt (a 24-word sentence, bounded by prompt_max_length), and measures per-call latency rather than per-step throughput. Like the training profilers it reads the same TrainPipelineConfig JSON as opentau-train, so you can point it at any policy config.

Because latency depends only on shapes / dtype / compile — not on weight values — the random-init numbers below equal what a fully-trained policy would produce on this same hardware and config (latency is weight- and data-independent). The harness forces IDENTITY normalization when pretrained_path is null (so the empty stats buffers are never read), builds a dummy N-image + prompt observation, optionally wraps policy.model.sample_actions in torch.compile, runs BENCH_N_WARMUP warmup calls followed by BENCH_N_TIMED timed calls with a cuda.synchronize() on both sides of each, and writes a JSON summary (mean / std / p50 / p95 / p99 / min / max, plus every raw sample and the full host / GPU / driver / library versions) to BENCH_OUTPUT_DIR.

Two scoping details matter for reading the numbers. First, the timed region is the full policy.sample_actions — it includes the CPU-side Qwen3-VL preprocessing (tokenize + image-process + chat-template) as well as the GPU forward, so it reflects true end-to-end call latency. Second, torch.compile is applied to policy.model.sample_actions (the backbone prefill + the flow-matching expert loop), not to that CPU preprocessing — so the compiled speedups below come entirely from the GPU forward.

Run it as a plain Python module — no accelerate launch needed (it is a single-process, single-GPU benchmark):

# Eager (no torch.compile)
BENCH_NO_COMPILE=1 \
    python -m opentau.scripts.benchmark_inference \
    --config_path=configs/benchmarks/cosmos3.json

# torch.compile, reduce-overhead mode (CUDA graphs) — the fastest setting
BENCH_COMPILE_MODE=reduce-overhead \
    python -m opentau.scripts.benchmark_inference \
    --config_path=configs/benchmarks/cosmos3.json

cosmos3 inference latency (B200)

Measured on 1× NVIDIA B200 (183 GB HBM, sm_100), batch size 1, bfloat16, torch 2.10.0+cu128 / CUDA 12.8 / cuDNN 91002 / driver 580.126.20 / transformers 4.57.6 / Python 3.10, using benchmark_inference.py with configs/benchmarks/cosmos3.json. The workload is the real pretrained TensorAuto/cosmos3-reason-32b reasoner backbone (~33B params, full depth, not truncated) with a randomly-initialized ~0.9B, full-64-layer action expert (condition_on_layer=None, so the expert reads a per-layer KV cache): 3 cameras at 224×224 + one 24-word prompt (capped at prompt_max_length=64 tokens) → a 50-action chunk (chunk_size = n_action_steps = 50), num_steps=10 Euler denoising steps, sdpa attention. Each number is the mean over BENCH_N_TIMED=50 timed calls after 5 warmup calls:

cosmos3 sample_actions latency — 3 img + prompt → 50-action chunk, num_steps=10 (1× B200, bf16, batch 1)
Configuration	mean (ms)	p50	p95	p99	speedup
Eager (no compile)	314.57	313.83	319.73	322.04	1.00×
`torch.compile` (default)	129.85	129.68	131.67	133.90	2.42×
`torch.compile` (reduce-overhead / CUDA graphs)	106.94	106.86	107.54	109.70	2.94×

Latency is very stable across the 50 timed calls: run-to-run std is 2.14 ms eager, 0.85 ms compiled (default), and 0.45 ms with CUDA graphs — the p99 sits within ~2–3 % of the mean in every case.

Where the time goes (num_steps decomposition). A sample_actions call is one 32B backbone prefill plus num_steps passes of the action expert, so latency is affine in num_steps: L(N) ≈ F + N·s, where F is a fixed cost — the single Qwen3-VL backbone prefill over the 3 images + prompt, the CPU-side Qwen3-VL preprocessing, and rope / noise setup — and s is one ~0.9B action-expert forward over the 51 suffix tokens (1 state + 50 action) cross-attending the cached backbone KV. Re-running at num_steps=1 (eager 82.13 ms, compiled-default 57.97 ms) and solving the two points gives, eager, s = (314.57−82.13)/9 = 25.83 ms/step and F ≈ 56.3 ms; compiled-default, s = (129.85−57.97)/9 = 7.99 ms/step and F ≈ 50.0 ms. At num_steps=10 the 10 expert passes dominate the call — 258 of 315 ms (82 %) eager, 80 of 130 ms (62 %) compiled. This is why torch.compile helps so much: its win is almost entirely on the expert loop (25.8 → 8.0 ms/step, ~3.2×), which is launch- and fusion-bound — many small kernels across 64 layers × 10 steps — exactly what Inductor fusion and CUDA graphs (reduce-overhead) collapse; the compute-bound 32B prefill barely moves (56 → 50 ms).

Note

The first warmup call pays the torch.compile cost: at num_steps=10 the initial call takes ~58 s in default mode and ~53 s in reduce-overhead (which logs “cudagraph partition into 2 partitions due to CUDAGraph-unsafe custom ops” — it still graphs most of the forward); num_steps=1 / default is ~15 s. Steady state is reached by warmup call 2, so BENCH_N_WARMUP=5 leaves ample margin. The harness asserts the call-1 / call-N ratio is ≥ 5× when compiling, to catch a silent compile bail.

Environment variables (all optional):

Variable	Default	Effect
`BENCH_NO_COMPILE`	0	When `1`, skip `torch.compile` entirely; `policy.model.sample_actions` runs eager. (The timed region is still the full `policy.sample_actions` call.) Use for the eager baseline.
`BENCH_COMPILE_MODE`	`max-autotune`	The `mode=` passed to `torch.compile` (ignored when `BENCH_NO_COMPILE=1`). The results above used `default` and `reduce-overhead` (the latter enables CUDA graphs and was fastest).
`BENCH_N_WARMUP`	5	Warmup calls before timing. Call 1 absorbs the compile cost; leave at ≥ 5 so the harness can assert compile actually fired.
`BENCH_N_TIMED`	50	Timed calls after warmup. The reported mean / std / p50 / p95 / p99 / min / max are computed over these.
`BENCH_OUTPUT_DIR`	`benchmark_results`	Directory for the per-run JSON summary (`<host>_<policy>_<timestamp>.json`), which stores the config snapshot, host / GPU / driver / library versions, the warmup series, every timed sample, and the summary statistics.