Dequantize the bitsandbytes base before the vLLM weight push

tests/test_vllm_client_server.py

@@ -15,19 +15,24 @@
 import os
 import subprocess
 from types import SimpleNamespace
+from unittest.mock import patch
 
 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.distributed import DistributedBackend
+from trl.generation.vllm_client import VLLMClient, _dense_param_data, parse_logprobs
+from trl.generation.vllm_generation import _check_quantization_supported, 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 +42,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 +135,132 @@ 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)
+
+    def test_dense_param_data_dequantizes_a_4bit_parameter_with_its_quant_state(self):
+        # The fix itself: a `Params4bit` must go through `dequantize_4bit` with its own `quant_state`, and the packed
+        # buffer must never be what comes back. The kernel needs a GPU, so it is replaced by a stand-in that records
+        # its arguments and returns a dense tensor of the unpacked shape.
+        packed = torch.zeros(16, 1, dtype=torch.uint8)
+        quant_state = object()
+        param = bnb.nn.Params4bit(packed, requires_grad=False, quant_state=quant_state)
+        dense = torch.randn(4, 8)
+        calls = []
+
+        def fake_dequantize_4bit(data, state):
+            calls.append((data.data_ptr(), state))
+            return dense
+
+        with patch.object(bnb.functional, "dequantize_4bit", fake_dequantize_4bit):
+            out = _dense_param_data(param)
+        assert out is dense
+        assert calls == [(packed.data_ptr(), quant_state)]
+
+    def test_update_model_params_sends_dense_tensors_with_dense_metadata(self):
+        # `update_model_params` must describe and send the dequantized weight, not the packed `(16, 1)` uint8 storage
+        # a `Params4bit` holds. The kernel is replaced by a stand-in as above; the server call is captured.
+        packed = torch.zeros(16, 1, dtype=torch.uint8)
+        dense = torch.randn(4, 8)
+        model = nn.Module()
+        model.weight = bnb.nn.Params4bit(packed, requires_grad=False, quant_state=object())
+        model.bias = nn.Parameter(torch.zeros(4))
+        client = VLLMClient.__new__(VLLMClient)
+        sent = {}
+
+        def fake_update_named_params(metadata, named_params):
+            sent["metadata"] = metadata
+            sent["params"] = list(named_params)
+
+        with (
+            patch.object(bnb.functional, "dequantize_4bit", lambda data, state: dense),
+            patch.object(client, "update_named_params", fake_update_named_params),
+        ):
+            client.update_model_params(model)
+
+        assert sent["metadata"] == [("weight", "float32", [4, 8]), ("bias", "float32", [4])]
+        assert sent["params"][0][0] == "weight" and sent["params"][0][1] is dense
+        assert sent["params"][1][0] == "bias" and sent["params"][1][1].data_ptr() == model.bias.data.data_ptr()
+
+    @pytest.mark.parametrize(
+        ("module_factory", "fsdp_version", "fsdp_use_orig_params", "pre_quantized", "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, None, False, "4-bit quantized base under FSDP2"),
+            # Under FSDP1, `summon_full_params` keeps the `Params4bit` and its `quant_state` only with
+            # `use_orig_params=True`; with Accelerate's default `False` (or the unresolved `None`) it exposes the packed
+            # storage as a plain tensor, and that is what the sync would push.
+            (lambda: bnb.nn.Linear4bit(8, 8), 1, False, False, "FSDP1 with `use_orig_params=False`"),
+            (lambda: bnb.nn.Linear4bit(8, 8), 1, None, False, "FSDP1 with `use_orig_params=False`"),
+            # vLLM has never supported in-flight 8-bit, independent of the sharding strategy.
+            (lambda: bnb.nn.Linear8bitLt(8, 8), 0, None, False, "8-bit quantization"),
+            (lambda: bnb.nn.Linear8bitLt(8, 8), 1, True, False, "8-bit quantization"),
+            (lambda: bnb.nn.Linear8bitLt(8, 8), 2, None, False, "8-bit quantization"),
+            # A checkpoint saved already quantized makes colocated vLLM read its `quantization_config` and allocate
+            # packed weights, so the dense push cannot fill them whatever the sharding is.
+            (lambda: bnb.nn.Linear4bit(8, 8), None, None, True, "saved already quantized"),
+            (lambda: bnb.nn.Linear4bit(8, 8), 1, True, True, "saved already quantized"),
+            # Negative controls: every other combination reaches the dequantizing push and must be accepted.
+            (lambda: bnb.nn.Linear4bit(8, 8), 1, True, False, None),  # FSDP1 with the original params exposed
+            (lambda: bnb.nn.Linear4bit(8, 8), 0, None, False, None),  # no FSDP at all
+            (lambda: nn.Linear(8, 8), 2, None, False, None),  # dense base under FSDP2
+            (lambda: nn.Linear(8, 8), 1, False, False, None),  # dense base under FSDP1, default setting
+            (lambda: nn.Linear(8, 8), 0, None, False, None),  # dense base, no FSDP
+            (lambda: nn.Linear(8, 8), None, None, True, None),  # no bitsandbytes layer, so nothing to refuse
+            # `DistributedBackend.fsdp_version` is `None`, not `0`, when FSDP is off, so these are the values the
+            # guard actually receives in production. The `0` cases above only cover the documented sentinel.
+            (lambda: bnb.nn.Linear8bitLt(8, 8), None, None, False, "8-bit quantization"),
+            (lambda: bnb.nn.Linear4bit(8, 8), None, None, False, None),
+            (lambda: nn.Linear(8, 8), None, None, False, None),
+        ],
+    )
+    def test_check_quantization_supported(
+        self, module_factory, fsdp_version, fsdp_use_orig_params, pre_quantized, expected_message
+    ):
+        model = nn.Sequential(module_factory())
+        if expected_message is None:
+            _check_quantization_supported(model, fsdp_version, fsdp_use_orig_params, pre_quantized)  # must not raise
+        else:
+            with pytest.raises(ValueError, match=expected_message):
+                _check_quantization_supported(model, fsdp_version, fsdp_use_orig_params, pre_quantized)
+
+
+class TestDistributedBackendFsdpVersion(TrlTestCase):
+    """`fsdp_version` was added to Accelerate's FSDP plugin in 1.6.0; older plugins can only configure FSDP1."""
+
+    @staticmethod
+    def _accelerator(fsdp_plugin):
+        state = SimpleNamespace(deepspeed_plugin=None, fsdp_plugin=fsdp_plugin)
+        return SimpleNamespace(state=state)
+
+    def test_reads_fsdp_version_from_a_current_plugin(self):
+        plugin = SimpleNamespace(fsdp_version=2, use_orig_params=None)
+        with patch("trl.distributed.accelerate.__version__", "1.6.0"):
+            backend = DistributedBackend(self._accelerator(plugin))
+        assert backend.fsdp_version == 2 and backend.is_fsdp
+
+    def test_reports_fsdp1_for_a_plugin_that_predates_the_attribute(self):
+        # Accelerate < 1.6.0 exposes no `fsdp_version`; reading it would raise, and treating its absence as
+        # "no FSDP" would let a 4-bit base slip past the guard under FSDP1.
+        plugin = SimpleNamespace(use_orig_params=False)
+        with patch("trl.distributed.accelerate.__version__", "1.5.0"):
+            backend = DistributedBackend(self._accelerator(plugin))
+        assert backend.fsdp_version == 1 and backend.fsdp_use_orig_params is False
+
+    def test_reports_no_fsdp_without_a_plugin(self):
+        backend = DistributedBackend(self._accelerator(None))
+        assert backend.fsdp_version is None and backend.fsdp_use_orig_params is None and not backend.is_fsdp
+
+
 @pytest.mark.slow
 @require_torch_multi_accelerator
 @require_vllm

trl/distributed.py

@@ -16,6 +16,9 @@
 
 from contextlib import contextmanager
 
+import accelerate
+from packaging.version import Version
+
 
 class DistributedBackend:
     """Abstracts distributed backend specifics (DeepSpeed ZeRO, FSDP) behind a uniform API.
@@ -42,7 +45,17 @@ def __init__(self, accelerator):
         ds_plugin = accelerator.state.deepspeed_plugin
         fsdp_plugin = getattr(accelerator.state, "fsdp_plugin", None)
         self.zero_stage = ds_plugin.zero_stage if ds_plugin else 0
-        self.fsdp_version = getattr(fsdp_plugin, "fsdp_version", None) if fsdp_plugin else None
+        if fsdp_plugin is None:
+            self.fsdp_version = None
+        elif Version(accelerate.__version__) < Version("1.6.0"):
+            # The plugin gained `fsdp_version` together with FSDP2 support in Accelerate 1.6.0; before that FSDP1 was
+            # the only version it could configure.
+            self.fsdp_version = 1
+        else:
+            self.fsdp_version = fsdp_plugin.fsdp_version
+        # FSDP1 only: whether `summon_full_params` exposes the original parameter objects (Accelerate's default is
+        # `False`, which yields plain tensors and drops the bitsandbytes `quant_state`).
+        self.fsdp_use_orig_params = fsdp_plugin.use_orig_params if fsdp_plugin else None
 
     @property
     def is_zero3(self) -> bool:

trl/experimental/online_dpo/online_dpo_trainer.py

@@ -41,7 +41,6 @@
     PreTrainedTokenizerBase,
     ProcessorMixin,
     TrainerCallback,
-    is_bitsandbytes_available,
 )
 from transformers.models.auto.modeling_auto import MODEL_FOR_IMAGE_TEXT_TO_TEXT_MAPPING_NAMES
 from transformers.trainer_utils import EvalPrediction
@@ -54,8 +53,10 @@
     maybe_apply_chat_template,
     prepare_multimodal_messages,
 )
+from ...distributed import DistributedBackend
 from ...extras.profiling import profiling_context
-from ...generation.vllm_client import VLLMClient
+from ...generation.vllm_client import VLLMClient, _dense_param_data
+from ...generation.vllm_generation import _check_quantization_supported
 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
@@ -68,10 +69,6 @@
     from transformers.trainer_pt_utils import nested_gather
 
 
-if is_bitsandbytes_available():
-    import bitsandbytes as bnb
-
-
 if is_peft_available():
     from peft import PeftConfig
 
@@ -430,6 +427,16 @@ def __init__(
                     "`pip install trl[vllm]` to use it."
                 )
 
+            # Both modes push dense weights at sync time (see `_dense_param_data`), so the check runs before either
+            # branch. Only colocate mode builds the engine from `model.name_or_path`, where a checkpoint saved already
+            # quantized makes vLLM allocate packed weights; `hf_quantizer.pre_quantized` records whether the checkpoint
+            # was one.
+            dist = DistributedBackend(self.accelerator)
+            pre_quantized = (
+                self.vllm_mode == "colocate" and model.hf_quantizer is not None and model.hf_quantizer.pre_quantized
+            )
+            _check_quantization_supported(model, dist.fsdp_version, dist.fsdp_use_orig_params, pre_quantized)
+
             if self.vllm_mode == "server":
                 if self.accelerator.is_main_process:
                     if args.vllm_server_base_url is not None:
@@ -453,14 +460,10 @@ 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
-                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):
-                            raise ValueError("vLLM does not support in-flight 8-bit quantization.")
+                # 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.
                 vllm_kwargs = {
                     "model": model.name_or_path,
                     "tensor_parallel_size": self.vllm_tensor_parallel_size,
@@ -474,7 +477,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.
@@ -753,7 +755,7 @@ def _sync_fsdp2_params_to_vllm(self, module: nn.Module):
         for name, param in module.state_dict().items():
             # When using PEFT, we need to recover the original parameter name
             name = name.removeprefix("base_model.model.").replace(".base_layer", "")
-            # Skip PEFT layers: they don’t exist in vLLM, and they are merged already.
+            # Skip PEFT layers: they don't exist in vLLM, and they are merged already.
             if is_peft_model(module) and module.prefix in name:
                 continue
             # When module to save, remove its prefix and discard the original module
@@ -763,6 +765,8 @@ def _sync_fsdp2_params_to_vllm(self, module: nn.Module):
 
             if param.is_cpu:
                 param = param.to(self.accelerator.device)
+            # No `_dense_param_data` here: `state_dict()` already dropped any `quant_state`, so a 4-bit base cannot be
+            # served on this path and `_check_quantization_supported` refuses it up front.
             param = param.full_tensor()
 
             if self.vllm_mode == "server" and self.accelerator.is_main_process:
@@ -822,7 +826,7 @@ def _move_model_to_vllm_inner(self):
                     for name, param in self.model.named_parameters():
                         # When using PEFT, we need to recover the original parameter name
                         name = name.removeprefix("base_model.model.").replace(".base_layer", "")
-                        # Skip PEFT layers: they don’t exist in vLLM, and they are merged already.
+                        # Skip PEFT layers: they don't exist in vLLM, and they are merged already.
                         if self.model.prefix in name:
                             continue
                         # When module to save, remove its prefix and discard the original module
@@ -831,10 +835,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 +856,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 +883,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"""