def test_load_quantized_checkpoint(tmp_path):
    """Test that a checkpoint saved from a quantized model can be loaded back into a quantized model."""

    class Model(torch.nn.Module):
        def __init__(self):
            super().__init__()
            self.linear = torch.nn.Linear(16, 16, bias=False)

        def forward(self, x):
            return self.linear(x)

    fabric = Fabric(accelerator="cuda", devices=1, plugins=BitsandbytesPrecision("nf4-dq"))
    model = Model()
    model = fabric.setup(model)
    model(torch.randn(2, 16, device=fabric.device))
    state_dict = model.state_dict()
    # The checkpoint contains quantized weights
    assert state_dict["linear.weight"].dtype == torch.uint8
    assert state_dict["linear.weight"].shape == (128, 1)
    torch.save(state_dict, tmp_path / "checkpoint.pt")

    fabric = Fabric(accelerator="cuda", devices=1, plugins=BitsandbytesPrecision("nf4-dq"))
    model = Model()
    model = fabric.setup(model)
    state_dict = torch.load(tmp_path / "checkpoint.pt", weights_only=True)
    model.load_state_dict(state_dict)
    assert model.linear.weight.dtype == torch.uint8
    assert model.linear.weight.shape == (128, 1)
    # Shapes match during forward (weight is being dequantized during forward)
    model(torch.randn(2, 16, device=fabric.device))

    # Test with lazy load (LitGPT uses this)
    # TODO: Replace `_lazy_load` with `torch.load(..., mmap=True)` in LitGPT
    state_dict = _lazy_load(tmp_path / "checkpoint.pt")
    model.load_state_dict(state_dict)
    assert model.linear.weight.dtype == torch.uint8
    assert model.linear.weight.shape == (128, 1)
Test that a checkpoint saved from a quantized model can be loaded back into a quantized model.
tests/tests_fabric/plugins/precision/test_bitsandbytes.py