Keep quantized embedding output dtype in sync with module conversions

bitsandbytes/nn/modules.py

@@ -830,7 +830,15 @@ def maybe_rearrange_weight(state_dict, prefix, local_metadata, strict, missing_k
         raise ValueError(f"Only 'row' weight format is supported, got {weight_format}")
 
 
-class Embedding8bit(nn.Embedding):
+class _EmbeddingWithOutputDtype(nn.Embedding):
+    def _apply(self, fn, recurse=True):
+        output_dtype = fn(torch.empty(0, dtype=self.dtype)).dtype
+        module = super()._apply(fn, recurse)
+        self.dtype = output_dtype
+        return module
+
+
+class Embedding8bit(_EmbeddingWithOutputDtype):
     """
     This class implements [LLM.int8()](https://arxiv.org/abs/2208.07339) algorithm for embedding layer
 
@@ -877,7 +885,7 @@ def forward(self, input: Tensor) -> Tensor:
         return output.to(self.dtype)
 
 
-class Embedding4bit(nn.Embedding):
+class Embedding4bit(_EmbeddingWithOutputDtype):
     """
     This is the base class similar to Linear4bit. It implements the 4-bit quantization algorithm presented in
     [QLoRA](https://arxiv.org/abs/2305.14314) for embeddings.

tests/test_modules.py

@@ -492,3 +492,70 @@ def test_embedding_not_implemented_error():
     with pytest.raises(NotImplementedError):
         emb = bnb.nn.Embedding8bit(32, 32)
         emb.state_dict()
+
+
+@pytest.mark.parametrize("device", get_available_devices())
+@pytest.mark.parametrize("embedding_class", [bnb.nn.Embedding8bit, bnb.nn.EmbeddingFP4, bnb.nn.EmbeddingNF4])
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
+@pytest.mark.parametrize("quantize_first", [False, True])
+def test_quantized_embedding_dtype_model_updates(device, embedding_class, dtype, quantize_first):
+    if device == "hpu" and embedding_class is not bnb.nn.Embedding8bit:
+        quant_type = "fp4" if embedding_class is bnb.nn.EmbeddingFP4 else "nf4"
+        if not is_supported_on_hpu(quant_type, dtype):
+            pytest.skip("This quantization and dtype configuration is not supported on HPU")
+    embedding = embedding_class(16, 64, dtype=torch.float32)
+    model = nn.Sequential(embedding, nn.Linear(64, 3))
+    packed_weight = None
+    if quantize_first:
+        model.to(device)
+        packed_weight = embedding.weight.detach().clone()
+    model.to(dtype=dtype)
+    if not quantize_first:
+        model.to(device)
+    if packed_weight is not None:
+        torch.testing.assert_close(embedding.weight, packed_weight, atol=0, rtol=0)
+    tokens = torch.tensor([[1, 3, 5], [0, 15, 4]], device=device)
+    assert embedding(tokens).dtype == dtype
+    assert embedding.dtype == dtype
+
+    with torch.no_grad():
+        dequantized = embedding(torch.arange(16, device=device))
+    reference = nn.Sequential(
+        nn.Embedding.from_pretrained(dequantized, freeze=True),
+        nn.Linear(64, 3, device=device, dtype=dtype),
+    )
+    reference[1].load_state_dict(model[1].state_dict())
+    optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9)
+    reference_optimizer = torch.optim.SGD(reference.parameters(), lr=0.1, momentum=0.9)
+    for _ in range(2):
+        optimizer.zero_grad()
+        reference_optimizer.zero_grad()
+        actual = model(tokens)
+        expected = reference(tokens)
+        torch.testing.assert_close(actual, expected, atol=0, rtol=0)
+        actual.float().square().mean().backward()
+        expected.float().square().mean().backward()
+        for parameter, reference_parameter in zip(model[1].parameters(), reference[1].parameters()):
+            torch.testing.assert_close(parameter.grad, reference_parameter.grad, atol=0, rtol=0)
+        optimizer.step()
+        reference_optimizer.step()
+        for parameter, reference_parameter in zip(model[1].parameters(), reference[1].parameters()):
+            torch.testing.assert_close(parameter, reference_parameter, atol=0, rtol=0)
+    assert not embedding.weight.requires_grad
+
+
+@pytest.mark.parametrize("device", get_available_devices())
+@pytest.mark.parametrize("embedding_class", [bnb.nn.Embedding8bit, bnb.nn.EmbeddingFP4, bnb.nn.EmbeddingNF4])
+def test_quantized_embedding_dtype_convenience_methods(device, embedding_class):
+    if device == "hpu" and embedding_class is not bnb.nn.Embedding8bit:
+        pytest.skip("FP16 4-bit embedding inference is not supported on HPU")
+    embedding = embedding_class(16, 64).to(device)
+    tokens = torch.tensor([0, 7, 15], device=device)
+    expected = embedding(tokens).detach()
+    packed_weight = embedding.weight.detach().clone()
+    for method, dtype in (("half", torch.float16), ("float", torch.float32), ("bfloat16", torch.bfloat16)):
+        assert getattr(embedding, method)() is embedding
+        assert embedding.dtype == dtype
+        actual = embedding(tokens)
+        torch.testing.assert_close(actual, expected.to(dtype), atol=0, rtol=0)
+        torch.testing.assert_close(embedding.weight, packed_weight, atol=0, rtol=0)