Skip to content

Support saving and restoring quantized embedding checkpoints - #2139

Draft
Afloat16 wants to merge 4 commits into
bitsandbytes-foundation:mainfrom
Afloat16:feat/quantized-embedding-checkpoints-20261113
Draft

Afloat16 wants to merge 4 commits into
bitsandbytes-foundation:mainfrom
Afloat16:feat/quantized-embedding-checkpoints-20261113

Conversation

@Afloat16

Copy link
Copy Markdown

Problem

Embedding8bit and Embedding4bit currently raise NotImplementedError from state_dict(). A model containing these frozen quantized embeddings therefore cannot use the normal model/optimizer checkpoint workflow, even when only its downstream layers are trainable.

Change

  • Save INT8 weights with their row scales and FP4/NF4 weights with the existing packed QuantState metadata.
  • Restore into fresh or already quantized embedding modules without requantizing the saved values.
  • Preserve the destination parameter object for ordinary copy-based loading, and support assign=True, including materializing a meta-initialized embedding.
  • Restore the weight/module quantization-state links, keep copied weights and metadata independent of the input checkpoint, and reject mismatched logical embedding shapes or missing INT8 scales.
  • Replace the previous unsupported-save test with the supported unquantized checkpoint contract.

Validation

Exact production commit: 68fef90cd796dc324d12abb5cbaa3fa7166fa0ae.

44 new native CPU cases cover INT8/FP4/NF4, both partial and block-aligned row widths, copy/assign modes, fresh and prequantized destinations, FP16/BF16 packed storage, and meta initialization. The model tests train a downstream layer, save model and optimizer states to disk, restore both, and compare the next two SGD steps exactly.

Original source: all 44 new cases fail. Patched source: all 44 pass. The existing embedding/Int8Params selection passes 64 tests with one existing unsupported-CPU skip. Ruff and formatting pass for all changed Python files.

Native validation: https://github.com/Afloat16/bitsandbytes/actions/runs/43405757953

The run compiles the repository's AVX2 CPU extension and uses PyTorch 2.14.0+cpu. The validation concerns checkpoint fidelity and resumed execution; it does not measure GPU throughput.

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 <ShuffleMonkey@163.com>
Cover fresh and meta embeddings with native assign=True restores.

Signed-off-by: Afloat <ShuffleMonkey@163.com>
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