Estimate model memory before downloading the weights

A model's parameter count is a useful starting point for a memory estimate. It is not the amount of memory the whole application will need.

I separate the budget into weights, working memory, cached state and the application around the model. That makes a failed load easier to explain than one large number labelled "model size".

The architecture and figures below are hypothetical. They are calculations, not hardware benchmarks or claims about a particular model.

Weights-only estimates exclude metadata, mixed-precision tensors, allocator overhead and runtime workspaces.

Start with the weights

For P parameters stored at b bits each:

weight_bytes = P * b / 8

Three billion parameters at sixteen bits need six billion bytes, about 5.59 GiB. At four bits, the simple estimate falls to 1.5 billion bytes, about 1.40 GiB.

That is a lower-level arithmetic estimate. A real quantised format may store scale factors, block metadata and some tensors at a higher precision. The file size and loaded memory can therefore be larger than parameters / 2 for four-bit weights.

Keep decimal GB and binary GiB separate when comparing the result with an operating system's memory display. One GiB is 2**30 bytes.

Budget for the context cache

For a conventional transformer key/value cache, a useful simplified estimate is:

cache_bytes = 2 * layers * kv_heads * head_dimension
              * sequence_length * bytes_per_element * batch_size

The leading two represents keys and values. Use the number of key/value heads, which can differ from the number of query heads.

With 28 layers, 8 key/value heads, a head dimension of 128, 4,096 cached tokens, two bytes per element and a batch size of one:

cache_bytes = 469,762,048
            = 448 MiB

Doubling the cached sequence length doubles this estimate to 896 MiB. Actual cache implementations can use different layouts, precisions or allocation strategies, so check the runtime rather than assuming the formula describes every architecture.

Put both calculations in one place

parameters = 3_000_000_000
weight_bits = 4
layers, kv_heads, head_dim = 28, 8, 128
tokens, element_bytes, batch_size = 4096, 2, 1

weights = parameters * weight_bits / 8
cache = (2 * layers * kv_heads * head_dim
         * tokens * element_bytes * batch_size)

print(round(weights / 2**30, 3))  # 1.397 GiB
print(cache / 2**20)             # 448 MiB

The sum is still not a complete capacity requirement. It leaves out temporary tensors, compute buffers, framework allocations and everything else the process is doing.

Measure the high point

Loading and inference can have different peaks. A loader may briefly hold a source representation and a converted representation together. A long prompt can require more working memory than the short prompt used in an initial test.

I would measure idle memory, loading peak, prompt-processing peak and generation with the intended context length. Include another request if the application supports concurrent inference. A batch size of one does not describe two simultaneous sessions.

Memory mapped weights complicate the reading further: virtual address space, resident memory and accelerator allocations are different measurements. Record which one the tool is reporting.

Treat the estimate as a filter

The calculation is good at ruling out an obviously unsuitable configuration before a large download. It is also useful for finding the variable worth reducing. A smaller context can change cached state without changing the weights; a different precision can change weight storage without fixing an unrelated application leak.

Once the estimate is close to the available capacity, run a representative load test and leave headroom for the actual process. A successful one-token demonstration is a narrow result.

Reference: Hugging Face's quantisation documentation describes implementation-specific weight compression and memory considerations.