def METHOD_NAME(ckpt_dir, model, optimizer=None, mode="train", load_recovery=None):
    nranks = dist.get_world_size()
    if nranks > 1:
        hcg = env.get_hcg()
        mp_rank = hcg.get_model_parallel_rank()
        pp_rank = hcg.get_stage_id()
        sharding_rank = hcg.get_sharding_parallel_rank()
    load_recovery = {} if load_recovery is None else load_recovery
    if ckpt_dir and isinstance(ckpt_dir, str):
        logger.info("Try to load checkpoint from %s " % ckpt_dir)
        if mode == "quant":
            load_dir = ckpt_dir
        else:
            load_dir = (
                "{}/mp_{:0>2d}_sharding_{:0>2d}_pp_{:0>2d}".format(ckpt_dir, mp_rank, sharding_rank, pp_rank)
                if nranks > 1
                else ckpt_dir
            )
        model_path = os.path.join(load_dir, "model.pdparams")
        opt_path = os.path.join(load_dir, "model_state.pdopt")
        meta_path = os.path.join(load_dir, "meta_state.pdopt")
        if os.path.exists(model_path):
            model_dict = paddle.METHOD_NAME(model_path)
            for name, param in model.state_dict().items():
                assert name in model_dict.keys(), "No param named `{}` was found in checkpoint file.".format(name)
                if param.dtype != model_dict[name].dtype:
                    model_dict[name] = model_dict[name].cast(param.dtype)
            model.set_state_dict(model_dict)
        else:
            raise ValueError("No model checkpoint file found in %s." % model_path)
        if mode == "train":
            if os.path.exists(opt_path):
                opt_dict = paddle.METHOD_NAME(opt_path)
                optimizer.set_state_dict(opt_dict)
            else:
                raise ValueError("No optimizer checkpoint file found in %s." % opt_path)
            if os.path.exists(meta_path):
                meta_dict = paddle.METHOD_NAME(meta_path)
                load_recovery.update(
                    {"step": meta_dict["step"], "epoch": meta_dict["epoch"], "rng_state": meta_dict["cuda_rng_state"]}
                )
            else:
                raise ValueError("No meta checkpoint file found in %s." % meta_path)
        logger.info("successfully load checkpoints")
    else:
        logger.warning("`load` requires a valid value of `ckpt_dir`.")
        raise TypeError("`load` requires a valid value of `ckpt_dir`.")