AI Training vs. Inference: Why They Need Different Kinds of Memory

Training a large language model and running that model for inference use the same underlying neural network, but they put very different pressure on memory. A training step must keep model parameters, gradients, optimizer state, and the activations needed for backpropagation. Inference does not need gradients or optimizer state, but it still needs the trained weights, temporary runtime memory, and — for autoregressive generation — a KV cache that grows with active context. That structural difference, not simply "training is bigger," is why the two workloads scale differently and why the common shorthand about compute-bound training and bandwidth-bound inference is useful only with important caveats.

What Training Actually Holds in Memory

Backpropagation needs more than the model's parameters. A full training step also keeps gradients, optimizer state, and activations from the forward pass that are needed during the backward pass. Parameters, gradients, optimizer states, and activations are therefore four major components of training memory, although the exact footprint depends on precision, optimizer choice, parallelism, checkpointing, and framework implementation.

One widely cited baseline comes from Microsoft's DeepSpeed/ZeRO work. For mixed-precision training with Adam, Microsoft estimated that a one-trillion-parameter model requires roughly 16 TB just for parameters, gradients, and optimizer states. That corresponds to about 16 bytes per parameter under the specific mixed-precision Adam assumptions used in that analysis.

The qualifier matters. The 16-bytes-per-parameter figure is not a universal law of training. Different optimizers, precisions, sharding strategies, and optimizer-state compression techniques change the number substantially. Adam keeps two moment estimates per parameter in addition to the parameter and gradient state, while simpler optimizers can require less optimizer memory. Any training-memory estimate is therefore much more useful when the optimizer and precision assumptions are stated explicitly.

Activation Memory Grows With Batch Size and Sequence Length, Not Just Model Size

For a fixed model, optimizer, and precision, the total logical model-state footprint is tied mainly to parameter count. Activation memory behaves differently: it changes with factors including batch size, sequence length, model architecture, and parallelism strategy, and it can become a major capacity constraint in long-context training.

NVIDIA researchers measured this directly in Reducing Activation Recomputation in Large Transformer Models. Their sequence-parallelism and selective-recomputation techniques reduced activation memory by about 5× and cut the execution-time overhead associated with activation recomputation by more than 90%. In one 530-billion-parameter GPT-3-style training setup on 2,240 A100 GPUs, Model FLOPs Utilization (MFU) rose from 42.1% with full recomputation to 54.2%. Those are results from that specific setup, not a universal multiplier for every training run.

Why Large Models Need Many GPUs at Once

Large-model training often exceeds the HBM capacity of a single accelerator, and fitting the model is only part of the problem: the required compute would also make training impractically slow without large-scale parallelism. Modern training therefore combines techniques such as tensor parallelism, pipeline parallelism, and data parallelism, with model states and activations partitioned or replicated in different ways depending on the system.

The Megatron-LM paper reported training iterations for a one-trillion-parameter model at 502 petaFLOP/s across 3,072 GPUs. Meta later said that Llama 3.1 405B was trained on more than 16,000 NVIDIA H100 GPUs and more than 15 trillion tokens. Those GPU counts reflect both memory requirements and the enormous amount of computation that must be completed within a practical training time.

A Real Example: Llama 3.1 405B's Training vs. Inference Memory Footprint

Training: About 6.48 TB Under a 16-Byte Adam Baseline

If the 16-bytes-per-parameter mixed-precision Adam baseline above is applied to Llama 3.1 405B, the arithmetic is straightforward: 405 billion parameters × 16 bytes ≈ 6.48 TB of aggregate model-state memory. That figure excludes activations and does not represent the actual per-GPU footprint after techniques such as sharding, replication, checkpointing, or offloading are applied.

For a rough capacity comparison, an NVIDIA H100 SXM has 80 GB of HBM3. Dividing 6.48 TB by 80 GB gives about 81 H100 SXM GPUs' worth of raw HBM capacity just for that hypothetical aggregate model-state footprint. That is a capacity-equivalent calculation, not a statement that 81 GPUs would be sufficient to train Llama 3.1 405B in practice.

Inference: About 810 GB Down to 203 GB for the Weights

Inference removes the training-only gradients and optimizer state, but it still requires more than just the checkpoint. The model weights must be loaded, and serving also consumes runtime workspace, temporary activations, and — during autoregressive generation — KV-cache memory.

Hugging Face's Llama 3.1 deployment guide estimates that the 405B checkpoint alone requires about 810 GB in FP16, 405 GB in FP8, or 203 GB in INT4. Hugging Face explicitly notes that these figures cover checkpoint memory only and exclude reserved space for kernels or CUDA graphs. It also separately estimates substantial KV-cache memory at long context lengths.

An eight-GPU H100 SXM node has about 640 GB of raw HBM capacity, so the 810 GB FP16 checkpoint cannot fit in that node even before runtime overhead is counted. The 405 GB FP8 checkpoint leaves much more room, but whether a real serving workload fits comfortably still depends on context length, batch/concurrency, KV-cache size, and framework overhead.

Put side by side, the hypothetical 6.48 TB training model-state estimate is about eight times the 810 GB FP16 weight-only inference estimate. That comparison is useful only as an illustration. The training number assumes the specific 16-byte mixed-precision Adam model-state baseline, while the inference number counts FP16 checkpoint weights only. Training activations and inference KV cache are excluded from that 8× ratio, and both can be substantial.

Batch Size Means Something Different in Training and Inference

Training does not have a user waiting for each individual result, so its batch configuration is usually chosen around a combination of optimization behavior, hardware efficiency, memory capacity, and overall time-to-train. Large global batches can be distributed across many GPUs and may also be built through gradient accumulation.

Online inference has a different objective. A serving system must balance throughput against response latency, especially when many requests arrive concurrently. Current MLPerf Inference datacenter benchmarks therefore separate latency-constrained server scenarios from throughput-oriented offline scenarios rather than treating "inference" as one workload.

Offline inference — bulk document processing is one example — is usually much less sensitive to per-request latency and can optimize more aggressively for throughput. Online serving, by contrast, has to keep time-to-first-token and per-token latency within acceptable limits while still using batching efficiently. Training generally optimizes time-to-train rather than interactive per-request latency.

Is Training Always Compute-Bound and Inference Always Bandwidth-Bound?

It is common shorthand to say training is compute-bound and inference is bandwidth-bound, but that compresses several different bottlenecks into one sentence. Training often demands enormous compute, yet memory capacity, communication, and activation storage can become limiting factors as model and context sizes grow.

Recent long-context training research makes the memory side concrete. MEMO reported training a 7B model with a one-million-token sequence length on eight A800 GPUs by managing activation memory through fine-grained swapping and recomputation. SlimPipe targets the accumulation of activation memory across microbatches in pipeline-parallel long-context training. These are research systems rather than universal production recipes, but both demonstrate that activation memory can become a first-order constraint.

Inference also has to be split into phases. NVIDIA's inference guidance describes the prefill phase, which processes the input prompt, as highly parallel and compute-intensive, while the autoregressive decode phase is typically memory-bound because weights and KV-cache data must be moved repeatedly while only one new token is generated at a time. That prefill-versus-decode distinction is more accurate than saying all inference is bandwidth-bound.

For the hardware side of this problem, see HBM Explained: Why AI GPUs Need High-Bandwidth Memory. For the broader system-level gap between AI compute growth and memory, see The AI Memory Wall Explained.

The next layer of the problem is capacity beyond local GPU HBM. Technologies such as CXL-attached memory are being explored for memory expansion and pooling, while inference systems are also increasingly looking at multi-tier KV-cache designs. Those are separate architectural questions, and each deserves its own treatment rather than being treated as a simple replacement for HBM.