Make sLSTM checkpoints portable across backends (#127)

The sLSTM cell stores its recurrent kernel and bias in a backend-specific
internal layout (sLSTMCell_vanilla vs sLSTMCell_cuda use different shapes and
orderings). state_dict serialized those internal parameters directly, so a
checkpoint saved with one backend could not be loaded into a cell using the
other: the recurrent kernel raised a size mismatch and the bias loaded silently
with a scrambled gate/head ordering -- even though the cell already exposes a
single canonical (external) layout via its ext2int/int2ext conversions and
documents cross-backend conversion as an intended workflow.

Register a state_dict hook that stores the recurrent kernel and bias in the
backend-agnostic external layout, and a load-state-dict pre-hook that converts
an external-layout checkpoint back into the current backend's internal layout.
External tensors are identified by rank (recurrent kernel 4D, bias 3D), so
legacy checkpoints saved in the internal layout still load into a cell of their
original backend.

Adds regression tests covering cross-backend load (both directions),
same-backend round-trips, and legacy internal-layout checkpoints.

Author: jaideeppyne

tests/test_slstm_backend_state_dict.py

@@ -0,0 +1,86 @@
+"""Regression tests for sLSTM cross-backend checkpoint loading (issue #127).
+
+The sLSTM cell stores its recurrent kernel and bias in a backend-specific
+internal layout. ``state_dict`` used to serialize those internal parameters
+directly, so a checkpoint saved with one backend could not be loaded into a
+cell using another backend. These tests run on CPU via ``skip_backend_init``.
+"""
+
+import pytest
+import torch
+
+from xlstm.blocks.slstm.cell import (
+    sLSTMCell_cuda,
+    sLSTMCell_vanilla,
+    sLSTMCellConfig,
+)
+
+
+def _external_weights(cell):
+    """Return the backend-agnostic (external) recurrent kernel and bias."""
+    return (
+        cell._recurrent_kernel_int2ext(cell._recurrent_kernel_),
+        cell._bias_int2ext(cell._bias_),
+    )
+
+
+def _randomize(cell):
+    with torch.no_grad():
+        cell._recurrent_kernel_.copy_(torch.randn_like(cell._recurrent_kernel_))
+        cell._bias_.copy_(torch.randn_like(cell._bias_))
+
+
+@pytest.mark.parametrize(
+    "src_cls,dst_cls",
+    [
+        (sLSTMCell_cuda, sLSTMCell_vanilla),
+        (sLSTMCell_vanilla, sLSTMCell_cuda),
+    ],
+)
+def test_state_dict_loads_across_backends(src_cls, dst_cls):
+    """A checkpoint saved with one backend loads into the other without error
+    and preserves the (canonical) weights."""
+    config = sLSTMCellConfig(hidden_size=16, num_heads=4)
+    src = src_cls(config, skip_backend_init=True)
+    _randomize(src)
+    dst = dst_cls(config, skip_backend_init=True)
+
+    missing, unexpected = dst.load_state_dict(src.state_dict(), strict=True)
+    assert missing == [] and unexpected == []
+
+    src_kernel, src_bias = _external_weights(src)
+    dst_kernel, dst_bias = _external_weights(dst)
+    assert torch.allclose(dst_kernel, src_kernel)
+    assert torch.allclose(dst_bias, src_bias)
+
+
+@pytest.mark.parametrize("cls", [sLSTMCell_vanilla, sLSTMCell_cuda])
+def test_state_dict_same_backend_round_trip(cls):
+    """Saving and reloading within the same backend keeps the internal weights
+    bit-for-bit identical."""
+    config = sLSTMCellConfig(hidden_size=16, num_heads=4)
+    src = cls(config, skip_backend_init=True)
+    _randomize(src)
+    dst = cls(config, skip_backend_init=True)
+
+    dst.load_state_dict(src.state_dict())
+    assert torch.allclose(dst._recurrent_kernel_, src._recurrent_kernel_)
+    assert torch.allclose(dst._bias_, src._bias_)
+
+
+@pytest.mark.parametrize("cls", [sLSTMCell_vanilla, sLSTMCell_cuda])
+def test_state_dict_loads_legacy_internal_checkpoint(cls):
+    """A legacy checkpoint stored in the internal layout still loads into a cell
+    of its original backend (backward compatibility)."""
+    config = sLSTMCellConfig(hidden_size=16, num_heads=4)
+    src = cls(config, skip_backend_init=True)
+    _randomize(src)
+    legacy = {
+        "_recurrent_kernel_": src._recurrent_kernel_.detach().clone(),
+        "_bias_": src._bias_.detach().clone(),
+    }
+
+    dst = cls(config, skip_backend_init=True)
+    dst.load_state_dict(legacy)
+    assert torch.allclose(dst._recurrent_kernel_, src._recurrent_kernel_)
+    assert torch.allclose(dst._bias_, src._bias_)

xlstm/blocks/slstm/cell.py

@@ -262,11 +262,63 @@ def __init__(self, config: sLSTMCellConfig):
 
         self.reset_parameters()
 
+        # Persist the backend-agnostic (external/canonical) layout in `state_dict` so
+        # checkpoints can be moved between backends (e.g. "cuda" <-> "vanilla"). The
+        # stored `_recurrent_kernel_`/`_bias_` parameters are kept in a backend-specific
+        # internal layout; without these hooks a checkpoint saved with one backend fails
+        # to load into another -- the recurrent kernel raises a size mismatch and the
+        # bias loads silently with a scrambled gate/head ordering. See issue #127.
+        self._register_state_dict_hook(sLSTMCellBase._state_dict_to_external_hook)
+        self._register_load_state_dict_pre_hook(
+            sLSTMCellBase._load_state_dict_to_internal_pre_hook, with_module=True
+        )
+
         if self.config.hidden_size % self.config.num_heads != 0:
             raise ValueError(
                 f"Hidden Size {self.config.hidden_size} must be divisible by head num {self.config.num_heads}"
             )
 
+    @staticmethod
+    def _state_dict_to_external_hook(module, state_dict, prefix, local_metadata):
+        """Store the recurrent kernel and bias in the backend-agnostic external layout."""
+        recurrent_kernel_key = prefix + "_recurrent_kernel_"
+        bias_key = prefix + "_bias_"
+        if recurrent_kernel_key in state_dict:
+            state_dict[recurrent_kernel_key] = module._recurrent_kernel_int2ext(
+                state_dict[recurrent_kernel_key]
+            )
+        if bias_key in state_dict:
+            state_dict[bias_key] = module._bias_int2ext(state_dict[bias_key])
+        return state_dict
+
+    @staticmethod
+    def _load_state_dict_to_internal_pre_hook(
+        module,
+        state_dict,
+        prefix,
+        local_metadata,
+        strict,
+        missing_keys,
+        unexpected_keys,
+        error_msgs,
+    ):
+        """Convert an external-layout checkpoint to this backend's internal layout.
+
+        External tensors are identified by their rank (recurrent kernel: 4D, bias: 3D).
+        Legacy checkpoints saved in the internal layout (recurrent kernel: 3D, bias: 1D)
+        are left untouched so they still load into a cell of their original backend.
+        """
+        recurrent_kernel_key = prefix + "_recurrent_kernel_"
+        bias_key = prefix + "_bias_"
+        recurrent_kernel = state_dict.get(recurrent_kernel_key)
+        if recurrent_kernel is not None and recurrent_kernel.ndim == 4:
+            state_dict[recurrent_kernel_key] = module._recurrent_kernel_ext2int(
+                recurrent_kernel
+            )
+        bias = state_dict.get(bias_key)
+        if bias is not None and bias.ndim == 3:
+            state_dict[bias_key] = module._bias_ext2int(bias)
+
     def __repr__(self):
         return (
             f"{self.__class__.__name__}(function={self.config.function}, "