From 165132966f6993af3a93bf0cf54b2212c0363f01 Mon Sep 17 00:00:00 2001 From: Afloat Date: Fri, 13 Nov 2026 17:46:05 +0800 Subject: [PATCH 1/4] Save and restore quantized embedding checkpoints Serialize INT8 row scales and packed 4-bit quantization metadata. Restore fresh or prequantized embeddings through native state_dict loading while preserving parameter identity for copy-based loads. Signed-off-by: Afloat --- bitsandbytes/nn/modules.py | 87 ++++++++++++++++++- tests/test_embedding_checkpoint.py | 134 +++++++++++++++++++++++++++++ 2 files changed, 219 insertions(+), 2 deletions(-) create mode 100644 tests/test_embedding_checkpoint.py diff --git a/bitsandbytes/nn/modules.py b/bitsandbytes/nn/modules.py index ebc0b0943..a1d4697c5 100644 --- a/bitsandbytes/nn/modules.py +++ b/bitsandbytes/nn/modules.py @@ -857,7 +857,41 @@ def __init__(self, num_embeddings, embedding_dim, device=None, dtype=None): self.weight = Int8Params(self.weight.data, has_fp16_weights=False, requires_grad=False) def _save_to_state_dict(self, destination, prefix, keep_vars): - raise NotImplementedError("Saving Embedding8bit module is not implemented") + super()._save_to_state_dict(destination, prefix, keep_vars) + if self.weight.SCB is not None: + destination[prefix + "SCB"] = self.weight.SCB if keep_vars else self.weight.SCB.detach() + + def _load_from_state_dict( + self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs + ): + weight_key = prefix + "weight" + scale_key = prefix + "SCB" + packed_weight = state_dict.get(weight_key) + scales = state_dict.get(scale_key) + quantized = packed_weight is not None and packed_weight.dtype == torch.int8 + if quantized: + if scales is None: + error_msgs.append(f"Missing quantization scales for {weight_key}") + return + if packed_weight.shape != (self.num_embeddings, self.embedding_dim) or scales.shape != ( + self.num_embeddings, + ): + error_msgs.append(f"Invalid quantized embedding shape for {weight_key}") + return + assign = local_metadata.get("assign_to_params_buffers", False) + device = packed_weight.device if assign else self.weight.device + self.weight.data = torch.empty_like(packed_weight, device=device) + scales = scales.to(device=device).clone() + state_dict.pop(scale_key) + + super()._load_from_state_dict( + state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs + ) + if quantized: + if not isinstance(self.weight, Int8Params): + self.weight = Int8Params(self.weight.data, requires_grad=False, has_fp16_weights=False) + self.weight.CB = self.weight.data + self.weight.SCB = scales def forward(self, input: Tensor) -> Tensor: if not hasattr(self.weight, "SCB"): @@ -961,7 +995,56 @@ def _forward_with_partial_dequantize(self, input: Tensor): return output.to(self.dtype) def _save_to_state_dict(self, destination, prefix, keep_vars): - raise NotImplementedError("Saving Embedding4bit module is not implemented") + super()._save_to_state_dict(destination, prefix, keep_vars) + if self.weight.quant_state is not None: + for name, value in self.weight.quant_state.as_dict(packed=True).items(): + destination[prefix + "weight." + name] = value if keep_vars else value.detach() + + def _load_from_state_dict( + self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs + ): + weight_key = prefix + "weight" + quantized_keys = [ + key + for key in state_dict + if key.startswith(weight_key + ".") + and key.split(".")[-1] in QuantState.valid_qs_keys + QuantState.valid_qs_type_keys + ] + restored = None + if quantized_keys: + if weight_key not in state_dict: + error_msgs.append(f"Missing packed weight for {weight_key}") + return + assign = local_metadata.get("assign_to_params_buffers", False) + packed_weight = state_dict[weight_key] + device = packed_weight.device if assign else self.weight.device + stats = {key: state_dict[key] for key in quantized_keys} + # Loading must not share mutable quantization metadata with the + # checkpoint: later device moves update these tensors in place. + restored = Params4bit.from_prequantized( + packed_weight, + copy.deepcopy(stats), + requires_grad=False, + device=device, + ) + if restored.quant_state.shape != (self.num_embeddings, self.embedding_dim): + error_msgs.append(f"Invalid quantized embedding shape for {weight_key}") + return + for key in quantized_keys: + state_dict.pop(key) + if assign: + state_dict[weight_key] = restored + else: + self.weight.data = torch.empty_like(restored.data, device=device) + self.weight.__dict__.update(restored.__dict__) + state_dict[weight_key] = restored.data + + super()._load_from_state_dict( + state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs + ) + if restored is not None: + self.weight.module = self + self.quant_state = self.weight.quant_state def forward(self, input: Tensor) -> Tensor: fix_4bit_weight_quant_state_from_module(self) diff --git a/tests/test_embedding_checkpoint.py b/tests/test_embedding_checkpoint.py new file mode 100644 index 000000000..f5e3ad6ea --- /dev/null +++ b/tests/test_embedding_checkpoint.py @@ -0,0 +1,134 @@ +import copy + +import pytest +from tests.helpers import get_available_devices +import torch + +import bitsandbytes as bnb + + +def make_embedding(kind, width=64, dtype=torch.float32, storage=torch.uint8): + if kind == "int8": + return bnb.nn.Embedding8bit(16, width, dtype=dtype) + return bnb.nn.Embedding4bit(16, width, dtype=dtype, quant_type=kind, quant_storage=storage) + + +@pytest.mark.parametrize("device", get_available_devices()) +@pytest.mark.parametrize("kind", ["int8", "fp4", "nf4"]) +@pytest.mark.parametrize("width", [32, 64]) +@pytest.mark.parametrize("assign", [False, True]) +@pytest.mark.parametrize("prequantized", [False, True]) +def test_embedding_checkpoint_roundtrip(device, kind, width, assign, prequantized, tmp_path): + torch.manual_seed(29) + source = torch.nn.Sequential( + make_embedding(kind, width).to(device), + torch.nn.Linear(width, 4, device=device), + ) + inputs = torch.tensor([[1, 3, 1], [7, 9, 11]], device=device) + optimizer = torch.optim.SGD(source[1].parameters(), lr=0.03) + optimizer.zero_grad() + source(inputs).square().mean().backward() + optimizer.step() + path = tmp_path / "embedding.pt" + torch.save({"model": source.state_dict(), "optimizer": optimizer.state_dict()}, path) + checkpoint = torch.load(path, map_location=device, weights_only=True) + original_checkpoint = copy.deepcopy(checkpoint["model"]) + + restored = ( + torch.nn.Sequential( + make_embedding(kind, width), + torch.nn.Linear(width, 4), + ).to(device) + if prequantized + else torch.nn.Sequential( + make_embedding(kind, width), + torch.nn.Linear(width, 4), + ) + ) + if not prequantized: + # Keep the embedding weight unquantized until load_state_dict while + # placing the destination tensors on the requested device. + restored[0].weight.data = restored[0].weight.data.to(device) + restored[1].to(device) + original_parameter = restored[0].weight + result = restored.load_state_dict(checkpoint["model"], strict=True, assign=assign) + assert not result.missing_keys and not result.unexpected_keys + if not assign: + assert restored[0].weight is original_parameter + assert restored[0].weight.data_ptr() != checkpoint["model"]["0.weight"].data_ptr() + expected_type = bnb.nn.Int8Params if kind == "int8" else bnb.nn.Params4bit + assert isinstance(restored[0].weight, expected_type) + assert not restored[0].weight.requires_grad + if kind == "int8": + assert restored[0].weight.CB.data_ptr() == restored[0].weight.data_ptr() + assert restored[0].weight.SCB.data_ptr() != checkpoint["model"]["0.SCB"].data_ptr() + else: + assert restored[0].weight.module is restored[0] + assert restored[0].quant_state is restored[0].weight.quant_state + assert restored[0].weight.quant_state.absmax.data_ptr() != checkpoint["model"]["0.weight.absmax"].data_ptr() + for key, value in original_checkpoint.items(): + torch.testing.assert_close(checkpoint["model"][key], value, rtol=0, atol=0) + + restored_optimizer = torch.optim.SGD(restored[1].parameters(), lr=0.03) + restored_optimizer.load_state_dict(checkpoint["optimizer"]) + for _ in range(2): + expected = source(inputs) + actual = restored(inputs) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + optimizer.zero_grad() + restored_optimizer.zero_grad() + expected.square().mean().backward() + actual.square().mean().backward() + for actual_parameter, expected_parameter in zip(restored[1].parameters(), source[1].parameters()): + torch.testing.assert_close(actual_parameter.grad, expected_parameter.grad, rtol=0, atol=0) + optimizer.step() + restored_optimizer.step() + for actual_parameter, expected_parameter in zip(restored[1].parameters(), source[1].parameters()): + torch.testing.assert_close(actual_parameter, expected_parameter, rtol=0, atol=0) + + +@pytest.mark.parametrize("kind", ["int8", "fp4", "nf4"]) +@pytest.mark.parametrize("keep_vars", [False, True]) +def test_embedding_checkpoint_unquantized_state(kind, keep_vars): + source = make_embedding(kind) + destination = make_embedding(kind) + state = source.state_dict(keep_vars=keep_vars) + assert list(state) == ["weight"] + destination.load_state_dict(state, strict=True) + torch.testing.assert_close(destination.weight, source.weight, rtol=0, atol=0) + + +@pytest.mark.parametrize("device", get_available_devices()) +@pytest.mark.parametrize("kind", ["fp4", "nf4"]) +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +def test_embedding_checkpoint_floating_storage(device, kind, dtype): + source = make_embedding(kind, dtype=dtype, storage=dtype).to(device) + destination = make_embedding(kind, dtype=dtype, storage=dtype) + destination.weight.data = destination.weight.data.to(device) + state = source.state_dict() + destination.load_state_dict(state, strict=True) + inputs = torch.tensor([1, 3, 7], device=device) + torch.testing.assert_close(destination(inputs), source(inputs), rtol=0, atol=0) + assert destination.weight.quant_storage == dtype + assert destination.weight.quant_state.dtype == dtype + + +@pytest.mark.parametrize("kind", ["int8", "fp4", "nf4"]) +def test_embedding_checkpoint_rejects_wrong_logical_shape(kind): + source = make_embedding(kind, width=64).to("cpu") + destination = make_embedding(kind, width=32) + original_weight = destination.weight.detach().clone() + with pytest.raises(RuntimeError, match="Invalid quantized embedding shape"): + destination.load_state_dict(source.state_dict()) + torch.testing.assert_close(destination.weight, original_weight, rtol=0, atol=0) + + +def test_embedding_checkpoint_requires_int8_scales(): + source = make_embedding("int8").to("cpu") + state = source.state_dict() + del state["SCB"] + destination = make_embedding("int8") + original_weight = destination.weight.detach().clone() + with pytest.raises(RuntimeError, match="Missing quantization scales"): + destination.load_state_dict(state) + torch.testing.assert_close(destination.weight, original_weight, rtol=0, atol=0) From 27c7605235adc8f1179edb0919a142a4daa27f0b Mon Sep 17 00:00:00 2001 From: Afloat Date: Fri, 13 Nov 2026 17:49:17 +0800 Subject: [PATCH 2/4] Materialize packed shapes before assignment-based checkpoint loading Cover fresh and meta embeddings with native assign=True restores. Signed-off-by: Afloat --- bitsandbytes/nn/modules.py | 6 +++++- tests/test_embedding_checkpoint.py | 13 +++++++++++++ 2 files changed, 18 insertions(+), 1 deletion(-) diff --git a/bitsandbytes/nn/modules.py b/bitsandbytes/nn/modules.py index a1d4697c5..75ca7bfaa 100644 --- a/bitsandbytes/nn/modules.py +++ b/bitsandbytes/nn/modules.py @@ -880,7 +880,10 @@ def _load_from_state_dict( return assign = local_metadata.get("assign_to_params_buffers", False) device = packed_weight.device if assign else self.weight.device - self.weight.data = torch.empty_like(packed_weight, device=device) + if assign: + self.weight = Int8Params(packed_weight, requires_grad=False, has_fp16_weights=False) + else: + self.weight.data = torch.empty_like(packed_weight, device=device) scales = scales.to(device=device).clone() state_dict.pop(scale_key) @@ -1033,6 +1036,7 @@ def _load_from_state_dict( for key in quantized_keys: state_dict.pop(key) if assign: + self.weight = restored state_dict[weight_key] = restored else: self.weight.data = torch.empty_like(restored.data, device=device) diff --git a/tests/test_embedding_checkpoint.py b/tests/test_embedding_checkpoint.py index f5e3ad6ea..38f26e06f 100644 --- a/tests/test_embedding_checkpoint.py +++ b/tests/test_embedding_checkpoint.py @@ -132,3 +132,16 @@ def test_embedding_checkpoint_requires_int8_scales(): with pytest.raises(RuntimeError, match="Missing quantization scales"): destination.load_state_dict(state) torch.testing.assert_close(destination.weight, original_weight, rtol=0, atol=0) + + +@pytest.mark.parametrize("kind", ["int8", "fp4", "nf4"]) +@pytest.mark.parametrize("width", [32, 64]) +def test_embedding_checkpoint_assign_materializes_meta(kind, width): + source = make_embedding(kind, width=width).to("cpu") + with torch.device("meta"): + destination = make_embedding(kind, width=width) + assert destination.weight.device.type == "meta" + destination.load_state_dict(source.state_dict(), assign=True, strict=True) + assert destination.weight.device.type == "cpu" + inputs = torch.tensor([[1, 7], [4, 9]]) + torch.testing.assert_close(destination(inputs), source(inputs), rtol=0, atol=0) From 3d5968010e26902badecb3113e1fed30a0ee9c4e Mon Sep 17 00:00:00 2001 From: Afloat Date: Fri, 13 Nov 2026 17:55:53 +0800 Subject: [PATCH 3/4] Update embedding checkpoint contract tests for supported save/load Signed-off-by: Afloat --- tests/test_modules.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/tests/test_modules.py b/tests/test_modules.py index 95f78b6d3..d9ae0e9f5 100644 --- a/tests/test_modules.py +++ b/tests/test_modules.py @@ -484,11 +484,11 @@ def test_4bit_linear_weight_fsdp_fix(device): assert module.weight.quant_state is not None -def test_embedding_not_implemented_error(): - with pytest.raises(NotImplementedError): - emb = bnb.nn.Embedding4bit(32, 32) - emb.state_dict() - - with pytest.raises(NotImplementedError): - emb = bnb.nn.Embedding8bit(32, 32) - emb.state_dict() +@pytest.mark.parametrize("embedding_cls", [bnb.nn.Embedding8bit, bnb.nn.EmbeddingFP4, bnb.nn.EmbeddingNF4]) +def test_embedding_unquantized_checkpoint(embedding_cls): + emb = embedding_cls(32, 32) + restored = embedding_cls(32, 32) + result = restored.load_state_dict(emb.state_dict(), strict=True) + assert not result.missing_keys + assert not result.unexpected_keys + torch.testing.assert_close(restored.weight, emb.weight, rtol=0, atol=0) From 68fef90cd796dc324d12abb5cbaa3fa7166fa0ae Mon Sep 17 00:00:00 2001 From: Afloat Date: Fri, 13 Nov 2026 17:59:41 +0800 Subject: [PATCH 4/4] Apply repository import grouping to regression tests Signed-off-by: Afloat --- tests/test_embedding_checkpoint.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_embedding_checkpoint.py b/tests/test_embedding_checkpoint.py index 38f26e06f..09e84d186 100644 --- a/tests/test_embedding_checkpoint.py +++ b/tests/test_embedding_checkpoint.py @@ -1,10 +1,10 @@ import copy import pytest -from tests.helpers import get_available_devices import torch import bitsandbytes as bnb +from tests.helpers import get_available_devices def make_embedding(kind, width=64, dtype=torch.float32, storage=torch.uint8):