Principle:FMInference FlexLLMGen Distributed Checkpoint Management
| Field | Value |
|---|---|
| Sources | Paper: FlexGen, DeepSpeed ZeRO Documentation |
| Domains | Checkpointing, Distributed_Training, Model_Parallelism |
| Last Updated | 2026-02-09 00:00 GMT |
Overview
A structured approach to saving and loading model checkpoints from distributed training that supports reshaping across different tensor-parallel, pipeline-parallel, and data-parallel configurations.
Description
Distributed checkpoint management solves the problem of persisting and restoring model state across training runs that may use different parallelism topologies. When a model is trained with TP=4, PP=2, DP=8, the checkpoint consists of many sharded files. To load this checkpoint with different parallelism degrees (e.g., TP=2, PP=4), the checkpoint manager must reshape the data.
Key principles:
- File-based sharding -- Each rank's model state is stored in a separate file with a standardized prefix (mp_rank_XX_model_states.pt for model files, layer_XX for layer files). This enables partial loading without reading the entire checkpoint.
- 2D parallel mapping -- Megatron-style 2D maps index checkpoint files by (PP_index, TP_index) pairs. When the parallelism degree changes, a reshape operation computes a new mapping that correctly distributes parameter slices to the new topology.
- Tensor merging and splitting -- When TP degree decreases, parameter tensors from multiple ranks must be concatenated along the correct dimension. Self-attention output weights concatenate along dim=1, while most other weights concatenate along dim=0. Sequential parameters (biases, layer norms) that are replicated across TP ranks are taken from a single rank.
- ZeRO state coordination -- ZeRO optimizer states (partitioned gradients and optimizer states) must also be reshaped to match the new topology. The ZeRO checkpoint tracks source TP/PP/DP degrees to enable this.
- Layer-to-pipeline mapping -- Transformer layers are evenly distributed across pipeline stages. The mapping from layer indices to PP indices is recomputed when PP degree changes.
Usage
Use distributed checkpoint management when:
- Changing the number of GPUs between training and inference.
- Converting a training checkpoint to an inference-optimized format.
- Resuming training with a different parallelism configuration after hardware changes.
- Loading Megatron-LM or DeepSpeed checkpoints into the FlexLLMGen inference pipeline.
Theoretical Basis
Distributed checkpointing views the model as a 3D tensor decomposition over TP, PP, and DP dimensions. Each parameter tensor is partitioned according to a split plan that assigns tensor slices to ranks. Reshaping between parallelism configurations is equivalent to computing a new split plan and performing the corresponding gather/scatter operations on the stored tensor slices.
For parameters split along TP:
W_original = concat(W_rank0, W_rank1, ..., W_rankN, dim=split_dim)
W_new_rank_i = split(W_original, new_tp_degree, dim=split_dim)[i]
The split_dim depends on the parameter type: column-parallel layers split along the output dimension, row-parallel layers split along the input dimension.