Skip to main content

Model Checkpoint

Author : Guangyang Lu


Example Code:

This function is experiential.


In this tutorial, you will learn how to save and load model checkpoints.

To leverage the power of parallel strategies in Colossal-AI, modifications to models and tensors are needed, for which you cannot directly use or torch.load to save or load model checkpoints. Therefore, we have provided you with the API to achieve the same thing.

Moreover, when loading, you are not demanded to use the same parallel strategy as saving.

How to use​


There are two ways to train a model in Colossal-AI, by engine or by trainer. Be aware that we only save the state_dict. Therefore, when loading the checkpoints, you need to define the model first.

Save when using engine​

from colossalai.utils import save_checkpoint
model = ...
engine, _, _, _ = colossalai.initialize(model=model, ...)
for epoch in range(num_epochs):
... # do some training
save_checkpoint('', epoch, model)

Save when using trainer​

from colossalai.trainer import Trainer, hooks
model = ...
engine, _, _, _ = colossalai.initialize(model=model, ...)
trainer = Trainer(engine, ...)
hook_list = [
hooks.SaveCheckpointHook(1, '', model)


from colossalai.utils import load_checkpoint
model = ...
load_checkpoint('', model)
... # train or test