feat: Bitpack-on-quantize mode for GPTQ and DBF - #43
Merged
aki916f merged 53 commits intoJul 30, 2026
Merged
Conversation
Copy link
Copy Markdown
Contributor
Sorry, something went wrong.
…, inference paths
…itpack_mode_merge Resolved conflicts in: - onecomp/quantizer/gptq/_gptq.py - tests/onecomp/quantizer/gptq/test_gptq.py
…eature/ysmz/bitpack_mode
Add a DBF-specific bitpack_on_quantize flag (default True) to the DBF quantizer dataclass and validate it as a bool in validate_params(). DBF packs arbitrary shapes via padding, so no bit-width/shape constraint is imposed; only the bool type is checked.
Add packed-state metadata (dbf_A_is_packed, dbf_B_is_packed, dbf_A_original_shape, dbf_B_original_shape) and a get_unpacked_binary_factors() helper to DBFResult so callers can obtain the unpacked +/-1 factors regardless of storage. The helper returns float16 for both packed and unpacked storage so the representation difference never leaks to callers. Add shape-aware unpack_binary_matrix() (and pack_binary_matrix alias) to dbf_layer so the pack/unpack shape-restore logic lives in one place. Old results without the metadata fields are treated as unpacked via getattr defaults.
In DBF.quantize_layer(), when bitpack_on_quantize is enabled and DBF produced both binary factors, pack dbf_A / dbf_B into uint8 CPU tensors immediately after run_dbf() and record the packed-state metadata on the DBFResult. The unpacked references are dropped early to reduce RAM held during quantization. On DBF failure / is_dbf_quantized=False the factors are left untouched, preserving the existing dequant error behavior.
compute_dequantized_weight() now recovers the +/-1 factors via get_unpacked_binary_factors(), so it returns the same weight whether dbf_A / dbf_B are stored packed or unpacked (pack/unpack is value-preserving). The reconstruction formula is unchanged. Route the existing in-layer unpack call sites (BitLinearPacked.forward and DoubleBinaryLinear._unpack_bp) through the shared unpack_binary_matrix helper to avoid duplicating the shape-restore logic.
DoubleBinaryLinear.__init__() / from_quantization_result() now accept already-packed dbf_A / dbf_B (with their original shapes) and register them straight into bp1 / bp3 without an unpack/re-pack round-trip; unpacked inputs are packed as before. GemLite initialization builds a short-lived unpacked copy for packed inputs only (not retained). The saved checkpoint format (bp1 / bp3) and from_saved_state() are unchanged.
Update the existing DBF quantizer test assertions that assumed unpacked float16 dbf_A / dbf_B: check_quantize_layer branches on the packed flags (uint8/1D when packed, float16/2D otherwise) and reconstructs the weight via get_unpacked_binary_factors(); check_equal_results and the apply-to-module helper compare/store factors through the same helper so they hold for both representations. No changes were needed in the blockwise post-processing optimizers or the vLLM DBF plugin, which work on the unchanged bp1 / bp3 packed buffers.
…allback Add a bitpack_on_quantize flag (default True) to AutoBitQuantizer and propagate it to candidate quantizers that support on-quantize bitpacking: validate_params() syncs the flag onto GPTQ/DBF child candidates (via _sync_child_bitpack_on_quantize) before their own validation, _assign_all_dbf() forwards it to the all-DBF quantizer, and the inject_dbf() fallback path forwards it to every DBF it creates. A non-bool flag raises a clear ValueError.
Add DBF quantizer tests covering both representations: default flag is True; packing stores uint8 1D factors with shape metadata; disabling keeps unpacked float16; packed and unpacked dequantized weights are bit-identical (both from quantize_layer and hand-built results); pack_binary -> unpack_binary_matrix is bit-exact across shapes; missing original shape raises; from_quantization_result registers bp1/bp3 with no re-pack and forward matches the dequantized linear; and a non-bool flag raises. Add AutoBit tests that the flag is synced onto DBF candidates and forwarded to inject_dbf fallbacks (1-bit GPTQ on a 64-wide layer to stay under the DBF threshold).
QEP, LPCD, chunked / multi-GPU quantization, and the cumulative-error analyzer all consume DBF results only via compute_dequantized_weight().to(device).to(dtype) and never touch dbf_A / dbf_B directly, so no functional change is needed for the on-quantize bitpack. Add a regression smoke test exercising that exact consumption path (analyzer._update_weights) with a packed DBFResult, confirming it succeeds and yields weights identical to the unpacked representation (no packed state leaks to callers).
- pre_process pipeline: add DBF quantized/dequantized save/load round-trip cases for TinyLlama and Qwen3 (5 cases: GPTQ q/dq, DBF q/dq, RTN dq) - tests/vllm_plugins/dbf/test_dbf_e2e.py: DBF quantize -> save -> config verification -> vLLM generation smoke (slow, CUDA/vLLM gated) - tests/vllm_plugins/dbf/test_dbf_config.py: DBF plugin config parse/dispatch smoke (DbfConfig)
Add DBF packed-result smoke coverage for QEP, LPCD, and chunked calibration calc_quant_error paths. The tests verify that each Runner path produces packed dbf_A / dbf_B results and consumes them through compute_dequantized_weight() without leaking storage details to callers.
register_online_hadamard_hooks -> get_hadK introspects module.in_features, but DoubleBinaryLinear only stored _bp1_shape/_bp3_shape, raising AttributeError when re-registering Hadamard hooks on saved DBF-quantized rotated models. Expose in_features/out_features like nn.Linear (derived from the original unpacked shapes; also set in from_saved_state).
Quantize a layer with bitpack_on_quantize=False vs True under identical seeds and assert the dequantized weights match bit-exactly (compute_dequantized_weight -> torch.equal). The corresponding DBF equivalence test is in tests/onecomp/quantizer/dbf/test_dbf_bitpack_equivalence.py.
Parallels example_gptq_vllm_inference.py for the DBF path: quantize TinyLlama with DBF, save, then load and generate with vLLM.
Add a bitpack_on_quantize row to the parameters table in docs/algorithms/dbf.md (DBF) and docs/algorithms/autobit.md (propagation to candidate / injected-DBF quantizers).
…smz/bitpack_mode Resolved conflicts in: - CHANGELOG.md
- Document the breaking change in CHANGELOG - Add a comment explaining qzeros reshape after pack/unpack
…ikoji/bitpack_dbf
… feature/ysmz/bitpack_mode Resolve conflicts in GPTQ/DBF bitpack-on-quantize integration: - Unify AutoBit propagation for GPTQ and DBF child quantizers and auto-created DBF fallbacks - Preserve GPTQ packable-width validation and DBF unrestricted 1-bit packing behavior
Equivalence is already covered by tests, and these examples do not demonstrate user-facing usage.
- document GPTQ quantize-time and save-time packing constraints - clarify AutoBit propagation, overrides, and fused-group restrictions - document DBF and JointQ bitpacking behavior - record these doc updates in the DBF and GPTQ bitpack changelog sections
…smz/bitpack_mode Resolve the CHANGELOG.md conflict and fix a silent DBF auto-merge breakage: - Preserve both bitpack-mode and upstream v1.3.0 changelog entries. - Drop upstream's premature in_features/out_features assignment in DoubleBinaryLinear.__init__, which indexed 1-D packed uint8 factors. - Keep the shape-derived assignment after packed/unpacked factor shapes are resolved, preserving the attributes required by rotation support.
- Remove references to the deleted DBF and GPTQ equivalence examples - Add the missing Qwen3.6 vLLM example to README and the inference guide
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
This PR introduces a bitpack_on_quantize mode that keeps quantization results in bitpacked form from the moment each layer is quantized, for both GPTQ and DBF.
Motivation. Until now, quantization results held unpacked tensors for the whole run — int32 qweight / qzeros for GPTQ, ±1 float16 dbf_A / dbf_B for DBF — and packing happened only later, when an inference layer was created or the model was saved. For large models, the accumulated unpacked results can account for a substantial share of the memory held during quantization. Packing them right after each layer is quantized reduces storage for qweight / qzeros by roughly 8x at 4-bit and for the DBF binary factors by roughly 16x, and lets the large unpacked tensors be freed early.
What changes. bitpack_on_quantize is added to the base Quantizer (default False) and enabled by default for GPTQ and DBF. AutoBitQuantizer propagates the setting to existing GPTQ and DBF child candidates before child validation and passes it to DBF fallback quantizers created by inject_dbf(). JointQ deliberately keeps its internal GPTQ initialization unpacked to preserve its existing optimization path.
Compatibility. Supported downstream paths include compute_dequantized_weight(), inference-layer construction (GPTQLinear and DoubleBinaryLinear), the covered QEP / LPCD / chunked-calibration paths, save/load round-trips, and DBF vLLM inference. Equivalence tests verify that packed storage reconstructs the same dequantized weights as unpacked storage for fixed GPTQ and DBF quantization results.
Breaking changes.
Bug fix. DoubleBinaryLinear now exposes in_features / out_features, fixing the AttributeError raised by register_online_hadamard_hooks() -> get_hadK() when re-registering Hadamard hooks on saved, DBF-quantized, rotation-preprocessed models.
Documentation, a new DBF + vLLM inference example, and unit / equivalence / runner smoke / vLLM e2e tests are included.
Changes
GPTQ
New Feature / Breaking Changes: GPTQ bitpack-on-quantize mode
Validation / compatibility tweaks
Documentation
Tests
DBF
New Feature / Breaking Changes: DBF bitpack-on-quantize mode
Validation / compatibility tweaks
Bug Fix
Examples
Documentation
Tests