From 32f5f088fd93cd013909b0fd1fb16ee4e8ef9601 Mon Sep 17 00:00:00 2001
From: ErenAta16 <erena6466@gmail.com>
Date: Wed, 29 Jul 2026 18:02:02 +0300
Subject: [PATCH] Run replace_linear's post_processing_function on the
 replacement module

`post_processing_function` is documented as "a function name of the replacement
linear class that is called after processing", but the lookup targets the module
being replaced:

    func = getattr(module, post_processing_function, None)
    if func is not None:
        func(module)

`module` at that point is still the original `torch.nn.Linear`, which never
carries the hook, so `getattr(..., None)` returns `None` and the block does
nothing. The parameter is a no-op for every caller.

    replace_linear(model, MyLinear, post_processing_function="post_init")
    -> fc replaced with MyLinear, post_init calls: []

Two things in three lines: the wrong object, and `func(module)` passing an extra
positional argument to what is already a bound method.

Look the hook up on `model._modules[name]` and call it with no argument. After:

    -> fc and block[0] replaced, post_init calls: ['4x8', '8x16']
       (lm_head skipped via skip_modules, as before)

A replacement class that does not define the method still passes through
untouched, so this cannot start raising for anyone.
---
 bitsandbytes/utils.py |  8 ++++--
 tests/test_utils.py   | 66 +++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 72 insertions(+), 2 deletions(-)
 create mode 100644 tests/test_utils.py

diff --git a/bitsandbytes/utils.py b/bitsandbytes/utils.py
index 513baceab..c3a3c4b43 100644
--- a/bitsandbytes/utils.py
+++ b/bitsandbytes/utils.py
@@ -157,9 +157,13 @@ def replace_linear(
                 model._modules[name].bias = old_module.bias
 
             if post_processing_function is not None:
-                func = getattr(module, post_processing_function, None)
+                # Look the hook up on the replacement, not on the module we just swapped out: the
+                # original is a plain `nn.Linear` and never carries it, so `getattr(..., None)`
+                # returned None and the hook silently never ran. It is already bound to the new
+                # instance, so it takes no argument.
+                func = getattr(model._modules[name], post_processing_function, None)
                 if func is not None:
-                    func(module)
+                    func()
     return model
 
 
diff --git a/tests/test_utils.py b/tests/test_utils.py
new file mode 100644
index 000000000..e810a40de
--- /dev/null
+++ b/tests/test_utils.py
@@ -0,0 +1,66 @@
+import torch
+from torch import nn
+
+from bitsandbytes.utils import replace_linear
+
+
+class _TrackingLinear(nn.Linear):
+    """Stands in for a quantized replacement that needs a step after construction."""
+
+    post_init_calls: list[str] = []
+
+    def post_init(self):
+        _TrackingLinear.post_init_calls.append(f"{self.in_features}x{self.out_features}")
+
+
+class _Model(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.fc = nn.Linear(4, 8)
+        self.block = nn.Sequential(nn.Linear(8, 16), nn.ReLU())
+        self.lm_head = nn.Linear(16, 4)
+
+
+def test_replace_linear_runs_the_post_processing_hook():
+    """`post_processing_function` is documented as a method of the replacement class.
+
+    It used to be looked up on the module being replaced, which is a plain `nn.Linear` and never
+    carries it, so `getattr(..., None)` returned `None` and the hook silently never ran.
+    """
+    _TrackingLinear.post_init_calls.clear()
+    model = _Model()
+
+    replace_linear(model, _TrackingLinear, post_processing_function="post_init")
+
+    assert isinstance(model.fc, _TrackingLinear)
+    assert isinstance(model.block[0], _TrackingLinear)
+    assert not isinstance(model.lm_head, _TrackingLinear)  # skip_modules default
+    assert _TrackingLinear.post_init_calls == ["4x8", "8x16"]
+
+
+def test_replace_linear_without_hook_is_unchanged():
+    _TrackingLinear.post_init_calls.clear()
+    model = _Model()
+
+    replace_linear(model, _TrackingLinear)
+
+    assert isinstance(model.fc, _TrackingLinear)
+    assert _TrackingLinear.post_init_calls == []
+
+
+def test_replace_linear_tolerates_a_replacement_without_the_hook():
+    """A replacement class that does not define the method must not raise."""
+    model = _Model()
+
+    replace_linear(model, nn.Linear, post_processing_function="post_init")
+
+    assert isinstance(model.fc, nn.Linear)
+
+
+def test_replace_linear_copy_weights():
+    model = _Model()
+    original = model.fc.weight.detach().clone()
+
+    replace_linear(model, _TrackingLinear, copy_weights=True)
+
+    assert torch.equal(model.fc.weight.detach(), original)