@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 8fe7fc8634

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 15, 04:45 UTC

[Quantization] Improve BitsAndBytesModelLoader (#20242)

Signed-off-by: Jee Jee Li <pandaleefree@gmail.com>

Jee Jee Li committed Jun 30, 2025 at 18:22 UTC 8fe7fc863481a7d48c6f5bcc7bb40b2c7ebb5476

1 file changed +72 -51

vllm/model_executor/model_loader/bitsandbytes_loader.py

+72 -51

@@ -20,8 +20,6 @@ from vllm.distributed import (get_tensor_model_parallel_rank,

20 get_tensor_model_parallel_world_size)

21 # yapf: enable

22 from vllm.logger import init_logger

23 -# yapf conflicts with isort for this block

24 -# yapf: disable

23 from vllm.model_executor.layers.linear import (LinearBase,

24 MergedColumnParallelLinear,

25 QKVParallelLinear,

@@ -39,6 +37,8 @@ from vllm.model_executor.utils import (get_packed_modules_mapping,

37 set_weight_attrs)

38 from vllm.platforms import current_platform

40 +# yapf conflicts with isort for this block

41 +

42 logger = init_logger(__name__)

@@ -54,11 +54,17 @@ class BitsAndBytesModelLoader(BaseModelLoader):

54 self.unsharded_weights_modules: list[str] = []

55 # Save the module names that are sharded by column.

56 self.column_sharded_weights_modules: list[str] = []

57 + # Modules whose weights might have fused on disk

58 + # we need their output_sizes to make shard in flight correctly with TP

59 + self.maybe_fused_weights_modules: dict[str, list[int]] = {}

60 # Store all module names (from transformers) that support

61 # BNB quantization.

62 self.target_modules: list[str] = []

63 # mapping weight names from transformers to vllm.

64 self.weight_mapper: Callable = lambda name: name

65 + self.pre_quant: bool = False

66 + self.load_8bit: bool = False

67 + self.is_pool_model: bool = False

69 def _get_weight_files(

70 self,

@@ -134,13 +140,14 @@ class BitsAndBytesModelLoader(BaseModelLoader):

140 return hf_weights_files, use_safetensors

141

142 def _hf_weight_iter(self, hf_weights_files, use_safetensors: bool):

137 - def _maybe_pool_model(module_name:str):

143 +

144 + def _maybe_pool_model(module_name: str):

145 # For pool model, we need to add the prefix `model.`

146 # for the weight name if possible.

147 if self.is_pool_model and self.target_modules[0]. \

148 startswith("model.") and not module_name.startswith(

149 "model."):

143 - return "model."+module_name

150 + return "model." + module_name

151

152 return module_name

153

@@ -159,8 +166,7 @@ class BitsAndBytesModelLoader(BaseModelLoader):

166 # mapping weight names from transformers to vllm while preserving

167 # original names.

168 mapped_name = self.weight_mapper(org_name)

162 - mapped_name=_maybe_pool_model(mapped_name)

163 -

169 + mapped_name = _maybe_pool_model(mapped_name)

170

171 yield org_name, mapped_name, param

172

@@ -168,8 +174,6 @@ class BitsAndBytesModelLoader(BaseModelLoader):

174 self,

175 model_name_or_path: str,

176 revision: Optional[str],

171 - pre_quant: bool,

172 - load_8bit: bool,

177 ) -> tuple[Generator[tuple[str, torch.Tensor], None, None], dict[str,

178 Any]]:

179 """Get an iterator to the model weights with bitsandbytes quantization,

@@ -192,8 +196,8 @@ class BitsAndBytesModelLoader(BaseModelLoader):

196

197 quant_state_dict: dict[str, Any] = {}

198

195 - if pre_quant:

196 - if load_8bit:

199 + if self.pre_quant:

200 + if self.load_8bit:

201 return self._quantized_8bit_generator(

202 hf_weights_files, use_safetensors,

203 quant_state_dict), quant_state_dict

@@ -390,10 +394,13 @@ class BitsAndBytesModelLoader(BaseModelLoader):

394 yield org_weight_name, processed_weight

395

396 def _get_bnb_target_modules(self, model: nn.Module) -> None:

393 -

397 + """

398 + Identify and collect all modules that support BitsAndBytes

399 + quantization.

400 + """

401 for name, module in model.named_modules():

395 - if (isinstance(module, LinearBase) and

396 - hasattr(module.quant_method, "quant_config")):

402 + if (isinstance(module, LinearBase)

403 + and hasattr(module.quant_method, "quant_config")):

404 if modules_info := self.modules_mapping.get_sub_modules(name):

405 # Map vllm's names to transformers's names.

406 rep_name, sub_modules = modules_info

@@ -409,29 +416,11 @@ class BitsAndBytesModelLoader(BaseModelLoader):

416 ), "vllm currently does not support BNB quantization for"

417 f" {type(model).__name__}"

418

412 - def load_weights(self, model: nn.Module, model_config: ModelConfig) -> None:

413 - if not hasattr(model, "load_weights"):

414 - raise AttributeError(

415 - "The required method 'load_weights' is not defined in class"

416 - f" {type(model).__name__}.")

417 -

418 - if not hasattr(model, "packed_modules_mapping"):

419 - raise AttributeError(

420 - f"Model {type(model).__name__} does not support BitsAndBytes "

421 - "quantization yet. No 'packed_modules_mapping' found.")

422 - self.is_pool_model=is_pooling_model(model)

423 -

424 - self.modules_mapping = ParamMapping(get_packed_modules_mapping(model))

425 -

426 - # For some models like Molmo, we need to use hf_to_vllm_mapper

427 - # to ensure correct loading of weights.

428 - if hf_to_vllm_mapper := getattr(model, "hf_to_vllm_mapper", None):

429 - self.weight_mapper = lambda name: hf_to_vllm_mapper._map_name(name)

430 -

431 - # Modules whose weights might have fused on disk

432 - # we need their output_sizes to make shard in flight correctly with TP

433 - self.maybe_fused_weights_modules: dict[str, list[int]] = {}

434 - self._get_bnb_target_modules(model)

419 + def _classify_module_sharding(self, model: nn.Module):

420 + """

421 + Categorize modules based on their weight sharding requirements

422 + for tensor parallelism.

423 + """

424 for name, module in model.named_modules():

425 # Some modules like `ReplicatedLinear` should not have their weights

426 # sharded. The reason for implementing it this way is to avoid new

@@ -449,19 +438,27 @@ class BitsAndBytesModelLoader(BaseModelLoader):

438 elif isinstance(module, (RowParallelLinear, )):

439 self.column_sharded_weights_modules.append(name)

440

452 - self.model_type = type(model).__name__

441 + def _verify_model_compatibility(self, model: nn.Module,

442 + model_config: ModelConfig) -> None:

443 + """

444 + Verify that the model is compatible with BitsAndBytes quantization.

445 + """

446 + if not hasattr(model, "load_weights"):

447 + raise AttributeError(

448 + "The required method 'load_weights' is not defined in class"

449 + f" {type(model).__name__}.")

450

454 - logger.info("Loading weights with BitsAndBytes quantization. "

455 - "May take a while ...")

451 + if not hasattr(model, "packed_modules_mapping"):

452 + raise AttributeError(

453 + f"Model {type(model).__name__} does not support BitsAndBytes "

454 + "quantization yet. No 'packed_modules_mapping' found.")

455

456 quant_config = getattr(model_config.hf_config, "quantization_config",

457 None)

459 -

460 - pre_quant = False

458 if quant_config is not None:

459 quant_method = quant_config.get("quant_method")

460 if quant_method == "bitsandbytes":

464 - pre_quant = True

461 + self.pre_quant = True

462 else:

463 raise ValueError(

464 f"BitsAndBytes loader does not support {quant_method} "

@@ -469,20 +466,43 @@ class BitsAndBytesModelLoader(BaseModelLoader):

466

467 # The quant_states in pre_quantized models cannot work with a split

468 # weight tensor. So TP does not work with pre_quantized bnb models.

472 - if pre_quant and get_tensor_model_parallel_world_size() > 1:

469 + if self.pre_quant and get_tensor_model_parallel_world_size() > 1:

470 raise ValueError(

471 "Prequant BitsAndBytes models with tensor parallelism is not "

472 "supported. Please try with pipeline parallelism.")

473 + if self.pre_quant:

474 + self.load_8bit = quant_config.get("load_in_8bit", False)

475 +

476 + def _initialize_loader_state(self, model: nn.Module,

477 + model_config: ModelConfig) -> None:

478 + """

479 + Initialize the loader's internal state based on the model and

480 + configuration.

481 + """

482 + self.is_pool_model = is_pooling_model(model)

483 + self.modules_mapping = ParamMapping(get_packed_modules_mapping(model))

484

477 - load_8bit = False

478 - if pre_quant:

479 - load_8bit = quant_config.get("load_in_8bit", False)

485 + # For some models like Molmo, we need to use hf_to_vllm_mapper

486 + # to ensure correct loading of weights.

487 + if hf_to_vllm_mapper := getattr(model, "hf_to_vllm_mapper", None):

488 + self.weight_mapper = lambda name: hf_to_vllm_mapper._map_name(name)

489

481 - qweight_iterator, quant_state_dict = (

482 - self._get_quantized_weights_iterator(model_config.model,

483 - model_config.revision,

484 - pre_quant, load_8bit))

490 + self._get_bnb_target_modules(model)

491 + self._classify_module_sharding(model)

492

493 + def load_weights(self, model: nn.Module,

494 + model_config: ModelConfig) -> None:

495 +

496 + self._verify_model_compatibility(model, model_config)

497 + self._initialize_loader_state(model, model_config)

498 +

499 + logger.info("Loading weights with BitsAndBytes quantization. "

500 + "May take a while ...")

501 + qweight_iterator, quant_state_dict = (

502 + self._get_quantized_weights_iterator(

503 + model_config.model,

504 + model_config.revision,

505 + ))

506 weights_to_load = {name for name, _ in model.named_parameters()}

507 loaded_weights = model.load_weights(qweight_iterator)

508 # Some models may have weights loading tracker unimplemented.

@@ -562,10 +582,11 @@ class BitsAndBytesModelLoader(BaseModelLoader):

582 offsets = torch.tensor(offsets).cpu()

583 set_weight_attrs(param, {"bnb_shard_offsets": offsets})

584

565 - if load_8bit:

585 + if self.load_8bit:

586 set_weight_attrs(

587 param, {"matmul_state": [None] * len(quant_states)})

588 torch.cuda.empty_cache()

589 +

590 def download_model(self, model_config: ModelConfig) -> None:

591 self._prepare_weights(model_config.model, model_config.revision)

592