[tgghzyh]

Linear8bitLt.to() moves its cached INT8 weights and row scales when called on the layer directly. Moving a parent module uses PyTorch's recursive _apply instead. After a forward pass, model.to("cpu") can therefore leave both tensors on the GPU even though the layer's weight reports CPU. The row scales returned by state_dict() are also still on the GPU.

import bitsandbytes as bnb
import torch

model = torch.nn.Sequential(bnb.nn.Linear8bitLt(128, 256, bias=False, has_fp16_weights=False))
model.eval().to("cuda")
with torch.inference_mode():
    model(torch.randn(4, 128, device="cuda", dtype=torch.float16))
model.to("cpu")
assert model[0].weight.device.type == "cpu"            # passes on main
assert model.state_dict()["0.SCB"].device.type == "cpu"  # fails on main

Handle cached state during _apply so direct and parent-module moves share the same path. Keep the INT8 cache aliased to the moved weight when it already holds the quantized data, preserve the scale dtype, and drop input-dependent outlier indices left on the previous device. This extends the direct-layer movement covered by #1769.

On an A10, a randomly initialized Qwen3-0.6B checkpoint loaded through Transformers with BitsAndBytesConfig(load_in_8bit=True, llm_int8_threshold=0.0) had 196 quantized projections. After moving the model to CPU, the base revision retained 441,778,176 bytes (421.3125 MiB) of cached INT8 weights and scales on the GPU. With the fix, all 196 caches and all checkpoint scales moved to CPU. After moving back to GPU, the logits matched the original GPU output exactly. A second CPU forward also exactly reproduced the first CPU result.

The regression test exercises a nested layer with both INT8 and retained floating-point weights. Both new cases failed on the base revision and passed on an RTX 3080 and an A10 with the fix. The full tests/test_linear8bitlt.py run on each GPU finished with 55 passed and 2 expected failures.

BNB_TEST_DEVICE=cuda python -m pytest -q tests/test_linear8bitlt.py -k device_movement