fix: Handle non-Int8Params weights in Linear8bitLt (tied weights)

When lm_head is converted to Linear8bitLt but its weight is tied to an
embedding layer, weight tying replaces the Int8Params with a regular
nn.Parameter that lacks SCB/CB attributes, causing AttributeError.

Add an isinstance check in Linear8bitLt.forward() to fall back to
F.linear when the weight is not Int8Params (e.g. due to weight tying).
Also make _save_to_state_dict and _load_from_state_dict robust to
non-Int8Params weights via getattr defaults.

Fixes #1634

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: TimDettmers

bitsandbytes/nn/modules.py

@@ -1007,9 +1007,9 @@ def _save_to_state_dict(self, destination, prefix, keep_vars):
         scb_name = "SCB"
 
         # case 1: .cuda was called, SCB is in self.weight
-        param_from_weight = getattr(self.weight, scb_name)
+        param_from_weight = getattr(self.weight, scb_name, None)
         # case 2: self.init_8bit_state was called, SCB is in self.state
-        param_from_state = getattr(self.state, scb_name)
+        param_from_state = getattr(self.state, scb_name, None)
 
         key_name = prefix + f"{scb_name}"
 
@@ -1048,18 +1048,19 @@ def _load_from_state_dict(
         for key in unexpected_copy:
             input_name = key[len(prefix) :]
             if input_name == "SCB":
-                if self.weight.SCB is None:
+                weight_scb = getattr(self.weight, "SCB", None)
+                if weight_scb is None:
                     # buffers not yet initialized, can't access them directly without quantizing first
                     raise RuntimeError(
                         "Loading a quantized checkpoint into non-quantized Linear8bitLt is "
                         "not supported. Please call module.cuda() before module.load_state_dict()",
                     )
 
                 input_param = state_dict[key]
-                self.weight.SCB.copy_(input_param)
+                weight_scb.copy_(input_param)
 
                 if self.state.SCB is not None:
-                    self.state.SCB = self.weight.SCB
+                    self.state.SCB = weight_scb
 
                 unexpected_keys.remove(key)
 
@@ -1085,6 +1086,11 @@ def to(self, *args, **kwargs):
         return result
 
     def forward(self, x: torch.Tensor):
+        # If weight is not Int8Params (e.g. due to weight tying with a non-quantized module
+        # like an embedding layer), fall back to regular linear. See issue #1634.
+        if not isinstance(self.weight, Int8Params):
+            return torch.nn.functional.linear(x, self.weight, self.bias)
+
         self.state.is_training = self.training
         if self.weight.CB is not None:
             self.init_8bit_state()

tests/test_linear8bitlt.py

@@ -335,3 +335,71 @@ 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)
+
+
+class TiedWeightModel(torch.nn.Module):
+    """A minimal model with tied weights between an embedding and lm_head, mimicking
+    architectures like OPT where lm_head.weight is shared with the embedding layer."""
+
+    def __init__(self, vocab_size, hidden_dim):
+        super().__init__()
+        self.embed_tokens = torch.nn.Embedding(vocab_size, hidden_dim)
+        self.q_proj = torch.nn.Linear(hidden_dim, hidden_dim)
+        self.v_proj = torch.nn.Linear(hidden_dim, hidden_dim)
+        self.out_proj = torch.nn.Linear(hidden_dim, hidden_dim)
+        self.lm_head = torch.nn.Linear(hidden_dim, vocab_size, bias=False)
+        # Tie weights
+        self.lm_head.weight = self.embed_tokens.weight
+
+    def forward(self, x):
+        h = self.embed_tokens(x)
+        h = self.out_proj(self.q_proj(h) + self.v_proj(h))
+        return self.lm_head(h)
+
+
+@pytest.mark.parametrize("device", get_available_devices())
+def test_linear8bitlt_tied_weights_no_crash(device):
+    """Test that Linear8bitLt gracefully handles tied weights (issue #1634).
+
+    When lm_head is replaced with Linear8bitLt but its weight is tied to
+    an embedding layer, the weight becomes a regular Parameter instead of
+    Int8Params. The forward pass should still work via F.linear fallback.
+    """
+    vocab_size, hidden_dim = 32, 64
+    model = TiedWeightModel(vocab_size, hidden_dim)
+
+    skip_modules = ["q_proj", "v_proj"]
+
+    # Replace non-skipped linear layers with Linear8bitLt (simulating what
+    # HuggingFace transformers does with llm_int8_skip_modules)
+    from bitsandbytes.utils import replace_linear
+
+    model = replace_linear(
+        model,
+        lambda inf, outf, bias: Linear8bitLt(inf, outf, bias=bias, has_fp16_weights=False),
+        skip_modules=skip_modules,
+        copy_weights=True,
+    )
+
+    # Re-tie weights (as transformers does after module replacement)
+    model.lm_head.weight = model.embed_tokens.weight
+
+    model = model.to(device)
+
+    # Verify: skipped modules remain nn.Linear
+    assert type(model.q_proj) is torch.nn.Linear, "q_proj should remain nn.Linear"
+    assert type(model.v_proj) is torch.nn.Linear, "v_proj should remain nn.Linear"
+
+    # Verify: non-skipped, non-tied modules are Linear8bitLt
+    assert isinstance(model.out_proj, Linear8bitLt), "out_proj should be Linear8bitLt"
+
+    # Verify: lm_head is Linear8bitLt but with a regular Parameter (from tying)
+    assert isinstance(model.lm_head, Linear8bitLt), "lm_head should be Linear8bitLt"
+    assert not isinstance(model.lm_head.weight, bnb.nn.Int8Params), (
+        "lm_head.weight should be a regular Parameter due to tying"
+    )
+
+    # Forward pass should NOT crash (this was the bug in issue #1634)
+    x = torch.randint(0, vocab_size, (2, 8), device=device)
+    output = model(x)
+    assert output.shape == (2, 8, vocab_size)