LLM Training Memory Calculator for ZeRO and FSDP

Executive Summary

An LLM training memory calculator should produce a per-rank range, not a binary promise. The lower bound is the sum of persistent model states divided only where the selected strategy actually shards them. The upper planning bound adds activations, the largest simultaneously gathered parameter unit, communication and optimizer temporaries, runtime allocations, and explicit headroom. This distinction matters because DeepSpeed's own state estimator says activation and intermediate memory must be added separately ^[1], while PyTorch exposes allocated and allocator-reserved memory as different quantities ^[2]. A procurement decision should therefore carry a versioned estimate, sensitivity range, and measured validation record.

For a common mixed-precision Adam configuration, one documented ledger is 6 bytes per parameter for FP16 weights plus an FP32 main copy, 4 bytes for FP32 gradients, and 8 bytes for two FP32 Adam states, or 18 bytes before activations and transients ^{[3] [4]}. That is a configuration, not a universal rule. FSDP2, for example, can retain high-precision sharded parameters without another high-precision optimizer copy ^[5]. The calculator must ask for the exact optimizer, precision policy, framework version, and implementation.

ZeRO stage 1 shards optimizer state, stage 2 also shards gradients, and stage 3 also partitions parameters and gathers them around computation ^{[6] [7]}. PyTorch FSDP FULL_SHARD likewise shards parameters, gradients, and optimizer states ^[8]. Yet average state divided by rank count is not peak HBM: an AllGather can materialize one or more layers in transient output buffers ^[9].

The resulting distributed LLM training GPU calculator should report three counts: the mathematical minimum, the topology-compatible count, and a validation range. Megatron Core expresses total ranks as TP x PP x CP x EP x DP ^[10]. Capacity should be shown in both decimal GB and binary GiB, where 1 GiB is 1,073,741,824 bytes and 1 GB is 1,000,000,000 bytes ^{[11] [12]}. Offload can change the capacity ledger, but it cannot establish acceptable throughput. Final approval requires a representative microbatch on the pinned stack, with peak allocated, peak reserved, non-framework allocations, and checkpoint staging observed.

Introduction and Background

The planning question is deceptively simple: can a proposed full-training or fine-tuning job fit, and how many accelerators does it require? A correct answer spans four coupled ledgers: GPU high-bandwidth memory (HBM), host random-access memory (RAM), local or shared non-volatile memory express (NVMe) storage, and checkpoint storage. It also spans time. Parameters may be sharded while idle, gathered for a layer, prefetched for the next layer, and consolidated again for a checkpoint.

The report uses per rank to mean the memory visible to one training process, ordinarily one accelerator. It separates persistent state from peak additions, and it treats all user inputs as part of a versioned record. This is necessary because ordinary data parallelism gives every GPU a complete model copy ^[13], while tensor parallelism (TP) splits selected layers and pipeline parallelism (PP) assigns different layers to stages ^{[14] [15]}. Neither operation is equivalent to dividing every byte by the total GPU count.

As of September 19, 2026, the framework and hardware version are first-class inputs. Current PyTorch documentation distinguishes FSDP1 from per-parameter FSDP2, and the migration tutorial maps FULL_SHARD to reshard_after_forward=True behavior ^[16]. Hardware capacity also changes the topology choice: an eight-GPU DGX H100 has 640 GB aggregate GPU memory, while an eight-GPU DGX H200 has 1,128 GB ^{[17] [18]}. Aggregate capacity is descriptive, however. Fit must hold on every rank, including the most heavily loaded pipeline stage.

GPU Smith is an adjacent infrastructure advisor, not a competing training framework. Its published process starts with workload modelling and a reference architecture and ends with documented acceptance testing ^{[19] [20]}. That posture is useful here: the calculator is an auditable requirements model, not a substitute for a measured acceptance test.

Key Changes

Change 1: Replace one bytes-per-parameter constant with a ledger

A planning model starts with named tensors, not folklore. Let P be total parameters and T be trainable parameters. For every state i, record count N_i, storage bytes b_i, replication factor r_i, sharding divisor s_i, and location. The persistent bytes on rank j are the sum of N_i x b_i x r_i / s_i, adjusted for stage placement and imbalance. This notation exposes whether a value applies to P, T, a local pipeline partition, or a selected tensor-parallel shard.

Table 1 defines the minimum input and output ledger for an enterprise calculator.

The ledger prevents two frequent category errors. First, an FP32 master copy should not be counted because a generic article once counted it. NVIDIA describes master weights as one solution for optimizer updates ^[27], while its performance guidance says different arrangements have different footprints ^[28]. Second, a decimal device label must not be mixed with a binary calculation. The output should show raw bytes, GB, and GiB side by side.

Change 2: Model average shards and peak gathers separately

The ZeRO memory calculator branch needs a strategy table plus a peak-event model. Stage 1 partitions optimizer states. Stage 2 additionally partitions gradients. Stage 3 partitions the 16-bit parameters and gathers them as computation requires ^[29]. DeepSpeed's published lower-bound expression for stage 3 without offload is largest-layer memory plus 18 x P / total GPUs for the documented configuration ^[30]. The important term is the largest layer: it survives even when persistent state is deeply sharded.

Table 2 maps the major strategies. “Peak addition” is deliberately qualitative because wrapper boundaries, prefetch policy, buckets, and release timing determine its size.

The table answers ZeRO stage 3 memory requirements and FSDP GPU memory requirements at the right level: persistent states divide, but live unsharded units do not. FSDP's all-gather limiter is intended to constrain memory to two consecutive FSDP instances ^[36]. DeepSpeed similarly exposes a prefetch bucket measured as the maximum number of elements fetched ahead ^[37]. Those configuration values belong in the calculator schema.

Peak control is also an implementation choice. DeepSpeed documents sequential all-gather as a way to avoid large temporary flattening buffers ^[38]. Red Hat's communication description confirms that sharded weights may be gathered on every GPU before each layer's forward and backward passes ^[39]. These controls should be inputs, never hidden constants.

Change 3: Treat parallelism as a matrix, not one divisor

The distributed LLM training GPU calculator should calculate a world-size lattice before it calculates a purchase quantity. The main dimensions are:

• Data parallelism (DP): replicas process different samples; the ZeRO or FSDP sharding group is commonly tied to this dimension.

• Tensor parallelism (TP): selected matrices and their work are divided within a layer. Parameters outside those plans may remain replicated.

• Pipeline parallelism (PP): layer groups are assigned to stages. The maximum stage, not the average stage, determines fit.

• Context parallelism (CP): sequence work is partitioned where the framework supports it; it does not imply optimizer-state division.

• Expert parallelism (EP): mixture-of-experts modules distribute experts, while shared dense tensors remain subject to their own plan. Expert parallelism also introduces token dispatch, for which Megatron recommends an all-to-all dispatcher ^[40].

• Sequence parallelism and replication: these are additional flags on selected tensors, not automatic divisors for the whole ledger.

The valid total is TP x PP x CP x EP x DP, subject to framework compatibility. The sharding divisor for model states is the group that actually shards them, not necessarily world size; AWS describes state partitioning across GPUs within a data-parallel group ^[41]. In a layout illustration (Hypothetical Example), a 64-GPU job with TP=8, PP=2, and DP=4 does not automatically give a 64-way optimizer-state reduction. If sharding is scoped to DP, the divisor is 4.

Topology changes the rounded answer. AWS recommends power-of-two TP degrees ^[42] and, for multi-node TP plus PP, recommends keeping TP within nodes and placing PP across nodes ^[43]. In a rounding illustration (Hypothetical Example), the arithmetic minimum may be 10 GPUs while the first legal or operationally sensible layout is 16. The output must retain both numbers and the rounding reason.

Full Training, LoRA, QLoRA, and Offload Branches

Full training and parameter-efficient adaptation

Full fine-tuning retrains every model parameter ^[44], so T=P for gradients and optimizer states. Low-Rank Adaptation (LoRA) constrains a dense update to a low-rank decomposition rather than learning the complete update matrix ^[45]. In the Hugging Face Parameter-Efficient Fine-Tuning (PEFT) implementation, only adapter parameters are updated, greatly narrowing optimizer and gradient scope ^[46].

The calculator must use separate branches:

• Full training: base weights, gradients, optimizer states, and any main weights apply across all trainable parameters.

• LoRA: base weights stay resident according to their storage policy, while adapter parameters create the trainable gradient and optimizer ledger.

• QLoRA: the frozen base is stored in 4-bit form and gradients pass through it into trainable LoRA adapters ^[47]. The base is dequantized to BFloat16 for matrix multiplication ^[48], so storage bits alone do not describe peak compute buffers.

• Custom partial tuning: the user supplies a trainable mask or count. No frozen-base assumption is inferred from a method label.

QLoRA's paper reports 0.373 bits per parameter less metadata under its double-quantization block assumptions ^[49]. That is useful for reproducing that configuration, but it is not a universal 4-bit overhead. More importantly, the paper says activation gradients account for most LLM fine-tuning memory ^[50]. Reducing trainable state does not remove the activation branch.

The implementation must agree with the method label. Hugging Face states that training with 8-bit or 4-bit base weights supports updating extra parameters rather than those base weights ^[51]. Its LoRA reference also distinguishes ordinary projection targeting from QLoRA-style adapters on every linear transformer layer ^[52]. Adapter target lists, ranks, and dtypes therefore belong in the serialized input.

CPU and NVMe offload

An LLM CPU offload calculator is three linked capacity ledgers plus a transfer path. It should expose:

• Device-resident bytes: shards, live gathered parameters, activation range, current collectives, and runtime reserve.

• Host-resident bytes: offloaded parameters, gradients or optimizer states, pinned transfer buffers, checkpoint consolidation, dataloader demand, and operating-system reserve.

• NVMe working set: offloaded state, queue or staging buffers, checkpoint shards, temporary full-state materialization, and retention policy.

• Transfer path: device to host, host to NVMe, peer collectives, and whether computation moves with the state.

• Performance constraints: buyer-supplied minimum samples per second, maximum step time, and acceptable checkpoint window, validated rather than inferred from capacity.

DeepSpeed parameter offload is valid only with stage 3 ^[53]. Its optimizer computation runs on CPU even when NVMe is the backing device ^[54]. FSDP2's CPU policy moves parameters, gradients, and optimizer states to CPU and consequently places the optimizer step there ^[55]. These facts describe placement, not throughput. A “fits in RAM” result must never be translated into a training-rate claim.

Implementation Considerations and Process Changes

Versioned input schema

Every run should serialize enough information to reproduce the estimate:

• Identity: calculator schema version, timestamp, model revision, architecture, parameter count, largest wrapped unit, and mixture-of-experts structure.

• Training method: pretraining, full fine-tuning, LoRA, QLoRA, or a custom trainable subset.

• Software: framework, DeepSpeed or PyTorch version, optimizer class and implementation, attention kernel, precision policy, compiler settings, and allocator configuration.

• Batch geometry: sequence length, packing distribution, microbatch per rank, gradient accumulation, and global batch.

• Memory controls: activation checkpoint policy, CPU or NVMe offload, collective bucket sizes, prefetch policy, wrapping boundaries, and reshard timing.

• Parallel dimensions: DP, TP, PP, CP, EP, sharding group, replication group, and rank-to-node placement.

• Hardware: usable HBM per rank, host RAM per node, NVMe free space, accelerators per node, fabric, and reserved capacity policy.

• Output policy: arithmetic fit margin, topology rounding, required confidence band, checkpoint retention, and acceptance thresholds.

Activation checkpointing is not a simple percentage toggle. It reduces saved tensors by recomputing them during backward ^[56]. Standard self-attention is quadratic in sequence length for both time and memory ^[57], while FlashAttention avoids storing the large intermediate attention matrix for backward ^[58]. Architecture, kernel, sequence packing, and checkpoint policy must therefore select or calibrate the activation formula.

GPU-count output and validation plan

The LLM training GPU count calculator should emit these outputs in order:

1. Persistent per-rank state: lower-bound arithmetic for the selected sharding group.

2. Peak event additions: activations, largest gathered unit, prefetch, reductions, optimizer temporaries, and runtime range.

3. Raw fit count: the smallest rank count whose high estimate is below usable HBM.

4. Legal parallel layout: the first TP x PP x CP x EP x DP product that satisfies divisibility and placement constraints.

5. Node-rounded procurement count: complete nodes or the explicitly supported partial-node shape.

6. Host and storage requirement: per node and aggregate, including checkpoint staging.

7. Validation range: at least the topology-rounded minimum and the next practical layout when uncertainty overlaps the HBM limit.

8. Reasons: dominant memory rows, unverified assumptions, and which measurement can narrow the range.

Validation begins before the whole model is allocated. FSDP2 documentation supports initialization after sharding, avoiding a full pre-sharding model allocation ^[16]. Then run a representative microbatch and record:

• Configuration hash: code, model revision, exact command, environment, and serialized calculator input.

• Peak allocated HBM: the maximum held by live tensors.

• Peak reserved HBM: the caching allocator's maximum managed amount.

• External device memory: libraries such as NCCL may allocate memory outside PyTorch's profiler visibility ^[59].

• Fragmentation evidence: memory snapshot, segment distribution, and any allocator changes. PyTorch documents split-size tuning as a possible aid for borderline fragmentation cases ^[60].

• Steady-state window: exclude warm-up allocations. TensorFlow Profiler likewise advises avoiding the first batches because initialization can distort results ^[61].

• Error decomposition: measured peak minus estimated midpoint, attributed to activations, collective overlap, optimizer temporaries, runtime, or untracked allocations.

• Checkpoint test: save, load, reshard, and stage a checkpoint under the intended topology. PyTorch Distributed Checkpoint can load-time reshard between cluster topologies ^[62].

Data Analysis and Evidence

Worked planning calculation (Hypothetical Example)

Consider a fictional 7 billion parameter dense model under full fine-tuning. This is a unit-checking example, not a benchmark. Assume the documented 18-byte mixed-precision Adam ledger, an eight-rank DP sharding group, a user-measured 0.6 GB largest wrapped unit, 12 to 20 GiB of activations, 2 to 4 GiB of collective and optimizer temporaries, 3 to 5 GiB of runtime allocations, and 15 percent safety headroom. No throughput is claimed.

The unsharded state is 7,000,000,000 x 18 = 126,000,000,000 bytes, or 126 GB and 117.35 GiB. Stage 3 or full sharding divides that persistent state by eight to 15.75 GB, or 14.67 GiB, then adds the 0.56 GiB largest unit. Before headroom, the range is 14.67 + 0.56 + 12 + 2 + 3 = 32.23 GiB to 14.67 + 0.56 + 20 + 4 + 5 = 44.23 GiB. Applying 15 percent headroom yields 37.07 to 50.86 GiB per rank.

Table 3 compares outputs from the same declared assumptions.

This comparison isolates the consequence of parameter replication. It does not prove that stage 2 or stage 3 is faster, because communication overlap, kernel choice, interconnect, and offload paths are outside a capacity-only calculation. Pipeline schedules also change activation peaks: every concurrently active microbatch carries activations and gradients ^[63].

Hardware capacity is necessary but not sufficient

Official system specifications illustrate why node layout belongs beside per-rank output. An eight-GPU DGX B200 supplies 1,440 GB aggregate GPU memory ^[64]. An eight-accelerator AMD MI300X baseboard supplies 1.5 TB HBM3, or 192 GB per accelerator ^{[65] [66]}. DGX H100 and H200 list 2 TB of host memory ^[67].

The physical fabrics differ as well. DGX H100 and H200 document 900 GB/s GPU-to-GPU bandwidth through fourth-generation NVLink ^[68], while DGX B200 documents 14.4 TB/s aggregate fifth-generation NVLink switch bandwidth ^[69]. AMD describes the eight-accelerator MI300X platform as fully meshed over Infinity Fabric ^[70]. These are topology inputs, not interchangeable throughput predictions.

Those specifications constrain the ledgers, but they do not replace them. Aggregate HBM cannot repair one oversized pipeline stage. Host capacity cannot be allocated entirely to optimizer state because the operating system, data pipeline, pinned buffers, and checkpoint consolidation also need space. Interconnect bandwidth does not establish end-to-end training throughput without message sizes and overlap. The calculator should preserve every capacity and performance constraint as a separately testable claim.

Offload expands the usable hierarchy rather than erasing movement costs. The ZeRO-Infinity paper explicitly spans GPU, CPU, and NVMe memory ^[71]. Likewise, an official server specification can confirm an eight-GPU H100 configuration ^[17], but only a workload trace can show whether the planned traffic meets the acceptance threshold.

Implications and Future Directions

For enterprise platform and finance teams, the most important implication is that GPU memory requirements for LLM training are not one scalar. They are a versioned distribution over ranks and training phases. A defensible procurement packet should include the input schema, formula version, low and high estimate by ledger row, topology diagram, configuration hash, validation logs, and the exception list.

Several process changes follow:

• Procure against a validated topology: buy complete supported units after TP, PP, and sharding groups are laid out.

• Retain uncertainty explicitly: use a range until activation and temporary allocations are measured.

• Separate capacity from performance: offload can make a state ledger fit while failing a buyer's step-time threshold.

• Treat checkpointing as a peak event: include save, load, consolidation, and topology change in host and storage capacity.

• Recalculate after software changes: framework releases, optimizer implementations, wrap policies, kernels, and allocator settings can move the peak.

• Test the heaviest rank: pipeline imbalance and expert placement can make the average misleading.

• Record both allocated and reserved memory: a single dashboard number cannot explain allocator headroom.

Future calculator versions should add architecture-specific activation plugins for grouped-query attention, mixture-of-experts routing, sequence packing distributions, and fused optimizers. They should also ingest a measured pilot trace and update uncertain rows instead of applying an opaque correction factor. This creates a useful feedback loop: arithmetic proposes the topology, a small-scale or full-rank test measures the peak, and the next version reconciles the difference.

Frequently Asked Questions (FAQs)

How many GPUs are needed to train an LLM?

Calculate per-rank high memory under each legal parallel layout, select the smallest layout below usable HBM, then round to a supported node topology. Report the raw minimum, topology-rounded count, and validation range. Do not divide total bytes by aggregate HBM alone.

What does a ZeRO memory calculator need beyond parameter count?

It needs optimizer states, gradient and parameter dtypes, main-weight policy, trainable count, DP sharding degree, largest gathered unit, collective buckets, prefetch, activations, runtime reserve, offload placement, and headroom. DeepSpeed's cold estimator can avoid loading the model after parameter counts are known ^[72], but its state estimate still does not replace the activation and temporary ledger.

What should an FSDP memory calculator report?

It should report persistent shards by FSDP strategy, the largest unsharded unit, prefetch overlap, reshard timing, activation and runtime ranges, host placement, and both allocated and reserved measured peaks. The calculation must also name whether it targets FSDP1 or FSDP2 because their state representation and controls differ.

Are ZeRO stage 3 and FSDP FULL_SHARD identical?

They share the broad policy of sharding parameters, gradients, and optimizer states. PyTorch explicitly calls FSDP its ZeRO-3 implementation ^[73]. They are not interchangeable calculator profiles because gather timing, wrapping, prefetch, optimizer representation, checkpoint APIs, and versions differ.

Does CPU or NVMe offload guarantee a job will run fast enough?

No. It can solve a device-capacity constraint by relocating state, but the result still depends on transfer volume, overlap, CPU optimizer work, storage behavior, and the workload's required step time. Capacity fit and throughput acceptance are separate gates.

Why can measured memory exceed the estimate?

Likely causes include activation assumptions, overlapping gathered units, optimizer temporaries, allocator reserve, fragmentation, CUDA context and libraries, and allocations outside the framework profiler. Short-lived temporaries can create an out-of-memory spike even when steady state appears safe ^[74].

Conclusion

A useful LLM training memory calculator is a transparent model of tensors, placement, time, and topology. It starts with a byte-accurate persistent-state ledger, applies only the sharding divisors that the pinned framework actually uses, and adds the peak events that average arithmetic hides. It handles full training, LoRA, and QLoRA as different branches. It reports GPU, host, NVMe, and checkpoint capacity independently.

The decision output is a range with three GPU counts: mathematical minimum, topology-compatible procurement count, and validation range. A result is ready for procurement only after a representative microbatch and checkpoint path have been measured on the intended software stack, with allocated, reserved, and external device memory recorded. That approach turns a convenient estimate into an auditable infrastructure requirement without pretending that capacity alone predicts training performance.