RFC: torch.cuda.estimate_memory

Ran into this again last week — spun up a run on a rented A100, model loaded fine, made it about 40 steps in and OOM'd. Would've been nice to know that before burning ~15 min of GPU time.

As far as I can tell there's no way to ask ahead of time "will this fit" — memory_allocated() and memory_reserved() only tell you what's already happened, not what's about to happen. So you just... run it and find out.

Not just me apparently, there's a package on PyPI (torch-kitsune) that exists purely to do this prediction outside of torch, and I found a 2025 paper (VeritasEst) doing basically the same thing with static/CPU analysis. Feels like something that should just live in core.

Rough idea for the API:

estimate = torch.cuda.estimate_memory(model, sample_input, optimizer=optimizer)

estimate.total_mb
estimate.fits(device=0) # bool

with a breakdown available too (params_mb, optimizer_mb, activations_mb) in case you want to know what's actually eating the memory.

Params and optimizer state are easy to get exactly right — just sum sizes, apply the known multiplier per optimizer type (Adam ~2x, etc). Activations are the annoying part. My best idea is running the forward pass on meta tensors to get shapes without actually allocating anything, but that's going to fall over on data-dependent control flow or custom ops that don't support meta. So probably has to ship as "our best estimate" rather than something exact — worth saying that up front so nobody's surprised later.

Two things I'd want opinions on before going further:

should this live under torch.cuda, or somewhere backend-agnostic since the idea isn't CUDA-specific
is it worth doing activations at all in a first version, or ship params+optimizer only and add activations once people are actually using it

Not trying to solve batch-size auto-tuning or checkpointing here, just the "will this OOM" question.

Start by reviewing the existing torch.cuda memory_allocated() and memory_reserved() APIs, then assess the proposed torch.cuda.estimate_memory interface and its meta-tensor activation approach. Resolve whether the first version should be CUDA-specific or backend-agnostic and whether it estimates parameters and optimizer state only or includes activations; done means the scope and API are agreed.