Aussie AI Blog

LLM Training versus Inference Efficiency

April 25, 2026

by David Spuler, Ph.D.

Training Versus Inference

LLM training and inference have different compute, memory, and network bottlenecks. The cost of performing large-model LLM training for frontier models is becoming a barrier to the advancement towards better models. There are a variety of techniques to speed up LLM training, and here we compare them to the inference situation.

Training and inference have fundamentally different roles in the LLM cost hierarchy. Here are some of the key points of difference and similarity:

Users — training has no users; inference costs are amortized over many user queries.
Inputs — user queries are the input for inference versus huge datasets of training data.
Outputs — inference outputs an answer to users; training doesn't output text.
Components — both have many overlapping components (e.g., attention, FFN/GLU), but some different ones.

Differences and Similarities

Some of the biggest technical differences between the two architectures include:

Backward pass — training has a whole phase called "backward propagation" or "backprop" that's not in inference.
Decoding phase — training doesn't have the inherently sequential "decoding phase" that outputs the tokens in inference.
Parallelism — LLM training compute is inherently parallelizable, and similar to the "prefill phase" of inference (but decoding is sequential).
KV cache — optimizing KV caching data is critical for inference, but training does not use a KV cache at all.

Transformer components that are used in both inference and training:

Tokenization and embeddings
Attention modules (e.g. Flash Attention)
FFNs or GLUs (matrix multiplications and activation function)
Layer Normalization (LayerNorm or RMSNorm)
Unembedding module (final Softmax, logits, and probability generation)

One of the fundamental similarities between training and inference is that they both need optimization of the same basic architectural features:

Compute
Memory usage
Network cost
Disk storage

Let's look at each of these attributes in turn.

Transformer Compute Components

Inference has these extra components:

Decoding algorithm (e.g. greedy, top-k, beam)
KV caching mechanisms (e.g., storing the KV cache, KV cache compression, etc.)
Attention optimizations of the KV cache (e.g., Paged Attention in vLLM, Radix Attention in SGLang)
Speculative decoding (not useful in training, which is already parallel)
Token output mechanism

Training has various unique components:

Backwards propagation ("backprop")
Gradient optimizer (usually AdamW)
Checkpointing (aggregate all model updates)

The optimization of compute has many of the same approaches in training:

Kernel optimizations
Vectorization (GPU)
Multi-GPU parallelization
Low-precision arithmetic (FP8 or FP4)

Sometimes, one of the best ways to speed up compute is to avoid having the GPU wait for data transfers from memory.

Memory Usage

Memory usage is often more important than compute, and memory access efficiency is critical for both inference and training. The memory profile differs in inference between the two phases:

Prefill — compute-bound with tokenwise parallelized steps.
Decoding — memory-bound with sequential algorithm.

The decoding phase is actually somewhat more nuanced:

Attention module (decoding phase) — memory-bound because of KV cache data for each layer.
FFN or GLU — compute-bound matrix multiplications that don't use the KV cache.

Note that the attention module in the prefill phase is more compute-bound. It does need to store the per-layer KV cache data, but this is cheaper than reading the KV cache data in the various self-attention patterns during decoding. Also, the size of the KV cache is linear with the sequence length, so it depends on the input tokens in prefill, but on the output length in the decoding phase.

Training is inherently parallel (prefill-like) and more compute-bound in the main forward and backward passes, especially in their matrix multiplications (GEMM kernels). Generally, the compute cost of the backward pass is the main bottleneck, because it involves more matrix operations than the forward pass. However, memory cost is significant in the gradient optimizer, which has to read and write not only the computed gradients, but also the internal optimizer state (extra AdamW data that must be stored), and the updated model weights. Fused AdamW kernels can minimize this memory usage, but not eliminate it. Note that the weight matrices are the same in forward and backward passes, with the backward pass GEMM kernels using the transpose of the matrices used in the forward pass.

The ways to reduce memory costs in LLM training include such techniques as:

Kernel fusion — fewer memory accesses of temporary data (e.g., fused AdamW optimizer kernel).
Memory-efficient attention kernels (did someone say Flash Attention?)
GPU sharding — less data to store per GPU (e.g., DDP or FSDP).
Recomputation — trading extra compute time for less storage (e.g., activation checkpointing).
Data size reductions — FP8 or FP4 training.

Networking

Network bandwidth is one of the lesser known optimization methods. Both inference and training require very fast networking interconnects, but they have very different profiles. Obviously, this assumes a multi-GPU data center scenario with inter-GPU communication, not local on-device execution of an LLM on your desktop PC.

Inference uses networking for this traffic in multi-GPU or multi-server architectures:

Handling incoming user queries and routing them.
Sending input text or tokens out to disaggregated prefill and decoding servers.
Sending KV cache data to those servers (e.g., prefix KV caching)
Receiving the output answer tokens back from the servers.
Routing the answer text back to the correct user network session.

Sending KV cache data back and forth is often the largest part of network load in inference. Hence, there are optimizations that specifically target these types of network transfers, such as "Mooncake", which is a distributed file system specialized for KV cache data.

Training has heavy networking traffic from:

Sending out batches of training data to the GPUs
Receiving back gradients and updated parameter values.

Training tends to have a bursty network traffic profile, since both of those incoming and outgoing loads are occurring somewhat in lockstep. Inference network loads are more dependent on the user request patterns.

Some of the specific network optimization techniques in LLM training include:

Faster network interconnects — buy more gear (e.g., NVLink, InfiniBand, RoCE).
GPU sharding — reduces the network traffic to parts of the model.
Reducing CPU-GPU traffic — avoid CPU coordination traffic on the PCIe bus (e.g., CUDA graphs, fused kernels).
Gradient compression — send less gradient data in all-reduce operations.
Network congestion management — addressiing the burstiness of traffic in training.

Disk Storage

Using a faster SSD with NVMe is one non-obvious way to speed things up. Both inference and training need to offload some data from the GPU to CPU RAM, or often to disk. Again, the reasons are different for the two algorithms:

Inference — needs to store KV cache data (temporarily in GPU VRAM or CPU RAM, but later stored for use in multi-query KV cache reuse optimizations), which grows unboundedly with long outputs.
Training — stores gradients and optimizer state information (e.g. AdamW m and v first and second-order "moment" data)

There are various other niche algorithms where some data might be offloaded to RAM or disk, but the above are the main reasons.