From ad3d43aeea7134c758ce45ac5153d6c10ee1ad2c Mon Sep 17 00:00:00 2001 From: Nathan Azrak <42650258+nathan-az@users.noreply.github.com> Date: Mon, 29 Jan 2024 19:59:15 +0900 Subject: [PATCH] Make peft `bnb_4bit_compute_dtype` consistent with `torch_dtype` (#107) Co-authored-by: Nathan Azrak --- src/alignment/model_utils.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/src/alignment/model_utils.py b/src/alignment/model_utils.py index 52bf00f..d538fb2 100644 --- a/src/alignment/model_utils.py +++ b/src/alignment/model_utils.py @@ -39,11 +39,15 @@ def get_kbit_device_map() -> Dict[str, int] | None: return {"": get_current_device()} if torch.cuda.is_available() else None -def get_quantization_config(model_args) -> BitsAndBytesConfig | None: +def get_quantization_config(model_args: ModelArguments) -> BitsAndBytesConfig | None: if model_args.load_in_4bit: + compute_dtype = torch.float16 + if model_args.torch_dtype not in {"auto", None}: + compute_dtype = getattr(torch, model_args.torch_dtype) + quantization_config = BitsAndBytesConfig( load_in_4bit=True, - bnb_4bit_compute_dtype=torch.float16, # For consistency with model weights, we use the same value as `torch_dtype` which is float16 for PEFT models + bnb_4bit_compute_dtype=compute_dtype, bnb_4bit_quant_type=model_args.bnb_4bit_quant_type, bnb_4bit_use_double_quant=model_args.use_bnb_nested_quant, )