def METHOD_NAME(self, tmpdir, inputs):
    args_defaults = {
        'num_layers': 2,
        'hidden_size': 128,
        'num_attention_heads': 8,
        'max_position_embeddings': 128,
    }
    model = get_gpt2_model(args_defaults, mp_size=2)
    model = get_deepspeed_model(model)
    model.eval()
    device_name = get_accelerator().device_name()
    baseline = model(inputs[0].to(device_name), inputs[1].to(device_name), inputs[2].to(device_name))
    tag = 'mp_2'
    state_dict = {}
    state_dict['checkpoint_version'] = get_megatron_version()
    model.save_checkpoint(tmpdir, tag=tag, client_state=state_dict)
    dist.barrier()
    model.load_checkpoint(tmpdir, tag=tag, load_optimizer_states=False, load_lr_scheduler_states=False)
    device_name = get_accelerator().device_name()
    test = model(inputs[0].to(device_name), inputs[1].to(device_name), inputs[2].to(device_name))
    assert torch.allclose(baseline, test, rtol=1.0,
                          atol=1e-07), f"Baseline output {baseline} is not equal to save-then-load output {test}"