Principle:FMInference FlexLLMGen Inference Checkpoint Loading
| Field | Value |
|---|---|
| Sources | Paper: FlexGen, DeepSpeed Inference Documentation |
| Domains | Checkpointing, Inference, Model_Parallelism |
| Last Updated | 2026-02-09 00:00 GMT |
Overview
A weight loading strategy for inference engines that maps pre-trained checkpoint parameters to optimized kernel parameter layouts, handling tensor-parallel sharding, on-the-fly quantization, and multiple checkpoint format variations in a single pass.
Description
Inference checkpoint loading bridges the gap between how model weights are stored in training checkpoints and how they are organized in inference-optimized kernels. Training checkpoints typically store weights per-layer with standard PyTorch naming, while inference kernels may use fused layouts (e.g., combined QKV weights) and different data formats (e.g., INT8).
Key principles:
- Parameter name mapping -- Inference kernels have fixed internal parameter names (attn_qkvw, attn_ow, inter_w, output_w, norm_w, etc.) that differ from the source checkpoint's naming. The loading process uses a param_names mapping to translate between the two. Different model architectures (GPT-2, BLOOM, OPT) have different mappings.
- QKV fusion -- Training checkpoints may store Q, K, V projection weights as three separate tensors or as a single combined tensor. The loader handles both cases, concatenating separate Q/K/V tensors when needed and applying model-parallel slicing to the combined result.
- Tensor-parallel weight distribution -- For TP inference, each GPU needs only a slice of each weight tensor. The loader computes the correct slice based on the current rank and total TP degree, using different split dimensions for different parameter types:
- Column-parallel weights (QKV, MLP input): split along the output dimension.
- Row-parallel weights (attention output, MLP output): split along the input dimension.
- Biases: split or replicated depending on the associated weight's parallelism.
- On-the-fly quantization -- When INT8 inference is requested but the source checkpoint is FP16, the loader applies weight quantization during the loading process. This avoids the need to maintain separate quantized checkpoint files. Per-group quantization scales are computed and attached to the weight tensors.
- Checkpoint merging and splitting -- When the source checkpoint's TP degree differs from the target:
- Source TP > Target TP (merging) -- Multiple source shards are concatenated along the appropriate dimension.
- Source TP < Target TP (splitting) -- Source tensors are split into smaller slices for each target rank.
- Weight transposition -- Some checkpoint formats store weights in column-major order while the inference kernel expects row-major. The loader performs in-place transposition when needed, using a custom transpose function that reshapes without copying.
- Layer-type dispatch -- Different layer types (Linear, Embedding, LayerNorm, DeepSpeedTransformerInference) require different loading strategies. A layer_policies dictionary maps each layer class to its loading function.
- ZeRO-3 parameter materialization -- For models initialized with ZeRO-3 (where parameters are partitioned across data-parallel ranks), the loader uses GatheredParameters to temporarily materialize full parameters before loading.
Usage
Use inference checkpoint loading when initializing a DeepSpeed inference engine from a pre-trained model checkpoint. The loading process is triggered automatically by the InferenceEngine when a checkpoint path is provided in the configuration.
Key considerations:
- Checkpoint format -- Determine whether the checkpoint uses TP-sharded or PP-sharded format, and whether Q/K/V are combined or separate.
- TP degree matching -- For best performance, use the same TP degree as the source checkpoint. Cross-TP-degree loading works but requires tensor merging/splitting.
- Quantization -- INT8 quantization during loading provides the same model quality as pre-quantized checkpoints while avoiding checkpoint management overhead.
Theoretical Basis
The loading process performs a tensor redistribution operation that can be modeled as:
For each parameter P in the inference kernel:
1. Identify source parameter(s) S in the checkpoint
2. Load S to CPU/GPU memory
3. If QKV fusion needed: S = concat(Q, K, V)
4. If TP redistribution needed: S = slice(S, tp_rank, tp_degree)
5. If quantization needed: S = quantize(S, num_bits, num_groups)
6. If transposition needed: S = transpose(S)
7. Copy S into P
The order of operations matters: transposition must happen before quantization (so quantization groups are computed on the correctly-oriented tensor), and TP slicing must happen before transposition (to minimize the size of the transposed data).