Implementation:FMInference FlexLLMGen DeepSpeed Load Checkpoint

Field	Value
Sources	Repo: FlexLLMGen
Domains	Checkpointing, Inference, Model_Parallelism
Last Updated	2026-02-09 00:00 GMT

Overview

Vendored DeepSpeed checkpoint loading module for inference that handles loading pre-trained model weights into DeepSpeed's optimized inference transformer layers, with support for tensor-parallel sharding, weight quantization, and multiple checkpoint formats (tp-sharded and pipeline-parallel).

Description

load_checkpoint.py provides the load_model_with_checkpoint function, which is the primary mechanism for loading model weights into DeepSpeed's inference-optimized transformer layers. It handles the complex mapping between source checkpoint parameter names and the DeepSpeed inference kernel's internal parameter names (e.g., attn_qkvw, attn_ow, inter_w, output_w).

Key capabilities:

• Two checkpoint modes -- Supports tp (tensor-parallel) and pp (pipeline-parallel) checkpoint types, each with different loading strategies:
  • TP mode -- Loads parameters directly from sharded checkpoint files. Handles splitting (when source TP degree < target TP degree) and merging (when source TP degree > target TP degree) of weight tensors. Supports INT8 quantized checkpoints with associated scale tensors.
  • PP mode -- Uses named parameter mappings to copy weights via maybe_copy and maybe_copy_qkv helper functions. Supports Megatron-v2 QKV transposition.
• Weight quantization during load -- Applies on-the-fly weight quantization via a weight_quantizer object. For FP16 weights, quantizes after transposition. For INT8 checkpoints, preserves existing quantization scales.
• QKV handling -- Supports three QKV parameter formats:
  • Combined QKV weight (14 params): single qkv_w and qkv_b tensors.
  • Separate Q/K/V weights with bias (18 params): individual q_w, k_w, v_w with biases.
  • Separate Q/K/V weights without bias (12 params): individual q_w, k_w, v_w without biases.
• Model-parallel slicing -- Uses mp_replace to copy weight slices based on the current TP rank. Handles both qkv_copy (three-way split for Q/K/V) and regular copy for other layers.
• Layer-type dispatch -- Uses a layer_policies dictionary to dispatch loading logic based on module type. Standard layers (nn.Linear, nn.Embedding, nn.LayerNorm) use simple copy, while DeepSpeedTransformerInference layers use the specialized transformer loader.
• ZeRO-3 awareness -- Detects ZeRO-3 partitioned parameters (numel() == 0 or is_meta) and uses GatheredParameters context manager to materialize them before loading.
• Embedding weight tying -- After loading, ties lm_head.weight to the embedding weight for BLOOM-style models.
• Memory management -- Aggressively calls gc.collect() and deletes state dict references after loading to minimize peak memory usage.

This is AUTO_KEEP vendored code from DeepSpeed.

Code Reference

Field	Value
Repository	FlexLLMGen
File	benchmark/third_party/DeepSpeed/deepspeed/module_inject/load_checkpoint.py
Lines	1-362

Key Function:

def load_model_with_checkpoint(r_module,
                               sd,
                               mp_replace,
                               ckpt_type,
                               weight_quantizer=None,
                               rank=0,
                               param_names=None,
                               transformer_config=None,
                               megatron_v2=False):
    ...

I/O Contract

Inputs

Parameter	Type	Required	Description
r_module	nn.Module	Yes	The DeepSpeed inference model to load weights into
sd	list of dicts	Yes	List of state dictionaries from checkpoint shards
mp_replace	ReplaceWithTensorSlicing	Yes	Helper for model-parallel weight slicing and copying
ckpt_type	str	Yes	"tp" for tensor-parallel or "pp" for pipeline-parallel checkpoints
weight_quantizer	WeightQuantization	No	Quantizer for on-the-fly INT8 quantization during load
rank	int	No	Current TP rank (default: 0)
param_names	list	No	Ordered list of parameter name mappings for the transformer layer
transformer_config	object	No	Transformer configuration (heads count, etc.) for QKV reshaping
megatron_v2	bool	No	Enable Megatron-v2 QKV transposition (default: False)

Outputs

The function modifies r_module in-place, loading weights from the state dictionaries into the corresponding module parameters. No explicit return value.

Related Pages

• Principle:FMInference_FlexLLMGen_Inference_Checkpoint_Loading