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.

Which combinations are refused at build time. 8-bit anywhere, since vLLM has no in-flight 8-bit path. A 4-bit base under FSDP2, which reads weights from state_dict() as plain tensors with the quant_state already gone. And a 4-bit base under FSDP1 unless use_orig_params is True: measured on one H100 with a bitsandbytes Linear4bit(8, 8) wrapped in FSDP1, Accelerate's default use_orig_params=False makes summon_full_params expose a plain Parameter of shape (16, 1) with no quant_state, and that packed storage is what the sync would push; with use_orig_params=True the summoned weight is the Params4bit with its state and the push is the dense (8, 8). DistributedBackend exposes fsdp_use_orig_params so the guard can tell the two apart, and the error names the setting to flip. And, in colocate mode, a checkpoint that was saved already quantized with bitsandbytes: vLLM reads quantization_config from the checkpoint at name_or_path and allocates packed weights the dense push cannot fill. transformers records that case as hf_quantizer.pre_quantized, derived from the same config field, so the guard takes it as an argument; server mode passes False because the trainer does not know what the server was started from, and the docstring says so.

Two more sites were out of step. VLLMClient.update_model_params still sent param.data and described the packed (N, 1) uint8 shape in its metadata; it now routes both passes through _dense_param_data, which moved to vllm_client.py so the client can use it without a circular import. DistributedBackend read fsdp_version with a getattr default of None, which reports "no FSDP" on an Accelerate before 1.6.0, where the plugin has no such attribute and can only configure FSDP1; the version is now checked explicitly and older plugins report 1. OnlineDPO read the attribute directly, which raised there, and now builds the backend like VLLMGeneration does.

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.
• The FSDP1 use_orig_params split was measured on one H100 MIG slice (world size 1): use_orig_params=False pushes (16, 1) for an (8, 8) layer, True pushes (8, 8). The guard's test matrix covers False, None and True; a copy of the tree with the FSDP1 branch removed fails exactly the two refused rows.
• Not evaluable here: the dtype of a PEFT merge_adapter result on a bitsandbytes 4-bit base. PEFT raised a shape mismatch inside the merge before any dtype could be read, so the helper's documented model-dtype output remains verified only for the unmerged path.
• Not yet done: an end-to-end QLoRA plus vLLM run on GPU. The reproduction needs two devices and I have not had them.
• Third round (review follow-ups): the guard matrix gains a pre_quantized column with two refused rows and a dense control, a client test checks that update_model_params sends the dequantized tensor with dense metadata, and three backend tests cover a current plugin, a pre-1.6.0 plugin, and no plugin. All 23 pass on the fixed tree; a copy with the pre-quantized branch removed fails exactly the two refused rows; the previous head cannot collect the file at all, since it has no _dense_param_data in the client.

---

Note

Medium Risk
Changes the critical training→vLLM weight sync path and distributed setup guards; wrong behavior would break rollouts or silently mis-load weights, though unsupported setups now error explicitly.

Overview
Fixes QLoRA + vLLM weight sync (#4973) by dequantizing bitsandbytes Params4bit weights before pushing them to vLLM, instead of sending the packed uint8 buffer (which drops quant_state and breaks shape checks).

_dense_param_data is the single path for tensors sent to vLLM (client update_model_params, VLLMGeneration sync, and experimental Online DPO). Colocated engines are initialized without quantization="bitsandbytes" so vLLM allocates dense weights that match the push.

_check_quantization_supported fails fast at vLLM setup for combinations that cannot work with dense sync: 8-bit layers, 4-bit under FSDP2, 4-bit under FSDP1 unless use_orig_params=True, and colocate loads from checkpoints already saved quantized. DistributedBackend now exposes fsdp_use_orig_params and treats Accelerate < 1.6.0 as FSDP1 when inferring fsdp_version.

Tests cover dequantization, metadata for packed weights, the guard matrix, and FSDP version detection.

^{Reviewed by Cursor Bugbot for commit e40372b. 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 e40372b. Configure here.}

[cursor]

Colocate init assumes hf_quantizer exists

High Severity

Colocate vLLM init reads model.hf_quantizer to decide pre_quantized. Transformers only sets that attribute when a quantizer ran during from_pretrained, so an unquantized model raises AttributeError and never builds the engine. Colocate is the default mode, so ordinary dense training hits this before any weight sync.

Additional Locations (1)

• trl/experimental/online_dpo/online_dpo_trainer.py#L434-L437

 

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