megatron.checkpointing.load_checkpoint

megatron.checkpointing.load_checkpoint(model, optimizer, opt_param_scheduler, load_arg='load', strict=True)

Load a model checkpoint and return the iteration. strict (bool): whether to strictly enforce that the keys in

state_dict of the checkpoint match the names of parameters and buffers in model.