Dequantize the bitsandbytes base before the vLLM weight push

tests/test_vllm_client_server.py

@@ -17,17 +17,20 @@
 from types import SimpleNamespace
 
 import pytest
-from transformers import AutoModelForCausalLM, AutoProcessor, AutoTokenizer
+import torch
+from torch import nn
+from transformers import AutoModelForCausalLM, AutoProcessor, AutoTokenizer, is_bitsandbytes_available
 from transformers.testing_utils import torch_device
 
 from trl.generation.vllm_client import VLLMClient, parse_logprobs
-from trl.generation.vllm_generation import extract_logprobs
+from trl.generation.vllm_generation import _check_quantization_supported, _dense_param_data, extract_logprobs
 from trl.import_utils import is_vllm_available
 
 from .testing_utils import (
     TrlTestCase,
     kill_process,
     require_3_accelerators,
+    require_bitsandbytes,
     require_torch_multi_accelerator,
     require_vision,
     require_vllm,
@@ -37,6 +40,9 @@
 if is_vllm_available():
     from vllm import LLM, SamplingParams
 
+if is_bitsandbytes_available():
+    import bitsandbytes as bnb
+
 
 class TestParseLogprobs(TrlTestCase):
     def test_completion_logprobs_sorted_by_probability(self):
@@ -127,6 +133,42 @@ def test_extract_logprobs_returns_none_token_ids_when_logprobs_missing(self):
         assert all_token_ids is None
 
 
+@require_bitsandbytes
+class TestQuantizedWeightSync(TrlTestCase):
+    # Pure checks on the two helpers that guard a quantized base, so they run without an accelerator or a vLLM engine.
+
+    def test_dense_param_data_passes_through_an_unquantized_parameter(self):
+        # The helper only intercepts 4-bit parameters; anything else must reach vLLM byte for byte.
+        param = nn.Parameter(torch.randn(4, 8))
+        out = _dense_param_data(param)
+        assert out.data_ptr() == param.data.data_ptr()
+        assert out.shape == (4, 8)
+
+    @pytest.mark.parametrize(
+        ("module_factory", "fsdp_version", "expected_message"),
+        [
+            # A 4-bit base under FSDP2 cannot be dequantized: FSDP2 reads weights from `state_dict()`, which returns
+            # plain tensors, so the `quant_state` holding the scales is already gone. Refuse it at build time.
+            (lambda: bnb.nn.Linear4bit(8, 8), 2, "4-bit quantized base under FSDP2"),
+            # vLLM has never supported in-flight 8-bit, independent of the sharding strategy.
+            (lambda: bnb.nn.Linear8bitLt(8, 8), 0, "8-bit quantization"),
+            (lambda: bnb.nn.Linear8bitLt(8, 8), 2, "8-bit quantization"),
+            # Negative controls: every other combination reaches the dequantizing push and must be accepted.
+            (lambda: bnb.nn.Linear4bit(8, 8), 1, None),  # FSDP1 gathers through `summon_full_params`
+            (lambda: bnb.nn.Linear4bit(8, 8), 0, None),  # no FSDP at all
+            (lambda: nn.Linear(8, 8), 2, None),  # dense base under FSDP2
+            (lambda: nn.Linear(8, 8), 0, None),  # dense base, no FSDP
+        ],
+    )
+    def test_check_quantization_supported(self, module_factory, fsdp_version, expected_message):
+        model = nn.Sequential(module_factory())
+        if expected_message is None:
+            _check_quantization_supported(model, fsdp_version)  # must not raise
+        else:
+            with pytest.raises(ValueError, match=expected_message):
+                _check_quantization_supported(model, fsdp_version)
+
+
 @pytest.mark.slow
 @require_torch_multi_accelerator
 @require_vllm

trl/experimental/online_dpo/online_dpo_trainer.py

@@ -56,6 +56,7 @@
 )
 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
@@ -453,14 +454,25 @@ 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():
+                    fsdp_plugin = getattr(self.accelerator.state, "fsdp_plugin", None)
+                    fsdp_version = getattr(fsdp_plugin, "fsdp_version", 1) if fsdp_plugin else 1
                     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.")
+                        # FSDP2 syncs weights from `state_dict()`, which returns plain tensors: the bitsandbytes
+                        # `quant_state` holding the scales is gone before `_dense_param_data` can read it, so the
+                        # base cannot be dequantized for the push. Fail here rather than send packed storage to a
+                        # dense engine.
+                        if isinstance(module, bnb.nn.Linear4bit) and fsdp_version == 2:
+                            raise ValueError(
+                                "vLLM weight sync does not support a 4-bit quantized base under FSDP2. Train "
+                                "with DeepSpeed or on a single device, or load the base in full precision."
+                            )
                 vllm_kwargs = {
                     "model": model.name_or_path,
                     "tensor_parallel_size": self.vllm_tensor_parallel_size,
@@ -474,7 +486,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.
@@ -831,10 +842,10 @@ def _move_model_to_vllm_inner(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
@@ -852,10 +863,10 @@ def _move_model_to_vllm_inner(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))])
 
     def _sync_fsdp1_params_to_vllm(self, module: nn.Module, prefix: str = "", visited=None):
         """Memory-efficient post-order traversal of FSDP modules to extract full parameters and sync with vLLM."""
@@ -879,10 +890,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,60 @@ 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
+
+
+def _check_quantization_supported(model: nn.Module, fsdp_version: int) -> None:
+    """
+    Raise when the model's quantization cannot survive the dense weight push into vLLM.
+
+    The engine is built dense on purpose: the base may be 4-bit, but every weight is dequantized to the model dtype on
+    the way out (see `_dense_param_data`), and building with `quantization="bitsandbytes"` would allocate packed
+    `[out_features, in_features // 2]` weights that reject that push. Two combinations cannot be served that way and
+    are refused here rather than at the first sync. See https://github.com/huggingface/trl/issues/4973.
+
+    Args:
+        model (`torch.nn.Module`):
+            Model whose modules are inspected for bitsandbytes layers.
+        fsdp_version (`int`):
+            FSDP major version in use, `0` when FSDP is not enabled. Only version 2 is refused with a 4-bit base,
+            because it reads weights from `state_dict()`, which returns plain tensors whose `quant_state` is already
+            gone by the time `_dense_param_data` runs. FSDP1 reads through `summon_full_params` and keeps it.
+
+    Raises:
+        `ValueError`: if the model holds 8-bit layers, or 4-bit layers while `fsdp_version` is 2.
+    """
+    if not is_bitsandbytes_available():
+        return
+    for _, module in model.named_modules():
+        if isinstance(module, bnb.nn.Linear8bitLt):
+            raise ValueError("vLLM does not support in-flight 8-bit quantization.")
+        if isinstance(module, bnb.nn.Linear4bit) and fsdp_version == 2:
+            raise ValueError(
+                "vLLM weight sync does not support a 4-bit quantized base under FSDP2. Train with DeepSpeed or on a "
+                "single device, or load the base in full precision."
+            )
+
+
 class VLLMGeneration:
     """Handles vLLM-based generation for trainers.
 
@@ -334,14 +388,7 @@ def _init_vllm(self):
             # Ensure distributed rendezvous variables are set without colliding across concurrent runs
             ensure_master_addr_port()
 
-            quantization = None
-            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):
-                        raise ValueError("vLLM does not support in-flight 8-bit quantization.")
+            _check_quantization_supported(model, self._dist.fsdp_version)
 
             # Build LLM initialization kwargs
             self.llm = LLM(
@@ -359,7 +406,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:
@@ -403,7 +449,7 @@ def _iter_fsdp1_params(self, module: nn.Module, prefix: str = "", visited: set[s
                         continue  # skip FSDP subtrees already traversed
                     visited.add(full_name)
 
-                    yield full_name, param.data
+                    yield full_name, _dense_param_data(param)
 
     def _iter_fsdp2_params(self, module: nn.Module):
         """FSDP2-specific parameter iteration."""
@@ -464,7 +510,7 @@ def _iter_named_params(self):
                             continue
                         name = self._fix_param_name_to_vllm(name, extra_prefixes=["modules_to_save.default."])
 
-                        yield name, param.data
+                        yield name, _dense_param_data(param)
                 # Unmerge adapters while parameters are still gathered
                 model.unmerge_adapter()
                 # Parameters will automatically be repartitioned when exiting the context
@@ -476,7 +522,7 @@ def _iter_named_params(self):
                 for name, param in model.named_parameters():
                     name = self._fix_param_name_to_vllm(name)
                     with self._dist.gather_params([param]):
-                        yield name, param.data
+                        yield name, _dense_param_data(param)
 
     def sync_weights(self):
         """Synchronize model weights to vLLM.