AI model memory and bandwidth sizer

How it works

Whether a model fits, and how fast it can run, is mostly a memory question. This tool adds up the weights, the KV cache, the activations, the gradients and the optimiser state, splits them across tensor, pipeline and data parallel devices with ZeRO sharding, and compares the result with the memory of one accelerator. For serving it also finds the memory bandwidth a target token rate needs, because each decoded token reads the weights and the cache.

1. Count the parameters. Enter the parameters in billions, or let the tool count them from the architecture: embeddings (twice if the output head is not tied), then for every layer the attention projections (queries use all heads, keys and values use the KV heads), the feed-forward block (three matrices when gated, two when plain) and the norms. For a mixture of experts the feed-forward block is repeated for every expert and a router is added.attention per layer = hidden x heads x head size + 2 x hidden x KV heads x head size + heads x head size x hidden; MLP per expert = 3 (or 2) x hidden x FFN size; total = embeddings + layers x (attention + norms + experts x MLP)
2. Treat a mixture of experts. All experts must stay in memory, so memory follows the total parameters. Each token uses only the shared parameters plus its chosen experts, so compute follows the active parameters. When you enter totals instead of an architecture, the tool separates the expert and shared parts from the total and the active figures.per-expert parameters = (total − active) / (experts − active experts); shared = total − experts x per-expert
3. Size the weights. Weights are the parameter count times the bytes of the chosen precision: 4 for FP32, 2 for BF16, 1 for FP8 or INT8, and 0.5 for INT4. Scales for quantised formats are not counted.weights = parameters x bytes per parameter
4. Size the KV cache. Every token of every sequence stores a key and a value in every layer. Grouped-query attention stores only the KV heads, which is why the cache is smaller than with one KV head per query head.KV cache = 2 x layers x KV heads x head size x context tokens x sequences x bytes
5. Size the training activations. For a standard transformer layer with 16-bit activations, Korthikanti et al. count 34 sbh bytes plus 5 a s^2 b bytes stored per layer, where s is the sequence length, b the micro-batch, h the hidden size and a the heads. Of the 34, layer norms take 4, attention 11 and a feed-forward block of width 4h takes 19. Here the attention share is reduced for grouped-query attention and the feed-forward share follows the real width, with a gated block storing about 8 times its width. Selective recomputation drops the s^2 term, full recomputation keeps only each layer's input (2 sbh), and sequence parallelism splits what tensor parallelism leaves whole. The first pipeline stage holds one micro-batch per stage, so it stores the activations of all layers.per layer = s x b x h x (4 + (7 + 4 x KV heads / heads) + MLP share + 5 x a x s / h) x width / 2, divided by tensor parallel degree with sequence parallelism; selective: without the 5 a s / h term; full: 2 s b h
6. Add gradients and optimiser state. Gradients take the weight width (never below 16 bits). Adam with FP32 master weights keeps a master copy and two moments, 12 bytes a parameter, which with 2-byte weights and gradients makes the 16 bytes a parameter that Rajbhandari et al. use. SGD with momentum keeps a master copy and one moment; plain SGD keeps only the master copy. A trainable share below 100% (an adapter) limits gradients and state to that share. The output logits and their gradient are held in FP32.gradients = trainable parameters x bytes; optimiser state = trainable parameters x (4 master + 8 Adam moments); logits = 8 x s x b x vocabulary / tensor parallel
7. Shard across devices. Tensor and pipeline parallelism split the weights, gradients and optimiser state by their product. ZeRO then splits across the data parallel ranks: stage 1 the optimiser state, stage 2 also the gradients, stage 3 also the weights. When serving, the cache splits across the KV heads and pipeline stages, so tensor parallel degrees above the KV head count copy it.per device = weights / (tp x pp x [dp if ZeRO 3]) + gradients / (tp x pp x [dp if ZeRO 2 or 3]) + state / (tp x pp x [dp if ZeRO 1 to 3]) + activations + logits
8. Add overhead and compare with the device. A share for fragmentation and framework buffers, and a fixed amount for the runtime, are added to the sum. The tool then reports the headroom, the smallest layout that fits, and for serving the largest number of sequences that fit at the context length.device total = sum x (1 + overhead share) + reserved memory; headroom = device memory − device total
9. Find the decode bandwidth. One decode step produces one token for every sequence in the batch. It reads the weights it uses once, and every sequence's cache. For a mixture of experts the weights read are the shared part plus the experts the batch touches, taking routing as uniform. Pipeline stages run one after another on a token, so only tensor parallelism cuts what each device reads.bytes per step = weights used x bytes + KV cache bytes; bandwidth per device = (weights / tp + cache / min(tp, KV heads)) x tokens per second; experts touched = E x (1 − (1 − k / E)^batch)

Sources

The methods and default figures rest on these public sources. The wording here is our own; read the source for its full text.

• Reducing Activation Recomputation in Large Transformer Models (opens in a new tab), arXiv:2205.05198 (Korthikanti, Casper, Lym, McAfee, Andersch, Shoeybi, Catanzaro; NVIDIA), 10 May 2022. Activation memory per layer, selective and full recomputation, sequence parallelism.
• ZeRO: Memory Optimizations Toward Training Trillion Parameter Models (opens in a new tab), SC20, International Conference for High Performance Computing, Networking, Storage and Analysis (Rajbhandari, Rasley, Ruwase, He; Microsoft), November 2020 (arXiv v3, 13 May 2020). Sixteen bytes a parameter for mixed-precision Adam and the three ZeRO stages.
• GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints (opens in a new tab), arXiv:2305.13245 (Ainslie, Lee-Thorp, de Jong, Zemlyanskiy, Lebron, Sanghai; Google Research), EMNLP 2023, 22 May 2023. Grouped-query attention and its smaller KV cache.
• The Llama 3 Herd of Models (opens in a new tab), Meta AI, arXiv:2407.21783, 31 July 2024. Architecture of the 8B and 70B models: layers, hidden size, feed-forward size, heads, KV heads, vocabulary.
• Mixtral of Experts (opens in a new tab), Mistral AI, arXiv:2401.04088 (Jiang and others), 8 January 2024. Eight experts with two active per token: about 47 billion total and 13 billion active parameters.
• NVIDIA H100 Tensor Core GPU (opens in a new tab), NVIDIA, page read 25 October 2026. 80 GB HBM3 at 3.35 TB/s.
• ThinkSystem NVIDIA HGX B200 180GB 1000W GPU Product Guide (LP2226) (opens in a new tab), Lenovo Press, last updated 13 January 2026. 180 GB HBM3e at 7.7 TB/s.
• AMD Instinct MI300X Accelerator data sheet (opens in a new tab), Advanced Micro Devices, 2025. 192 GB HBM3 at 5.3 TB/s.
• TPU v5p (opens in a new tab), Google Cloud documentation, page updated 22 October 2026, read 25 October 2026. 95 GiB HBM at 2,765 GB/s per chip.
• Intel Gaudi 3 AI Accelerator White Paper (opens in a new tab), Intel, March 2025 (V1 Rev. 2). 128 GB HBM2e at 3.7 TB/s.