def test_load_full_state_dict_into_sharded_model(tmp_path):
    """Test that the strategy can load a full-state checkpoint into a FSDP sharded model."""
    from torch.distributed.fsdp import FullyShardedDataParallel as FSDP

    fabric = Fabric(accelerator="cuda", devices=1)
    fabric.seed_everything(0)
    trainer = BasicTrainer(fabric)
    trainer.run()

    # Save a full-state-dict checkpoint
    checkpoint_path = Path(fabric.broadcast(str(tmp_path / "full-checkpoint.pt")))
    state = {"model": trainer.model, "optimizer": trainer.optimizer, "steps": 1}
    fabric.save(checkpoint_path, state)

    # Gather all weights and store a copy manually
    with FSDP.summon_full_params(trainer.model, writeback=False, rank0_only=False):
        params_before = torch.cat([p.cpu().view(-1) for p in trainer.model.parameters()])

    # Create a FSDP sharded model
    fabric = Fabric(
        accelerator="cuda",
        strategy=FSDPStrategy(auto_wrap_policy=always_wrap_policy),
        devices=2,
    )
    fabric.launch()
    trainer = BasicTrainer(fabric)
    trainer.run()

    state = {"model": trainer.model, "optimizer": trainer.optimizer, "steps": 44}
    fabric.load(checkpoint_path, state)
    assert state["steps"] == 1

    # Gather all weights and compare
    with FSDP.summon_full_params(trainer.model, writeback=False, rank0_only=False):
        params_after = torch.cat([p.cpu().view(-1) for p in trainer.model.parameters()])
    assert torch.equal(params_before, params_after)

    # Create a raw state-dict checkpoint to test `Fabric.load_raw` too
    raw_checkpoint_path = checkpoint_path.with_name("model-state-dict")
    if fabric.global_rank == 0:
        checkpoint = torch.load(checkpoint_path, weights_only=True)
        torch.save(checkpoint["model"], raw_checkpoint_path)
    fabric.barrier()

    trainer.run()
    fabric.load_raw(raw_checkpoint_path, trainer.model)

    # Gather all weights and compare
    with FSDP.summon_full_params(trainer.model, writeback=False, rank0_only=False):
        params_after = torch.cat([p.cpu().view(-1) for p in trainer.model.parameters()])
    assert torch.equal(params_before, params_after)
Test that the strategy can load a full-state checkpoint into a FSDP sharded model.
tests/tests_fabric/strategies/test_fsdp_integration.py