Function load_checkpoint

codegeex/megatron/checkpointing.py:301–505

Load a model checkpoint and return the iteration. strict (bool): whether to strictly enforce that the keys in :attr:`state_dict` of the checkpoint match the names of parameters and buffers in model.

(
    model, 
    optimizer, 
    lr_scheduler, 
    load_arg="load", 
    strict=True,
)

def load_checkpoint(
    model, 
    optimizer, 
    lr_scheduler, 
    load_arg="load", 
    strict=True,
):
    """Load a model checkpoint and return the iteration.
    strict (bool): whether to strictly enforce that the keys in
        :attr:`state_dict` of the checkpoint match the names of
        parameters and buffers in model.
    """
    args = get_args()
    load_dir = getattr(args, load_arg)

    if args.deepspeed:
        loaded_dir, state_dict = model[0].load_checkpoint(load_dir)
        if loaded_dir is None:
            print_rank_0(
                "WARNING: could not find the metadata file {} ".format(load_dir)
            )
            print_rank_0(
                "    will not load any checkpoints and will start from " "random"
            )
            return 0
        release = False
    else:
        model = utils.unwrap_model(model)

        if load_dir.endswith(".pt"):
            checkpoint_name = load_dir
            release = True
        else:
            # Read the tracker file and set the iteration.
            tracker_filename = get_checkpoint_tracker_filename(load_dir)

            # If no tracker file, return iretation zero.
            if not os.path.isfile(tracker_filename):
                print_rank_0(
                    "WARNING: could not find the metadata file {} ".format(tracker_filename)
                )
                iteration = 0
                release = True
                checkpoint_name = get_checkpoint_name(load_dir, iteration, release)
                if not os.path.isfile(checkpoint_name):
                    print_rank_0(
                        "    will not load any checkpoints and will start from random"
                    )
                    return 0
            else:
                # Otherwise, read the tracker file and either set the iteration or
                # mark it as a release checkpoint.
                iteration = 0
                release = False
                with open(tracker_filename, "r") as f:
                    metastring = f.read().strip()
                    try:
                        iteration = int(metastring)