Principle:Deepspeedai DeepSpeed AutoTP Model Saving

Overview

Saving tensor-parallel model checkpoints by gathering partitioned parameters from all TP ranks into complete weight tensors before writing to disk.

Detailed Description

When saving a model trained with AutoTP, the checkpoint system must gather partitioned weight matrices from all TP ranks back into full tensors. DeepSpeed's save_checkpoint() orchestrates this process, and the underlying module_state_dict() mechanism uses GatherReplacedLayerParams to temporarily reconstruct full parameters before serialization. This ensures checkpoints are portable -- they can be loaded on a different number of GPUs or without tensor parallelism entirely.

The checkpoint saving process involves:

1. Tag creation: If no tag is provided, the global step count is used (e.g., global_step5000).
2. Directory creation: Rank 0 creates the checkpoint directory; all ranks synchronize via dist.barrier().
3. Parameter gathering: The get_layer_state_dict() helper function walks the model tree recursively. For each module, it enters a GatherReplacedLayerParams context which performs an AllGather to reconstruct full weight tensors from TP shards. Only rank 0 collects the state dict entries.
4. State serialization: The gathered state dict (containing full, un-partitioned weights) is serialized to disk along with optimizer states and client state.
5. Synchronization: A final dist.barrier() ensures all ranks complete before proceeding.

Key properties of TP checkpoint saving:

• Portability: Saved checkpoints contain the full model weights, not TP shards. They can be loaded with any TP configuration.
• AllGather-based reconstruction: For LinearAllreduce (row-parallel), the weight is transposed, AllGathered along the first dimension, then transposed back. For LinearLayer (column-parallel), the weight is AllGathered along the first dimension directly.
• Temporary gathering: The GatherReplacedLayerParams context manager stores the original partitioned data in data_partition, gathers the full tensor, and upon exit re-partitions by restoring from data_partition. This ensures training can continue after saving.
• Rank 0 only: Only rank 0 writes the model state dict to disk, but all ranks must participate in the AllGather communication.

Theoretical Basis

Checkpoint portability requires saving complete (un-partitioned) model weights. With tensor parallelism, each rank holds W_i, a shard of the full weight W. Gathering requires an AllGather collective operation:

• For column-parallel (weight split along dim 0): W = concat(W_0, W_1, ..., W_{tp_size-1}) along dimension 0.
• For row-parallel (weight split along dim 1): W = concat(W_0, W_1, ..., W_{tp_size-1}) along dimension 1 (implemented via transpose, AllGather, transpose).

This AllGather-then-serialize approach trades checkpoint-time communication cost for runtime flexibility. The alternative -- saving TP shards directly -- would require the same TP configuration at load time, severely limiting checkpoint reuse.

Knowledge Sources

• https://github.com/deepspeedai/DeepSpeed
• https://www.deepspeed.ai/
• https://arxiv.org/abs/1909.08053

Relationships

Implementation:Deepspeedai_DeepSpeed_DeepSpeedEngine_Save_For_TP

Metadata

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