Model Serving › Model Loading and Initialization

Before any request can be served, the model must be loaded into memory and prepared for inference. This sounds straightforward but involves several design decisions that affect startup time, memory usage, and steady-state serving performance.

The most basic approach is loading all model weights into GPU memory at startup. For a 7-billion-parameter model in 16-bit floating point, this requires approximately 14 GB of GPU memory just for weights, before accounting for activations and KV caches during inference. Larger models like 70B or 405B parameters require distributed loading across multiple GPUs.

The weight loading process involves several stages. First, the checkpoint files are read from disk. PyTorch checkpoints are typically stored as sharded files that are memory-mapped to avoid loading everything into CPU RAM simultaneously. Memory mapping lets the operating system handle the actual data movement on demand, paging in portions of the weight file only when the corresponding tensor is accessed. Then weights are transferred to GPU memory, which involves PCIe bandwidth if the GPU is discrete. For large models, this transfer alone can take tens of seconds.

Lazy loading defers this work, loading tensor shards only when they are first accessed. This reduces startup time but introduces latency spikes on first access, which is unacceptable in production where the first user request should not bear the cost of initializing the system. In practice, inference servers perform eager loading with warmup requests that pre-populate caches and ensure all code paths have been compiled before serving real traffic.

Why warmup requests? Many deep learning frameworks use just-in-time compilation: the first time a particular sequence length or batch size is encountered, the framework compiles a specialized CUDA kernel for that shape. Subsequent requests with the same shape reuse the compiled kernel and are much faster. Running warmup requests before opening the server to production traffic forces this compilation to happen in advance, preventing the first real users from experiencing artificially slow responses.

These estimates reveal why model serving infrastructure is designed to avoid frequent restarts. A 70B model takes over a minute just to load from disk and transfer to GPU memory, not counting model initialization and warmup. Production serving systems aim for zero-downtime deployments that gradually shift traffic to new instances rather than restarting existing ones. A cold start causes seconds of unavailability and delays the point when a new instance can absorb its share of production traffic. Other instances must carry additional load during that window.

Safetensors is an increasingly common checkpoint format that addresses some of these loading bottlenecks. Unlike PyTorch's native format, safetensors files store tensor data with a simple binary layout and a JSON header describing tensor shapes and offsets. This allows memory-mapped loading without any deserialization overhead, and different tensors can be loaded by different threads simultaneously because their file offsets are independent. For large models, safetensors reduces load time by up to 40% compared to traditional PyTorch checkpoints.

When a model is too large for a single GPU, tensor parallelism splits individual weight matrices across multiple GPUs. Each GPU holds a shard of every layer's weights, and all-reduce operations synchronize intermediate results between GPUs during each forward pass.

Consider a linear layer with weight matrix $\mathbf{W} \in \mathbb{R}^{d \times h}$. With tensor parallelism across $N$ GPUs, each GPU holds a column shard $\mathbf{W}_i \in \mathbb{R}^{d \times h/N}$. For an input $\mathbf{x}$, each GPU computes its partial output $\mathbf{y}_i = \mathbf{x} \mathbf{W}_i$, then the results are concatenated or summed across GPUs using collective operations.

where:

• $d$: the input dimension of the linear layer
• $h$: the output dimension of the linear layer
• $N$: the number of GPUs participating in tensor parallelism
• $\mathbf{W}_i$: the weight shard held on GPU $i$, containing columns $[(i-1)h/N, ih/N)$
• $\mathbf{y}_i$: the partial output computed on GPU $i$

The trade-off is communication overhead. Every layer requires GPU-to-GPU communication, which adds latency proportional to the number of layers and the bandwidth between GPUs. NVLink (on server-class hardware) provides much higher bandwidth than PCIe, making tensor parallelism more viable on GPU clusters connected with high-speed interconnects. On PCIe-only systems, the communication overhead for 8-way tensor parallelism can add tens of milliseconds per forward pass, enough to substantially increase response latency for interactive applications.

Pipeline parallelism is an alternative that assigns different layers to different GPUs rather than splitting individual layers. This reduces inter-GPU communication to just the activations passed between stages, but introduces pipeline bubbles where GPUs wait idle while upstream stages complete. Practical systems often combine tensor and pipeline parallelism for very large models: tensor parallelism within a node (where GPUs share NVLink) and pipeline parallelism across nodes (where communication is limited to Ethernet or InfiniBand).

The difference between NVLink and PCIe interconnects becomes critical for large-scale tensor parallelism. On PCIe, the communication overhead for an 8-GPU configuration with 80 layers adds tens of milliseconds per forward pass, substantially increasing latency. NVLink reduces this to a few milliseconds, making 8-way tensor parallelism practical for serving. This is one reason that server-class GPU machines with NVLink (such as DGX systems) command a substantial price premium over workstation-class hardware with the same number of GPUs but only PCIe connectivity.

Planning memory allocation before deployment prevents out-of-memory failures under load. GPU memory serves three main consumers during LLM inference.

Model weights occupy a fixed amount of memory determined by parameter count and data type. A 7B-parameter model in float16 uses 14 GB. Using int8 quantization (as discussed in Part XXXI) halves this to 7 GB, and int4 quantization halves it again to 3.5 GB, at the cost of some quality degradation.

Activations are the intermediate tensors computed during the forward pass. Their size depends on batch size and sequence length. For autoregressive decoding, activations are relatively small because only one token is processed per step. During prefill, however, the entire prompt is processed in parallel, and activation memory scales with sequence length. A system that accepts very long prompts must plan for the activation spike that occurs when a large prefill batch is processed.

The KV cache stores the key and value tensors for all tokens in currently active sequences. This is typically the most dynamic component: more concurrent requests means more KV cache consumption. Understanding the KV cache formula is essential for capacity planning. The KV cache size per request is:

$$\text{KV cache bytes} = 2 \times n\_layers \times n\_heads \times d\_head \times \text{seq\_len} \times \text{bytes\_per\_element}$$

where:

• $2$: one tensor for keys, one for values
• $n\_layers$: number of transformer layers in the model
• $n\_heads$: number of attention heads per layer
• $d\_head$: dimensionality of each attention head
• $\text{seq\_len}$: current sequence length for the request
• $\text{bytes\_per\_element}$: 2 for float16, 1 for int8

Notice that $n\_heads \times d\_head = d\_model$, the model's hidden dimension. So the formula simplifies to $2 \times n\_layers \times d\_model \times \text{seq\_len} \times \text{bytes}$, which makes intuitive sense: each of the $n\_layers$ layers stores a full-hidden-dimension key and value vector for every token position. As sequences grow longer or as more requests are served concurrently, KV cache consumption grows proportionally.

Traditional serving systems allocated a contiguous memory block for each request's KV cache at the start of the request, sized to the maximum allowed sequence length. This caused two problems. First, most sequences end long before reaching the maximum length, leaving most of the pre-allocated memory wasted. Second, memory fragmentation gradually reduced the effective memory available as many small gaps between allocations accumulated, making it impossible to serve new requests even when there seemed to be enough free memory in aggregate.

PagedAttention, introduced by the vLLM team in 2023, solved both problems by borrowing the virtual memory paging model from operating systems. Instead of allocating one contiguous block per request, PagedAttention divides KV cache memory into fixed-size blocks (pages), typically holding 16 or 32 token positions each. Each request is assigned pages on demand as new tokens are generated. When a request finishes, its pages are returned to a free pool and immediately available for new requests.

This design has three practical consequences. First, memory waste from over-allocation nearly disappears because pages are allocated incrementally. Second, fragmentation is eliminated because pages are interchangeable units of fixed size. Third, pages from different requests can share physical memory blocks if they have identical prefixes (prompt caching), because two requests starting with the same system prompt can point their early KV cache pages at the same physical memory rather than each maintaining their own copy. This prefix caching trick is particularly valuable for chat applications where most requests begin with the same assistant persona and instructions.