Move cached int8 linear weights when moving parent modules

bitsandbytes/nn/modules.py

@@ -1162,20 +1162,26 @@ def init_8bit_state(self):
         self.weight.CB = None
         self.weight.SCB = None
 
-    def to(self, *args, **kwargs):
-        # Call the parent to() method to handle standard parameter/buffer movement
-        result = super().to(*args, **kwargs)
-
-        device, _, _, _ = torch._C._nn._parse_to(*args, **kwargs)
-
-        # Handle state tensors if needed.
-        if device is not None:
-            if result.state.CB is not None:
-                result.state.CB = result.state.CB.to(device)
-            if result.state.SCB is not None:
-                result.state.SCB = result.state.SCB.to(device)
-
-        return result
+    def _apply(self, fn, recurse=True):
+        super()._apply(fn, recurse=recurse)
+        weight = self.weight
+
+        # The weight owns the row scales until the first forward transfers them to the matmul state.
+        for cache in (weight, self.state):
+            scales = getattr(cache, "SCB", None)
+            if scales is not None:
+                if scales.device != weight.device:
+                    cache.SCB = fn(scales) if scales.is_meta else scales.to(weight.device)
+                if weight.dtype == torch.int8:
+                    cache.CB = weight.data
+                elif cache.CB is not None and cache.CB.device != weight.device:
+                    cache.CB = fn(cache.CB)
+        if getattr(weight, "CB", None) is not None and weight.SCB is None:
+            # Moving an already-quantized Int8Params can reattach CB after init_8bit_state.
+            weight.CB = None
+        if self.state.idx is not None and self.state.idx.device != weight.device:
+            self.state.idx = None
+        return self
 
     def forward(self, x: torch.Tensor):
         self.state.is_training = self.training

tests/test_linear8bitlt.py

@@ -364,3 +364,29 @@ def test_linear8bitlt_device_movement(device):
 
     # Accelerator outputs should match both times.
     torch.testing.assert_close(out_accelerator_2, out_accelerator, rtol=1e-8, atol=1e-8)
+
+
+@pytest.mark.parametrize("device", get_available_devices(no_cpu=True))
+@pytest.mark.parametrize("has_fp16_weights", TRUE_FALSE)
+@torch.inference_mode()
+def test_nested_linear8bitlt_device_movement(device, has_fp16_weights):
+    layer = bnb.nn.Linear8bitLt(32, 128, bias=False, has_fp16_weights=has_fp16_weights)
+    model = torch.nn.Sequential(layer).eval().requires_grad_(False).to(device)
+    x = torch.randn(4, 32, dtype=torch.float16)
+    if not has_fp16_weights:
+        model.cpu()
+        assert layer.weight.SCB.device == layer.weight.device == x.device
+        assert layer.weight.CB.data_ptr() == layer.weight.data_ptr()
+        model.to(device)
+    reference = model(x.to(device)).cpu()
+
+    model.cpu()
+    assert layer.state.CB.device == layer.state.SCB.device == layer.weight.device == x.device
+    if not has_fp16_weights:
+        assert layer.state.CB.data_ptr() == layer.weight.data_ptr()
+    cpu_reference = model(x)
+
+    model.to(device)
+    torch.testing.assert_close(model(x.to(device)).cpu(), reference, rtol=0, atol=0)
+    model.to("cpu")
+    torch.testing.assert_close(model(x), cpu_reference, rtol=0, atol=0)