diff --git a/bitsandbytes/nn/modules.py b/bitsandbytes/nn/modules.py index ebc0b0943..6581f170c 100644 --- a/bitsandbytes/nn/modules.py +++ b/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 diff --git a/tests/test_linear8bitlt.py b/tests/test_linear8bitlt.py index b078e82b7..29a885f36 100644 --- a/tests/test_linear8bitlt.py +++ b/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)