Repository: vllm-project/vllm
File: vllm/model_executor/model_loader/bitsandbytes_loader.py (Python)

Code before:
        Verify that the model is compatible with BitsAndBytes quantization.
        """
        if not hasattr(model, "load_weights"):
            raise AttributeError(
                "The required method 'load_weights' is not defined in class"
                f" {type(model).__name__}."
            )

        if not hasattr(model, "packed_modules_mapping"):
            raise AttributeError(
                f"Model {type(model).__name__} does not support BitsAndBytes "
                "quantization yet. No 'packed_modules_mapping' found."
            )

        quant_config = getattr(model_config.hf_config, "quantization_config", None)
        if quant_config and (quant_method := quant_config.get("quant_method")):
            if quant_method == "bitsandbytes":
                self.pre_quant = True
            else:
                raise ValueError(
                    f"BitsAndBytes loader does not support {quant_method} quantization"
                )

        # The quant_states in bnb models cannot work with a split
        # weight tensor. So TP does not work with bnb models.
        if get_tensor_model_parallel_world_size() > 1:
            raise ValueError(
                "BitsAndBytes models with tensor parallelism is not "
                "supported. Please try with pipeline parallelism."
            )
        if quant_config and self.pre_quant:
            self.load_8bit = quant_config.get("load_in_8bit", False)

    def _initialize_loader_state(
        self, model: nn.Module, model_config: ModelConfig
    ) -> None:
        """
        Initialize the loader's internal state based on the model and
        configuration.
        """
        self.is_pool_model = is_pooling_model(model)
        self.modules_mapping = ParamMapping(get_packed_modules_mapping(model))

        if is_moe_model(model):
            self.expert_params_mapping = get_moe_expert_mapping(model)
            if not self.expert_params_mapping:
                raise AttributeError(
                    f"MoE Model {type(model).__name__} does not support "
                    "BitsAndBytes quantization yet. Ensure this model has "
                    "'get_expert_mapping' method."
                )

Review comment (performance):
Wait, pre-quant means loading from pre-quantized checkpoint like `unsloth/Qwen3-0.6B-unsloth-bnb-4bit`. This change will totally disable tp support with BNB online quantization:
```
python examples/offline_inference/basic/generate.py --model /mnt/data0/LLM/Qwen3-0.6B --language-model-only --quantization bitsandbytes -tp 2
```
**Main branch**
```
(EngineCore_DP0 pid=1411621) (Worker pid=1411627) (Worker_TP0 pid=1411627) INFO 03-03 18:06:24 [gpu_model_runner.py:4333] Model loading took 0.27 GiB memory and 1.173136 seconds
(EngineCore_DP0 pid=1411621) (Worker pid=1411627) (Worker_TP0 pid=1411627) INFO 03-03 18:06:31 [backends.py:916] Using cache directory: /home/mozf/.cache/vllm/torch_compile_cache/4dc805e8cc/rank_0_0/backbone for vLLM's torch.compile
(EngineCore_DP0 pid=1411621) (Worker pid=1411627) (Worker_TP0 pid=1411627) INFO 03-03 18:06:31 [backends.py:976] Dynamo bytecode transform time: 6.64 s
(EngineCore_DP0 pid=1411621) (Worker pid=1411630) (Worker_TP1 pid=1411630) INFO 03-03 18:06:38 [backends.py:350] Cache the graph of compile range (1, 8192) for later use
(EngineCore_DP0 pid=1411621) (Worker pid=1411627) (Worker_TP0 pid=1411627) INFO 03-03 18:06:38 [backends.py:350] Cache the graph of compile range (1, 8192) for later use
(EngineCore_DP0 pid=1411621) (Worker pid=1411627) (Worker_TP0 pid=1411627) INFO 03-03 18:06:44 [backends.py:366] Compiling a graph for compile range (1, 8192) takes 11.35 s
(EngineCore_DP0 pid=1411621) (Worker pid=1411627) (Worker_TP0 pid=1411627) INFO 03-03 18:06:44 [monitor.py:35] torch.compile takes 19.17 s in total
(EngineCore_DP0 pid=1411621) (Worker pid=1411627) (Worker_TP0 pid=1411627) INFO 03-03 18:06:44 [decorators.py:580] saving AOT compiled function to /home/mozf/.cache/vllm/torch_compile_cache/torch_aot_compile/5168fd1d484c8eb2809870fa3edce742932257929fa1cef1e3303c41dd36f8aa/rank_0_0/model
(EngineCore_DP0 pid=1411621) (Worker pid=1411630) (Worker_TP1 pid=1411630) INFO 03-03 18:06:44 [decorators.py:580] saving AOT compiled function to /home/mozf/.cache/vllm/torch_compile_cache/torch_aot_compile/5168fd1d484c8eb2809870fa3edce742932257929fa1cef1e3303c41dd36f8aa/rank_1_0/model
(EngineCore_DP0 pid=1411621) (Worker pid=1411627) (Worker_TP0 pid=1411627) INFO 03-03 18:06:44 [decorators.py:588] saved AOT compiled function to /home/mozf/.cache/vllm/torch_compile_cache/torch_aot_compile/5168fd1d484c8eb2809870fa3edce742932257929fa1cef1e3303c41dd36f8aa/rank_0_0/model
(EngineCore_DP0 pid=1411621) (Worker pid=1411630) (Worker_TP1 pid=1411630) INFO 03-03 18:06:45 [decorators.py:588] saved AOT compiled function to /home/mozf/.cache/vllm/torch_compile_cache/torch_aot_compile/5168fd1d484c8eb2809870fa3edce742932257929fa1cef1e3303c41dd36f8aa/rank_1_0/model
(EngineCore_DP0 pid=1411621) (Worker pid=1411627) (Worker_TP0 pid=1411627) INFO 03-03 18:06:46 [gpu_worker.py:423] Available KV cache memory: 20.3 GiB
(EngineCore_DP0 pid=1411621) INFO 03-03 18:06:54 [kv_cache_utils.py:1314] GPU KV cache size: 380,192 tokens
(EngineCore_DP0 pid=1411621) INFO 03-03 18:06:54 [kv_cache_utils.py:1319] Maximum concurrency for 40,960 tokens per request: 9.28x
Capturing CUDA graphs (mixed prefill-decode, PIECEWISE): 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████| 51/51 [00:05<00:00,  9.50it/s]
Capturing CUDA graphs (decode, FULL): 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 35/35 [00:03<00:00, 10.97it/s]
(EngineCore_DP0 pid=1411621) (Worker pid=1411627) (Worker_TP0 pid=1411627) INFO 03-03 18:07:04 [gpu_model_runner.py:5365] Graph capturing finished in 10 secs, took 1.35 GiB
(EngineCore_DP0 pid=1411621) INFO 03-03 18:07:04 [core.py:282] init engine (profile, create kv cache, warmup model) took 39.76 seconds
INFO 03-03 18:07:05 [llm.py:382] Supported tasks: ['generate']
WARNING 03-03 18:07:05 [model.py:1354] Default vLLM sampling parameters have been overridden by the model's `generation_config.json`: `{'temperature': 0.6, 'top_k': 20, 'top_p': 0.95}`. If this is not intended, please relaunch vLLM instance with `--generation-config vllm`.
Rendering prompts: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 4/4 [00:00<00:00, 157.40it/s]
Processed prompts: 100%|█████████████████████████████████████████████████████████████████████████████████████████████| 4/4 [00:00<00:00, 18.52it/s, est. speed input: 101.96 toks/s, output: 296.60 toks/s]
--------------------------------------------------
Prompt: 'Hello, my name is'
Generated text: " Lina. I'm 16. I'm in the 2nd"
--------------------------------------------------
Prompt: 'The president of the United States is'
Generated text: ' the president of the United States, and the president of the United States is the'
--------------------------------------------------
Prompt: 'The capital of France is'
Generated text: ' _______.\n\nOptions: A. Paris B. Rome C. Athens D.'
--------------------------------------------------
Prompt: 'The future of AI is'
Generated text: " not just about the technology itself; it's about how we use it. That"
--------------------------------------------------
```
**PR**
```
(EngineCore_DP0 pid=1409343) (Worker pid=1409352) (Worker_TP1 pid=1409352) ERROR 03-03 18:04:38 [multiproc_executor.py:800] ValueError: BitsAndBytes models with tensor parallelism is not supported. Please try with pipeline parallelism.
```

Code after:
        Verify that the model is compatible with BitsAndBytes quantization.
        """
        if not hasattr(model, "load_weights"):
            raise AttributeError(
                "The required method 'load_weights' is not defined in class"
                f" {type(model).__name__}."
            )

        if not hasattr(model, "packed_modules_mapping"):
            raise AttributeError(
                f"Model {type(model).__name__} does not support BitsAndBytes "
                "quantization yet. No 'packed_modules_mapping' found."
            )

        quant_config = getattr(model_config.hf_config, "quantization_config", None)
        if quant_config and (quant_method := quant_config.get("quant_method")):
            if quant_method == "bitsandbytes":
                self.pre_quant = True
            else:
                raise ValueError(
                    f"BitsAndBytes loader does not support {quant_method} quantization"
                )

        # The quant_states in pre_quantized models cannot work with a split
        # weight tensor. So TP does not work with pre_quantized bnb models.
        if self.pre_quant and get_tensor_model_parallel_world_size() > 1:
            raise ValueError(
                "Prequant BitsAndBytes models with tensor parallelism is not "
                "supported. Please try with pipeline parallelism."
            )
        if quant_config and self.pre_quant:
            self.load_8bit = quant_config.get("load_in_8bit", False)

    def _initialize_loader_state(
        self, model: nn.Module, model_config: ModelConfig
    ) -> None:
        """
        Initialize the loader's internal state based on the model and
        configuration.
        """
        self.is_pool_model = is_pooling_model(model)
        self.modules_mapping = ParamMapping(get_packed_modules_mapping(model))

        if is_moe_model(model):
            self.expert_params_mapping = get_moe_expert_mapping(model)
            if not self.expert_params_mapping:
                raise AttributeError(
                    f"MoE Model {type(model).__name__} does not support "
                    "BitsAndBytes quantization yet. Ensure this model has "
                    "'get_expert_mapping' method."
                )