def METHOD_NAME(self):
    CHECKPOINT_DIR = self.temp_dir
    model = FSDP(MyTestModule().cuda())
    optim = torch.optim.Adam(model.parameters(), lr=0.1)
    model(torch.rand(8, 8, device=dist.get_rank())).sum().backward()
    optim.step()
    with FSDP.state_dict_type(model, StateDictType.SHARDED_STATE_DICT):
        state_dict = {
            "model": model.state_dict(),
            "optim": FSDP.optim_state_dict(model, optim),
        }
        dcp.save_state_dict(
            state_dict=state_dict,
            storage_writer=FsspecWriter(CHECKPOINT_DIR),
            planner=dcp.DefaultSavePlanner(),
        )
    model_2 = FSDP(MyTestModule().cuda())
    optim_2 = torch.optim.Adam(model_2.parameters(), lr=0.1)
    with FSDP.summon_full_params(model):
        with FSDP.summon_full_params(model_2):
            for n_p1, n_p2 in zip(
                model.named_parameters(), model_2.named_parameters()
            ):
                self.assertNotEqual(n_p1[1], n_p2[1])
    # now load the model and ensure the values are the same
    with FSDP.state_dict_type(model_2, StateDictType.SHARDED_STATE_DICT):
        state_dict = {
            "model": model_2.state_dict(),
        }
        dcp.load_state_dict(
            state_dict=state_dict,
            storage_reader=FsspecReader(CHECKPOINT_DIR),
            planner=dcp.DefaultLoadPlanner(),
        )
        model_2.load_state_dict(state_dict["model"])
        optim_state = load_sharded_optimizer_state_dict(
            model_state_dict=state_dict["model"],
            optimizer_key="optim",
            storage_reader=FsspecReader(CHECKPOINT_DIR),
        )
        flattened_osd = FSDP.optim_state_dict_to_load(
            model_2, optim_2, optim_state["optim"]
        )
        optim_2.load_state_dict(flattened_osd)
    with FSDP.summon_full_params(model):
        with FSDP.summon_full_params(model_2):
            for n_p1, n_p2 in zip(
                model.named_parameters(), model_2.named_parameters()
            ):
                self.assertEqual(n_p1[1], n_p2[1])
    def opt_at(opt, idx):
        return list(iter(opt.state.values()))[idx]
    # Adam lazily creates its state
    self.assertEqual(
        opt_at(optim, 0)["exp_avg"], opt_at(optim_2, 0)["exp_avg"]
    )
    self.assertEqual(
        opt_at(optim, 0)["exp_avg_sq"], opt_at(optim_2, 0)["exp_avg_sq"]
    )