diff --git a/tests/test_vllm_client_server.py b/tests/test_vllm_client_server.py index d697692ca91..e00ee724e94 100644 --- a/tests/test_vllm_client_server.py +++ b/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,47 @@ 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 + # `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, "8-bit quantization"), + (lambda: bnb.nn.Linear4bit(8, 8), None, None), + (lambda: nn.Linear(8, 8), None, None), + ], + ) + 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 diff --git a/trl/experimental/online_dpo/online_dpo_trainer.py b/trl/experimental/online_dpo/online_dpo_trainer.py index 134670a49e8..4eca5e67a99 100644 --- a/trl/experimental/online_dpo/online_dpo_trainer.py +++ b/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""" diff --git a/trl/generation/vllm_generation.py b/trl/generation/vllm_generation.py index 31d21929b50..68d949a68ea 100644 --- a/trl/generation/vllm_generation.py +++ b/trl/generation/vllm_generation.py @@ -108,6 +108,61 @@ 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) -> 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`, *optional*): + FSDP major version in use, `None` when FSDP is not enabled, which is what `DistributedBackend` reports. + 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 +389,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 +407,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 +450,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 +511,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 +523,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.