Implementation:Deepspeedai DeepSpeed DeepSpeedEngine Save For TP

Overview

Concrete tool for saving tensor-parallel model checkpoints with parameter gathering provided by the DeepSpeed library.

Implementation Type

Method (instance method on DeepSpeedEngine)

Detailed Description

DeepSpeedEngine.save_checkpoint() with AutoTP-trained models uses GatherReplacedLayerParams internally to gather all TP-partitioned parameters before writing the checkpoint. The saved checkpoint contains the full (un-partitioned) model weights, making it portable across different TP configurations.

The checkpoint saving flow for TP models involves:

1. save_checkpoint(save_dir, tag, client_state, save_latest, exclude_frozen_parameters): The top-level entry point (L3695-3789 in engine.py). It creates the checkpoint directory, validates tag consistency across ranks, then delegates to _save_checkpoint().
2. _save_checkpoint(): Calls module_state_dict() to obtain the model state dict.
3. module_state_dict(): Returns the standard PyTorch state_dict() of the underlying module. For TP models, this returns partitioned weights by default.
4. Parameter gathering via get_layer_state_dict(): For consolidated checkpoints (e.g., when calling _consolidated_16bit_state_dict()), the engine uses a recursive get_layer_state_dict() function (L4170-4184) that wraps each module's parameters in a GatherReplacedLayerParams context. This performs AllGather across TP ranks to reconstruct the full weight tensors. Only rank 0 collects the gathered parameters into the state dict.
5. Synchronization: After saving, get_accelerator().synchronize() ensures all GPU communication is complete, followed by dist.barrier() for cross-rank coordination.

The GatherReplacedLayerParams context manager (L327-387 in layers.py) handles the gather/repartition cycle:

• On enter: Calls gather_params(params_list) on the first TP parameter, which stores the current partitioned data in data_partition and performs dist.all_gather_into_tensor() to reconstruct the full tensor.
• On exit: Calls _tp_partition(params_list) to re-slice the parameters back to their TP-partitioned form, restoring the training state.

Code Reference

• Repository: https://github.com/deepspeedai/DeepSpeed
• File: deepspeed/runtime/engine.py
• Lines: L3695-3789 (save_checkpoint), L3962+ (_save_checkpoint), L4165-4188 (get_layer_state_dict with GatherReplacedLayerParams)
• Signature: def save_checkpoint(self, save_dir, tag=None, client_state={}, save_latest=True, exclude_frozen_parameters=False) -> bool
• GatherReplacedLayerParams: deepspeed/module_inject/layers.py (L327-387)
• Import: Accessed via the DeepSpeed engine instance

Parameters

Parameter	Type	Required	Default	Description
save_dir	str	Yes	—	Directory path for saving the checkpoint
tag	str	No	None	Unique checkpoint tag; defaults to global_step{N}
client_state	dict	No	{}	Custom training state to include in the checkpoint
save_latest	bool	No	True	Write a "latest" file pointing to this checkpoint
exclude_frozen_parameters	bool	No	False	Exclude frozen (non-trainable) parameters from checkpoint

I/O

Direction	Name	Type	Description
Input	save_dir	str	Target directory for checkpoint files
Input	tag	str	Checkpoint identifier
Input	client_state	dict	User-defined state dictionary
Output	checkpoint files	files on disk	Full (gathered) model weights, optimizer states, and metadata
Output	return value	bool	True on success

Usage Example

import deepspeed
import torch
from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf")

engine, _, _, _ = deepspeed.initialize(
    model=model,
    config={
        "tensor_parallel": {"autotp_size": 4},
        "train_batch_size": 8,
        "bf16": {"enabled": True}
    }
)

# Training loop
for step, batch in enumerate(dataloader):
    loss = engine(batch["input_ids"], labels=batch["labels"]).loss
    engine.backward(loss)
    engine.step()

    if step % 1000 == 0:
        # Save TP checkpoint (parameters are auto-gathered internally)
        engine.save_checkpoint("checkpoints/", tag=f"step_{step}")

# Load checkpoint on a different TP configuration
model2 = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf")
engine2, _, _, _ = deepspeed.initialize(
    model=model2,
    config={
        "tensor_parallel": {"autotp_size": 2},  # different TP size
        "train_batch_size": 8,
        "bf16": {"enabled": True}
    }
)
engine2.load_checkpoint("checkpoints/", tag="step_5000")

Knowledge Sources

• https://github.com/deepspeedai/DeepSpeed
• https://www.deepspeed.ai/

Relationships

Principle:Deepspeedai_DeepSpeed_AutoTP_Model_Saving

Metadata

• Workflow: AutoTP_Training
• Type: Implementation
• Last Updated: 2026-02-09 00:00 GMT