def test_load_full_state_checkpoint_into_regular_model(distributed, tmp_path):
    """Test that a full-state checkpoint saved from a distributed model can be loaded back into a regular model."""

    # Save a regular full-state checkpoint from a distributed model
    model = FSDP2Model()
    strategy = ModelParallelStrategy(save_distributed_checkpoint=False)
    trainer = Trainer(
        default_root_dir=tmp_path,
        accelerator="gpu",
        devices=2,
        strategy=strategy,
        max_epochs=1,
        barebones=True,
    )
    trainer.fit(model)
    model_path = tmp_path / "last.ckpt"
    model_path = trainer.strategy.broadcast(model_path)
    trainer.save_checkpoint(model_path)
    model_state_dict = trainer.strategy.lightning_module_state_dict()
    optimizer_state_dict = trainer.strategy.optimizer_state(model.optimizers())

    if trainer.global_rank != 0:
        assert len(model_state_dict) == 0
        assert len(optimizer_state_dict) == 0

    # Create a regular model and load the checkpoint into it
    model = TemplateModel()
    trainer = Trainer(default_root_dir=tmp_path, accelerator="gpu", devices=2, strategy="ddp", max_epochs=1)
    trainer.fit(model, ckpt_path=model_path)
    restored_model_state_dict = trainer.strategy.lightning_module_state_dict()
    restored_optimizer_state_dict = trainer.strategy.optimizer_state(model.optimizers())

    if trainer.global_rank == 0:
        assert len(model_state_dict) == len(restored_model_state_dict)
        assert len(optimizer_state_dict) == len(restored_optimizer_state_dict)
        torch.testing.assert_close(model_state_dict, restored_model_state_dict, atol=0, rtol=0)
        torch.testing.assert_close(optimizer_state_dict, restored_optimizer_state_dict, atol=0, rtol=0)
    trainer.strategy.barrier()
Test that a full-state checkpoint saved from a distributed model can be loaded back into a regular model.
tests/tests_pytorch/strategies/test_model_parallel_integration.py