Author: Shenggui Li, Siqi Mai
In this tutorial, we will cover the use of
colossalai.initialize which injects features into your training components
(e.g. model, optimizer, dataloader) seamlessly. Calling
colossalai.initialize is the standard procedure before you run
into your training loops.
In the section below, I will cover how
colossalai.initialize works and what we should take note of.
In a typical workflow, we will launch distributed environment at the beginning of our training script.
Afterwards, we will instantiate our objects such as model, optimizer, loss function, dataloader etc. At this moment,
can come in to inject features into these objects. A pseudo-code example is like below:
# launch distributed environment
# create your objects
model = MyModel()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
criterion = torch.nn.CrossEntropyLoss()
train_dataloader = MyTrainDataloader()
test_dataloader = MyTrainDataloader()
# initialize features
engine, train_dataloader, test_dataloader, _ = colossalai.initialize(model,
colossalai.initialize function will return an
Engine object. The engine object is a wrapper
for model, optimizer and loss function. The engine object will run with features specified in the config file.
More details about the engine can be found in the Use Engine and Trainer in Training.