wassname/alignment-handbook
1
0
Fork 0
mirror of https://github.com/wassname/alignment-handbook.git synced 2026-08-01 12:20:05 +08:00
Code Issues Packages Projects Releases Wiki Activity

Make peft bnb_4bit_compute_dtype consistent with torch_dtype (#107)

Browse Source 
Co-authored-by: Nathan Azrak <nazrak@atlassian.com>
This commit is contained in:
Nathan Azrak
2024-01-29 11:59:15 +01:00
committed by GitHub
co-authored by Nathan Azrak
parent de7d8883cd
commit ad3d43aeea
1 changed files with 6 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/alignment/model_utils.py
+6 -2
View File
@@ -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,
)