Implementation:FMInference FlexLLMGen DeepSpeed Checkpoint

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

Overview

Vendored DeepSpeed checkpoint management class that loads, reshapes, and provides access to distributed model checkpoints across tensor-parallel (TP), pipeline-parallel (PP), and data-parallel (DP) dimensions.

Description

DeepSpeedCheckpoint is the central class for managing DeepSpeed checkpoint files. It handles checkpoint folders containing model files (mp_rank_*), layer files (layer_*), and ZeRO optimizer state files. The class supports reshaping checkpoints when the target parallelism degrees (TP, PP, DP) differ from the source.

Key capabilities:

• Parallelism detection -- Reads source TP, PP, and DP degrees from the ZeRO checkpoint metadata and compares them with the target degrees.
• 2D parallel mapping -- Maintains old and new 2D parallel maps (Megatron-style) for mapping rank files to (PP, TP) indices. When parallelism degrees change, the checkpoint is reshaped using reshape_meg_2d_parallel.
• Layer management -- Identifies embedding layers, transformer layers, and final layer norms from file prefixes. Builds mappings from PP/TP indices to the corresponding layer files.
• State dict merging -- For tensors that need to be concatenated across TP ranks, uses LAYER_CONCAT_DIM to determine the correct concatenation dimension (e.g., dim=1 for attention output weights, dim=0 for others). Sequential layers (biases, layer norms) are taken from rank 0 only.
• Global state -- Extracts iteration count and training arguments from rank 0's model file.

This is AUTO_KEEP vendored code from DeepSpeed.

Code Reference

Field	Value
Repository	FlexLLMGen
File	benchmark/third_party/DeepSpeed/deepspeed/checkpoint/deepspeed_checkpoint.py
Lines	1-315

Key Class:

class DeepSpeedCheckpoint(object):
    def __init__(self, dir, tp_degree=None, pp_degree=None, dp_degree=None):
        ...

    def get_transformer_state(self, tp_index: int, pp_index: int) -> list: ...
    def get_embedding_state(self, tp_index: int) -> Dict: ...
    def get_final_norm_state(self, tp_index: int) -> Dict: ...
    def get_2d_parallel_state(self, tp_index: int, pp_index: int) -> dict: ...
    def get_zero_checkpoint_state(self, pp_index, tp_index, dp_index) -> dict: ...
    def get_iteration(self): ...
    def get_args(self): ...

I/O Contract

Inputs

Parameter	Type	Required	Description
dir	str	Yes	Path to DeepSpeed checkpoint directory
tp_degree	int	No	Target tensor parallelism degree (auto-detected if None)
pp_degree	int	No	Target pipeline parallelism degree (auto-detected if None)
dp_degree	int	No	Target data parallelism degree (auto-detected if None)

Outputs

Method	Return Type	Description
get_transformer_state	list of dicts	State dicts for transformer layers at given TP/PP indices
get_embedding_state	dict	Merged embedding layer state dict for given TP index
get_final_norm_state	dict	Final layer norm state dict for given TP index
get_2d_parallel_state	dict	Merged model state dict for given TP/PP indices
get_iteration	int	Training iteration number from checkpoint

Related Pages

• Principle:FMInference_FlexLLMGen_Distributed_Checkpoint_Management