Dequantize the bitsandbytes base before the vLLM weight push

trl/experimental/online_dpo/online_dpo_trainer.py

@@ -51,6 +51,7 @@
 from ...data_utils import apply_chat_template, is_conversational, maybe_apply_chat_template
 from ...extras.profiling import profiling_context
 from ...generation.vllm_client import VLLMClient
+from ...generation.vllm_generation import _dense_param_data
 from ...import_utils import is_vllm_available
 from ...models.utils import prepare_deepspeed, prepare_fsdp, unwrap_model_for_generation
 from ...trainer.base_trainer import _BaseTrainer
@@ -448,13 +449,13 @@ def __init__(
                 # after the first optimizer step and remain in GPU memory throughout training. So we must reserve enough
                 # space for them.
                 # Configure vLLM parameters
-                vllm_quantization = None
+                # The engine is built dense on purpose. The base may be 4-bit, but every weight pushed at sync
+                # time is dequantized to the model dtype (see `_dense_param_data`), and building the engine with
+                # `quantization="bitsandbytes"` would allocate packed `[out_features, in_features // 2]` weights
+                # that reject that dense push. See https://github.com/huggingface/trl/issues/4973.
                 if is_bitsandbytes_available():
                     for _, module in model.named_modules():
-                        if isinstance(module, bnb.nn.Linear4bit):
-                            vllm_quantization = "bitsandbytes"
-                            break
-                        elif isinstance(module, bnb.nn.Linear8bitLt):
+                        if isinstance(module, bnb.nn.Linear8bitLt):
                             raise ValueError("vLLM does not support in-flight 8-bit quantization.")
                 vllm_kwargs = {
                     "model": model.name_or_path,
@@ -469,7 +470,6 @@ def __init__(
                     # Latest vLLM v1 memory profiler is misled by the high default value (i.e., 32768)
                     "max_num_batched_tokens": 4096,
                     "enable_sleep_mode": self.args.vllm_enable_sleep_mode,
-                    "quantization": vllm_quantization,
                 }
 
                 # vLLM requires the environment variables to be set for distributed training.
@@ -806,10 +806,10 @@ def _move_model_to_vllm(self):
                         name = self._fix_param_name_to_vllm(name, extra_prefixes=["modules_to_save.default."])
 
                         if self.vllm_mode == "server" and self.accelerator.is_main_process:
-                            self.vllm_client.update_named_param(name, param.data)
+                            self.vllm_client.update_named_param(name, _dense_param_data(param))
                         elif self.vllm_mode == "colocate":
                             llm_model = self.llm.llm_engine.model_executor.driver_worker.model_runner.model
-                            llm_model.load_weights([(name, param.data)])
+                            llm_model.load_weights([(name, _dense_param_data(param))])
                 # Unmerge adapters while parameters are still gathered
                 self.model.unmerge_adapter()
                 # Parameters will automatically be repartitioned when exiting the context
@@ -827,10 +827,10 @@ def _move_model_to_vllm(self):
                     name = self._fix_param_name_to_vllm(name)
                     with gather_if_zero3([param]):
                         if self.vllm_mode == "server" and self.accelerator.is_main_process:
-                            self.vllm_client.update_named_param(name, param.data)
+                            self.vllm_client.update_named_param(name, _dense_param_data(param))
                         elif self.vllm_mode == "colocate":
                             llm_model = self.llm.llm_engine.model_executor.driver_worker.model_runner.model
-                            llm_model.load_weights([(name, param.data)])
+                            llm_model.load_weights([(name, _dense_param_data(param))])
 
         # Reset cache on vLLM
         if self.vllm_mode == "server" and self.accelerator.is_main_process:
@@ -860,10 +860,10 @@ def _sync_fsdp1_params_to_vllm(self, module: nn.Module, prefix: str = "", visite
                     visited.add(full_name)
 
                     if self.vllm_mode == "server" and self.accelerator.is_main_process:
-                        self.vllm_client.update_named_param(full_name, param.data)
+                        self.vllm_client.update_named_param(full_name, _dense_param_data(param))
                     elif self.vllm_mode == "colocate":
                         llm_model = self.llm.llm_engine.model_executor.driver_worker.model_runner.model
-                        llm_model.load_weights([(full_name, param.data)])
+                        llm_model.load_weights([(full_name, _dense_param_data(param))])
 
     def _fix_param_name_to_vllm(self, name, extra_prefixes: list[str] | None = None):
         """Clean parameter names for vLLM compatibility"""

trl/generation/vllm_generation.py

@@ -108,6 +108,28 @@ def extract_logprobs(all_outputs: list["RequestOutput"]):
     import bitsandbytes as bnb
 
 
+def _dense_param_data(param: nn.Parameter) -> torch.Tensor:
+    """
+    Return a parameter's data as a dense tensor, dequantizing bitsandbytes 4-bit weights.
+
+    vLLM's weight loader expects a dense `[out_features, in_features]` tensor in the model dtype. A bitsandbytes 4-bit
+    base instead stores a flat packed `uint8` buffer whose scales live in `quant_state`, and reading `.data` drops that
+    `quant_state`. Pushing it unchanged therefore both mismatches the shape vLLM allocated and loses the scales
+    entirely. See https://github.com/huggingface/trl/issues/4973.
+
+    Args:
+        param (`torch.nn.Parameter`):
+            Parameter to read. Must be the parameter object rather than its `.data`, since `.data` has already
+            discarded the quantization state.
+
+    Returns:
+        `torch.Tensor`: the dense weight, dequantized when the parameter is 4-bit.
+    """
+    if is_bitsandbytes_available() and isinstance(param, bnb.nn.Params4bit):
+        return bnb.functional.dequantize_4bit(param.data, param.quant_state)
+    return param.data
+
+
 class VLLMGeneration:
     """Handles vLLM-based generation for trainers.
 
@@ -338,13 +360,13 @@ def _init_vllm(self):
             # Ensure distributed rendezvous variables are set without colliding across concurrent runs
             ensure_master_addr_port()
 
-            quantization = None
+            # The engine is built dense on purpose. The base may be 4-bit, but every weight pushed at sync time
+            # is dequantized to the model dtype (see `_dense_param_data`), and building the engine with
+            # `quantization="bitsandbytes"` would allocate packed `[out_features, in_features // 2]` weights that
+            # reject that dense push. See https://github.com/huggingface/trl/issues/4973.
             if is_bitsandbytes_available():
                 for _, module in model.named_modules():
-                    if isinstance(module, bnb.nn.Linear4bit):
-                        quantization = "bitsandbytes"
-                        break
-                    elif isinstance(module, bnb.nn.Linear8bitLt):
+                    if isinstance(module, bnb.nn.Linear8bitLt):
                         raise ValueError("vLLM does not support in-flight 8-bit quantization.")
 
             # Build LLM initialization kwargs
@@ -363,7 +385,6 @@ def _init_vllm(self):
                 max_num_batched_tokens=4096,
                 # Important so temperature scaling/logit tweaking affects the TIS log probs
                 logprobs_mode="processed_logprobs",
-                quantization=quantization,
                 trust_remote_code=self.trust_remote_code,
             )
             if self.enable_sleep_mode:
@@ -414,7 +435,7 @@ def _sync_fsdp1_params_to_vllm(self, module: nn.Module, prefix: str = "", visite
                         continue  # skip FSDP subtrees already traversed
                     visited.add(full_name)
 
-                    self._push_param_to_vllm(full_name, param.data)
+                    self._push_param_to_vllm(full_name, _dense_param_data(param))
 
     def _sync_fsdp2_params_to_vllm(self, module: nn.Module):
         """FSDP2-specific parameter synchronization."""
@@ -483,7 +504,7 @@ def sync_weights(self):
                             continue
                         name = self._fix_param_name_to_vllm(name, extra_prefixes=["modules_to_save.default."])
 
-                        self._push_param_to_vllm(name, param.data)
+                        self._push_param_to_vllm(name, _dense_param_data(param))
                 # Unmerge adapters while parameters are still gathered
                 model.unmerge_adapter()
                 # Parameters will automatically be repartitioned when exiting the context
@@ -495,7 +516,7 @@ def sync_weights(self):
                 for name, param in model.named_parameters():
                     name = self._fix_param_name_to_vllm(name)
                     with self._dist.gather_params([param]):
-                        self._push_param_to_vllm(name, param.data)
+                        self._push_param_to_vllm(name, _dense_param_data(param))
 
         # Reset cache on vLLM
         if self.mode == "server" and accelerator.is_main_process: