def test_load_full_state_dict_into_sharded_model(distributed, tmp_path):
    """Test that the strategy can load a full-state checkpoint into a distributed model."""
    fabric = Fabric(accelerator="cuda", devices=1)
    fabric.seed_everything(0)
    model, optimizer = _train(fabric)

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

    # Gather all weights and store a copy manually
    params_before = torch.cat([p.cpu().view(-1) for p in model.parameters()])

    # Create a FSDP sharded model
    strategy = ModelParallelStrategy(_parallelize_feed_forward_fsdp2, data_parallel_size=2, tensor_parallel_size=1)
    fabric = Fabric(accelerator="cuda", strategy=strategy, devices=2)
    fabric.launch()
    model, optimizer = _train(fabric)

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

    # Gather all weights and compare
    params_after = torch.cat([p.full_tensor().cpu().view(-1) for p in 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()

    _train(fabric, model, optimizer)
    fabric.load_raw(raw_checkpoint_path, model)

    # Gather all weights and compare
    params_after = torch.cat([p.full_tensor().cpu().view(-1) for p in model.parameters()])
    assert torch.equal(params_before, params_after)
Test that the strategy can load a full-state checkpoint into a distributed model.
tests/tests_fabric/strategies/test_model_parallel_integration.py