Skip to content

Keep quantized embedding output dtype in sync with module conversions - #2137

Draft
Afloat16 wants to merge 2 commits into
bitsandbytes-foundation:mainfrom
Afloat16:fix/quantized-embedding-output-dtype-20261112
Draft

Afloat16 wants to merge 2 commits into
bitsandbytes-foundation:mainfrom
Afloat16:fix/quantized-embedding-output-dtype-20261112

Conversation

@Afloat16

Copy link
Copy Markdown

Summary

Quantized embeddings store their output dtype when they are constructed. Converting a containing model with .half(), .bfloat16(), .float(), or .to(dtype=...) currently leaves that attribute unchanged. An embedding constructed in float32 consequently still returns float32 after the next linear layer has been converted, and the model fails with a dtype mismatch.

Add a shared embedding base that updates the output dtype after a successful module conversion. The existing 8-bit and 4-bit dequantization paths consume the updated dtype. Packed int8/uint8 weights retain their storage dtype, and the state dictionary format is unchanged.

Validation

  • 15 new native tests cover Embedding8bit, EmbeddingFP4, and EmbeddingNF4, FP16/BF16, conversion before and after quantization, and chained convenience conversions.
  • Real Sequential(embedding, linear) models take two SGD steps; outputs, linear gradients, and updated parameters match a floating-point embedding reference built from the dequantized weights.
  • Original source: all 15 new tests fail.
  • Patched source: 77 embedding tests pass, with one existing CPU-unsupported FSDP test skipped.
  • Ruff checks and formatting pass.

The repository's actual CPU extension was compiled with its CMake CPU backend (AVX2), and loaded with PyTorch 2.14.0+cpu. These results cover default packed int8/uint8 storage; floating-point packed 4-bit storage conversions are a separate existing issue.

Exact production head: 879f512b0d014cb53c3cacb6817803c92e09e75d.
https://github.com/Afloat16/bitsandbytes/actions/runs/43307725418

Make dtype conversion of parent models update the activation dtype for
8-bit and 4-bit embeddings, including already quantized weights.

Signed-off-by: Afloat <ShuffleMonkey@163.com>
Signed-off-by: Afloat <ShuffleMonkey@163.com>

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant