[BugFix] Fix queue checkpoint resume with lazy-initialized scalar buf…

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``).

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:

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)