[Afloat16]

Problem

Params4bit.__deepcopy__ copies its tensor and quantization state without sharing the caller's deepcopy memo, while leaving module pointing to the original layer. A copied Linear4bit therefore retains the source module in memory. Quantizing a copy made before quantization writes its new QuantState into the original module instead of its own. Copying an already quantized module also duplicates the QuantState independently for the module and its weight.

Change

Register the new parameter in the memo before recursively copying its state with that same memo. This preserves the module/weight cycle and shared QuantState references inside the copy, while isolating it from the original object graph.

Reproducer

import copy
import gc
import weakref
from bitsandbytes.nn import Linear4bit

original = Linear4bit(64, 32)
reference = weakref.ref(original)
cloned = copy.deepcopy(original)
cloned.to("cpu")
assert original.quant_state is None
assert cloned.quant_state is cloned.weight.quant_state
del original
gc.collect()
assert reference() is None

On unchanged main, quantizing the copy assigns the original module's QuantState and the copy keeps that original module alive.

Validation

Base a616886cacf86701514f0cb1d6be760f967481c0; patch 921db3aba1095030e4a608984bca0b1b31d731ad.

• The same 14 new cases: 14 failed before; 14 passed after.
• All new and existing copy/deepcopy cases: 38 passed.
• Related Linear4bit and module tests, excluding compile/FSDP state-dict tests: 279 passed, 2 skipped, 133 deselected.
• Native CPU NF4/FP4, nested and non-nested quantization, copying before/after quantization, shared references, independent storage, and source-object release are covered.
• Source and copied layers produce identical outputs, input gradients, bias gradients and bias parameters through three momentum-SGD steps.
• Repository-pinned Ruff 0.14.3 checks/formatting and git diff --check pass.

BNB_TEST_DEVICE=cpu OMP_NUM_THREADS=2 python -m pytest tests/test_linear4bit.py -k copy -q
BNB_TEST_DEVICE=cpu OMP_NUM_THREADS=2 python -m pytest tests/test_linear4bit.py tests/test_modules.py -k 'not torch_compile and not fsdp_state_dict' -q

Python 3.12.14 / PyTorch 2.15.0+cpu, with the current CPU native library built from source. GPU, distributed execution and the complete repository suite were not run.

Checked the module history, current PR catalog and all-state related issues/PRs. #1060 introduces the current standalone-parameter copy implementation; its tests do not cover owning modules, shared references, or source-module retention.