def METHOD_NAME(output_dir, model, optimizer=None, step=0, epoch=0, sharding_stage=2):
    """
    save the state dicts of model and optimizer into an checkpoint.
    """
    nranks = dist.get_world_size()
    if nranks > 1:
        hcg = env.get_hcg()
        dp_rank = hcg.get_data_parallel_rank()
        mp_rank = hcg.get_model_parallel_rank()
        pp_rank = hcg.get_stage_id()
        sharding_rank = hcg.get_sharding_parallel_rank()
    else:
        dp_rank = 0
    if dp_rank != 0:
        logger.info("DP_Rank %d doesn't save model" % dp_rank)
        return
    if output_dir and isinstance(output_dir, str):
        output_dir = os.path.join(output_dir, "epoch_%d_step_%d" % (epoch, step))
        if not os.path.exists(output_dir):
            os.makedirs(output_dir, exist_ok=True)
        logger.info("Save model to %s" % output_dir)
        save_dir = (
            "{}/mp_{:0>2d}_sharding_{:0>2d}_pp_{:0>2d}".format(output_dir, mp_rank, sharding_rank, pp_rank)
            if nranks > 1
            else output_dir
        )
        if sharding_stage == 3:
            model.get_all_parameters(convert2cpu=False)
        paddle.METHOD_NAME(model.state_dict(), os.path.join(save_dir, "model.pdparams"))
        if optimizer is not None:
            paddle.METHOD_NAME(optimizer.state_dict(), os.path.join(save_dir, "model_state.pdopt"))
        meta_dict = {"epoch": epoch, "step": step, "cuda_rng_state": paddle.get_cuda_rng_state()}
        paddle.METHOD_NAME(meta_dict, os.path.join(save_dir, "meta_state.pdopt"))
        save_auto_dir = os.path.join(output_dir, "auto_infer")
        save_for_auto_inference(os.path.join(save_auto_dir, "auto"), model)
    else:
        raise TypeError("`save` requires a valid value of `output_dir`.")