Training API¶
deepspeed.initialize()
returns a training engine in its first argument
of type DeepSpeedEngine
. This engine is used to progress training:
for step, batch in enumerate(data_loader):
#forward() method
loss = model_engine(batch)
#runs backpropagation
model_engine.backward(loss)
#weight update
model_engine.step()
Forward Propagation¶
Backward Propagation¶
Optimizer Step¶
Gradient Accumulation¶
Model Saving¶
Additionally when a DeepSpeed checkpoint is created, a script zero_to_fp32.py
is added there which can be used to reconstruct fp32 master weights into a single pytorch state_dict
file.