From 762439d36e0f920b23d5ac4a5b0277f43c4c7c6a Mon Sep 17 00:00:00 2001 From: mohsinm-dev Date: Tue, 17 Mar 2026 16:22:49 +0500 Subject: [PATCH] [BugFix] Fix queue checkpoint resume with lazy-initialized scalar buffers --- RELEASES.rst | 1 + .../tests/unit/test_nn_modules.py | 113 ++++++++++++++++++ stable_pretraining/utils/nn_modules.py | 50 ++++++++ 3 files changed, 164 insertions(+) diff --git a/RELEASES.rst b/RELEASES.rst index 494fbd50f..ff62b1448 100644 --- a/RELEASES.rst +++ b/RELEASES.rst @@ -15,3 +15,4 @@ Version 0.1 - RankMe, LiDAR metrics to monitor training. - Examples of extracting run data from WandB and utilizing it to create figures. - Fixed a bug in the logging functionality. +- Fixed ``OrderedQueue`` and ``UnsortedQueue`` checkpoint resume failing on scalar-label buffers when using lazy initialization (``shape=None``). diff --git a/stable_pretraining/tests/unit/test_nn_modules.py b/stable_pretraining/tests/unit/test_nn_modules.py index c6509cb9d..c3aca57e9 100644 --- a/stable_pretraining/tests/unit/test_nn_modules.py +++ b/stable_pretraining/tests/unit/test_nn_modules.py @@ -149,6 +149,62 @@ def test_multiple_large_batches(self): expected2 = torch.tensor([20, 30, 40], dtype=torch.float32) assert torch.allclose(result2, expected2) + def test_state_dict_roundtrip_lazy_scalar(self): + """Test state_dict save/load with lazy init and scalar labels.""" + q = UnsortedQueue(8, shape=None) + # Append scalar data (B,) to trigger lazy init to (max_length,) + q.append(torch.tensor([1.0, 2.0, 3.0])) + q.append(torch.tensor([4.0, 5.0])) + state = q.state_dict() + + # Fresh queue with lazy init — placeholder is (8, 1) + q2 = UnsortedQueue(8, shape=None) + q2.load_state_dict(state) + + assert torch.equal(q2.out, q.out) + assert torch.equal(q2.pointer, q.pointer) + assert torch.equal(q2.filled, q.filled) + + def test_state_dict_roundtrip_lazy_2d(self): + """Test state_dict save/load with lazy init and 2D data.""" + q = UnsortedQueue(6, shape=None) + q.append(torch.randn(2, 4)) + q.append(torch.randn(3, 4)) + state = q.state_dict() + + q2 = UnsortedQueue(6, shape=None) + q2.load_state_dict(state) + + assert torch.equal(q2.out, q.out) + assert torch.equal(q2.pointer, q.pointer) + + def test_state_dict_roundtrip_eager(self): + """Test state_dict save/load with explicit shape (no lazy init).""" + q = UnsortedQueue(5, shape=3) + q.append(torch.randn(2, 3)) + state = q.state_dict() + + q2 = UnsortedQueue(5, shape=3) + q2.load_state_dict(state) + + assert torch.equal(q2.out, q.out) + assert torch.equal(q2.pointer, q.pointer) + + def test_state_dict_roundtrip_lazy_scalar_recursive(self): + """Test that _load_from_state_dict works during recursive nn.Module loading.""" + # Simulate Lightning's recursive load by wrapping in a parent module + parent = torch.nn.Module() + parent.queue = UnsortedQueue(8, shape=None) + parent.queue.append(torch.tensor([1.0, 2.0, 3.0])) + state = parent.state_dict() + + parent2 = torch.nn.Module() + parent2.queue = UnsortedQueue(8, shape=None) + parent2.load_state_dict(state) + + assert torch.equal(parent2.queue.out, parent.queue.out) + assert torch.equal(parent2.queue.pointer, parent.queue.pointer) + @pytest.mark.unit class TestOrderedQueue: @@ -320,6 +376,63 @@ def test_multiple_large_batches_with_order(self): assert torch.allclose(result2, expected2) assert q.global_counter.item() == 9 + def test_state_dict_roundtrip_lazy_scalar(self): + """Test state_dict save/load with lazy init and scalar labels.""" + q = OrderedQueue(8, shape=None) + q.append(torch.tensor([1.0, 2.0, 3.0])) + q.append(torch.tensor([4.0, 5.0])) + state = q.state_dict() + + q2 = OrderedQueue(8, shape=None) + q2.load_state_dict(state) + + assert torch.equal(q2.out, q.out) + assert torch.equal(q2.pointer, q.pointer) + assert torch.equal(q2.filled, q.filled) + assert torch.equal(q2.global_counter, q.global_counter) + + def test_state_dict_roundtrip_lazy_2d(self): + """Test state_dict save/load with lazy init and 2D data.""" + q = OrderedQueue(6, shape=None) + q.append(torch.randn(2, 4)) + q.append(torch.randn(3, 4)) + state = q.state_dict() + + q2 = OrderedQueue(6, shape=None) + q2.load_state_dict(state) + + assert torch.equal(q2.out, q.out) + assert torch.equal(q2.pointer, q.pointer) + assert torch.equal(q2.global_counter, q.global_counter) + + def test_state_dict_roundtrip_eager(self): + """Test state_dict save/load with explicit shape (no lazy init).""" + q = OrderedQueue(5, shape=3) + q.append(torch.randn(2, 3)) + state = q.state_dict() + + q2 = OrderedQueue(5, shape=3) + q2.load_state_dict(state) + + assert torch.equal(q2.out, q.out) + assert torch.equal(q2.pointer, q.pointer) + assert torch.equal(q2.global_counter, q.global_counter) + + def test_state_dict_roundtrip_lazy_scalar_recursive(self): + """Test that _load_from_state_dict works during recursive nn.Module loading.""" + parent = torch.nn.Module() + parent.queue = OrderedQueue(8, shape=None) + parent.queue.append(torch.tensor([1.0, 2.0, 3.0])) + state = parent.state_dict() + + parent2 = torch.nn.Module() + parent2.queue = OrderedQueue(8, shape=None) + parent2.load_state_dict(state) + + assert torch.equal(parent2.queue.out, parent.queue.out) + assert torch.equal(parent2.queue.pointer, parent.queue.pointer) + assert torch.equal(parent2.queue.global_counter, parent.queue.global_counter) + @pytest.mark.unit class TestEMA: diff --git a/stable_pretraining/utils/nn_modules.py b/stable_pretraining/utils/nn_modules.py index c6d9c03f1..744389d6b 100644 --- a/stable_pretraining/utils/nn_modules.py +++ b/stable_pretraining/utils/nn_modules.py @@ -212,6 +212,33 @@ def _test(): assert 0 not in v.numpy() return True + def _load_from_state_dict( + self, + state_dict, + prefix, + local_metadata, + strict, + missing_keys, + unexpected_keys, + error_msgs, + ): + out_key = prefix + "out" + if out_key in state_dict: + self.out.resize_(state_dict[out_key].shape) + super()._load_from_state_dict( + state_dict, + prefix, + local_metadata, + strict, + missing_keys, + unexpected_keys, + error_msgs, + ) + + def load_state_dict(self, state_dict, strict=True, assign=False): + self.out.resize_(state_dict["out"].shape) + super().load_state_dict(state_dict, strict, assign) + class OrderedQueue(torch.nn.Module): """A queue that maintains insertion order of elements. @@ -410,6 +437,29 @@ def _test(): return True + def _load_from_state_dict( + self, + state_dict, + prefix, + local_metadata, + strict, + missing_keys, + unexpected_keys, + error_msgs, + ): + out_key = prefix + "out" + if out_key in state_dict: + self.out.resize_(state_dict[out_key].shape) + super()._load_from_state_dict( + state_dict, + prefix, + local_metadata, + strict, + missing_keys, + unexpected_keys, + error_msgs, + ) + def load_state_dict(self, state_dict, strict=True, assign=False): self.out.resize_(state_dict["out"].shape) super().load_state_dict(state_dict, strict, assign)