Performance and Benchmarking

Patterns for writing performance-sensitive code, and the two measurement tools that back any performance claim: the benchmark suite (how fast is it?) and the profiler (where does the time go?).

The four hot spots

1. Sampler steps: iterative, run 100-1000x, dominate wall time.
2. Score / energy gradients: autograd.grad calls are frequent and stack.
3. Loss forward + backward: called every training batch.
4. Host-device traffic: .item(), .cpu(), repeated .to() stall the GPU.

Optimise in that order. Everything else is noise.

Vectorise, don't loop

Work in batch dimensions; avoid Python-level iteration over samples.

# bad
energies = torch.stack([energy_fn(x[i]) for i in range(batch)])

# good
energies = energy_fn(x)                       # (batch,)

Sample many chains in parallel by putting the chain index in the leading dim.

Stay on device

Keep tensors on one device and dtype for the whole pipeline. Use self.device / self.dtype from TorchEBMModule inside the library; never hard-code "cuda".

.item(), .cpu(), .tolist(), and Python if tensor > 0: all trigger a full GPU sync; defer them until after the hot loop:

# bad: syncs every step
for step in range(n_steps):
    loss = loss_fn(x)
    if loss.item() < threshold:               # GPU stall
        break

# good: one sync per epoch
losses = []
for step in range(n_steps):
    losses.append(loss_fn(x).detach())
avg = torch.stack(losses).mean().item()       # single sync

The GPU-first contract

Every new sampler, loss, or trainer honours these rules; the existing components already do:

• Normalise transfers once at entry, never per step. Move conditioning and inputs to the device a single time when a call begins, then reuse them. TorchEBMModule._prepare_model_kwargs is the reference: it device-aligns model_kwargs once and the per-step drift closures capture the result.
• No .item() / .cpu() / .tolist() / .numpy() inside a step loop. Diagnostics and metrics stay device tensors and are stacked, then synced once at a logging boundary (BaseTrainer.train_step returns device scalars; train_epoch syncs once per epoch). Preallocate trajectories/diagnostics on device and fill them in place.
• Host scalars enter only through schedulers. get_scheduled_value returns a Python float computed on the host, so it folds into the next kernel launch with no device→host sync. Do not read step sizes or temperatures back off a device tensor.
• Respect the input dtype. BaseModel.gradient computes in the input's dtype (byte-identical for float32, no fp64 downcast); set model.force_fp32_gradient = True only when a low-precision model needs fp32-precision gradients.
• Adaptive/implicit integrators carry one bounded sync per iteration (the step-acceptance and solver-convergence checks are data-dependent and must reach the host). This is inherent; fixed-step integrators are the sync-free path. Greedy OT coupling is likewise host-bound by nature - use Sinkhorn in the training loop.

To catch a regression, wrap a hot loop in torch.cuda.set_sync_debug_mode("error"): any accidental host sync raises with the offending op named.

Reuse memory

Pre-allocate buffers once and fill them in place inside loops; for trajectories, write into a pre-allocated tensor instead of appending to a list:

traj = torch.empty(n_steps + 1, *x.shape, device=device, dtype=x.dtype)
traj[0] = x
for i in range(n_steps):
    x = sampler.step(x)
    traj[i + 1] = x

In-place ops (x.add_, x.mul_) are safe outside autograd-tracked paths.

Mixed precision and compilation

Inside the library, wrap large matmul blocks with self.autocast_context() rather than calling torch.autocast directly; this honours the user's configured dtype. At the benchmark/application layer, --amp and --compile apply the same transforms via benchmarks/registry.py::apply_mode, so eager, compiled, and mixed-precision results stay comparable.

Common pitfalls

• Implicit host-device copies: torch.tensor(x_numpy, device=...) inside a loop.
• Redundant .to() calls: BaseLoss.__call__ already moves inputs; subclass forward() must not.
• Missing torch.no_grad(): interpolation targets, momentum init, and random projections need no grad tracking.
• Tiny batches on GPU: prefer one big step over many small ones.
• isinstance checks inside the inner loop: resolve once before the loop.

Benchmarks: detecting change

The suite (pytest-benchmark under benchmarks/) auto-discovers every component exported from torchebm.*.__init__ and times its standard workload at three scales. Regular pytest tests/ never runs them.

bash benchmarks/run.sh --quick               # smoke run, small scale
bash benchmarks/run.sh --module losses       # one module
bash benchmarks/run.sh --filter "ScoreMatching"
bash benchmarks/run.sh --baseline            # save current results as baseline
bash benchmarks/run.sh --compare             # compare latest two runs
bash benchmarks/run.sh --ci                  # fail if geo-mean speedup < 0.95x

Components needing non-default construction get an entry in benchmarks/registry.py::COMPONENT_OVERRIDES; existing entries are the best reference. Modules and individual benchmarks can be excluded in benchmarks/benchmark.toml.

Publishing. Results and the dashboard live in the separate torchebm-benchmarks repository, deployed at soran-ghaderi.github.io/torchebm-benchmarks. After a run: copy the autosaved JSON from benchmarks/results/Linux-CPython-*/ into that repo and run bash scripts/publish.sh <path-to-json>; GitHub Pages auto-deploys.

Profiling: explaining change

benchmarks/profiler.py wraps torch.profiler around the same registry callables the benchmarks time. Profile only when an optimisation is non-trivial and evidence-driven: a dashboard regression to localise, a hot path rewrite to justify, or a suspected memory issue. Skip it for one-line fixes and cleanups.

The whole workflow is one before/after pair plus a diff:

python benchmarks/profiler.py run --component LangevinDynamics --scale medium --top 20 --label before
# make the change
python benchmarks/profiler.py run --component LangevinDynamics --scale medium --top 20 --label after
python benchmarks/profiler.py diff benchmarks/profiles/*_before benchmarks/profiles/*_after

Add --trace only when the top-N table is not enough (open in ui.perfetto.dev), --memory for allocator work (view at pytorch.org/memory_viz), --nvtx for Nsight Systems. Arbitrary callables profile via --callable module:factory. Outputs land under benchmarks/profiles/ (gitignored; profiles are local by design).

Division of labour: benchmarks detect a regression and track it across releases; the profiler explains it op by op on one run.