Dequantize the bitsandbytes base before the vLLM weight push

[behroozazarkhalili]

What does this PR do?

Fixes #4973. This is direction A from the discussion on that issue; direction B is still open, see below.

The report attributed the failure to merge_adapter() dequantizing the base. That is not what happens: PEFT's Linear4bit.merge leaves the weight as Params4bit, so nothing is upcast.

The measured cause is different. bitsandbytes stores a 4-bit weight as a flat packed uint8 buffer whose scales live in quant_state. The sync path read param.data, which drops quant_state, and handed that buffer to vLLM. On Qwen2.5-3B, 14 of 27 tensors failed the shape assert in vLLM's weight loader, and 0 quant_state tensors were transferred, so the scales never crossed even for the tensors whose shapes did fit.

The change

Every push site now goes through _dense_param_data, which dequantizes a Params4bit back to the model dtype and returns any other parameter's data unchanged. The dequantization acts on the parameter object rather than on .data, because .data has already discarded the quantization state by the time the push sees it.

The colocate engine is built dense to match. It previously received quantization="bitsandbytes" whenever a Linear4bit was found, which allocates packed [out_features, in_features // 2] weights and would reject the dense push with the same assert. The two halves only work together.

trl/experimental/online_dpo/online_dpo_trainer.py duplicates this block and is updated in lockstep.

Cost

The vLLM engine now holds a full-precision copy of the base, so the QLoRA memory saving covers training but not the rollout engine.

The alternative

Direction B, sending the packed weights together with quant_state and adding a --quantization flag, keeps the engine quantized and avoids that cost. It is a larger change and I have not built it. If you prefer B, say so and I will close this.

Verification

• ruff check, ruff format --check and the pinned doc-builder style pass on both files, each run alongside a control so a silent no-op would show up.
• Both branches of the helper are covered: a plain parameter passes its storage through, and a Params4bit routes to dequantize_4bit with quant_state forwarded and the packed shape intact.
• Not yet done: an end-to-end QLoRA plus vLLM run on GPU. The reproduction needs two devices and I have not had them.

---

Note

Medium Risk
Touches every vLLM weight-sync path and changes colocate engine quantization defaults; behavior is well-guarded for unsupported combos but increases GPU memory for QLoRA rollouts.

Overview
Fixes QLoRA + vLLM weight sync by dequantizing bitsandbytes 4-bit parameters before they are pushed to vLLM, instead of sending packed param.data (which drops quant_state and breaks shape checks).

Adds _dense_param_data (4-bit → dense via dequantize_4bit) and routes all server/colocate sync sites through it in VLLMGeneration and OnlineDPOTrainer. Colocate LLM init no longer sets quantization="bitsandbytes" so the engine expects dense weights matching the push.

Adds _check_quantization_supported: still rejects 8-bit bases; fails fast on 4-bit + FSDP2 (weights from state_dict() cannot be dequantized). Unit tests cover both helpers under @require_bitsandbytes.

Trade-off: the vLLM rollout engine holds a full-precision copy of the base; QLoRA memory savings apply to training, not the colocated engine.

^{Reviewed by Cursor Bugbot for commit 8da2109. Bugbot is set up for automated code reviews on this repo. Configure here.}

[bot-ci-comment]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[cursor]

Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

^{Reviewed by Cursor Bugbot for commit ea30fad. Configure here.}