🚀[FEA]: VRAM-aware autobatching for inference

[hholb]

Is this a new feature, an improvement, or a change to existing functionality?

Improvement

How would you describe the priority of this feature request

Low (would be nice)

Please provide a clear description of problem you would like to solve.

Disclaimer: I am new to AI weather forecasting and earth2studio. Please let me know if I am misunderstanding something or am missing something obvious.

Summary

run.deterministic / run.ensemble currently leave batch sizing to the user.
For deterministic runs the only batch lever is the number of initialization
times passed in time, which batch_func flattens into a single forward pass —
there is no memory-aware chunking, so oversizing simply OOMs. I'd like an
opt-in autobatcher that measures/estimates per-sample GPU memory and picks the
largest safe batch automatically, modeled on
torch-sim's autobatching.

Motivation / problem

We roll a global prognostic model across a full season of initialization dates
(one deterministic rollout per issue date, ~45-day lead) on an
A100-80GB. At batch size 1 the GPU sits at ~10% utilization for the entire
~650s rollout

Batching multiple initialization dates into one deterministic(time=[...]) call
is the fix (it turns the time dim into the model batch dim via batch_func),
but today that means:

1. Manual guesswork. We hard-code a batch size (SEASON_ROLLOUT_BATCH) and
   hope it fits.
2. OOM-driven iteration. Too high → CUDA OOM mid-rollout → wasted GPU
   minutes → lower it → rerun. Repeat per model, since the safe batch differs
   across models and with output-variable selection.
3. Non-portable constants. A value tuned for A100-80GB is wrong on
   40GB/L40S, and wrong again when we change lead time, output_coords, or grid.
4. Left-on-the-table utilization. To stay safe we under-batch, so we never
   actually saturate the GPU.

This is exactly the class of problem torch-sim solved for MD/relaxation
workloads, and prognostic weather rollouts have the same shape: many
independent states through one model, batch bounded by VRAM.

Current behavior (for reference)

• deterministic(time, nsteps, prognostic, data, io, output_coords=..., ...)
  has no batch-size parameter; all time entries become one batch via the
  batch_func decorator (earth2studio/models/batch.py), which flattens
  leading dims into the model batch dim and runs a single forward — no VRAM
  accounting, no sub-batching.
• ensemble(..., batch_size: int | None = None, ...) does chunk the
  ensemble dim, but the size is user-supplied and static — same manual-tuning /
  OOM problem, just on a different dimension.

Proposed feature

An opt-in, VRAM-aware autobatcher usable by the built-in workflows and directly
by users, e.g.:

# opt in, let earth2studio size the batch (not married to this API)
io = deterministic(time=all_init_dates, nsteps=180, prognostic=model,
                   data=data, io=io, autobatch=True)

Desired behavior:

• Estimate or measure peak per-sample forward memory once (a probe forward,
  or an analytic scaler keyed on grid size / #variables / #levels / lead
  window), then compute the largest batch that fits under a memory budget.
• Sub-batch transparently across the batchable dimension (init time for
  deterministic, ensemble members for ensemble) and stream results to the IO
  backend, so callers pass "all N" and never think about chunk size.
• Configurable budget + safety margin (target free-memory fraction, padding
  factor) with a conservative default.
• Graceful fallback / OOM backoff: on OOM, halve and retry the offending
  sub-batch rather than failing the whole run.
• Applies uniformly to deterministic, ensemble, and ideally diagnostic.

Prior art: torch-sim

torch-sim implements this pattern and is a good reference for the API and the
memory-estimation approach:

• Repo: https://github.com/TorchSim/torch-sim
• Docs: https://torchsim.github.io/torch-sim/
• Autobatching tutorial:
  https://torchsim.github.io/torch-sim/tutorials/autobatching_tutorial.html
• Implementation:
  https://github.com/TorchSim/torch-sim/blob/main/torch_sim/autobatching.py

Relevant pieces to mirror:

• measure_model_memory_forward(...) — probes peak GPU memory of one forward
  pass to ground the estimate empirically (uses
  torch.cuda.max_memory_allocated()).
• determine_max_batch_size(...) / estimate_max_memory_scaler(...) — derive
  the largest safe batch from measured/estimated per-sample cost.
• calculate_memory_scalers(..., memory_scales_with="n_atoms"|"n_atoms_x_density"|"n_edges")
  — a pluggable "what does memory scale with" metric. The weather analogue
  would be something like grid-points × #variables × #levels (× lead-window for
  the accumulated IO), which earth2studio can compute from the model's
  coords.
• BinningAutoBatcher (fixed set of states → memory-binned batches) and
  InFlightAutoBatcher (refills the batch as members finish — useful for
  variable-length / convergence workloads). The BinningAutoBatcher shape maps
  cleanly onto "N initialization dates through one prognostic model."
• Ergonomics worth copying: autobatcher=True as a one-line opt-in, plus
  max_memory_scaler / max_memory_padding knobs for users who want to skip
  the probe or add headroom.

Alternatives considered

• Manual batch_size on deterministic (mirroring ensemble): better than
  nothing and a reasonable first step, but still leaves users tuning per
  model/GPU/config and watching for OOM
• User-side chunking (what we do now): works, but every downstream project
  reinvents memory estimation and OOM backoff

Additional context

• Hardware: A100-80GB, single GPU per rollout.
• Observed: ~10% GPU utilization at batch 1; batch is currently VRAM-bound to
  ~2 for a 45-day global rollout, tuned by hand.
• We already restrict output_coords to the single variable we need, which both
  shrinks IO and raises the achievable batch — an autobatcher should account for
  the active output_coords when estimating memory.

Happy to help test against global prognostic models if useful.