Honor optimizer checkpoint load hooks for state migration

bitsandbytes/optim/optimizer.py

@@ -198,7 +198,13 @@ def load_state_dict(self, state_dict, move_to_device=True):
             move_to_device (`bool`, defaults to `True`):
                 Whether to move the optimizer's state to the device.
         """
-        # deepcopy, to be consistent with module API
+        state_dict = state_dict.copy()
+        for pre_hook in self._optimizer_load_state_dict_pre_hooks.values():
+            hook_result = pre_hook(self, state_dict)
+            if hook_result is not None:
+                state_dict = hook_result
+
+        # Do not mutate checkpoint data while unwrapping and moving quantization state.
         state_dict = deepcopy(state_dict)
 
         # Unwrap quantization states that were wrapped for FSDP compatibility
@@ -270,6 +276,9 @@ def update_group(group, new_group):
         param_groups = [update_group(g, ng) for g, ng in zip(groups, saved_groups)]
         self.__setstate__({"state": state, "param_groups": param_groups})
 
+        for post_hook in self._optimizer_load_state_dict_post_hooks.values():
+            post_hook(self)
+
     def to_gpu(self):
         for gindex, group in enumerate(self.param_groups):
             for pindex, p in enumerate(group["params"]):

tests/test_optim.py

@@ -1,3 +1,4 @@
+import copy
 import io
 import sys
 import time
@@ -16,6 +17,104 @@
 k = 20
 
 
+@pytest.mark.parametrize(
+    "optimizer_cls", [bnb.optim.Adam8bit, bnb.optim.Adam32bit, bnb.optim.SGD8bit, bnb.optim.SGD32bit]
+)
+@pytest.mark.parametrize("return_replacement", [False, True])
+@pytest.mark.parametrize("device", get_available_devices())
+def test_load_state_dict_hooks_migrate_groups_and_resume(optimizer_cls, return_replacement, device):
+    torch.manual_seed(43)
+    model = torch.nn.Sequential(
+        torch.nn.Linear(64, 64, bias=False), torch.nn.ReLU(), torch.nn.Linear(64, 64, bias=False)
+    ).to(device)
+    kwargs = {"momentum": 0.9} if "SGD" in optimizer_cls.__name__ else {}
+    optimizer = optimizer_cls(model.parameters(), lr=0.01, **kwargs)
+    inputs = torch.randn(4, 64, device=device)
+    for _ in range(3):
+        optimizer.zero_grad()
+        model(inputs).square().mean().backward()
+        optimizer.step()
+    checkpoint_buffer = io.BytesIO()
+    torch.save(optimizer.state_dict(), checkpoint_buffer)
+    checkpoint_buffer.seek(0)
+    checkpoint = torch.load(checkpoint_buffer)
+    expected_checkpoint = copy.deepcopy(checkpoint)
+    restored = copy.deepcopy(model)
+    restored_optimizer = optimizer_cls(
+        [{"params": restored[0].parameters()}, {"params": restored[2].parameters()}], lr=0.5, **kwargs
+    )
+    calls = []
+
+    def migrate_groups(opt, state_dict):
+        assert opt is restored_optimizer
+        calls.append("pre")
+        group = state_dict["param_groups"][0]
+        groups = [dict(group, params=[parameter_id]) for parameter_id in group["params"]]
+        if return_replacement:
+            return dict(state_dict, param_groups=groups)
+        state_dict["param_groups"] = groups
+
+    def check_loaded_state(opt):
+        assert opt is restored_optimizer
+        calls.append("post")
+        assert len(opt.param_groups) == 2
+        assert all(group["lr"] == 0.01 for group in opt.param_groups)
+        for parameter in restored.parameters():
+            assert opt.state[parameter]["step"] == 3
+            expected_dtype = torch.uint8 if "8bit" in optimizer_cls.__name__ else torch.float32
+            assert opt.state[parameter]["state1"].dtype == expected_dtype
+
+    restored_optimizer.register_load_state_dict_pre_hook(migrate_groups)
+    restored_optimizer.register_load_state_dict_post_hook(check_loaded_state)
+    restored_optimizer.load_state_dict(checkpoint)
+    assert calls == ["pre", "post"]
+    torch.testing.assert_close(checkpoint, expected_checkpoint, rtol=0, atol=0)
+    for _ in range(3):
+        optimizer.zero_grad()
+        restored_optimizer.zero_grad()
+        expected = model(inputs)
+        actual = restored(inputs)
+        torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+        expected.square().mean().backward()
+        actual.square().mean().backward()
+        optimizer.step()
+        restored_optimizer.step()
+        for original_parameter, restored_parameter in zip(model.parameters(), restored.parameters()):
+            torch.testing.assert_close(restored_parameter, original_parameter, rtol=0, atol=0)
+            torch.testing.assert_close(
+                restored_optimizer.state[restored_parameter], optimizer.state[original_parameter], rtol=0, atol=0
+            )
+
+
+@pytest.mark.parametrize("move_to_device", [False, True])
+def test_load_state_dict_hook_order_and_removal(move_to_device):
+    parameter = torch.nn.Parameter(torch.ones(2))
+    optimizer = bnb.optim.Adam32bit([parameter])
+    checkpoint = optimizer.state_dict()
+    calls = []
+    pre = optimizer.register_load_state_dict_pre_hook(lambda opt, state: calls.append("pre"))
+    pre_first = optimizer.register_load_state_dict_pre_hook(lambda opt, state: calls.append("pre_first"), prepend=True)
+    post = optimizer.register_load_state_dict_post_hook(lambda opt: calls.append("post"))
+    post_first = optimizer.register_load_state_dict_post_hook(lambda opt: calls.append("post_first"), prepend=True)
+    optimizer.load_state_dict(checkpoint, move_to_device=move_to_device)
+    assert calls == ["pre_first", "pre", "post_first", "post"]
+    for handle in (pre, pre_first, post, post_first):
+        handle.remove()
+    calls.clear()
+    optimizer.load_state_dict(checkpoint, move_to_device=move_to_device)
+    assert calls == []
+
+
+def test_load_state_dict_post_hook_requires_successful_load():
+    optimizer = bnb.optim.Adam32bit([torch.nn.Parameter(torch.ones(2))])
+    calls = []
+    optimizer.register_load_state_dict_pre_hook(lambda opt, state: calls.append("pre"))
+    optimizer.register_load_state_dict_post_hook(lambda opt: calls.append("post"))
+    with pytest.raises(ValueError, match="different number of parameter groups"):
+        optimizer.load_state_dict({"state": {}, "param_groups": []})
+    assert calls == ["pre"]
+
+
 def assert_most_approx_close(a, b, rtol=1e-3, atol=1e-3, max_error_count=0):
     idx = torch.isclose(a, b, rtol=rtol, atol=atol)
     error_count = (idx == 0).sum().item()