Overview

As we are developing new features for the Colossal-AI system, we find it difficult to save/load model/optimizer checkpoints. It is because that different features need to handle the save/load logic on its own without a common protocol. This is not a good idea as we don't want it to be that the model trained with one feature can only be loaded with the same feature. This limits the usage of the checkpoint and hinders the integration with the community.

Therefore, it is important to design a unified checkpoint system as a protocol. Some important factors should be considered.

1. Checkpoint format should be feature/parallelism agnostic, which means that the checkpoint saved by one feature can be loaded for another feature. This makes sure that the checkpoint does not contain any feature-related information, but only the model parameters.
2. Support large-scale model saving/loading. We need to consider cases such as sharded model and sharded parameters.
3. Integration with the AI ecosystem, e.g. supports huggingface model weight format

Background

First of all, we should understand which use cases this unified system will cater to. Let's assume we have a model like this:

When we train this model, this model can have different ways of placement over GPUs. Let's assume we are training with 2 GPUs.

• Data Parallel (the model is replicated over GPUs)

• Tensor Parallel (the model weights are sharded over GPUs)

• Pipeline Parallel (the model is sharded by layer over GPUs)

• Auto Parallel (some parameters are sharded, some are not, depending on the search results)

Therefore, the unified checkpoint system should support at least the features mentioned above.

Currently, there are mainly two ways to save/load the model checkpoints:

• PyTorch native way using state_dict(). For example, PyTorch saves the whole model into a single file called model.pth

• Hugging Face way by sharding the model by parameters. There are many weight files, each weight files stores a collection of the parameters. A index.json file is used to specify which parameter goes into which file. An example can be https://huggingface.co/facebook/opt-66b/blob/main/pytorch_model.bin.index.json.

Design

According to the information mentioned above, what we need to support can be expressed as a matrix:

Therefore, we will have the following APIs to cover these usages:

class CheckpointIO:

    def load_model(self,
                   model: Union[nn.Module, ModelWrapper],
                   checkpoint: str,
                   strict: bool = True) -> Union[nn.Module, ModelWrapper]:
          ...

    def save_model(self,
                               model: Union[nn.Module, ModelWrapper],
                               checkpoint: str,
                               shard: bool = False,
                               gather_dtensor: bool = True,
                               prefix: str = None,
                               size_per_shard: int = 1024):
          ...

    def load_optimizer(self, optimizer: Optimizer, checkpoint: str):
           ...

    def save_optimizer(self,
                                     optimizer: Optimizer,
                                     checkpoint: str,
                                     shard: bool = False,
                                     gather_dtensor: bool = True,
                                     prefix: str = None,
                                     size_per_shard: int = 1024):
        ...

We can focus on the load_model and save_model methods.

save_model has two arguments to define its checkpoint format:

• shard (bool): Will save in the HF sharded model format with a index.json file if True. Otherwise, the model weight will be saved in a single file.
• gather_dtensor (bool): will gather the distributed tensor back to its global tensor if True. Otherwise, save it as a tensor shard.

To better explain the outcome of different cases, I will use file structure to illustrate:

1. shard = False, gather_dtensor = True

- checkpoints
   - model.pth

2. shard = True, gather_dtensor = True

- checkpoints
   - model.bin.index.json
   - model_0001-of-0003.bin
   - model_0002-of-0003.bin
   - model_0003-of-0003.bin

3. shard = False, gather_dtensor = False. The dtensors will be stored in an individual folder and each tensor sharded is numbered. The saved tensor format can be:

{
    "shard_info": Any
    "tensor": torch.Tensor
}

The file structure will look like:

- checkpoints
   - model.bin.index.json
   - model.pth
   - dtensor
      - linear_weight.bin.1
      - linear_weight.bin.2
      - linear_bias.bin.1
      - linear_bias.bin.2

The weight map key-value pair for dtensor in the index.json file will look like linear_weight: linear_weight.bin.*

4. shard = True, gather_dtensor = False

- checkpoints
   - model.bin.index.json
   - model_0001-of-0003.bin
   - model_0002-of-0003.bin
   - model_0003-of-0003.bin
   - dtensor
      - linear_weight.bin.1
      - linear_weight.bin.2
      - linear_bias.bin.1
      - linear_bias.bin.2

When loading models, we only support cases of shard = True, gather_detnsor = False and shard = False, gather_dtensor = False. Therefore, merging dtensors to a global tensor can be done offline via our CLI.